From 04e4eb25ce77ce2c46cc7cae5c23562d36805324 Mon Sep 17 00:00:00 2001 From: "Pramodith (via Claude Code)" Date: Thu, 18 Jun 2026 14:46:22 +0000 Subject: [PATCH 01/10] feat(calib): top-k/top-p LM-head shortlist calibration harness Instrument test_dflash (DFLASH_TOPK_CALIB=1, chain verify path) to measure how large a draft-derived candidate shortlist M must be to cover the target's greedy argmax / top-k / top-p nucleus, for a candidate-restricted LM head (ggml_get_rows over the 248k-vocab head). Reports set-coverage@M and top-p nucleus MASS-coverage@M (the quality-relevant metric for sampling) on an M grid up to 65536. Findings (5 prompts, 2550 positions, Qwen3.6-27B Q4 target + Q4 3.6 draft): greedy viable (M~1k -> ~97%); sampling set-coverage poor at small M but nucleus MASS-coverage ~99.6% at M=4096 (large M is cheap: head ~6% of step traffic). Adds calib/{prep_prompts.py,run_calib.sh,aggregate.py} + docs writeup for scale-up (incl. B200 sm_100 build). Co-Authored-By: Claude Opus 4.8 (1M context) --- calib/.gitignore | 2 + calib/aggregate.py | 47 +++++++ calib/prep_prompts.py | 59 +++++++++ calib/run_calib.sh | 23 ++++ docs/topk-head-optimization.md | 223 +++++++++++++++++++++++++++++++++ server/test/test_dflash.cpp | 175 ++++++++++++++++++++++++++ 6 files changed, 529 insertions(+) create mode 100644 calib/.gitignore create mode 100644 calib/aggregate.py create mode 100644 calib/prep_prompts.py create mode 100755 calib/run_calib.sh create mode 100644 docs/topk-head-optimization.md diff --git a/calib/.gitignore b/calib/.gitignore new file mode 100644 index 000000000..709a3903a --- /dev/null +++ b/calib/.gitignore @@ -0,0 +1,2 @@ +*.bin +out/ diff --git a/calib/aggregate.py b/calib/aggregate.py new file mode 100644 index 000000000..4a0416cbc --- /dev/null +++ b/calib/aggregate.py @@ -0,0 +1,47 @@ +#!/usr/bin/env python3 +"""Position-weighted aggregate of per-prompt [topk-calib] tables. + + python3 calib/aggregate.py calib/out/calib.log +""" +import re, sys +from pathlib import Path + +def main(): + txt = Path(sys.argv[1]).read_text() + blocks = re.split(r'\[topk-calib\] positions=', txt)[1:] + set_metrics = ["greedy (target argmax)", "sampling top_k=8", + "sampling top_k=20", "sampling top_p=0.95 nucleus"] + set_agg = {m: {} for m in set_metrics} # metric -> {M: weighted%} + mass_agg = {} # M -> weighted% + totpos = 0; nblk = 0 + for b in blocks: + pos = int(re.match(r'(\d+)', b).group(1)); totpos += pos; nblk += 1 + cur = None + for line in b.splitlines(): + if "set-coverage" in line: + for m in set_metrics: + if m in line: cur = m + if "MASS coverage" in line: + cur = "__mass__" + cm = re.search(r'coverage@M=(\d+)\s*:\s*([\d.]+)%', line) + if cm and cur and cur != "__mass__": + M, v = int(cm.group(1)), float(cm.group(2)) + set_agg[cur][M] = set_agg[cur].get(M, 0.0) + pos * v + mm = re.search(r'mass@M=(\d+)\s*:\s*([\d.]+)%', line) + if mm: + M, v = int(mm.group(1)), float(mm.group(2)) + mass_agg[M] = mass_agg.get(M, 0.0) + pos * v + print(f"prompts={nblk} total_positions={totpos}\n") + Ms = sorted({M for d in set_agg.values() for M in d} | set(mass_agg)) + hdr = "M".ljust(8) + "".join(f"{M:>9}" for M in Ms) + print(hdr) + for m in set_metrics: + row = m[:26].ljust(8) + "".join( + f"{set_agg[m].get(M,0)/totpos:9.1f}" for M in Ms) + print(row) + if mass_agg: + print("\n-- top_p=0.95 nucleus MASS coverage (mean %) --") + print("mass".ljust(8) + "".join(f"{mass_agg.get(M,0)/totpos:9.2f}" for M in Ms)) + +if __name__ == "__main__": + main() diff --git a/calib/prep_prompts.py b/calib/prep_prompts.py new file mode 100644 index 000000000..3ffb9332a --- /dev/null +++ b/calib/prep_prompts.py @@ -0,0 +1,59 @@ +#!/usr/bin/env python3 +"""Tokenize calibration prompts to int32 .bin files for test_dflash. + +Pulls prompts from the same datasets the benchmark uses (HumanEval / GSM8K / +MATH-500), applies the Qwen3.6 chat template, and writes one packed-int32 file +per prompt into the output dir. Each file is a prompt test_dflash can decode. + + python3 calib/prep_prompts.py --n 200 --out calib --thinking off + +Deps: transformers, datasets, jinja2 (CPU only). Network access to HF for the +tokenizer + datasets (set HF_TOKEN to avoid rate limits). +""" +import argparse, struct, sys +from pathlib import Path + +DATASETS = [ + # (label, hf_id, config, split, field-extractor) + ("he", "openai/openai_humaneval", None, "test", lambda x: x["prompt"]), + ("gsm", "openai/gsm8k", "main", "test", lambda x: f"Question: {x['question']}\nAnswer: "), + ("math", "HuggingFaceH4/MATH-500", None, "test", lambda x: f"Problem: {x['problem']}\nSolution: "), +] + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--n", type=int, default=200, help="total prompts (split across datasets)") + ap.add_argument("--out", default="calib", help="output dir for .bin files") + ap.add_argument("--tokenizer", default="Qwen/Qwen3.6-27B") + ap.add_argument("--thinking", choices=["on", "off"], default="off") + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--max-prompt-tok", type=int, default=1024, + help="skip prompts longer than this (keeps verify ctx small)") + args = ap.parse_args() + + from datasets import load_dataset + from transformers import AutoTokenizer + tok = AutoTokenizer.from_pretrained(args.tokenizer, trust_remote_code=True) + out = Path(args.out); out.mkdir(parents=True, exist_ok=True) + + per = max(1, args.n // len(DATASETS)) + written = 0 + for label, hf_id, cfg, split, extract in DATASETS: + ds = load_dataset(hf_id, cfg, split=split).shuffle(seed=args.seed) + take = min(per, len(ds)) + for i in range(take): + raw = extract(ds[i]) + text = tok.apply_chat_template( + [{"role": "user", "content": raw}], tokenize=False, + add_generation_prompt=True, enable_thinking=(args.thinking == "on")) + ids = tok.encode(text, add_special_tokens=False) + if not ids or len(ids) > args.max_prompt_tok: + continue + with open(out / f"{label}_{i:04d}.bin", "wb") as f: + for t in ids: + f.write(struct.pack("&1 \ + | grep -E "topk-calib|out of memory|failed|CUDA error" +done diff --git a/docs/topk-head-optimization.md b/docs/topk-head-optimization.md new file mode 100644 index 000000000..5058af103 --- /dev/null +++ b/docs/topk-head-optimization.md @@ -0,0 +1,223 @@ +# Candidate-restricted LM head (top-k logits) — design, calibration, status + +**Branch:** `feat/topk-head-calib` · **Worktree:** `/workspace/lucebox-topk` +**Status (2026-06-18):** premise validated on a smoke sample; instrumentation extended to +sampling; full calibration run pending. No production kernel implemented yet. + +--- + +## 1. Motivation / the idea + +The target LM head is the matmul `logits = output.weight @ hidden` over the **full vocab**. +For Qwen3.6-27B the head is `output.weight = [n_embd=5120 × n_vocab=248320]`, quantized **Q6_K +≈ 1.0 GB**, i.e. **~6% of the 16.8 GB** of weights streamed per target forward. The 248k vocab +is unusually large, which makes the head a bigger slice than on a typical 32k-vocab model. + +**Idea:** instead of computing logits over all 248k vocab entries, compute them only over a +small **candidate shortlist** `S` of vocab indices, then argmax/sample within `S`. In ggml this +is clean because vocab is the *row* dimension (`ne1`) of `output.weight`: a restricted head is +`ggml_get_rows(output.weight, candidate_ids)` (gathers k quantized rows) followed by a small +matmul. Restricting to k≈256–1024 shrinks the head matmul to <0.5% of its size. + +**Where the shortlist comes from:** the DFlash draft already predicts a per-position +distribution (its hidden states projected through the target head, `project_hidden_to_topk`). +Its **top-M** tokens are the natural candidate set `S`. Calibration decides M. + +--- + +## 2. Greedy formulation + the correctness caveat + +Greedy spec-decode today is **exact / lossless**: it only commits a draft token when it equals +the target's true full-vocab argmax, and the rejection "bonus" is the true argmax. Output == +pure target-greedy decoding. + +A candidate-restricted head makes argmax run over `S` only. This stays exact **iff the true +argmax ∈ S**. Calibration ("greedy token always in draft top-M") bounds the failure rate, but +"always on a calibration set" ≠ "always at inference" — so this becomes **approximate greedy**, +not lossless. Must validate output quality, not just speed. (You cannot cheaply *verify* +exactness: confirming no out-of-`S` token beats the in-`S` max needs the full matmul.) + +### Realistic ceiling +- Head ≈ 6% of weight traffic; restricting it ≈ removes that 6% on the **verify** side. +- But the shortlist itself is produced by `project_hidden_to_topk`, which already runs one full + target-head matmul over the draft block. So you eliminate the *verify-side* head only. +- **Net realistic decode speedup ≈ ~5%** here (more than a typical model because vocab=248k). + Modest but real. Worth it only if calibration shows small M suffices. + +--- + +## 3. Extension to sampling (top-k / top-p) — and why it can be *exact* + +Generalize the condition: `S` must contain the verifier's **post-filter support**. The sampler +chain order (`server/src/common/sampler.h:6`) is: + +``` +rep_penalty → freq/pres → top_k → softmax(temp) → top_p → draw +``` + +The post-filter support is the target's **top-K_s** tokens (then the top-p nucleus within them). +So calibrate M such that **draft-top-M ⊇ target-top-K_s**. Greedy is the K_s=1 case. + +**Key insight — normalization:** `top_k` is applied *before* `softmax`. So if `top_k=K_s` is +finite, `softmax` normalizes over just those K_s kept tokens **in both the full and restricted +pipelines** — the global denominator never enters. Therefore, as long as `S ⊇ target-top-K_s`: +- restricted-head `top_k` keeps exactly the true top-K_s (they're the K_s globally-highest and + all in `S`), `softmax`/`top_p`/draw are identical ⇒ **sampling is EXACT (lossless)**. +- This covers Qwen3.6 defaults (`top_k=20, top_p=0.95`): just need `S ⊇ target top-20`. + +**Exceptions:** +- `top_k = 0` (pure top-p): `softmax` normalizes over the kept set; the missing tail inflates + probabilities and shrinks the nucleus ⇒ **approximate**. Over-cover (`S` carries ≥0.999 mass) + or correct the threshold with an estimate of the uncovered tail mass. +- pure temperature (no truncation): every token has nonzero prob ⇒ can't shortlist; needs full + head. (Fine — the optimization only pays off when truncation already discards the tail.) +- penalties run before `top_k`, so strictly the condition is on the *penalized* top-K_s; raw-logit + draft-top-M is a good superset proxy. + +--- + +## 4. Calibration methodology (the make-or-break test, no kernel needed) + +Measure whether the premise holds before building any kernel. Per chain-verify position +(exact alignment: `target_tok[i]` ↔ draft row `i+1`): +- **greedy:** rank of the target argmax within the draft's ranked candidate list. `coverage@M` + = fraction of positions with rank < M. +- **sampling:** the MAX draft-rank over the target's top-K_s set / top-p nucleus → the smallest + M with `draft-top-M ⊇ target-top-K_s`. `coverage@M` = fraction fully covered. + +**Decision rule:** if `coverage@M` reaches ~99.9% at small M (≈64–256 greedy, somewhat larger +for top_k=20), build the restricted head. If it only saturates at large M, the speedup +evaporates — report and stop. + +--- + +## 5. What's implemented (this branch) + +Instrumentation only — gated by env `DFLASH_TOPK_CALIB=1`, in the **chain** verify path of +`server/test/test_dflash.cpp` (run WITHOUT `--ddtree`; chain gives exact position alignment): +1. accumulators declared before the decode loop (`calib_hist`, `samp_hist_k8/k20/p95`); +2. transfer the draft's full per-position logits (`draft_sg.logits → draft_logits_buf`) — the + fast path otherwise only reads GPU argmax; +3. transfer the target's full logits (`sg.logits → verify_logits_buf`) in the batched path; +4. per-position rank histograms (greedy argmax rank; max draft-rank over target top-8 / top-20 / + top-p=0.95 nucleus at temp 1.0); +5. print `coverage@M` curves at end of run. + +Build: `cmake -B server/build -S server -G Ninja -DCMAKE_BUILD_TYPE=Release +-DCMAKE_CUDA_ARCHITECTURES=86 -DGGML_CUDA_NCCL=OFF` then `cmake --build server/build --target +test_dflash -j`. (NCCL OFF is required — the fork's `ncclCommInitAll` aborts on this box's driver.) + +Run: `calib/run_calib.sh` (uses `GPU=${GPU:-1}` — GPU 0 was occupied by another session; +6 tokenized prompts in `calib/*.bin`, generated with the Qwen3.6 tokenizer, thinking off). + +### Results so far (smoke: 1 prompt, 210 positions, greedy metric) +``` +coverage@k=1 : 74.3% (draft top-1 == target greedy ≈ per-token accept rate) +coverage@k=64 : 97.6% +coverage@k=256 : 99.5% +coverage@k=512 : 100% (mean rank 4.9, max rank 354) +``` +⇒ Greedy premise holds strongly: target greedy token within draft top-512 (~0.14% of vocab) on +this sample; a safe M≈1024 is likely lossless and still makes the head matmul ~free. Sampling +coverage numbers pending the full multi-prompt run with the extended binary. + +--- + +## 5b. Full calibration results (2026-06-18) + +5 prompts (code1/code2/gen1/gen2/math1; math2 dropped on a transient GPU OOM), +**2550 positions**, n_gen=256, chain path, Qwen3.6-27B Q4_K_M target + Q4_K_M 3.6 draft, +position-weighted `coverage@M` (= fraction of positions whose set is fully inside draft top-M): + +| coverage@M | greedy (argmax) | top_k=8 | top_k=20 (Qwen def) | top_p=0.95 | +|-----------:|----------------:|--------:|--------------------:|-----------:| +| 128 | 93.7% | 29.8% | 3.3% | 62.5% | +| 512 | 96.5% | 57.6% | 19.1% | 68.5% | +| 1024 | 97.5% | 69.3% | 35.0% | 71.8% | +| 4096 | 99.4% | 87.0% | **64.3%** | 78.4% | + +Per-prompt greedy@M=128 ranged 80%–99% (code/math high, "general" prompts lower); +tails are long (max greedy rank up to ~133k on a few low-confidence positions). + +**Verdict.** +- **Greedy: viable.** M≈1024 → ~97.5%, M≈128 → ~94%. A restricted head with M≈1k–2k is + near-lossless and still <1% of the 248k-wide matmul. Build it. +- **Sampling (top_k / top_p): NOT viable with this draft.** Exact `top_k` needs the draft + shortlist to contain the target's *entire* top-K_s; even M=4096 covers the full top-20 only + 64% of the time. The draft matches the target's **#1** token well but ranks the rest of the + nucleus very differently (the target's #15 can sit at draft-rank thousands). Covering the + nucleus needs an impractically large M, erasing the speedup. +- **Caveat:** this is the weak Q4 3.6 draft (same one that underperformed the benchmark). A + BF16/stronger draft would likely rank the nucleus more faithfully and could move the sampling + numbers; re-run the calibration with it before concluding. The structural asymmetry + (argmax easy, full-nucleus hard) will persist to some degree. + +## 5c. Large-M is cheap; MASS-coverage reopens sampling (2026-06-18) + +The head scales linearly in M and is only ~6% of step traffic, so large M stays cheap: +M=8192 → head ~0.2% of step (~5.9% speedup), M=32768 → ~0.8% (~5.2%), M=65536 → ~1.6% (~4.4%). +So we can afford M far larger than first assumed. + +More importantly, **set-coverage is the wrong metric for sampling quality** — missing a low-prob +nucleus token barely distorts the draw. The right metric is **covered probability MASS**. The +instrumentation now reports, for the top_p=0.95 nucleus, the mean fraction of nucleus *mass* +inside draft-top-M, plus the count of positions with <99% mass covered. Extended M grid to 65536. + +Smoke (1 prompt, math1, 225 positions) — mass coverage: +``` +mass@M=256 : 98.93% (32/225 positions <99% covered) +mass@M=4096 : 99.60% ( 3/225 positions <99%) +mass@M=65536 : 99.9993% ( 0/225 positions <99%) +``` +⇒ At M≈4096–8192 (still ~5.6–5.9% speedup) the per-position sampling distortion is <0.5% of +mass; M=65536 is lossless-in-practice. **Sampling is likely viable as an approximate-but-faithful +mode after all** — pending the large-prompt run to confirm the worst-case tail. (Set-coverage of +the full top-20 also reaches ~100% by M=65536 on this prompt.) + +### Tooling for scale-up +- `calib/prep_prompts.py --n N` — tokenize N prompts from HumanEval/GSM8K/MATH-500 (Qwen3.6 chat + template) into `calib/*.bin`. Deps: transformers, datasets, jinja2. +- `calib/run_calib.sh` — `GPU=0 NGEN=256 bash calib/run_calib.sh | tee calib/out/calib.log` + (chain path, no `--ddtree`; one table per prompt). +- `calib/aggregate.py calib/out/calib.log` — position-weighted aggregate (set + mass coverage). +- CPU note: each position does 2× nth_element+sort over the 248k vocab (host-side), so very large + prompt counts are CPU-bound regardless of GPU; parallelize across prompts if needed. + +### B200 build (Blackwell sm_100, CUDA ≥12.8) +``` +git fetch && git checkout feat/topk-head-calib +git submodule update --init --recursive +cmake -B server/build -S server -G Ninja -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_CUDA_ARCHITECTURES=100 -DGGML_CUDA_NCCL=OFF +cmake --build server/build --target test_dflash -j +# place target + draft GGUFs under server/models/ and server/models/draft/ +python3 calib/prep_prompts.py --n 300 +GPU=0 NGEN=256 bash calib/run_calib.sh | tee calib/out/calib.log +python3 calib/aggregate.py calib/out/calib.log +``` + +## 6. Next steps +1. **[done]** Full calibration — see §5b. Greedy validated; sampling needs a better draft. +2. Implement the restricted head for the **greedy** path in `build_qwen35_graph` / the verify graph + (`server/src/qwen35/graph_builders.cpp`, LM head at ~line 360): plumb candidate ids + (draft top-M) → `ggml_get_rows(output.weight, ids)` → small matmul → argmax/sample over the + k logits → map index back to vocab id. Reuse the draft's `project_hidden_to_topk` output as + the candidate set (it's already computed). +3. Validate quality vs full-head greedy (token-exact match rate) and measure real decode speedup. +4. Sampling is **on hold** pending a better draft (§5b): the Q4 3.6 draft can't cover the + target's top-k/top-p nucleus at a useful M. Before revisiting, re-run §5b's calibration with + a BF16/stronger draft. If coverage improves: keep `top_k` finite for exactness (chain order + makes it lossless when `draft-top-M ⊇ target-top-K_s`), and gate pure-top_p / pure-temperature + to the full head (or over-cover with a tail-mass correction). + +## 7. Key code references +- LM head matmul + argmax: `server/src/qwen35/graph_builders.cpp:360-365` (chain), `:432` (tree), + `build_lm_head_projection_step` `:446`. +- Candidate source: `Qwen35DFlashTarget::project_hidden_to_topk` `server/src/qwen35/qwen35_dflash_target.cpp:614`; + `extract_draft_topk` `server/src/common/ddtree.cpp:12`. +- Greedy chain verify + acceptance: `server/test/test_dflash.cpp` (`!seq_verify` path ~3692, accept ~3800); + generic loop `server/src/common/dflash_spec_decode.cpp:196-207` (layer-split). +- Sampler chain: `server/src/common/sampler.{h,cpp}` (`needs_logit_processing` `sampler.h:38`, + `sample_logits`). +- Sampled-verify (server, single-GPU): `server/src/qwen35/qwen35_backend.cpp:1741-1779` (gate), + `:2033-2051` (sampled acceptance walk). diff --git a/server/test/test_dflash.cpp b/server/test/test_dflash.cpp index 25f885225..da54394cc 100644 --- a/server/test/test_dflash.cpp +++ b/server/test/test_dflash.cpp @@ -3073,6 +3073,37 @@ int main(int argc, char ** argv) { return ok; }; + // ── top-k LM-head calibration (DFLASH_TOPK_CALIB=1) ────────────────── + // Premise test for the candidate-restricted LM head: at each chain verify + // position, record the rank of the target's full-vocab greedy argmax token + // within the draft's ranked candidate list for the SAME position (draft row + // i+1 predicts the token that target_tok[i] predicts). coverage@k = fraction + // of positions whose target greedy token falls in the draft's top-k. Chain + // path only (exact position alignment); run WITHOUT --ddtree. + static const bool g_topk_calib = [](){ + const char * e = std::getenv("DFLASH_TOPK_CALIB"); + return e != nullptr && std::strcmp(e, "0") != 0; + }(); + std::vector calib_hist; // greedy: rank (0 = top-1) -> count + int64_t calib_total = 0, calib_rank_sum = 0; + int calib_max_rank = 0; + // Sampling-coverage: histogram of the MAX draft-rank over the target's + // top-Ks set (and top-p=0.95 nucleus). coverage@M = fraction of positions + // whose entire target top-Ks / nucleus is within draft top-M. With the + // top_k -> softmax -> top_p chain order, draft-top-M ⊇ target-top-Ks makes + // top_k(+top_p) sampling exact under the restricted head. + std::vector samp_hist_k8, samp_hist_k20, samp_hist_p95; + // M grid for coverage curves (also the cap MCAP = last entry = max rank we + // resolve; tokens beyond MCAP count as "not covered" for every M). + static const std::vector calib_Mgrid = { + 1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768, 65536}; + const int calib_MCAP = calib_Mgrid.back(); + // top-p=0.95 nucleus MASS coverage: per-M sum of (covered nucleus prob mass) + // and count of positions where covered mass < 0.99 (the quality-relevant + // metric: missing a low-prob nucleus token barely distorts sampling). + std::vector massp95_sum(calib_Mgrid.size(), 0.0); + std::vector massp95_bad(calib_Mgrid.size(), 0); + while (n_generated < n_gen) { const int need_commit_budget = n_gen - n_generated; @@ -3243,6 +3274,15 @@ int main(int argc, char ** argv) { // replay see the correct prefix. draft_tok[0] = last_tok; + // Calibration: the chain fast path leaves draft_logits_buf unpopulated + // (it reads GPU argmax). Pull the draft's full per-position logits so we + // can rank the target's greedy token against them after verify. The + // bridge path already filled draft_logits_buf above. + if (g_topk_calib && !draft_hidden_bridge) { + ggml_backend_tensor_get(draft_sg.logits, draft_logits_buf.data(), 0, + sizeof(float) * (size_t)vocab * q_len); + } + // DDTree top-K extraction. Positions 1..q_len-1 of the draft output // are the per-position distributions for the block-diffusion tree // (position 0 is the root/bonus slot and is fixed to last_tok). Only @@ -3745,6 +3785,14 @@ int main(int argc, char ** argv) { ggml_backend_tensor_get(sg.argmax_tokens, target_tok.data(), 0, sizeof(int32_t) * q_len); + // Calibration: the batched path normally reads only the GPU argmax. + // Pull the full target logits so we can find the target's top-K_s / + // top-p nucleus for the sampling-coverage metric. (seq_verify path + // already fills verify_logits_buf per position.) + if (g_topk_calib) { + ggml_backend_tensor_get(sg.logits, verify_logits_buf.data(), 0, + sizeof(float) * (size_t)vocab * q_len); + } } else { // Sequential verify: q_len independent single-token decodes. // Each call writes K/V at slot committed+i and advances SSM by 1. @@ -3785,6 +3833,101 @@ int main(int argc, char ** argv) { auto T_verify_logits = sync_us(); tt_verify_logits += std::chrono::duration(T_verify_logits - T_verify_compute).count(); + // Calibration: for each chain verify position, rank the target's tokens + // within draft row (i+1)'s ranked logits (rank 0 == draft top-1). + // - greedy : rank of the target argmax -> calib_hist + // - sampling : MAX draft-rank over the target's top-Ks / top-p nucleus + // -> samp_hist_* (the M needed so draft-top-M covers it) + if (g_topk_calib) { + // Stamped top-MCAP ranking avoids an O(vocab log vocab) full sort per + // position: nth_element+sort only the top-MCAP, and a per-token stamp + // marks validity so we never O(vocab)-reset draft_rank. Tokens beyond + // MCAP are treated as rank == MCAP (uncovered for every M in the grid). + static std::vector draft_rank, d_idx, t_idx; + static std::vector draft_stamp; + static int64_t calib_stamp = 0; + if ((int)draft_rank.size() != vocab) { + draft_rank.assign(vocab, 0); draft_stamp.assign(vocab, 0); + d_idx.resize(vocab); t_idx.resize(vocab); + } + const int mcap = std::min(calib_MCAP, vocab); + auto bump = [](std::vector & h, int r) { + if ((int)h.size() <= r) h.resize(r + 1, 0); + h[r]++; + }; + static std::vector> nucleus; // (draft_rank, prob) + for (int i = 0; i < q_len - 1; i++) { + const float * drow = draft_logits_buf.data() + (size_t)(i + 1) * vocab; + const float * trow = verify_logits_buf.data() + (size_t)i * vocab; + + // ── draft top-MCAP ranks (stamped) ── + for (int v = 0; v < vocab; v++) d_idx[v] = v; + if (mcap < vocab) + std::nth_element(d_idx.begin(), d_idx.begin() + mcap, d_idx.end(), + [&](int a, int b){ return drow[a] > drow[b]; }); + std::sort(d_idx.begin(), d_idx.begin() + mcap, + [&](int a, int b){ return drow[a] > drow[b]; }); + ++calib_stamp; + for (int r = 0; r < mcap; r++) { + draft_rank[d_idx[r]] = r; draft_stamp[d_idx[r]] = calib_stamp; + } + auto rank_of = [&](int tok)->int { + return (draft_stamp[tok] == calib_stamp) ? draft_rank[tok] : calib_MCAP; + }; + + // ── greedy: rank of target argmax ── + const int grank = rank_of(target_tok[i]); + bump(calib_hist, grank); + calib_total++; calib_rank_sum += grank; + if (grank > calib_max_rank) calib_max_rank = grank; + + // ── target softmax(temp=1) over full vocab + sorted top-MCAP ── + double maxz = trow[0]; + for (int v = 1; v < vocab; v++) if (trow[v] > maxz) maxz = trow[v]; + double Z = 0.0; + for (int v = 0; v < vocab; v++) Z += std::exp((double)trow[v] - maxz); + for (int v = 0; v < vocab; v++) t_idx[v] = v; + if (mcap < vocab) + std::nth_element(t_idx.begin(), t_idx.begin() + mcap, t_idx.end(), + [&](int a, int b){ return trow[a] > trow[b]; }); + std::sort(t_idx.begin(), t_idx.begin() + mcap, + [&](int a, int b){ return trow[a] > trow[b]; }); + + // ── top-k set coverage (max draft-rank over target top-Ks) ── + int mr8 = 0, mr20 = 0; + for (int j = 0; j < 20 && j < mcap; j++) { + const int dr = rank_of(t_idx[j]); + if (j < 8 && dr > mr8) mr8 = dr; + if (dr > mr20) mr20 = dr; + } + bump(samp_hist_k8, mr8); + bump(samp_hist_k20, mr20); + + // ── top-p=0.95 nucleus: set coverage + MASS coverage per M ── + nucleus.clear(); + double cum = 0.0; int mrp = 0; + for (int j = 0; j < mcap; j++) { + const int tok = t_idx[j]; + const double prob = std::exp((double)trow[tok] - maxz) / Z; + const int dr = rank_of(tok); + nucleus.emplace_back(dr, prob); + if (dr > mrp) mrp = dr; + cum += prob; + if (cum >= 0.95) break; + } + bump(samp_hist_p95, mrp); + const double tot = cum > 0 ? cum : 1.0; + for (size_t mi = 0; mi < calib_Mgrid.size(); mi++) { + const int M = calib_Mgrid[mi]; + double covered = 0.0; + for (const auto & nt : nucleus) if (nt.first < M) covered += nt.second; + const double frac = covered / tot; + massp95_sum[mi] += frac; + if (frac < 0.99) massp95_bad[mi]++; + } + } + } + std::printf("[step %d] committed=%d last_tok=%d\n", n_draft_steps, committed, last_tok); // 5) Greedy longest-prefix accept with standard spec-decoding comparison. @@ -4027,6 +4170,38 @@ int main(int argc, char ** argv) { n_draft_steps++; } + if (g_topk_calib && calib_total > 0) { + std::printf("\n[topk-calib] positions=%lld mean_rank=%.3f max_rank=%d (vocab=%d)\n", + (long long)calib_total, + (double)calib_rank_sum / (double)calib_total, + calib_max_rank, vocab); + // Set-coverage: fraction of positions whose ENTIRE set is within draft top-M. + auto cov = [&](const char * label, const std::vector & h) { + std::printf("[topk-calib] --- %s: set-coverage (entire set in draft top-M) ---\n", label); + for (int k : calib_Mgrid) { + int64_t c = 0; + for (int r = 0; r < k && r < (int)h.size(); r++) c += h[r]; + std::printf("[topk-calib] coverage@M=%-6d : %8.4f%%\n", + k, 100.0 * (double)c / (double)calib_total); + } + }; + cov("greedy (target argmax)", calib_hist); + cov("sampling top_k=8", samp_hist_k8); + cov("sampling top_k=20 (Qwen def)", samp_hist_k20); + cov("sampling top_p=0.95 nucleus", samp_hist_p95); + // Mass-coverage: mean fraction of the top-p=0.95 nucleus PROBABILITY MASS + // captured by draft top-M (the quality-relevant metric for sampling — a + // missed low-prob nucleus token barely distorts the draw), plus the count + // of positions where <99% of the mass is covered. + std::printf("[topk-calib] --- top_p=0.95 nucleus: MASS coverage (mean over positions) ---\n"); + for (size_t mi = 0; mi < calib_Mgrid.size(); mi++) { + std::printf("[topk-calib] mass@M=%-6d : %8.4f%% (positions <99%% covered: %lld/%lld)\n", + calib_Mgrid[mi], + 100.0 * massp95_sum[mi] / (double)calib_total, + (long long)massp95_bad[mi], (long long)calib_total); + } + } + auto t_gen1 = std::chrono::steady_clock::now(); double gen_s = std::chrono::duration(t_gen1 - t_gen0).count(); double tps = n_generated / std::max(1e-9, gen_s); From cda8468829dc60f05d46541e157207fec5455caf Mon Sep 17 00:00:00 2001 From: "Pramodith (via Claude Code)" Date: Thu, 18 Jun 2026 15:21:06 +0000 Subject: [PATCH 02/10] avoid redundant logits read --- server/scripts/bench_llm.py | 6 +-- server/src/qwen35/qwen35_dflash_target.cpp | 31 ++++++++++-- server/test/test_dflash.cpp | 56 +++++++++++++++++----- 3 files changed, 75 insertions(+), 18 deletions(-) diff --git a/server/scripts/bench_llm.py b/server/scripts/bench_llm.py index 4722a7f3f..b8b31c91d 100644 --- a/server/scripts/bench_llm.py +++ b/server/scripts/bench_llm.py @@ -35,7 +35,7 @@ DRAFT = None TEST_DFLASH = os.environ.get("DFLASH_BIN", str(ROOT / "build" / f"test_dflash{BIN_SUFFIX}")) TEST_GENERATE = os.environ.get("DFLASH_BIN_AR", str(ROOT / "build" / f"test_generate{BIN_SUFFIX}")) -TOKENIZER = os.environ.get("DFLASH_TOKENIZER", "Qwen/Qwen3.5-27B") +TOKENIZER = os.environ.get("DFLASH_TOKENIZER", "Qwen/Qwen3.6-27B") TMPDIR = Path(tempfile.gettempdir()) / "dflash_bench" TMPDIR.mkdir(parents=True, exist_ok=True) @@ -53,8 +53,8 @@ def _gsm_gold(x): BENCHES = [ - ("HumanEval", "openai_humaneval", None, "test", lambda x: x["prompt"], None, N_GEN), - ("GSM8K", "gsm8k", "main", "test", lambda x: f"Question: {x['question']}\nAnswer: ", _gsm_gold, 1024), + ("HumanEval", "openai/openai_humaneval", None, "test", lambda x: x["prompt"], None, N_GEN), + ("GSM8K", "openai/gsm8k", "main", "test", lambda x: f"Question: {x['question']}\nAnswer: ", _gsm_gold, 1024), ("Math500", "HuggingFaceH4/MATH-500", None, "test", lambda x: f"Problem: {x['problem']}\nSolution: Put your final answer in \\boxed{{}}.\n", lambda x: x["answer"], 2048), ] diff --git a/server/src/qwen35/qwen35_dflash_target.cpp b/server/src/qwen35/qwen35_dflash_target.cpp index 88be8c979..0a2657683 100644 --- a/server/src/qwen35/qwen35_dflash_target.cpp +++ b/server/src/qwen35/qwen35_dflash_target.cpp @@ -278,13 +278,38 @@ bool Qwen35DFlashTarget::verify_tree( return false; } - // Posterior = per-node target argmax. Tree-shaped graphs can return -1 from - // the GPU argmax shortcut, so compute argmax on CPU from full logits. + // Posterior = per-node target argmax. + // + // The verify graph already computes a batched per-node GPU argmax + // (sg_.argmax_tokens, built by build_target_step_tree). When the caller does + // not need the full logits (greedy decode, logits_out == nullptr) we read + // those N_actual int32s directly and skip the vocab×N_actual D2H + CPU + // argmax entirely — eliminates the verify-logits transfer hotspot. + // + // Historically the GPU argmax shortcut has returned -1 for tree-shaped + // verify graphs on some builds; guard against that by validating every row + // and falling back to the CPU path for the step if any index is bad. + // Escape hatch: DFLASH_GPU_VERIFY_ARGMAX=0 forces the legacy CPU path. + static const bool kGpuVerifyArgmax = []() { + const char * v = std::getenv("DFLASH_GPU_VERIFY_ARGMAX"); + return v == nullptr || v[0] != '0'; + }(); const int vocab = (int)sg_.logits->ne[0]; + posterior_out.resize(N_actual); + + if (kGpuVerifyArgmax && !logits_out && sg_.argmax_tokens) { + ggml_backend_tensor_get(sg_.argmax_tokens, posterior_out.data(), 0, + sizeof(int32_t) * N_actual); + bool ok = true; + for (int i = 0; i < N_actual; i++) { + if (posterior_out[i] < 0 || posterior_out[i] >= vocab) { ok = false; break; } + } + if (ok) return true; // fast path; otherwise fall through to CPU argmax + } + std::vector logits((size_t)vocab * N_actual); ggml_backend_tensor_get(sg_.logits, logits.data(), 0, sizeof(float) * (size_t)vocab * N_actual); - posterior_out.resize(N_actual); for (int i = 0; i < N_actual; i++) { const float * row = logits.data() + (size_t)i * vocab; int am = 0; float best = row[0]; diff --git a/server/test/test_dflash.cpp b/server/test/test_dflash.cpp index 25f885225..b4b4a3186 100644 --- a/server/test/test_dflash.cpp +++ b/server/test/test_dflash.cpp @@ -3417,16 +3417,40 @@ int main(int argc, char ** argv) { T_verify_compute = sync_us(); tt_verify_compute += std::chrono::duration(T_verify_compute - T_verify_set).count(); - // DDTree test mode reads full logits and computes posterior on - // CPU. The GPU argmax shortcut has returned -1 for tree-shaped - // verify graphs on some builds, which makes the harness stop - // after the root even though logits are valid. This is test-only; - // server decode paths are unaffected. + // DDTree posterior: per-node argmax over the verify logits. + // default : full vocab×N D2H + CPU argmax (legacy). + // GPU_VERIFY_ARGMAX=1: read the in-graph batched GPU argmax + // (tree_verify_argmax) — N int32s, no bulk D2H. + // GPU_VERIFY_ARGMAX=2: run BOTH and report per-step mismatches + // (validates the historical "-1 / tie" concern). + static const int kGpuVerifyArgmax = [](){ + const char * v = std::getenv("DFLASH_GPU_VERIFY_ARGMAX"); + return v ? std::atoi(v) : 0; + }(); std::vector posterior(N_actual); - ggml_backend_tensor_get(sg.logits, verify_logits_buf.data(), 0, - sizeof(float) * (size_t)vocab * N_actual); - for (int i = 0; i < N_actual; i++) { - posterior[i] = argmax_f32(verify_logits_buf.data() + (size_t)i * vocab, vocab); + bool logits_resident = false; // verify_logits_buf populated this step? + if (kGpuVerifyArgmax == 1 && sg.argmax_tokens) { + ggml_backend_tensor_get(sg.argmax_tokens, posterior.data(), 0, + sizeof(int32_t) * N_actual); + } else { + ggml_backend_tensor_get(sg.logits, verify_logits_buf.data(), 0, + sizeof(float) * (size_t)vocab * N_actual); + logits_resident = true; + for (int i = 0; i < N_actual; i++) { + posterior[i] = argmax_f32(verify_logits_buf.data() + (size_t)i * vocab, vocab); + } + if (kGpuVerifyArgmax == 2 && sg.argmax_tokens) { + std::vector gpu_post(N_actual); + ggml_backend_tensor_get(sg.argmax_tokens, gpu_post.data(), 0, + sizeof(int32_t) * N_actual); + int mism = 0, first = -1; + for (int i = 0; i < N_actual; i++) + if (gpu_post[i] != posterior[i]) { mism++; if (first < 0) first = i; } + if (mism) + std::fprintf(stderr, "[verify_argmax cmp] step=%d N=%d mismatches=%d " + "first@%d gpu=%d cpu=%d\n", n_draft_steps, N_actual, mism, + first, gpu_post[first], posterior[first]); + } } auto T_verify_logits_ddtree = sync_us(); tt_verify_logits += std::chrono::duration( @@ -3437,10 +3461,18 @@ int main(int argc, char ** argv) { int bonus_node_idx = 0; std::vector accepted = follow_verified_tree(tree, posterior.data(), next_token, &bonus_node_idx); if (g_sampler.temp > 0.0f) { + // Sampling needs the bonus node's full logit row. If we took the + // GPU-argmax path we skipped the bulk D2H, so fetch just that row. std::vector bonus_logits(vocab); - std::memcpy(bonus_logits.data(), - verify_logits_buf.data() + (size_t)bonus_node_idx * vocab, - (size_t)vocab * sizeof(float)); + if (logits_resident) { + std::memcpy(bonus_logits.data(), + verify_logits_buf.data() + (size_t)bonus_node_idx * vocab, + (size_t)vocab * sizeof(float)); + } else { + ggml_backend_tensor_get(sg.logits, bonus_logits.data(), + (size_t)bonus_node_idx * vocab * sizeof(float), + (size_t)vocab * sizeof(float)); + } next_token = sample_logits(bonus_logits.data(), vocab, g_sampler, out_all, g_sampler_rng); } const int accept_depth = (int)accepted.size(); // includes root From 252b47b3e0d8004cf452c7192828b79b4f6b49f5 Mon Sep 17 00:00:00 2001 From: "Pramodith (via Claude Code)" Date: Thu, 18 Jun 2026 15:56:11 +0000 Subject: [PATCH 03/10] topk kernels --- server/CMakeLists.txt | 6 +- server/src/common/draft_topk_cuda.cu | 205 +++++++++++++++++++++ server/src/common/draft_topk_cuda.h | 34 ++++ server/src/qwen35/qwen35_dflash_target.cpp | 23 ++- server/test/test_dflash.cpp | 38 +++- 5 files changed, 293 insertions(+), 13 deletions(-) create mode 100644 server/src/common/draft_topk_cuda.cu create mode 100644 server/src/common/draft_topk_cuda.h diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 5bfcb5ea7..a89db888c 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -395,7 +395,11 @@ if(DFLASH27B_GPU_BACKEND STREQUAL "hip") elseif(DFLASH27B_GPU_BACKEND STREQUAL "cuda") target_sources(dflash_common PRIVATE src/flashprefill_select.cpp - src/flashprefill.cpp) + src/flashprefill.cpp + src/common/draft_topk_cuda.cu) + # PUBLIC so consumers (e.g. the test_dflash executable) also see the macro + # and take the GPU draft top-K path instead of the CPU fallback. + target_compile_definitions(dflash_common PUBLIC DFLASH27B_HAVE_DRAFT_TOPK_CUDA=1) # Multi-arch: scan all resolved arches and compile every applicable # flashprefill kernel variant. This lets a single binary run on mixed # GPUs (e.g. Volta sm_70 + Pascal sm_61) without "no kernel image" errors. diff --git a/server/src/common/draft_topk_cuda.cu b/server/src/common/draft_topk_cuda.cu new file mode 100644 index 000000000..8ebfdd2cd --- /dev/null +++ b/server/src/common/draft_topk_cuda.cu @@ -0,0 +1,205 @@ +// GPU top-K + log-prob extraction for DDTree draft distributions. +// See draft_topk_cuda.h for the contract. Mirrors extract_draft_topk (ddtree.cpp). + +#include "draft_topk_cuda.h" + +#include +#include +#include +#include +#include + +namespace dflash::common { + +namespace { + +constexpr int kMaxK = 8; // ddtree_K is 8 in practice; K>kMaxK → CPU fallback +constexpr int kBlock = 1024; // threads per position (power of two for the reduction) +// With only n_positions (~15) blocks, occupancy is capped by blocks, not +// threads — so we want many warps per block to hide the vocab read latency. +// 1024 threads × kMaxK × (float+int32) = 64 KiB dynamic shared, which exceeds +// the 48 KiB default and needs an opt-in (see ensure_smem_optin below). + +// One block per draft position. A single strided pass over the vocabulary +// accumulates a per-thread online logsumexp (running max + sum) and a +// per-thread sorted-descending top-K; a shared-memory tree reduction then +// merges both across threads. log_prob[k] = scaled_logit[k] - log_z. +__global__ void draft_topk_kernel(const float * __restrict__ logits, + int vocab, int K, float inv_t, + float * __restrict__ out_lp, + int32_t * __restrict__ out_ids) { + const int row = blockIdx.x; + const int tid = threadIdx.x; + const float * __restrict__ li = logits + (size_t)row * vocab; + + float lmax = -FLT_MAX; + float lsum = 0.0f; + float topv[kMaxK]; + int topi[kMaxK]; +#pragma unroll + for (int k = 0; k < kMaxK; k++) { topv[k] = -FLT_MAX; topi[k] = -1; } + + // ---- single strided pass: logsumexp + local top-K ------------------ + // j ascends within a thread, so a strict-greater insert keeps the lower + // id on ties (matches the CPU heap's first-wins behaviour). + for (int j = tid; j < vocab; j += kBlock) { + const float l = li[j] * inv_t; + if (l > lmax) { lsum = lsum * __expf(lmax - l) + 1.0f; lmax = l; } + else { lsum += __expf(l - lmax); } + if (l > topv[K - 1]) { + int p = K - 1; + while (p > 0 && topv[p - 1] < l) { + topv[p] = topv[p - 1]; topi[p] = topi[p - 1]; --p; + } + topv[p] = l; topi[p] = j; + } + } + + // ---- block reduction -------------------------------------------------- + __shared__ float s_max[kBlock]; + __shared__ float s_sum[kBlock]; + extern __shared__ char s_raw[]; + float * s_topv = reinterpret_cast(s_raw); + int32_t * s_topi = reinterpret_cast(s_topv + (size_t)kBlock * K); + + s_max[tid] = lmax; + s_sum[tid] = lsum; + for (int k = 0; k < K; k++) { + s_topv[(size_t)tid * K + k] = topv[k]; + s_topi[(size_t)tid * K + k] = topi[k]; + } + __syncthreads(); + + for (int stride = kBlock / 2; stride > 0; stride >>= 1) { + if (tid < stride) { + // logsumexp merge of (max,sum) pairs + const float am = s_max[tid], as = s_sum[tid]; + const float bm = s_max[tid + stride], bs = s_sum[tid + stride]; + const float m = fmaxf(am, bm); + s_sum[tid] = as * __expf(am - m) + bs * __expf(bm - m); + s_max[tid] = m; + + // merge two sorted-desc top-K lists, keep the K largest (lower id on tie) + float av[kMaxK]; int ai[kMaxK]; + float bv[kMaxK]; int bi[kMaxK]; + for (int k = 0; k < K; k++) { + av[k] = s_topv[(size_t)tid * K + k]; + ai[k] = s_topi[(size_t)tid * K + k]; + bv[k] = s_topv[(size_t)(tid + stride) * K + k]; + bi[k] = s_topi[(size_t)(tid + stride) * K + k]; + } + int ia = 0, ib = 0; + for (int k = 0; k < K; k++) { + bool takeA; + if (ib >= K) takeA = true; + else if (ia >= K) takeA = false; + else if (av[ia] > bv[ib]) takeA = true; + else if (av[ia] < bv[ib]) takeA = false; + else takeA = (ai[ia] <= bi[ib]); // tie → lower id + if (takeA) { s_topv[(size_t)tid * K + k] = av[ia]; s_topi[(size_t)tid * K + k] = ai[ia]; ia++; } + else { s_topv[(size_t)tid * K + k] = bv[ib]; s_topi[(size_t)tid * K + k] = bi[ib]; ib++; } + } + } + __syncthreads(); + } + + if (tid == 0) { + const float log_z = s_max[0] + logf(s_sum[0]); + for (int k = 0; k < K; k++) { + out_lp[(size_t)row * K + k] = s_topv[k] - log_z; + out_ids[(size_t)row * K + k] = s_topi[k]; + } + } +} + +// Per-device scratch for the [n_positions × K] outputs, grown as needed. The +// decode loop is single-threaded, so a plain static cache is safe and avoids a +// cudaMalloc/cudaFree on every step. +struct Scratch { + int device = -1; + size_t cap = 0; // elements + float * d_lp = nullptr; + int32_t * d_ids = nullptr; +}; +Scratch g_scratch; + +bool ensure_scratch(int device, size_t n) { + if (g_scratch.device == device && g_scratch.cap >= n) return true; + if (g_scratch.d_lp) cudaFree(g_scratch.d_lp); + if (g_scratch.d_ids) cudaFree(g_scratch.d_ids); + g_scratch = Scratch{}; + if (cudaMalloc(&g_scratch.d_lp, n * sizeof(float)) != cudaSuccess) return false; + if (cudaMalloc(&g_scratch.d_ids, n * sizeof(int32_t)) != cudaSuccess) { + cudaFree(g_scratch.d_lp); g_scratch.d_lp = nullptr; return false; + } + g_scratch.device = device; + g_scratch.cap = n; + return true; +} + +} // namespace + +bool extract_draft_topk_cuda(const void * d_logits, + int n_positions, int vocab, int K, + float * out_log_probs, + int32_t * out_token_ids, + float temperature) { + if (!d_logits || n_positions <= 0 || vocab <= 0 || K <= 0 || K > kMaxK) return false; + + cudaPointerAttributes attr{}; + if (cudaPointerGetAttributes(&attr, d_logits) != cudaSuccess) { + cudaGetLastError(); // clear the error so we don't poison the next CUDA call + return false; + } + if (attr.type != cudaMemoryTypeDevice) return false; + + int prev = 0; + cudaGetDevice(&prev); + const int dev = attr.device; + if (dev != prev) cudaSetDevice(dev); + + static const bool kProfile = std::getenv("DFLASH_TOPK_PROFILE") != nullptr; + bool ok = false; + const size_t n = (size_t)n_positions * K; + if (ensure_scratch(dev, n)) { + const float inv_t = 1.0f / fmaxf(1e-3f, temperature); + const size_t smem = (size_t)kBlock * K * (sizeof(float) + sizeof(int32_t)); + // Opt in to >48 KiB dynamic shared (once per device). If it fails the + // launch below will error and we fall back to the CPU path. + static int s_optin_dev = -1; + if (s_optin_dev != dev) { + constexpr int kMaxSmem = (int)((size_t)kBlock * kMaxK * + (sizeof(float) + sizeof(int32_t))); + cudaFuncSetAttribute(draft_topk_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, kMaxSmem); + s_optin_dev = dev; + } + cudaEvent_t e_k0, e_k1, e_c1; + if (kProfile) { cudaEventCreate(&e_k0); cudaEventCreate(&e_k1); cudaEventCreate(&e_c1); cudaEventRecord(e_k0); } + draft_topk_kernel<<>>( + static_cast(d_logits), vocab, K, inv_t, + g_scratch.d_lp, g_scratch.d_ids); + if (kProfile) cudaEventRecord(e_k1); + if (cudaGetLastError() == cudaSuccess && cudaDeviceSynchronize() == cudaSuccess) { + const cudaError_t e1 = cudaMemcpy(out_log_probs, g_scratch.d_lp, + n * sizeof(float), cudaMemcpyDeviceToHost); + const cudaError_t e2 = cudaMemcpy(out_token_ids, g_scratch.d_ids, + n * sizeof(int32_t), cudaMemcpyDeviceToHost); + ok = (e1 == cudaSuccess && e2 == cudaSuccess); + } + if (kProfile) { + cudaEventRecord(e_c1); cudaEventSynchronize(e_c1); + float k_ms = 0, c_ms = 0; + cudaEventElapsedTime(&k_ms, e_k0, e_k1); + cudaEventElapsedTime(&c_ms, e_k1, e_c1); + std::fprintf(stderr, "[topk] kernel=%.3f ms sync+copy=%.3f ms (n_pos=%d vocab=%d)\n", + k_ms, c_ms, n_positions, vocab); + cudaEventDestroy(e_k0); cudaEventDestroy(e_k1); cudaEventDestroy(e_c1); + } + } + if (!ok) cudaGetLastError(); + if (dev != prev) cudaSetDevice(prev); + return ok; +} + +} // namespace dflash::common diff --git a/server/src/common/draft_topk_cuda.h b/server/src/common/draft_topk_cuda.h new file mode 100644 index 000000000..36ec0c8fe --- /dev/null +++ b/server/src/common/draft_topk_cuda.h @@ -0,0 +1,34 @@ +// GPU top-K + log-prob extraction for DDTree draft distributions. +// +// Drop-in device-side replacement for extract_draft_topk (ddtree.cpp): instead +// of D2H-ing the full [vocab × n_positions] draft logits and running an OpenMP +// heap top-K + online-logsumexp on the CPU, this runs the whole thing on the +// GPU directly on the logits tensor's device buffer and copies back only the +// [n_positions × K] results. +// +// Semantics match extract_draft_topk exactly: +// scaled = logit * (1 / max(1e-3, temperature)) +// log_z = logsumexp_j(scaled_j) (per position) +// out_log_probs = top-K scaled logits (desc), minus log_z +// out_token_ids = matching vocab ids; ties broken toward the lower id +// +// Returns false (caller must fall back to the CPU path) when CUDA is +// unavailable, the pointer is not device memory, K is out of range, or any +// CUDA call fails. Only compiled into CUDA builds; see CMakeLists.txt. + +#pragma once + +#include + +namespace dflash::common { + +// d_logits: device pointer to row-major [n_positions][vocab] f32 logits (the +// position stride is `vocab` floats — pass an offset pointer to skip +// leading positions). out_* are HOST buffers of size n_positions*K. +bool extract_draft_topk_cuda(const void * d_logits, + int n_positions, int vocab, int K, + float * out_log_probs, + int32_t * out_token_ids, + float temperature); + +} // namespace dflash::common diff --git a/server/src/qwen35/qwen35_dflash_target.cpp b/server/src/qwen35/qwen35_dflash_target.cpp index 0a2657683..65e181928 100644 --- a/server/src/qwen35/qwen35_dflash_target.cpp +++ b/server/src/qwen35/qwen35_dflash_target.cpp @@ -5,6 +5,7 @@ #include "graph_builders.h" #include "step_graph.h" #include "attn_masks.h" +#include "common/draft_topk_cuda.h" // gpu_runtime_compat.h maps the raw cudaStream_t / cudaMemcpy* symbols used // below (rollback_to / rollback_to_tree) onto their HIP equivalents. Without // it the file only compiles on CUDA via a transitive ; HIP @@ -656,12 +657,28 @@ bool Qwen35DFlashTarget::project_hidden_to_topk( if (st != GGML_STATUS_SUCCESS) return false; const int vocab = (int)proj_sg_.logits->ne[0]; + top_log_probs.assign((size_t)n_tokens * K, 0.0f); + top_token_ids.assign((size_t)n_tokens * K, 0); + +#ifdef DFLASH27B_HAVE_DRAFT_TOPK_CUDA + // GPU path: top-K + logsumexp directly on the logits device buffer, skipping + // the vocab×n_tokens D2H and the CPU heap extract. Falls back to the CPU path + // on any failure. Escape hatch: DFLASH_GPU_DRAFT_TOPK=0. + static const bool kGpuDraftTopk = []() { + const char * v = std::getenv("DFLASH_GPU_DRAFT_TOPK"); + return v == nullptr || v[0] != '0'; + }(); + if (kGpuDraftTopk && + extract_draft_topk_cuda(proj_sg_.logits->data, n_tokens, vocab, K, + top_log_probs.data(), top_token_ids.data(), + temperature)) { + return true; + } +#endif + std::vector logits((size_t)vocab * n_tokens); ggml_backend_tensor_get(proj_sg_.logits, logits.data(), 0, sizeof(float) * (size_t)vocab * n_tokens); - - top_log_probs.assign((size_t)n_tokens * K, 0.0f); - top_token_ids.assign((size_t)n_tokens * K, 0); extract_draft_topk(logits.data(), n_tokens, vocab, K, top_log_probs.data(), top_token_ids.data(), temperature); return true; diff --git a/server/test/test_dflash.cpp b/server/test/test_dflash.cpp index b4b4a3186..c68de6713 100644 --- a/server/test/test_dflash.cpp +++ b/server/test/test_dflash.cpp @@ -270,6 +270,7 @@ using dflash::common::free_qwen35_layer_split_shards; #include "qwen35_layer_split_dflash_target.h" #include "common/dflash_spec_decode.h" #include "common/gguf_mmap.h" +#include "common/draft_topk_cuda.h" using dflash::common::is_eos_tok; // ─── Layer-split daemon — extracted to src/qwen35/layer_split_daemon.{h,cpp} ─ @@ -3268,16 +3269,35 @@ int main(int argc, char ** argv) { } } else { // DDTree K>1: need real log-probs for best-first tree scoring. - // Transfer full logits for positions 1..q_len-1. - if (!draft_hidden_bridge) { - ggml_backend_tensor_get(draft_sg.logits, draft_logits_buf.data(), 0, - sizeof(float) * vocab * q_len); + bool topk_done = false; +#ifdef DFLASH27B_HAVE_DRAFT_TOPK_CUDA + // GPU path: top-K + logsumexp on the draft logits device buffer + // (positions 1..q_len-1), no full-vocab D2H. Escape: DFLASH_GPU_DRAFT_TOPK=0. + static const bool kGpuDraftTopk = [](){ + const char * v = std::getenv("DFLASH_GPU_DRAFT_TOPK"); + return v == nullptr || v[0] != '0'; + }(); + if (kGpuDraftTopk && !draft_hidden_bridge) { + topk_done = dflash::common::extract_draft_topk_cuda( + (const float *)draft_sg.logits->data + (size_t)vocab, + L, vocab, ddtree_K, + ddtree_top_log_probs.data(), + ddtree_top_token_ids.data(), + ddtree_temp); + } +#endif + if (!topk_done) { + // Transfer full logits for positions 1..q_len-1, extract on CPU. + if (!draft_hidden_bridge) { + ggml_backend_tensor_get(draft_sg.logits, draft_logits_buf.data(), 0, + sizeof(float) * vocab * q_len); + } + extract_draft_topk(draft_logits_buf.data() + (size_t)vocab, + L, vocab, ddtree_K, + ddtree_top_log_probs.data(), + ddtree_top_token_ids.data(), + ddtree_temp); } - extract_draft_topk(draft_logits_buf.data() + (size_t)vocab, - L, vocab, ddtree_K, - ddtree_top_log_probs.data(), - ddtree_top_token_ids.data(), - ddtree_temp); } } auto T_draft_logits = sync_us(); From 0aa84b29ac124315aed04064fba8b26ef4d6b901 Mon Sep 17 00:00:00 2001 From: pramodith Date: Thu, 18 Jun 2026 16:59:04 +0000 Subject: [PATCH 04/10] calibration script --- calib/calibrate.py | 330 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 330 insertions(+) create mode 100644 calib/calibrate.py diff --git a/calib/calibrate.py b/calib/calibrate.py new file mode 100644 index 000000000..add7875a3 --- /dev/null +++ b/calib/calibrate.py @@ -0,0 +1,330 @@ +#!/usr/bin/env python3 +"""End-to-end top-k/top-p LM-head calibration driver (GPU). + +Single entry point that replaces the run_calib.sh + aggregate.py pair. It can +download the models, tokenize prompts, build the CUDA `test_dflash` binary, +run the calibration per-prompt on a chosen GPU, and aggregate the +`[topk-calib]` coverage/mass tables into a report (stdout + JSON). + +Why a *driver* and not a pure-Python re-implementation: the DFlash draft is an +EAGLE-style head that consumes the target's mid-stack hidden states and shares +the target's LM head (see server/src/common/dflash_draft_graph.h). Stock +llama.cpp / llama-cpp-python cannot load or run it, so the only faithful way to +produce the draft candidate set is this repo's C++ `test_dflash`. This script +drives that binary on the GPU (Blackwell sm_100 by default) and does all the +orchestration + analysis in Python. + +Typical use (Blackwell B200, GPU 0): + + # one-shot full pipeline (download ~16GB target + draft, prep, build, run) + python3 calib/calibrate.py --all --prep 200 + + # later re-runs once models/bins/binary exist + python3 calib/calibrate.py --gpu 0 --ngen 256 --json calib/out/calib.json + +By default (no stage flags) it only runs+aggregates and tells you which flag to +add if a prerequisite is missing. Stage flags (--download/--prep/--build) opt +in to the heavy steps; --all enables them all. + +Deps for --prep: transformers, datasets, jinja2. Deps for --download: the `hf` +CLI (huggingface_hub). The run/aggregate stages are stdlib-only. +""" +import argparse +import json +import os +import re +import shutil +import subprocess +import sys +from pathlib import Path + +REPO = Path(__file__).resolve().parent.parent +DEF_TARGET = REPO / "server/models/Qwen3.6-27B-Q4_K_M.gguf" +DEF_DRAFT = REPO / "server/models/draft/dflash-draft-3.6-q4_k_m.gguf" +DEF_BIN = REPO / "server/build/test_dflash" +DEF_PROMPTS = REPO / "calib" +DEF_OUTDIR = REPO / "calib/out" + +# HF repos (see README.md): target + DFlash draft. +TARGET_REPO = "unsloth/Qwen3.6-27B-GGUF" +TARGET_FILE = "Qwen3.6-27B-Q4_K_M.gguf" +DRAFT_REPO = "Lucebox/Qwen3.6-27B-DFlash-GGUF" +DRAFT_FILE = "dflash-draft-3.6-q4_k_m.gguf" + +SET_METRICS = [ + "greedy (target argmax)", + "sampling top_k=8", + "sampling top_k=20 (Qwen def)", + "sampling top_p=0.95 nucleus", +] + + +def run(cmd, **kw): + """Echo + run a command, raising on failure.""" + print(" $ " + " ".join(str(c) for c in cmd), file=sys.stderr) + return subprocess.run([str(c) for c in cmd], check=True, **kw) + + +# ── stage: download ────────────────────────────────────────────────────────── +def stage_download(target: Path, draft: Path): + if not shutil.which("hf"): + sys.exit("error: `hf` CLI not found. `pip install huggingface_hub` (set " + "HF_TOKEN to avoid rate limits), or download the GGUFs manually.") + if not target.exists(): + target.parent.mkdir(parents=True, exist_ok=True) + print(f"[download] target -> {target.parent}", file=sys.stderr) + run(["hf", "download", TARGET_REPO, TARGET_FILE, "--local-dir", target.parent]) + else: + print(f"[download] target present, skip ({target})", file=sys.stderr) + if not draft.exists(): + draft.parent.mkdir(parents=True, exist_ok=True) + print(f"[download] draft -> {draft.parent}", file=sys.stderr) + run(["hf", "download", DRAFT_REPO, DRAFT_FILE, "--local-dir", draft.parent]) + else: + print(f"[download] draft present, skip ({draft})", file=sys.stderr) + + +# ── stage: prep prompts ────────────────────────────────────────────────────── +def stage_prep(prompts_dir: Path, n: int, thinking: str): + print(f"[prep] tokenizing {n} prompts -> {prompts_dir}", file=sys.stderr) + run([sys.executable, str(REPO / "calib/prep_prompts.py"), + "--n", n, "--out", prompts_dir, "--thinking", thinking]) + + +# ── stage: build ───────────────────────────────────────────────────────────── +def stage_build(binary: Path, cuda_arch: int): + build_dir = REPO / "server/build" + print("[build] ensuring git submodules (deps/llama.cpp, ...)", file=sys.stderr) + run(["git", "-C", REPO, "submodule", "update", "--init", "--recursive"]) + print(f"[build] configuring test_dflash (sm_{cuda_arch}, NCCL OFF)", file=sys.stderr) + run(["cmake", "-B", build_dir, "-S", REPO / "server", "-G", "Ninja", + "-DCMAKE_BUILD_TYPE=Release", + f"-DCMAKE_CUDA_ARCHITECTURES={cuda_arch}", + "-DGGML_CUDA_NCCL=OFF"]) + run(["cmake", "--build", build_dir, "--target", "test_dflash", "-j"]) + if not binary.exists(): + sys.exit(f"error: build finished but {binary} not found") + + +# ── stage: run calibration over all prompt bins ────────────────────────────── +def stage_run(args, bins, outdir: Path): + """Run test_dflash per prompt; return {name: full_log_text} for parsing.""" + outdir.mkdir(parents=True, exist_ok=True) + logs = {} + for i, p in enumerate(bins, 1): + name = p.stem + out_bin = outdir / f"{name}_out.bin" + log_path = outdir / f"{name}.log" + if args.skip_existing and log_path.exists(): + print(f"[run {i}/{len(bins)}] {name}: cached log, skip", file=sys.stderr) + logs[name] = log_path.read_text() + continue + print(f"[run {i}/{len(bins)}] {name} (gpu {args.gpu})", file=sys.stderr) + # DFLASH27B_PREFILL_UBATCH=1: token-by-token prefill. The batched prefill + # attention path emits all-NaN logits for some (n_tokens, kv_start) configs, + # which makes the greedy prefill argmax fall to token 0 and the run abort + # silently (~76% of prompts, biased toward code). Token-by-token prefill is + # numerically correct and only marginally slower (prefill ≪ the 256-tok gen). + env = {"CUDA_VISIBLE_DEVICES": str(args.gpu), "DFLASH_TOPK_CALIB": "1", + "DFLASH27B_PREFILL_UBATCH": str(args.prefill_ubatch)} + cmd = [args.bin, args.target, args.draft, p, args.ngen, out_bin, + "--fast-rollback", f"--max-ctx={args.maxctx}"] + # Inherit current env, then override the two calib vars. + full_env = {**os.environ, **env} + proc = subprocess.run([str(c) for c in cmd], env=full_env, + stdout=subprocess.PIPE, stderr=subprocess.STDOUT, + text=True) + log_path.write_text(proc.stdout) + logs[name] = proc.stdout + if proc.returncode != 0 or "[topk-calib]" not in proc.stdout: + tail = "\n".join(proc.stdout.splitlines()[-8:]) + print(f" ! {name} produced no calib table (rc={proc.returncode}). " + f"tail:\n{tail}", file=sys.stderr) + return logs + + +# ── parse one prompt's [topk-calib] block ──────────────────────────────────── +def parse_block(text: str): + """Return dict for one calib run, or None if no table present. + + {positions, mean_rank, max_rank, vocab, + set: {metric: {M: pct}}, mass: {M: pct}, mass_bad: {M: (bad, total)}} + """ + m = re.search(r"\[topk-calib\] positions=(\d+)\s+mean_rank=([\d.]+)\s+" + r"max_rank=(\d+)\s+\(vocab=(\d+)\)", text) + if not m: + return None + res = { + "positions": int(m.group(1)), + "mean_rank": float(m.group(2)), + "max_rank": int(m.group(3)), + "vocab": int(m.group(4)), + "set": {k: {} for k in SET_METRICS}, + "mass": {}, + "mass_bad": {}, + } + cur = None + for line in text.splitlines(): + sc = re.search(r"--- (.+?): set-coverage", line) + if sc: + label = sc.group(1) + cur = label if label in res["set"] else None + continue + if "MASS coverage" in line: + cur = "__mass__" + continue + cm = re.search(r"coverage@M=(\d+)\s*:\s*([\d.]+)%", line) + if cm and cur and cur != "__mass__": + res["set"][cur][int(cm.group(1))] = float(cm.group(2)) + mm = re.search(r"mass@M=(\d+)\s*:\s*([\d.]+)%.*?<99% covered:\s*(\d+)/(\d+)", + line) + if mm: + M = int(mm.group(1)) + res["mass"][M] = float(mm.group(2)) + res["mass_bad"][M] = (int(mm.group(3)), int(mm.group(4))) + return res + + +# ── aggregate + report ─────────────────────────────────────────────────────── +def aggregate(per_prompt: dict): + """Position-weighted aggregate across prompts.""" + ok = {n: b for n, b in per_prompt.items() if b} + totpos = sum(b["positions"] for b in ok.values()) + if totpos == 0: + return None + Ms = sorted({M for b in ok.values() for M in b["mass"]} + | {M for b in ok.values() for d in b["set"].values() for M in d}) + + # Accumulate sums + per-cell weight so cells no prompt reported stay None + # (rather than a misleading 0.0). In real output every metric shares the + # same M-grid, so all cells are populated. + set_sum = {m: {M: 0.0 for M in Ms} for m in SET_METRICS} + set_w = {m: {M: 0.0 for M in Ms} for m in SET_METRICS} + mass_sum = {M: 0.0 for M in Ms} + mass_w = {M: 0.0 for M in Ms} + mass_bad = {M: 0 for M in Ms} + rank_w = 0.0 + max_rank = 0 + for b in ok.values(): + w = b["positions"] + rank_w += b["mean_rank"] * w + max_rank = max(max_rank, b["max_rank"]) + for m in SET_METRICS: + for M, v in b["set"][m].items(): + set_sum[m][M] += v * w + set_w[m][M] += w + for M, v in b["mass"].items(): + mass_sum[M] += v * w + mass_w[M] += w + for M, (bad, _) in b["mass_bad"].items(): + mass_bad[M] += bad + set_agg = {m: {M: (set_sum[m][M] / set_w[m][M] if set_w[m][M] else None) + for M in Ms} for m in SET_METRICS} + mass_agg = {M: (mass_sum[M] / mass_w[M] if mass_w[M] else None) for M in Ms} + return { + "prompts": len(ok), "failed": len(per_prompt) - len(ok), + "total_positions": totpos, "mean_rank": rank_w / totpos, + "max_rank": max_rank, "Ms": Ms, + "set": set_agg, "mass": mass_agg, "mass_bad": mass_bad, + "per_prompt": {n: {"positions": b["positions"], "mean_rank": b["mean_rank"], + "max_rank": b["max_rank"]} for n, b in ok.items()}, + } + + +def print_report(agg): + Ms = agg["Ms"] + print(f"\n=== top-k head calibration aggregate ===") + print(f"prompts={agg['prompts']} (failed={agg['failed']}) " + f"total_positions={agg['total_positions']} " + f"mean_rank={agg['mean_rank']:.2f} max_rank={agg['max_rank']}\n") + def cell(v, fmt): + return f"{'-':>10}" if v is None else f"{v:10.{fmt}f}" + hdr = "metric".ljust(28) + "".join(f"{M:>10}" for M in Ms) + print("-- set-coverage (entire set within draft top-M, position-weighted %) --") + print(hdr) + for m in SET_METRICS: + row = m[:27].ljust(28) + "".join(cell(agg['set'][m][M], 2) for M in Ms) + print(row) + print("\n-- top_p=0.95 nucleus MASS coverage (mean %, + positions <99% covered) --") + print("mass %".ljust(28) + "".join(cell(agg['mass'][M], 3) for M in Ms)) + print("bad(<99%)".ljust(28) + "".join(f"{agg['mass_bad'][M]:10d}" for M in Ms)) + print("\n-- per-prompt --") + for n, s in sorted(agg["per_prompt"].items()): + print(f" {n:<16} pos={s['positions']:<5} mean_rank={s['mean_rank']:7.2f} " + f"max_rank={s['max_rank']}") + + +# ── main ───────────────────────────────────────────────────────────────────── +def main(): + ap = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--target", type=Path, default=DEF_TARGET) + ap.add_argument("--draft", type=Path, default=DEF_DRAFT) + ap.add_argument("--bin", dest="bin", type=Path, default=DEF_BIN) + ap.add_argument("--prompts", type=Path, default=DEF_PROMPTS, + help="dir holding tokenized *.bin prompts") + ap.add_argument("--out-dir", type=Path, default=DEF_OUTDIR) + ap.add_argument("--gpu", type=int, default=0, help="CUDA_VISIBLE_DEVICES index") + ap.add_argument("--ngen", type=int, default=256) + ap.add_argument("--maxctx", type=int, default=1024) + ap.add_argument("--prefill-ubatch", type=int, default=1, + help="DFLASH27B_PREFILL_UBATCH. Default 1 (token-by-token) avoids " + "the batched-prefill NaN bug that aborts ~76%% of prompts. Set " + ">1 to use batched prefill (faster, but reintroduces the bug).") + ap.add_argument("--cuda-arch", type=int, default=100, + help="CMAKE_CUDA_ARCHITECTURES for --build (Blackwell B200=100)") + ap.add_argument("--json", type=Path, default=None, help="write aggregate JSON here") + ap.add_argument("--skip-existing", action="store_true", + help="reuse a prompt's cached .log instead of re-running it") + # stage opt-ins + ap.add_argument("--download", action="store_true", help="hf download missing GGUFs") + ap.add_argument("--prep", type=int, metavar="N", default=0, + help="tokenize N prompts via prep_prompts.py before running") + ap.add_argument("--thinking", choices=["on", "off"], default="off") + ap.add_argument("--build", action="store_true", help="cmake build test_dflash") + ap.add_argument("--all", action="store_true", + help="enable --download + --build (use with --prep N)") + ap.add_argument("--no-run", action="store_true", + help="only run enabled prep/build/download stages, skip calibration") + args = ap.parse_args() + + if args.all: + args.download = True + args.build = True + + if args.download: + stage_download(args.target, args.draft) + if args.prep: + stage_prep(args.prompts, args.prep, args.thinking) + if args.build: + stage_build(args.bin, args.cuda_arch) + + if args.no_run: + return + + # prerequisite checks with actionable hints + if not args.target.exists(): + sys.exit(f"error: target not found: {args.target}\n add --download") + if not args.draft.exists(): + sys.exit(f"error: draft not found: {args.draft}\n add --download") + if not args.bin.exists(): + sys.exit(f"error: binary not found: {args.bin}\n add --build") + bins = sorted(args.prompts.glob("*.bin")) + if not bins: + sys.exit(f"error: no *.bin prompts in {args.prompts}\n add --prep N") + print(f"[run] {len(bins)} prompts, ngen={args.ngen}, gpu={args.gpu}", file=sys.stderr) + + logs = stage_run(args, bins, args.out_dir) + per_prompt = {n: parse_block(t) for n, t in logs.items()} + agg = aggregate(per_prompt) + if not agg: + sys.exit("error: no calibration tables parsed from any prompt run") + print_report(agg) + if args.json: + args.json.parent.mkdir(parents=True, exist_ok=True) + args.json.write_text(json.dumps(agg, indent=2)) + print(f"\n[json] wrote {args.json}", file=sys.stderr) + + +if __name__ == "__main__": + main() From 823070ad75062fa0e96c0e5942349d18eaa7bd58 Mon Sep 17 00:00:00 2001 From: pramodith Date: Thu, 18 Jun 2026 17:00:31 +0000 Subject: [PATCH 05/10] docs(calib): top-k head calibration results (B200, 102k positions) Full top-k/top-p LM-head shortlist calibration on Qwen3.6-27B Q4_K_M target + Q4_K_M 3.6 draft, run on a single B200 via calib/calibrate.py over 198 balanced prompts (HumanEval/GSM8K/MATH-500), 102,045 chain-verify positions. Verdict: greedy viable (target argmax in draft top-1024 at 99.46%); sampling exact-set-coverage impractical but mass-coverage faithful (top_p=0.95 nucleus mass 99.8% at M=8192). Also documents the batched-prefill NaN bug and the DFLASH27B_PREFILL_UBATCH=1 workaround used to capture all prompts. Co-Authored-By: Claude Opus 4.8 --- docs/topk-head-calibration-results.md | 115 ++++++++++++++++++++++++++ 1 file changed, 115 insertions(+) create mode 100644 docs/topk-head-calibration-results.md diff --git a/docs/topk-head-calibration-results.md b/docs/topk-head-calibration-results.md new file mode 100644 index 000000000..e76ddd524 --- /dev/null +++ b/docs/topk-head-calibration-results.md @@ -0,0 +1,115 @@ +# Candidate-restricted LM head — calibration results + +**Date:** 2026-06-18 · **Branch:** `feat/topk-head-calib` +**Companion:** design + methodology in [`topk-head-optimization.md`](topk-head-optimization.md) + +Full top-k/top-p shortlist calibration for the candidate-restricted LM head, run on +a single **NVIDIA B200** (Blackwell, sm_100) via `calib/calibrate.py` driving the +`test_dflash` chain-verify path (exact position alignment, `DFLASH_TOPK_CALIB=1`). + +## Setup + +| | | +|---|---| +| Target | Qwen3.6-27B **Q4_K_M** (`server/models/Qwen3.6-27B-Q4_K_M.gguf`, vocab 248,320) | +| Draft | DFlash 3.6 **Q4_K_M** (`dflash-draft-3.6-q4_k_m.gguf`) | +| Prompts | 198 — 66 each from HumanEval / GSM8K / MATH-500 (Qwen3.6 chat template, thinking off, ≤1024 tok) | +| Decode | greedy (`temp=0`), `n_gen=256`, chain path (no `--ddtree`), `--max-ctx=1024` | +| Positions | **102,045** chain-verify positions | +| GPU | B200, GPU 0, `DFLASH27B_PREFILL_UBATCH=1` (see caveat) | + +Sampling metrics (top_k=8 / top_k=20 / top_p=0.95) are computed at fixed internal +temperature 1.0 per position, independent of the greedy decode trajectory. + +> **Run quality:** 198/198 prompts succeeded (0 dropped), balanced across all three +> datasets — gsm 37,905 · he 22,935 · math 41,205 positions. This required the +> `DFLASH27B_PREFILL_UBATCH=1` workaround below; without it ~76% of prompts abort and +> the survivors skew heavily away from code (HumanEval). + +## Set-coverage — fraction of positions whose entire set ⊆ draft top-M (%) + +| M | greedy | top_k=8 | top_k=20 (Qwen def) | top_p=0.95 | +|------:|-------:|--------:|--------------------:|-----------:| +| 1 | 57.30 | 0.00 | 0.00 | 47.02 | +| 2 | 70.24 | 0.00 | 0.00 | 56.23 | +| 4 | 79.88 | 0.00 | 0.00 | 62.19 | +| 8 | 87.50 | 0.06 | 0.00 | 67.00 | +| 16 | 92.68 | 2.68 | 0.00 | 71.08 | +| 32 | 95.94 | 10.53 | 0.02 | 74.72 | +| 64 | 97.64 | 24.51 | 0.66 | 78.10 | +| 128 | 98.56 | 41.73 | 4.11 | 81.10 | +| 256 | 99.02 | 58.15 | 12.77 | 83.59 | +| 512 | 99.27 | 71.41 | 26.01 | 85.97 | +| 1024 | 99.46 | 81.44 | 41.79 | 88.14 | +| 2048 | 99.63 | 88.51 | 57.68 | 90.39 | +| 4096 | 99.78 | 93.46 | 71.72 | 92.68 | +| 8192 | 99.86 | 96.54 | 82.94 | 94.89 | +| 16384 | 99.92 | 98.23 | 90.45 | 96.56 | +| 32768 | 99.95 | 99.11 | 94.93 | 97.73 | +| 65536 | 99.97 | 99.55 | 97.44 | 98.61 | + +`mean_rank = 61.6`, `max_rank = 65536` (grid ceiling; a long low-confidence tail, +dominated by code prompts, extends beyond it). + +## top_p=0.95 nucleus — MASS coverage + +Set-coverage is the wrong quality metric for sampling (missing a low-prob nucleus +token barely distorts the draw). The right one is **covered probability mass**. + +| M | mean nucleus mass covered (%) | positions <99% covered (of 102,045) | +|------:|------:|------:| +| 256 | 98.318 | 13,639 | +| 512 | 98.827 | 10,421 | +| 1024 | 99.176 | 7,499 | +| 2048 | 99.460 | 4,924 | +| 4096 | 99.682 | 2,827 | +| 8192 | 99.816 | 1,435 | +| 16384 | 99.895 | 702 | +| 32768 | 99.939 | 340 | +| 65536 | 99.963 | 157 | + +## Verdict + +- **Greedy: viable.** Target argmax ∈ draft top-M at **99.46% for M=1024** and + 99.78% at M=4096 — near-lossless, and the restricted head is <0.5% of the + 248k-wide matmul. Build the restricted head for the greedy verify path. +- **Sampling, exact set-coverage: not practical.** Exact `top_k` needs the draft + shortlist to contain the target's *entire* top-K_s; the Q4 draft covers the full + top-20 only ~72% of the time even at M=4096 (97% at M=65536). The draft ranks the + target's #1 well but the rest of the nucleus poorly. +- **Sampling, approximate-but-faithful: viable.** By covered mass, M≈8192 captures + 99.8% of the nucleus (only ~1.4% of positions <99% covered) at ~5.9% step speedup; + M=65536 is lossless-in-practice (99.963%, 0.15% of positions <99%). Gate pure-top_p + / pure-temperature to the full head, or over-cover with a tail-mass correction. +- A stronger (BF16) draft would likely improve the sampling numbers; the structural + asymmetry (argmax easy, full nucleus hard) will persist to some degree. + +## Caveat — batched-prefill NaN bug (worked around, not fixed) + +The batched prefill attention in `test_dflash` emits **all-NaN logits** for certain +`(n_tokens, kv_start)` shapes (e.g. a single 51-token ubatch, or any 2nd+ ubatch with +`kv_start>0`). The greedy `argmax` of an all-NaN vector returns index 0, so the prefill +reports `last_tok=0` and the first draft step aborts with a silent `return 1` right +after `[migrate]`. This affected ~76% of prompts and biased survivors away from code. + +Root cause is the batched FA / KQ-mask path — it reproduces on `main` with no calib +env, independent of this work. Token-by-token prefill is numerically correct, so the +calibration runs with `DFLASH27B_PREFILL_UBATCH=1` (the `calib/calibrate.py` default, +`--prefill-ubatch 1`). The underlying kernel bug still needs a real fix; it would +silently corrupt any batched-prefill decode, not just calibration. + +## Reproduce + +```bash +# toolchain (uv): cmake/ninja + tokenizer deps +uv venv .venv-calib && uv pip install --python .venv-calib/bin/python \ + cmake ninja transformers datasets jinja2 huggingface_hub + +# download models (target -> server/models, draft -> server/models/draft) +# build test_dflash (sm_100), prep 200 prompts, run on GPU 0, aggregate -> JSON +PATH="$PWD/.venv-calib/bin:$PATH" .venv-calib/bin/python calib/calibrate.py \ + --all --prep 200 --gpu 0 --ngen 256 --prefill-ubatch 1 \ + --json calib/out/calib200_fixed.json +``` + +Raw per-prompt tables + the position-weighted aggregate: `calib/out/calib200_fixed.json`. From 39be91454dd9bbeb3ddf2f8def489524861543fa Mon Sep 17 00:00:00 2001 From: "Pramodith (via Claude Code)" Date: Thu, 18 Jun 2026 17:18:56 +0000 Subject: [PATCH 06/10] perf(draft_topk_cuda): split-K + register top-K, ~10x faster kernel The draft top-K + logsumexp kernel launched only n_positions (~15) blocks, leaving most of the GPU's SMs idle, and kept its per-thread top-K in a data-dependent insertion index that the compiler spilled to local memory and re-read on every vocab element. Both made the ~9 MB vocab scan run at a small fraction of peak DRAM bandwidth. Rework into a split-K two-pass design: - pass 1 (draft_topk_partial) splits each position's vocab scan across many blocks (2D grid n_positions x split) so all SMs stay busy; - pass 2 (draft_topk_combine) merges the per-split partials per position. Template both kernels on K (compile-time) so the top-K stays register-resident via a branchless unrolled bubble instead of spilling, and read logits as float4 (one coalesced 16-byte transaction per 4 logits) with a scalar fallback when a row base is not 16-byte aligned (vocab % 4 != 0), preserving any-vocab correctness. split is auto-tuned (env override DFLASH_TOPK_SPLIT). Measured on an RTX 3090 (n=15, vocab=151936, K=8): - GPU kernel time: 392 us -> 36.3 us (30.6 partial + 5.75 combine), 10.8x - full call (kernel+sync+D2H): 0.407 ms -> 0.053 ms, 7.7x Full-call speedup is 5.9-8.4x across n in {7,15,31,63}. Output is bit-for-bit equivalent to the CPU reference (id_mismatches=0) across K in {1,2,4,8} and both aligned and odd vocab; compute-sanitizer memcheck clean on both paths. Adds bench_topk.cu, a standalone microbenchmark + CPU-reference correctness harness (not wired into the build) used to profile and A/B this change. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01T55hNb5cgyCwNYnNAE1hun --- server/src/common/bench_topk.cu | 100 +++++++ server/src/common/draft_topk_cuda.cu | 388 ++++++++++++++++++++------- 2 files changed, 398 insertions(+), 90 deletions(-) create mode 100644 server/src/common/bench_topk.cu diff --git a/server/src/common/bench_topk.cu b/server/src/common/bench_topk.cu new file mode 100644 index 000000000..ff6531b7e --- /dev/null +++ b/server/src/common/bench_topk.cu @@ -0,0 +1,100 @@ +// Standalone microbenchmark + correctness check for extract_draft_topk_cuda. +// Build: +// nvcc -O3 -arch=sm_86 -o /tmp/bench_topk \ +// server/src/common/bench_topk.cu server/src/common/draft_topk_cuda.cu +// Run: +// /tmp/bench_topk [n_positions] [vocab] [K] [iters] +#include "draft_topk_cuda.h" +#include +#include +#include +#include +#include +#include +#include +#include + +using dflash::common::extract_draft_topk_cuda; + +// CPU reference matching the documented semantics. +static void cpu_ref(const float* logits, int n, int vocab, int K, float temp, + std::vector& lp, std::vector& ids) { + const float inv_t = 1.0f / std::fmax(1e-3f, temp); + lp.assign((size_t)n * K, 0.f); + ids.assign((size_t)n * K, -1); + for (int r = 0; r < n; r++) { + const float* li = logits + (size_t)r * vocab; + // online logsumexp + float lmax = -FLT_MAX, lsum = 0.f; + for (int j = 0; j < vocab; j++) { + float l = li[j] * inv_t; + if (l > lmax) { lsum = lsum * std::exp(lmax - l) + 1.f; lmax = l; } + else lsum += std::exp(l - lmax); + } + float log_z = lmax + std::log(lsum); + // top-K with lower-id-wins on ties + std::vector idx(vocab); + for (int j = 0; j < vocab; j++) idx[j] = j; + std::partial_sort(idx.begin(), idx.begin() + K, idx.end(), + [&](int a, int b){ + float la = li[a]*inv_t, lb = li[b]*inv_t; + if (la != lb) return la > lb; + return a < b; + }); + for (int k = 0; k < K; k++) { + int j = idx[k]; + lp[(size_t)r*K+k] = li[j]*inv_t - log_z; + ids[(size_t)r*K+k] = j; + } + } +} + +int main(int argc, char** argv) { + int n = argc > 1 ? atoi(argv[1]) : 15; + int vocab = argc > 2 ? atoi(argv[2]) : 151936; + int K = argc > 3 ? atoi(argv[3]) : 8; + int iters = argc > 4 ? atoi(argv[4]) : 200; + float temp = 1.0f; + + printf("n=%d vocab=%d K=%d iters=%d\n", n, vocab, K, iters); + + std::vector h_logits((size_t)n * vocab); + std::mt19937 rng(1234); + std::normal_distribution dist(0.f, 4.f); + for (auto& x : h_logits) x = dist(rng); + + float* d_logits = nullptr; + cudaMalloc(&d_logits, h_logits.size() * sizeof(float)); + cudaMemcpy(d_logits, h_logits.data(), h_logits.size()*sizeof(float), cudaMemcpyHostToDevice); + + std::vector lp((size_t)n*K); + std::vector ids((size_t)n*K); + + // correctness + bool ok = extract_draft_topk_cuda(d_logits, n, vocab, K, lp.data(), ids.data(), temp); + if (!ok) { printf("FAIL: kernel returned false\n"); return 1; } + + std::vector rlp; std::vector rids; + cpu_ref(h_logits.data(), n, vocab, K, temp, rlp, rids); + int mism = 0; float maxerr = 0.f; + for (size_t i = 0; i < lp.size(); i++) { + if (ids[i] != rids[i]) { if (mism < 10) printf(" id mismatch @%zu gpu=%d cpu=%d\n", i, ids[i], rids[i]); mism++; } + maxerr = std::fmax(maxerr, std::fabs(lp[i]-rlp[i])); + } + printf("correctness: id_mismatches=%d max_lp_err=%.3e\n", mism, maxerr); + + // warmup + for (int i = 0; i < 10; i++) extract_draft_topk_cuda(d_logits, n, vocab, K, lp.data(), ids.data(), temp); + cudaDeviceSynchronize(); + + cudaEvent_t e0, e1; cudaEventCreate(&e0); cudaEventCreate(&e1); + cudaEventRecord(e0); + for (int i = 0; i < iters; i++) + extract_draft_topk_cuda(d_logits, n, vocab, K, lp.data(), ids.data(), temp); + cudaEventRecord(e1); cudaEventSynchronize(e1); + float ms = 0; cudaEventElapsedTime(&ms, e0, e1); + printf("avg full call (kernel+sync+copy): %.4f ms\n", ms / iters); + + cudaFree(d_logits); + return 0; +} diff --git a/server/src/common/draft_topk_cuda.cu b/server/src/common/draft_topk_cuda.cu index 8ebfdd2cd..7adc5045a 100644 --- a/server/src/common/draft_topk_cuda.cu +++ b/server/src/common/draft_topk_cuda.cu @@ -13,91 +13,228 @@ namespace dflash::common { namespace { -constexpr int kMaxK = 8; // ddtree_K is 8 in practice; K>kMaxK → CPU fallback -constexpr int kBlock = 1024; // threads per position (power of two for the reduction) -// With only n_positions (~15) blocks, occupancy is capped by blocks, not -// threads — so we want many warps per block to hide the vocab read latency. -// 1024 threads × kMaxK × (float+int32) = 64 KiB dynamic shared, which exceeds -// the 48 KiB default and needs an opt-in (see ensure_smem_optin below). - -// One block per draft position. A single strided pass over the vocabulary -// accumulates a per-thread online logsumexp (running max + sum) and a -// per-thread sorted-descending top-K; a shared-memory tree reduction then -// merges both across threads. log_prob[k] = scaled_logit[k] - log_z. -__global__ void draft_topk_kernel(const float * __restrict__ logits, - int vocab, int K, float inv_t, - float * __restrict__ out_lp, - int32_t * __restrict__ out_ids) { +constexpr int kMaxK = 8; // ddtree_K is 8 in practice; K>kMaxK → CPU fallback +constexpr int kBlock = 256; // threads per block (power of two for the reduction) +constexpr int kMaxSplit = 128; // max vocab splits per position (combine-block cap) + +// The workload is tiny in rows (n_positions ~ 15) but huge in vocab (~152k), +// so it is purely DRAM-bandwidth bound: ~9 MB to read per call. The original +// one-block-per-position layout launched only ~15 blocks, leaving most of the +// GPU's SMs idle and hitting only a fraction of peak bandwidth. We instead +// split each position's vocab scan across `split` blocks (a 2D grid of +// n_positions × split) so the whole device stays busy, then a cheap second +// kernel merges the `split` partials per position into the final top-K. + +// Merge two sorted-descending top-K lists (av/ai, bv/bi) into dst (dv/di), +// keeping the K largest; ties resolve toward the lower id. dst may alias av/ai +// safely only via the temporary below, so callers pass a distinct dst when the +// inputs come from shared memory. +template +__device__ __forceinline__ void merge_topk(const float * av, const int * ai, + const float * bv, const int * bi, + float * dv, int * di) { + int ia = 0, ib = 0; +#pragma unroll + for (int k = 0; k < K; k++) { + bool takeA; + if (ib >= K) takeA = true; + else if (ia >= K) takeA = false; + else if (av[ia] > bv[ib]) takeA = true; + else if (av[ia] < bv[ib]) takeA = false; + else takeA = (ai[ia] <= bi[ib]); // tie → lower id + if (takeA) { dv[k] = av[ia]; di[k] = ai[ia]; ia++; } + else { dv[k] = bv[ib]; di[k] = bi[ib]; ib++; } + } +} + +// Fold one logit (raw value `LRAW`, vocab id `ID`) into a thread's running +// online-logsumexp (lmax,lsum) and its register-resident sorted-descending +// top-K (topv,topi). K is a compile-time constant so the unrolled bubble uses +// only fixed indices — the arrays stay in registers instead of spilling to +// local memory the way a data-dependent insertion index would. The new value +// enters at slot K-1 and bubbles up only past strictly-smaller entries, so on +// ties the earlier (lower-id, since ids ascend within a thread) entry wins — +// matching the CPU heap's first-wins behaviour. +#define DFLASH_TOPK_CONSUME(LRAW, ID) \ + do { \ + const float _l = (LRAW) * inv_t; \ + if (_l > lmax) { lsum = lsum * __expf(lmax - _l) + 1.0f; lmax = _l; } \ + else { lsum += __expf(_l - lmax); } \ + if (_l > topv[K - 1]) { \ + topv[K - 1] = _l; topi[K - 1] = (ID); \ + _Pragma("unroll") \ + for (int _p = K - 1; _p > 0; --_p) { \ + if (topv[_p] > topv[_p - 1]) { \ + const float _tv = topv[_p]; \ + topv[_p] = topv[_p - 1]; topv[_p - 1] = _tv; \ + const int _ti = topi[_p]; \ + topi[_p] = topi[_p - 1]; topi[_p - 1] = _ti; \ + } \ + } \ + } \ + } while (0) + +// ---- pass 1: per-(position, split) partial logsumexp + local top-K --------- +// Grid: (n_positions, split). Block s of row `row` scans its contiguous vocab +// chunk and writes one partial (max, sum, sorted top-K) to part_* at index +// row*split + s. Contiguous chunks keep the strided per-warp reads coalesced. +// VEC selects 16-byte float4 loads (one coalesced transaction per 4 logits) — +// only used when every row's base pointer is 16-byte aligned (vocab % 4 == 0 +// and an aligned tensor); otherwise the scalar path is used. +template +__global__ void draft_topk_partial(const float * __restrict__ logits, + int vocab, float inv_t, int split, + float * __restrict__ part_max, + float * __restrict__ part_sum, + float * __restrict__ part_v, + int32_t * __restrict__ part_i) { const int row = blockIdx.x; + const int s = blockIdx.y; const int tid = threadIdx.x; const float * __restrict__ li = logits + (size_t)row * vocab; float lmax = -FLT_MAX; float lsum = 0.0f; - float topv[kMaxK]; - int topi[kMaxK]; + float topv[K]; + int topi[K]; #pragma unroll - for (int k = 0; k < kMaxK; k++) { topv[k] = -FLT_MAX; topi[k] = -1; } - - // ---- single strided pass: logsumexp + local top-K ------------------ - // j ascends within a thread, so a strict-greater insert keeps the lower - // id on ties (matches the CPU heap's first-wins behaviour). - for (int j = tid; j < vocab; j += kBlock) { - const float l = li[j] * inv_t; - if (l > lmax) { lsum = lsum * __expf(lmax - l) + 1.0f; lmax = l; } - else { lsum += __expf(l - lmax); } - if (l > topv[K - 1]) { - int p = K - 1; - while (p > 0 && topv[p - 1] < l) { - topv[p] = topv[p - 1]; topi[p] = topi[p - 1]; --p; - } - topv[p] = l; topi[p] = j; + for (int k = 0; k < K; k++) { topv[k] = -FLT_MAX; topi[k] = -1; } + + if (VEC) { + // Partition the float4s of the row into `split` contiguous chunks. + const int vocab4 = vocab >> 2; // # of full float4s + const int chunk4 = (vocab4 + split - 1) / split; + const int b4 = s * chunk4; + const int e4 = min(b4 + chunk4, vocab4); + const float4 * __restrict__ li4 = reinterpret_cast(li); + for (int j4 = b4 + tid; j4 < e4; j4 += kBlock) { + const float4 f = li4[j4]; + const int base = j4 << 2; + DFLASH_TOPK_CONSUME(f.x, base + 0); + DFLASH_TOPK_CONSUME(f.y, base + 1); + DFLASH_TOPK_CONSUME(f.z, base + 2); + DFLASH_TOPK_CONSUME(f.w, base + 3); } + // Tail elements past the last full float4 (only when vocab % 4 != 0); + // the last split owns them so no id is scanned twice. + if (s == split - 1) { + for (int j = (vocab4 << 2) + tid; j < vocab; j += kBlock) + DFLASH_TOPK_CONSUME(li[j], j); + } + } else { + const int chunk = (vocab + split - 1) / split; + const int begin = s * chunk; + const int end = min(begin + chunk, vocab); + for (int j = begin + tid; j < end; j += kBlock) + DFLASH_TOPK_CONSUME(li[j], j); } - // ---- block reduction -------------------------------------------------- + // ---- block reduction over kBlock threads ------------------------------ __shared__ float s_max[kBlock]; __shared__ float s_sum[kBlock]; - extern __shared__ char s_raw[]; - float * s_topv = reinterpret_cast(s_raw); - int32_t * s_topi = reinterpret_cast(s_topv + (size_t)kBlock * K); + __shared__ float s_topv[kBlock * K]; + __shared__ int32_t s_topi[kBlock * K]; s_max[tid] = lmax; s_sum[tid] = lsum; +#pragma unroll for (int k = 0; k < K; k++) { - s_topv[(size_t)tid * K + k] = topv[k]; - s_topi[(size_t)tid * K + k] = topi[k]; + s_topv[tid * K + k] = topv[k]; + s_topi[tid * K + k] = topi[k]; } __syncthreads(); for (int stride = kBlock / 2; stride > 0; stride >>= 1) { if (tid < stride) { - // logsumexp merge of (max,sum) pairs const float am = s_max[tid], as = s_sum[tid]; const float bm = s_max[tid + stride], bs = s_sum[tid + stride]; const float m = fmaxf(am, bm); s_sum[tid] = as * __expf(am - m) + bs * __expf(bm - m); s_max[tid] = m; - // merge two sorted-desc top-K lists, keep the K largest (lower id on tie) - float av[kMaxK]; int ai[kMaxK]; - float bv[kMaxK]; int bi[kMaxK]; + float dv[K]; int di[K]; + merge_topk(&s_topv[tid * K], &s_topi[tid * K], + &s_topv[(tid + stride) * K], &s_topi[(tid + stride) * K], + dv, di); +#pragma unroll for (int k = 0; k < K; k++) { - av[k] = s_topv[(size_t)tid * K + k]; - ai[k] = s_topi[(size_t)tid * K + k]; - bv[k] = s_topv[(size_t)(tid + stride) * K + k]; - bi[k] = s_topi[(size_t)(tid + stride) * K + k]; + s_topv[tid * K + k] = dv[k]; + s_topi[tid * K + k] = di[k]; } - int ia = 0, ib = 0; + } + __syncthreads(); + } + + if (tid == 0) { + const int idx = row * split + s; + part_max[idx] = s_max[0]; + part_sum[idx] = s_sum[0]; +#pragma unroll + for (int k = 0; k < K; k++) { + part_v[(size_t)idx * K + k] = s_topv[k]; + part_i[(size_t)idx * K + k] = s_topi[k]; + } + } +} + +#undef DFLASH_TOPK_CONSUME + +// ---- pass 2: merge the `split` partials per position into the final top-K -- +// Grid: n_positions blocks of blockDim = pow2_ceil(split) threads. Thread t +// loads partial t (or an identity when t >= split), then a shared-memory tree +// reduction merges all partials. log_prob[k] = top_v[k] - log_z. +template +__global__ void draft_topk_combine(const float * __restrict__ part_max, + const float * __restrict__ part_sum, + const float * __restrict__ part_v, + const int32_t * __restrict__ part_i, + int split, + float * __restrict__ out_lp, + int32_t * __restrict__ out_ids) { + const int row = blockIdx.x; + const int tid = threadIdx.x; // blockDim is pow2_ceil(split) + + __shared__ float s_max[kMaxSplit]; + __shared__ float s_sum[kMaxSplit]; + __shared__ float s_topv[kMaxSplit * K]; + __shared__ int32_t s_topi[kMaxSplit * K]; + + if (tid < split) { + const int idx = row * split + tid; + s_max[tid] = part_max[idx]; + s_sum[tid] = part_sum[idx]; +#pragma unroll + for (int k = 0; k < K; k++) { + s_topv[tid * K + k] = part_v[(size_t)idx * K + k]; + s_topi[tid * K + k] = part_i[(size_t)idx * K + k]; + } + } else { + s_max[tid] = -FLT_MAX; + s_sum[tid] = 0.0f; +#pragma unroll + for (int k = 0; k < K; k++) { + s_topv[tid * K + k] = -FLT_MAX; + s_topi[tid * K + k] = -1; + } + } + __syncthreads(); + + for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (tid < stride) { + const float am = s_max[tid], as = s_sum[tid]; + const float bm = s_max[tid + stride], bs = s_sum[tid + stride]; + const float m = fmaxf(am, bm); + s_sum[tid] = as * __expf(am - m) + bs * __expf(bm - m); + s_max[tid] = m; + + float dv[K]; int di[K]; + merge_topk(&s_topv[tid * K], &s_topi[tid * K], + &s_topv[(tid + stride) * K], &s_topi[(tid + stride) * K], + dv, di); +#pragma unroll for (int k = 0; k < K; k++) { - bool takeA; - if (ib >= K) takeA = true; - else if (ia >= K) takeA = false; - else if (av[ia] > bv[ib]) takeA = true; - else if (av[ia] < bv[ib]) takeA = false; - else takeA = (ai[ia] <= bi[ib]); // tie → lower id - if (takeA) { s_topv[(size_t)tid * K + k] = av[ia]; s_topi[(size_t)tid * K + k] = ai[ia]; ia++; } - else { s_topv[(size_t)tid * K + k] = bv[ib]; s_topi[(size_t)tid * K + k] = bi[ib]; ib++; } + s_topv[tid * K + k] = dv[k]; + s_topi[tid * K + k] = di[k]; } } __syncthreads(); @@ -105,6 +242,7 @@ __global__ void draft_topk_kernel(const float * __restrict__ logits, if (tid == 0) { const float log_z = s_max[0] + logf(s_sum[0]); +#pragma unroll for (int k = 0; k < K; k++) { out_lp[(size_t)row * K + k] = s_topv[k] - log_z; out_ids[(size_t)row * K + k] = s_topi[k]; @@ -112,29 +250,80 @@ __global__ void draft_topk_kernel(const float * __restrict__ logits, } } -// Per-device scratch for the [n_positions × K] outputs, grown as needed. The -// decode loop is single-threaded, so a plain static cache is safe and avoids a +// Per-device scratch for the [n_positions × K] outputs plus the +// [n_positions × split × K] pass-1 partials, grown as needed. The decode loop +// is single-threaded, so a plain static cache is safe and avoids a // cudaMalloc/cudaFree on every step. struct Scratch { - int device = -1; - size_t cap = 0; // elements - float * d_lp = nullptr; - int32_t * d_ids = nullptr; + int device = -1; + size_t cap = 0; // output elements (n_positions*K) + size_t part_cap = 0; // partial-list elements (n_positions*split*kMaxK) + float * d_lp = nullptr; + int32_t * d_ids = nullptr; + float * d_pmax = nullptr; // [n_positions*split] + float * d_psum = nullptr; // [n_positions*split] + float * d_pv = nullptr; // [n_positions*split*kMaxK] + int32_t * d_pi = nullptr; // [n_positions*split*kMaxK] }; Scratch g_scratch; -bool ensure_scratch(int device, size_t n) { - if (g_scratch.device == device && g_scratch.cap >= n) return true; - if (g_scratch.d_lp) cudaFree(g_scratch.d_lp); - if (g_scratch.d_ids) cudaFree(g_scratch.d_ids); +void free_scratch() { + if (g_scratch.d_lp) cudaFree(g_scratch.d_lp); + if (g_scratch.d_ids) cudaFree(g_scratch.d_ids); + if (g_scratch.d_pmax) cudaFree(g_scratch.d_pmax); + if (g_scratch.d_psum) cudaFree(g_scratch.d_psum); + if (g_scratch.d_pv) cudaFree(g_scratch.d_pv); + if (g_scratch.d_pi) cudaFree(g_scratch.d_pi); g_scratch = Scratch{}; - if (cudaMalloc(&g_scratch.d_lp, n * sizeof(float)) != cudaSuccess) return false; - if (cudaMalloc(&g_scratch.d_ids, n * sizeof(int32_t)) != cudaSuccess) { - cudaFree(g_scratch.d_lp); g_scratch.d_lp = nullptr; return false; - } - g_scratch.device = device; - g_scratch.cap = n; +} + +// Allocate output + partial buffers. n = n_positions*K outputs; +// n_parts = n_positions*split partials (each carrying K entries). +bool ensure_scratch(int device, size_t n, size_t n_parts) { + const size_t n_part_lists = n_parts * (size_t)kMaxK; // upper bound on K + if (g_scratch.device == device && g_scratch.cap >= n && + g_scratch.part_cap >= n_part_lists) + return true; + free_scratch(); + if (cudaMalloc(&g_scratch.d_lp, n * sizeof(float)) != cudaSuccess) goto fail; + if (cudaMalloc(&g_scratch.d_ids, n * sizeof(int32_t)) != cudaSuccess) goto fail; + if (cudaMalloc(&g_scratch.d_pmax, n_parts * sizeof(float)) != cudaSuccess) goto fail; + if (cudaMalloc(&g_scratch.d_psum, n_parts * sizeof(float)) != cudaSuccess) goto fail; + if (cudaMalloc(&g_scratch.d_pv, n_part_lists * sizeof(float)) != cudaSuccess) goto fail; + if (cudaMalloc(&g_scratch.d_pi, n_part_lists * sizeof(int32_t)) != cudaSuccess) goto fail; + g_scratch.device = device; + g_scratch.cap = n; + g_scratch.part_cap = n_part_lists; return true; +fail: + free_scratch(); + return false; +} + +inline int pow2_ceil(int x) { + int p = 1; + while (p < x) p <<= 1; + return p; +} + +// Choose how many blocks to split each position's vocab scan across. The work +// is bandwidth bound, so we want enough total blocks (n_positions × split) to +// saturate the device while keeping each chunk large enough that the strided +// reads stay coalesced and per-block overhead stays amortized. +int pick_split(int vocab, int n_positions) { + if (const char * v = std::getenv("DFLASH_TOPK_SPLIT")) { + int s = std::atoi(v); + if (s >= 1 && s <= kMaxSplit) return s; + } + // Aim for ~240 total blocks (≈3 waves on a ~80-SM device, measured sweet + // spot for vocab~150k), but cap the chunk floor at ~2k elements so a block + // never scans too little to be worth launching. + int by_blocks = (240 + n_positions - 1) / n_positions; + int by_chunk = vocab / 2048; + int split = by_blocks < by_chunk ? by_blocks : by_chunk; + if (split < 1) split = 1; + if (split > kMaxSplit) split = kMaxSplit; + return split; } } // namespace @@ -160,25 +349,44 @@ bool extract_draft_topk_cuda(const void * d_logits, static const bool kProfile = std::getenv("DFLASH_TOPK_PROFILE") != nullptr; bool ok = false; - const size_t n = (size_t)n_positions * K; - if (ensure_scratch(dev, n)) { - const float inv_t = 1.0f / fmaxf(1e-3f, temperature); - const size_t smem = (size_t)kBlock * K * (sizeof(float) + sizeof(int32_t)); - // Opt in to >48 KiB dynamic shared (once per device). If it fails the - // launch below will error and we fall back to the CPU path. - static int s_optin_dev = -1; - if (s_optin_dev != dev) { - constexpr int kMaxSmem = (int)((size_t)kBlock * kMaxK * - (sizeof(float) + sizeof(int32_t))); - cudaFuncSetAttribute(draft_topk_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, kMaxSmem); - s_optin_dev = dev; - } + const int split = pick_split(vocab, n_positions); + const size_t n = (size_t)n_positions * K; + const size_t n_parts = (size_t)n_positions * split; + if (ensure_scratch(dev, n, n_parts)) { + const float inv_t = 1.0f / fmaxf(1e-3f, temperature); cudaEvent_t e_k0, e_k1, e_c1; if (kProfile) { cudaEventCreate(&e_k0); cudaEventCreate(&e_k1); cudaEventCreate(&e_c1); cudaEventRecord(e_k0); } - draft_topk_kernel<<>>( - static_cast(d_logits), vocab, K, inv_t, - g_scratch.d_lp, g_scratch.d_ids); + + const dim3 grid1(n_positions, split); + const int comb_block = pow2_ceil(split); + const float * lp_in = static_cast(d_logits); + // float4 loads are safe only when every row base is 16-byte aligned: + // the tensor base aligned and a vocab stride that is a multiple of 4. + const bool use_vec = (vocab % 4 == 0) && + (reinterpret_cast(lp_in) % 16 == 0); + // K (and the vectorization flag) are compile-time template parameters + // so the per-thread/per-partial top-K stays register-resident; dispatch + // the runtime K to its instantiation. K>kMaxK is already rejected above. +#define DFLASH_TOPK_LAUNCH(KV, VEC) \ + draft_topk_partial<<>>( \ + lp_in, vocab, inv_t, split, \ + g_scratch.d_pmax, g_scratch.d_psum, g_scratch.d_pv, g_scratch.d_pi); \ + draft_topk_combine<<>>( \ + g_scratch.d_pmax, g_scratch.d_psum, g_scratch.d_pv, g_scratch.d_pi, \ + split, g_scratch.d_lp, g_scratch.d_ids); +#define DFLASH_TOPK_CASE(KV) \ + case KV: \ + if (use_vec) { DFLASH_TOPK_LAUNCH(KV, true) } \ + else { DFLASH_TOPK_LAUNCH(KV, false) } \ + break; + switch (K) { + DFLASH_TOPK_CASE(1) DFLASH_TOPK_CASE(2) DFLASH_TOPK_CASE(3) DFLASH_TOPK_CASE(4) + DFLASH_TOPK_CASE(5) DFLASH_TOPK_CASE(6) DFLASH_TOPK_CASE(7) DFLASH_TOPK_CASE(8) + default: break; + } +#undef DFLASH_TOPK_CASE +#undef DFLASH_TOPK_LAUNCH + if (kProfile) cudaEventRecord(e_k1); if (cudaGetLastError() == cudaSuccess && cudaDeviceSynchronize() == cudaSuccess) { const cudaError_t e1 = cudaMemcpy(out_log_probs, g_scratch.d_lp, @@ -192,8 +400,8 @@ bool extract_draft_topk_cuda(const void * d_logits, float k_ms = 0, c_ms = 0; cudaEventElapsedTime(&k_ms, e_k0, e_k1); cudaEventElapsedTime(&c_ms, e_k1, e_c1); - std::fprintf(stderr, "[topk] kernel=%.3f ms sync+copy=%.3f ms (n_pos=%d vocab=%d)\n", - k_ms, c_ms, n_positions, vocab); + std::fprintf(stderr, "[topk] kernels=%.3f ms sync+copy=%.3f ms (n_pos=%d vocab=%d split=%d)\n", + k_ms, c_ms, n_positions, vocab, split); cudaEventDestroy(e_k0); cudaEventDestroy(e_k1); cudaEventDestroy(e_c1); } } From c2289b741ffd903467b50e76536ae4d356d1c569 Mon Sep 17 00:00:00 2001 From: "Pramodith (via Claude Code)" Date: Fri, 19 Jun 2026 08:52:44 +0000 Subject: [PATCH 07/10] topm kernels. --- server/models | 1 + server/scripts/bench_llm.py | 4 +- server/src/common/bench_restricted.cu | 146 ++++++++++++++ server/src/common/bench_topm.cu | 72 +++++++ server/src/common/restricted_lm_head_cuda.cu | 193 +++++++++++++++++++ server/src/common/restricted_lm_head_cuda.h | 33 ++++ server/src/common/topm_extract_cuda.cu | 172 +++++++++++++++++ server/src/common/topm_extract_cuda.h | 22 +++ server/src/internal.h | 5 + server/src/qwen35/graph_builders.cpp | 55 ++++-- server/src/qwen35/graph_builders.h | 6 +- server/src/qwen35/qwen35_target_graph.cpp | 26 +-- server/test/test_dflash.cpp | 113 ++++++++++- 13 files changed, 815 insertions(+), 33 deletions(-) create mode 120000 server/models create mode 100644 server/src/common/bench_restricted.cu create mode 100644 server/src/common/bench_topm.cu create mode 100644 server/src/common/restricted_lm_head_cuda.cu create mode 100644 server/src/common/restricted_lm_head_cuda.h create mode 100644 server/src/common/topm_extract_cuda.cu create mode 100644 server/src/common/topm_extract_cuda.h diff --git a/server/models b/server/models new file mode 120000 index 000000000..77851bf17 --- /dev/null +++ b/server/models @@ -0,0 +1 @@ +/workspace/lucebox-hub/server/models \ No newline at end of file diff --git a/server/scripts/bench_llm.py b/server/scripts/bench_llm.py index 4722a7f3f..001783257 100644 --- a/server/scripts/bench_llm.py +++ b/server/scripts/bench_llm.py @@ -53,8 +53,8 @@ def _gsm_gold(x): BENCHES = [ - ("HumanEval", "openai_humaneval", None, "test", lambda x: x["prompt"], None, N_GEN), - ("GSM8K", "gsm8k", "main", "test", lambda x: f"Question: {x['question']}\nAnswer: ", _gsm_gold, 1024), + ("HumanEval", "openai/openai_humaneval", None, "test", lambda x: x["prompt"], None, N_GEN), + ("GSM8K", "openai/gsm8k", "main", "test", lambda x: f"Question: {x['question']}\nAnswer: ", _gsm_gold, 1024), ("Math500", "HuggingFaceH4/MATH-500", None, "test", lambda x: f"Problem: {x['problem']}\nSolution: Put your final answer in \\boxed{{}}.\n", lambda x: x["answer"], 2048), ] diff --git a/server/src/common/bench_restricted.cu b/server/src/common/bench_restricted.cu new file mode 100644 index 000000000..c03932748 --- /dev/null +++ b/server/src/common/bench_restricted.cu @@ -0,0 +1,146 @@ +// Standalone correctness + timing for restricted_lm_head_q6k. +// nvcc -O3 -arch=sm_86 -o /tmp/bench_restricted \ +// server/src/common/bench_restricted.cu server/src/common/restricted_lm_head_cuda.cu +// /tmp/bench_restricted +#include "restricted_lm_head_cuda.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using dflash::common::restricted_lm_head_q6k; + +static constexpr int QK = 256; +struct BlockQ6K { uint8_t ql[128]; uint8_t qh[64]; int8_t scales[16]; uint16_t d; }; +static_assert(sizeof(BlockQ6K) == 210, "layout"); + +// CPU dequant of weight j∈[0,256) in block b — mirrors dequantize_row_q6_K. +static float cpu_dequant(const BlockQ6K& b, int j) { + float d = __half2float(*reinterpret_cast(&b.d)); + int half = j >> 7, jj = j & 127, group = jj >> 5, l = jj & 31, is = l >> 4; + int qlb = half * 64, qhb = half * 32, scb = half * 8; + uint8_t qh = b.qh[qhb + l]; + int q, si; + switch (group) { + case 0: q = (b.ql[qlb + l] & 0xF) | (((qh >> 0) & 3) << 4); si = scb + is + 0; break; + case 1: q = (b.ql[qlb + l + 32] & 0xF) | (((qh >> 2) & 3) << 4); si = scb + is + 2; break; + case 2: q = (b.ql[qlb + l] >> 4) | (((qh >> 4) & 3) << 4); si = scb + is + 4; break; + default:q = (b.ql[qlb + l + 32] >> 4) | (((qh >> 6) & 3) << 4); si = scb + is + 6; break; + } + return d * (float)b.scales[si] * (float)(q - 32); +} + +int main(int argc, char** argv) { + // ---- correctness on small dims ---- + { + const int n_embd = 5120, n_vocab = 3000, n_tokens = 5, M = 64; + const int bpr = n_embd / QK; + std::mt19937 rng(7); + std::uniform_int_distribution byte(0, 255); + std::vector head((size_t)n_vocab * bpr); + for (auto& b : head) { + for (auto& x : b.ql) x = byte(rng); + for (auto& x : b.qh) x = byte(rng); + for (auto& x : b.scales) x = (int8_t)byte(rng); + b.d = __half_as_ushort(__float2half(0.005f + 0.001f * (byte(rng) / 255.0f))); + } + std::normal_distribution nf(0, 1); + std::vector hidden((size_t)n_embd * n_tokens); + for (auto& x : hidden) x = nf(rng); + std::vector cand((size_t)M * n_tokens); + for (int p = 0; p < n_tokens; p++) { + // distinct random ids per position + std::vector ids(n_vocab); for (int i = 0; i < n_vocab; i++) ids[i] = i; + std::shuffle(ids.begin(), ids.end(), rng); + for (int c = 0; c < M; c++) cand[(size_t)p * M + c] = ids[c]; + } + // CPU reference argmax over candidates + std::vector ref(n_tokens); + for (int p = 0; p < n_tokens; p++) { + float best = -1e30f; int bestid = -1, bestc = -1; + for (int c = 0; c < M; c++) { + int id = cand[(size_t)p * M + c]; + const BlockQ6K* row = head.data() + (size_t)id * bpr; + double acc = 0; + for (int e = 0; e < n_embd; e++) + acc += (double)cpu_dequant(row[e >> 8], e & 255) * hidden[(size_t)p * n_embd + e]; + if ((float)acc > best || ((float)acc == best && c < bestc)) { best = (float)acc; bestid = id; bestc = c; } + } + ref[p] = bestid; + } + // GPU + void *d_head, *d_keys; float* d_hidden; int32_t *d_cand, *d_out; + cudaMalloc(&d_head, head.size() * sizeof(BlockQ6K)); + cudaMalloc(&d_hidden, hidden.size() * sizeof(float)); + cudaMalloc(&d_cand, cand.size() * sizeof(int32_t)); + cudaMalloc(&d_out, n_tokens * sizeof(int32_t)); + cudaMalloc(&d_keys, n_tokens * sizeof(uint64_t)); + cudaMemcpy(d_head, head.data(), head.size() * sizeof(BlockQ6K), cudaMemcpyHostToDevice); + cudaMemcpy(d_hidden, hidden.data(), hidden.size() * sizeof(float), cudaMemcpyHostToDevice); + cudaMemcpy(d_cand, cand.data(), cand.size() * sizeof(int32_t), cudaMemcpyHostToDevice); + bool ok = restricted_lm_head_q6k(d_head, n_embd, n_vocab, d_hidden, n_tokens, + d_cand, M, d_out, d_keys, 0); + cudaDeviceSynchronize(); + std::vector got(n_tokens); + cudaMemcpy(got.data(), d_out, n_tokens * sizeof(int32_t), cudaMemcpyDeviceToHost); + int mism = 0; + for (int p = 0; p < n_tokens; p++) { + if (got[p] != ref[p]) { printf(" MISMATCH p=%d gpu=%d cpu=%d\n", p, got[p], ref[p]); mism++; } + } + printf("correctness (n_embd=%d n_vocab=%d n_tok=%d M=%d): ok=%d mismatches=%d\n", + n_embd, n_vocab, n_tokens, M, (int)ok, mism); + cudaFree(d_head); cudaFree(d_hidden); cudaFree(d_cand); cudaFree(d_out); cudaFree(d_keys); + } + + // ---- timing at realistic dims: restricted-M vs full-vocab via same kernel ---- + { + const int n_embd = 5120, n_vocab = 248320, n_tokens = 16; + const int bpr = n_embd / QK; + std::mt19937 rng(1); + std::uniform_int_distribution byte(0, 255); + std::vector head((size_t)n_vocab * bpr); + for (auto& b : head) { + for (auto& x : b.ql) x = byte(rng); + for (auto& x : b.qh) x = byte(rng); + for (auto& x : b.scales) x = (int8_t)byte(rng); + b.d = __half_as_ushort(__float2half(0.005f)); + } + std::vector hidden((size_t)n_embd * n_tokens, 0.01f); + void *d_head, *d_keys; float* d_hidden; int32_t *d_cand, *d_out; + cudaMalloc(&d_head, head.size() * sizeof(BlockQ6K)); + cudaMalloc(&d_hidden, hidden.size() * sizeof(float)); + cudaMalloc(&d_out, n_tokens * sizeof(int32_t)); + cudaMalloc(&d_keys, n_tokens * sizeof(uint64_t)); + cudaMemcpy(d_head, head.data(), head.size() * sizeof(BlockQ6K), cudaMemcpyHostToDevice); + + auto time_M = [&](int M, int iters) -> float { + std::vector cand((size_t)M * n_tokens); + for (int p = 0; p < n_tokens; p++) + for (int c = 0; c < M; c++) cand[(size_t)p * M + c] = (c * 7 + p) % n_vocab; + cudaMalloc(&d_cand, cand.size() * sizeof(int32_t)); + cudaMemcpy(d_cand, cand.data(), cand.size() * sizeof(int32_t), cudaMemcpyHostToDevice); + for (int i = 0; i < 5; i++) + restricted_lm_head_q6k(d_head, n_embd, n_vocab, d_hidden, n_tokens, d_cand, M, d_out, d_keys, 0); + cudaDeviceSynchronize(); + cudaEvent_t e0, e1; cudaEventCreate(&e0); cudaEventCreate(&e1); + cudaEventRecord(e0); + for (int i = 0; i < iters; i++) + restricted_lm_head_q6k(d_head, n_embd, n_vocab, d_hidden, n_tokens, d_cand, M, d_out, d_keys, 0); + cudaEventRecord(e1); cudaEventSynchronize(e1); + float ms = 0; cudaEventElapsedTime(&ms, e0, e1); + cudaFree(d_cand); + return ms / iters; + }; + printf("\ntiming (n_embd=%d n_vocab=%d n_tok=%d):\n", n_embd, n_vocab, n_tokens); + printf(" full-vocab head (M=%d) : %.3f ms\n", n_vocab, time_M(n_vocab, 30)); + for (int M : {256, 1024, 2048, 8192}) + printf(" restricted M=%-6d : %.3f ms\n", M, time_M(M, 100)); + cudaFree(d_head); cudaFree(d_hidden); cudaFree(d_out); cudaFree(d_keys); + } + return 0; +} diff --git a/server/src/common/bench_topm.cu b/server/src/common/bench_topm.cu new file mode 100644 index 000000000..f2765f695 --- /dev/null +++ b/server/src/common/bench_topm.cu @@ -0,0 +1,72 @@ +// Correctness + timing for extract_topm_cuda. +// nvcc -O3 -arch=sm_86 -o /tmp/bench_topm \ +// server/src/common/bench_topm.cu server/src/common/topm_extract_cuda.cu +#include "topm_extract_cuda.h" +#include +#include +#include +#include +#include +#include +#include + +using dflash::common::extract_topm_cuda; +using dflash::common::extract_topm_scratch_bytes; + +int main(int argc, char** argv) { + const int vocab = argc > 1 ? atoi(argv[1]) : 248320; + const int n = argc > 2 ? atoi(argv[2]) : 16; + const int M = argc > 3 ? atoi(argv[3]) : 1024; + + std::vector logits((size_t)vocab * n); + std::mt19937 rng(123); + std::normal_distribution d(0, 4); + for (auto& x : logits) x = d(rng); + + float* d_logits; int32_t* d_cand; void* d_scr; + cudaMalloc(&d_logits, logits.size() * sizeof(float)); + cudaMalloc(&d_cand, (size_t)M * n * sizeof(int32_t)); + cudaMalloc(&d_scr, extract_topm_scratch_bytes(n)); + cudaMemcpy(d_logits, logits.data(), logits.size() * sizeof(float), cudaMemcpyHostToDevice); + + bool ok = extract_topm_cuda(d_logits, vocab, n, M, d_cand, d_scr, 0); + cudaDeviceSynchronize(); + std::vector cand((size_t)M * n); + cudaMemcpy(cand.data(), d_cand, cand.size() * sizeof(int32_t), cudaMemcpyDeviceToHost); + + // verify per position + int tot_argmax_in = 0, tot_overlap = 0, tot_distinct_ok = 0; + for (int p = 0; p < n; p++) { + const float* col = logits.data() + (size_t)p * vocab; + // true top-M + std::vector idx(vocab); + for (int i = 0; i < vocab; i++) idx[i] = i; + std::nth_element(idx.begin(), idx.begin() + M, idx.end(), + [&](int a, int b){ return col[a] > col[b]; }); + std::unordered_set trueM(idx.begin(), idx.begin() + M); + int argmax = (int)(std::max_element(col, col + vocab) - col); + + std::unordered_set got; + for (int c = 0; c < M; c++) got.insert(cand[(size_t)p * M + c]); + if (got.count(argmax)) tot_argmax_in++; + if ((int)got.size() == M) tot_distinct_ok++; + int ov = 0; for (int id : got) if (trueM.count(id)) ov++; + tot_overlap += ov; + } + printf("extract_topm vocab=%d n=%d M=%d ok=%d:\n", vocab, n, M, (int)ok); + printf(" argmax-in-candidates : %d/%d positions\n", tot_argmax_in, n); + printf(" exactly-M-distinct : %d/%d positions\n", tot_distinct_ok, n); + printf(" overlap with true top-M: %.2f%% (avg)\n", 100.0 * tot_overlap / ((double)M * n)); + + // timing + for (int i = 0; i < 10; i++) extract_topm_cuda(d_logits, vocab, n, M, d_cand, d_scr, 0); + cudaDeviceSynchronize(); + cudaEvent_t e0, e1; cudaEventCreate(&e0); cudaEventCreate(&e1); + cudaEventRecord(e0); + const int it = 200; + for (int i = 0; i < it; i++) extract_topm_cuda(d_logits, vocab, n, M, d_cand, d_scr, 0); + cudaEventRecord(e1); cudaEventSynchronize(e1); + float ms = 0; cudaEventElapsedTime(&ms, e0, e1); + printf(" time: %.3f ms/call\n", ms / it); + return 0; +} diff --git a/server/src/common/restricted_lm_head_cuda.cu b/server/src/common/restricted_lm_head_cuda.cu new file mode 100644 index 000000000..981773f48 --- /dev/null +++ b/server/src/common/restricted_lm_head_cuda.cu @@ -0,0 +1,193 @@ +// Candidate-restricted LM head for the greedy verify path (Q6_K weight). +// +// Instead of the full-vocab head matmul (output.weight [n_embd × n_vocab] Q6_K @ +// hidden), compute logits only over a per-position candidate shortlist S (the +// draft's top-M tokens), then argmax within S. See draft top-k calibration: +// the target greedy token lies in the draft top-M with ~94% (M=128) … ~97.5% +// (M=1024) coverage, so a restricted head is near-lossless for greedy and reads +// only M rows of the head instead of all 248k. +// +// This is a SINGLE fused kernel: each (candidate, position) block gathers the +// candidate's Q6_K row, dequantizes it in registers, dots it with that +// position's hidden vector, and folds the result into a per-position argmax via +// a packed 64-bit atomicMax — no intermediate gathered-weights tensor and no +// separate logits buffer. Mirrors the contract in restricted_lm_head_cuda.h. + +#include "restricted_lm_head_cuda.h" + +#include +#include +#include +#include +#include +#include + +namespace dflash::common { + +namespace { + +// Q6_K block layout (must match ggml-common.h block_q6_K exactly): 256 weights +// per super-block, 210 bytes. +constexpr int kQK = 256; +struct __align__(2) BlockQ6K { + uint8_t ql[kQK / 2]; // 128 B: low 4 bits + uint8_t qh[kQK / 4]; // 64 B: high 2 bits + int8_t scales[kQK / 16]; // 16 B: per-16 sub-block scales + uint16_t d; // 2 B: super-block scale (f16 bits) +}; +static_assert(sizeof(BlockQ6K) == 210, "block_q6_K layout mismatch"); + +constexpr int kWarp = 32; +constexpr int kWarpsPerBlock = 4; // 128-thread blocks, 1 warp per candidate +constexpr int kBlock = kWarp * kWarpsPerBlock; +constexpr int kCandPerWarp = 1; // candidates each warp pipelines (>1 ⇒ double-buffered) +constexpr int kNBuf = (kCandPerWarp > 1) ? 2 : 1; // shared-mem row buffers per warp + +// Dequantize one Q6_K weight at local index j∈[0,256) of super-block `b`. +// Mirrors dequantize_row_q6_K (ggml-quants.c): the 256 weights split into two +// 128-wide halves; within a half, weight jj uses group=jj/32, l=jj%32. +__device__ __forceinline__ float dequant_q6k(const BlockQ6K & b, int j, float d) { + const int half = j >> 7; // 0/1 + const int jj = j & 127; + const int group = jj >> 5; // 0..3 + const int l = jj & 31; + const int is = l >> 4; // 0/1 + const int ql_base = half * 64; + const int qh_base = half * 32; + const int sc_base = half * 8; + const uint8_t qhb = b.qh[qh_base + l]; + int q, sc_idx; + switch (group) { + case 0: q = (b.ql[ql_base + l] & 0xF) | (((qhb >> 0) & 3) << 4); sc_idx = sc_base + is + 0; break; + case 1: q = (b.ql[ql_base + l + 32] & 0xF) | (((qhb >> 2) & 3) << 4); sc_idx = sc_base + is + 2; break; + case 2: q = (b.ql[ql_base + l] >> 4) | (((qhb >> 4) & 3) << 4); sc_idx = sc_base + is + 4; break; + default:q = (b.ql[ql_base + l + 32] >> 4) | (((qhb >> 6) & 3) << 4); sc_idx = sc_base + is + 6; break; + } + return d * (float)b.scales[sc_idx] * (float)(q - 32); +} + +// Pack (logit, index) into one 64-bit key whose unsigned ordering matches the +// desired argmax (larger logit wins; ties → lower index). The float bits go in +// the high 32 bits (monotonic-mapped so the int compare matches float order), +// the *complemented* index in the low 32 (so lower index ranks higher on ties). +__device__ __forceinline__ uint64_t pack_key(float logit, int idx) { + uint32_t u = __float_as_uint(logit); + // Map IEEE float bits to a monotonically-increasing unsigned ordering. + u = (u & 0x80000000u) ? ~u : (u | 0x80000000u); + return ((uint64_t)u << 32) | (uint32_t)(0xFFFFFFFFu - (uint32_t)idx); +} +__device__ __forceinline__ int unpack_idx(uint64_t key) { + return (int)(0xFFFFFFFFu - (uint32_t)(key & 0xFFFFFFFFu)); +} + +// Decode a smem-resident Q6_K row (rw 32-bit words) and dot with hidden h. +__device__ __forceinline__ float row_dot(const uint32_t * sbuf, int blocks_per_row, + const float * __restrict__ h, int lane) { + const BlockQ6K * row = reinterpret_cast(sbuf); + float acc = 0.0f; + for (int sb = 0; sb < blocks_per_row; sb++) { + const BlockQ6K & b = row[sb]; + const float d = __half2float(*reinterpret_cast(&b.d)); + const int base = sb << 8; +#pragma unroll + for (int j = lane; j < kQK; j += kWarp) acc += dequant_q6k(b, j, d) * h[base + j]; + } +#pragma unroll + for (int off = kWarp / 2; off > 0; off >>= 1) + acc += __shfl_down_sync(0xffffffffu, acc, off); + return acc; +} + +// One warp per group of kCandPerWarp candidates for position p. Each warp +// software-pipelines its candidate rows: it issues a cp.async bulk copy of the +// next Q6_K row into shared memory while it decodes/dots the current one (the +// rows are the only sizeable global traffic, so overlapping their loads with +// the tiny dequant+dot compute hides the load latency). logit folds into +// best_key[p] via a packed atomicMax. Each row's 4200 bytes are 4-aligned, so +// the whole row is copied in one coalesced async pass (per-super-block copies +// would hit 2-mod-4 offsets that cp.async can't express). +__global__ void restricted_head_kernel(const BlockQ6K * __restrict__ head, + int n_embd, int blocks_per_row, + const float * __restrict__ hidden, + const int32_t * __restrict__ cand_ids, + int M, + unsigned long long * __restrict__ best_key) { + extern __shared__ uint32_t smem[]; // [kWarpsPerBlock][2][rw] + const int warp = threadIdx.x >> 5; + const int lane = threadIdx.x & 31; + const int p = blockIdx.y; + const int c0 = (blockIdx.x * kWarpsPerBlock + warp) * kCandPerWarp; + if (c0 >= M) return; + + const int rw = (blocks_per_row * (int)sizeof(BlockQ6K) + 3) / 4; // row words + uint32_t * buf[2] = { smem + (size_t)(warp * kNBuf + 0) * rw, + smem + (size_t)(warp * kNBuf + (kNBuf>1?1:0)) * rw }; + const float * __restrict__ h = hidden + (size_t)p * n_embd; + const int kc = min(kCandPerWarp, M - c0); // valid candidates for this warp + + auto issue = [&](int slot, int c) { + const int32_t id = cand_ids[(size_t)p * M + c]; + const uint32_t * src = reinterpret_cast(head + (size_t)id * blocks_per_row); + for (int w = lane; w < rw; w += kWarp) + __pipeline_memcpy_async(&buf[slot][w], &src[w], sizeof(uint32_t)); + __pipeline_commit(); + }; + + issue(0, c0); // prefetch first row + for (int k = 0; k < kc; k++) { + const bool has_next = (k + 1 < kc); + if (has_next) issue((k + 1) & 1, c0 + k + 1); + __pipeline_wait_prior(has_next ? 1 : 0); // current row ready (next still in flight) + __syncwarp(); + const float acc = row_dot(buf[k & 1], blocks_per_row, h, lane); + if (lane == 0) + atomicMax(&best_key[p], (unsigned long long)pack_key(acc, c0 + k)); + __syncwarp(); + } +} + +// Resolve the packed argmax keys to vocab ids: out_tokens[p] = cand_ids[p*M + c*] +// where c* is the winning candidate index for position p. +__global__ void resolve_kernel(const unsigned long long * __restrict__ best_key, + const int32_t * __restrict__ cand_ids, int M, + int n_tokens, int32_t * __restrict__ out_tokens) { + const int p = blockIdx.x * blockDim.x + threadIdx.x; + if (p >= n_tokens) return; + const int c = unpack_idx(best_key[p]); + out_tokens[p] = cand_ids[(size_t)p * M + c]; +} + +} // namespace + +bool restricted_lm_head_q6k(const void * d_head_q6k, + int n_embd, int n_vocab, + const float * d_hidden, int n_tokens, + const int32_t * d_cand_ids, int M, + int32_t * d_out_tokens, + void * d_scratch_keys, + cudaStream_t stream) { + if (n_embd <= 0 || n_embd % kQK != 0 || n_tokens <= 0 || M <= 0) return false; + (void)n_vocab; + const int blocks_per_row = n_embd / kQK; + auto * best_key = static_cast(d_scratch_keys); + + // Init per-position best keys to the minimum (all bits 0 → most-negative). + if (cudaMemsetAsync(best_key, 0, (size_t)n_tokens * sizeof(unsigned long long), + stream) != cudaSuccess) return false; + + const int cands_per_block = kWarpsPerBlock * kCandPerWarp; + const int rw = (blocks_per_row * (int)sizeof(BlockQ6K) + 3) / 4; // row words + const size_t smem = (size_t)kWarpsPerBlock * kNBuf * rw * sizeof(uint32_t); + dim3 grid((M + cands_per_block - 1) / cands_per_block, n_tokens); + restricted_head_kernel<<>>( + static_cast(d_head_q6k), n_embd, blocks_per_row, + d_hidden, d_cand_ids, M, best_key); + + const int rb = 128; + resolve_kernel<<<(n_tokens + rb - 1) / rb, rb, 0, stream>>>( + best_key, d_cand_ids, M, n_tokens, d_out_tokens); + + return cudaGetLastError() == cudaSuccess; +} + +} // namespace dflash::common diff --git a/server/src/common/restricted_lm_head_cuda.h b/server/src/common/restricted_lm_head_cuda.h new file mode 100644 index 000000000..39ab89bc7 --- /dev/null +++ b/server/src/common/restricted_lm_head_cuda.h @@ -0,0 +1,33 @@ +// Candidate-restricted greedy LM head (Q6_K weight). See the .cu for the idea. +// +// Computes, for each of n_tokens positions, the argmax vocab id over that +// position's M candidate ids — using a fused gather + Q6_K dequant + dot + +// argmax kernel that reads only the M candidate rows of the head, not all +// n_vocab. Approximate-greedy: exact iff the true argmax ∈ the candidate set +// (calibrated; see docs/topk-head-optimization.md). + +#pragma once + +#include +#include + +namespace dflash::common { + +// d_head_q6k : device pointer to the Q6_K output.weight, row-major +// [n_embd × n_vocab] (n_embd is the contiguous/block dim; one +// vocab row is n_embd/256 Q6_K blocks). +// d_hidden : device [n_embd × n_tokens] f32 pre-head hidden states. +// d_cand_ids : device [M × n_tokens] i32, column p = position p's M candidate +// vocab ids. +// d_out_tokens : device [n_tokens] i32, written with the per-position argmax id. +// d_scratch_keys : device scratch of >= n_tokens * sizeof(uint64_t) bytes. +// Returns false on bad args or a launch error (caller falls back to full head). +bool restricted_lm_head_q6k(const void * d_head_q6k, + int n_embd, int n_vocab, + const float * d_hidden, int n_tokens, + const int32_t * d_cand_ids, int M, + int32_t * d_out_tokens, + void * d_scratch_keys, + cudaStream_t stream); + +} // namespace dflash::common diff --git a/server/src/common/topm_extract_cuda.cu b/server/src/common/topm_extract_cuda.cu new file mode 100644 index 000000000..f1ade0498 --- /dev/null +++ b/server/src/common/topm_extract_cuda.cu @@ -0,0 +1,172 @@ +// Fast GPU top-M candidate extractor for the candidate-restricted LM head. +// +// Given the draft logits [vocab × n_tokens] (column p = position p's vocab +// logits, contiguous), produce cand_ids [M × n_tokens]: M candidate vocab ids +// per position that contain (with high probability) the target's argmax. Used +// to feed restricted_lm_head_q6k. The draft top-M is far cheaper to find +// approximately than exactly: we radix-threshold on the order-preserving uint +// mapping of the float logits via one histogram pass, pick the threshold bin +// whose cumulative top-down count first reaches M, then gather. Bins strictly +// above the threshold are always kept; the threshold bin fills the remainder +// (its intra-bin order is irrelevant — the bin spans <0.05% of the value +// range, so coverage of the true top-M is preserved). Reads the logits twice. + +#include "topm_extract_cuda.h" + +#include +#include +#include + +namespace dflash::common { + +namespace { + +constexpr int kBits = 16; // histogram key bits (16 ⇒ ~exact top-M) +constexpr int kBins = 1 << kBits; +constexpr int kBlock = 256; +constexpr int kThrThreads = 256; // threads for the parallel threshold scan +static_assert(kBins % kThrThreads == 0, "kBins must divide threshold threads"); + +// Order-preserving map float→uint32: larger float ⇒ larger uint (radix-sort key). +__device__ __forceinline__ uint32_t order_key(float f) { + uint32_t u = __float_as_uint(f); + return (u & 0x80000000u) ? ~u : (u | 0x80000000u); +} + +// Pass 1: per-position histogram of the top kBits of order_key over the vocab. +__global__ void histogram_kernel(const float * __restrict__ logits, + int vocab, int n_tokens, + unsigned int * __restrict__ hist) { // [n_tokens][kBins] + const int p = blockIdx.y; + unsigned int * h = hist + (size_t)p * kBins; + const float * __restrict__ col = logits + (size_t)p * vocab; + for (int v = blockIdx.x * blockDim.x + threadIdx.x; v < vocab; + v += gridDim.x * blockDim.x) { + atomicAdd(&h[order_key(col[v]) >> (32 - kBits)], 1u); + } +} + +// Pass 2: find threshold_bin[p] = the bin where the top-down cumulative count +// first reaches M, and above_cnt[p] = count strictly above it (< M). One block +// per position: each thread owns a contiguous chunk of bins, computes its chunk +// total, thread 0 builds the per-chunk "count above" prefix, the single thread +// whose chunk straddles the M-crossing rescans just its chunk. O(kBins/T + T). +__global__ void threshold_kernel(const unsigned int * __restrict__ hist, + int n_tokens, int M, + int * __restrict__ threshold_bin, + unsigned int * __restrict__ above_cnt) { + const int p = blockIdx.x; + if (p >= n_tokens) return; + const unsigned int * h = hist + (size_t)p * kBins; + const int t = threadIdx.x; + const int chunk = kBins / kThrThreads; + const int lo = t * chunk, hi = lo + chunk; // bins [lo, hi) + + __shared__ unsigned int part[kThrThreads]; // chunk totals + __shared__ unsigned int above[kThrThreads]; // count strictly above chunk t + unsigned int s = 0; + for (int b = lo; b < hi; b++) s += h[b]; + part[t] = s; + __syncthreads(); + + if (t == 0) { // suffix sums over chunks + unsigned int acc = 0; + for (int c = kThrThreads - 1; c >= 0; c--) { above[c] = acc; acc += part[c]; } + } + __syncthreads(); + + // The crossing chunk: count above it < M, but adding its total reaches M. + if (above[t] < (unsigned)M && (unsigned)M <= above[t] + part[t]) { + unsigned int cum = above[t]; + int bb = lo; + for (int b = hi - 1; b >= lo; b--) { + if (cum + h[b] >= (unsigned)M) { bb = b; break; } + cum += h[b]; + } + threshold_bin[p] = bb; + above_cnt[p] = cum; // count strictly above bb (< M) + } +} + +// Pass 3a: emit all ids in bins strictly above the threshold bin (guaranteed +// in the top-M). Uses a per-position fill counter. +__global__ void gather_above_kernel(const float * __restrict__ logits, + int vocab, int n_tokens, int M, + const int * __restrict__ threshold_bin, + int32_t * __restrict__ cand_ids, + unsigned int * __restrict__ fill) { + const int p = blockIdx.y; + const int tb = threshold_bin[p]; + const float * __restrict__ col = logits + (size_t)p * vocab; + int32_t * out = cand_ids + (size_t)p * M; + for (int v = blockIdx.x * blockDim.x + threadIdx.x; v < vocab; + v += gridDim.x * blockDim.x) { + if ((int)(order_key(col[v]) >> (32 - kBits)) > tb) { + const unsigned int slot = atomicAdd(&fill[p], 1u); + if (slot < (unsigned)M) out[slot] = v; + } + } +} + +// Pass 3b: fill the remaining [above_cnt, M) slots from the threshold bin. +__global__ void gather_fill_kernel(const float * __restrict__ logits, + int vocab, int n_tokens, int M, + const int * __restrict__ threshold_bin, + int32_t * __restrict__ cand_ids, + unsigned int * __restrict__ fill) { + const int p = blockIdx.y; + const int tb = threshold_bin[p]; + const float * __restrict__ col = logits + (size_t)p * vocab; + int32_t * out = cand_ids + (size_t)p * M; + for (int v = blockIdx.x * blockDim.x + threadIdx.x; v < vocab; + v += gridDim.x * blockDim.x) { + if ((int)(order_key(col[v]) >> (32 - kBits)) == tb) { + const unsigned int slot = atomicAdd(&fill[p], 1u); + if (slot < (unsigned)M) out[slot] = v; + else return; // this position is full + } + } +} + +} // namespace + +bool extract_topm_cuda(const float * d_logits, int vocab, int n_tokens, int M, + int32_t * d_cand_ids, void * d_scratch, cudaStream_t stream) { + if (vocab <= 0 || n_tokens <= 0 || M <= 0 || M > vocab) return false; + + // scratch layout: hist[n_tokens*kBins] u32 | threshold_bin[n_tokens] i32 | + // above_cnt[n_tokens] u32 | fill[n_tokens] u32 + auto * base = static_cast(d_scratch); + auto * hist = reinterpret_cast(base); + auto * threshold_bin = reinterpret_cast(hist + (size_t)n_tokens * kBins); + auto * above_cnt = reinterpret_cast(threshold_bin + n_tokens); + auto * fill = above_cnt + n_tokens; + + if (cudaMemsetAsync(hist, 0, (size_t)n_tokens * kBins * sizeof(unsigned int), stream) != cudaSuccess) + return false; + cudaMemsetAsync(fill, 0, (size_t)n_tokens * sizeof(unsigned int), stream); + + const int nsplit = 64; // blocks per position for the vocab scan + dim3 grid(nsplit, n_tokens); + histogram_kernel<<>>(d_logits, vocab, n_tokens, hist); + threshold_kernel<<>>( + hist, n_tokens, M, threshold_bin, above_cnt); + gather_above_kernel<<>>( + d_logits, vocab, n_tokens, M, threshold_bin, d_cand_ids, fill); + // Reset fill to above_cnt so 3b continues filling from the right slot. + cudaMemcpyAsync(fill, above_cnt, (size_t)n_tokens * sizeof(unsigned int), + cudaMemcpyDeviceToDevice, stream); + gather_fill_kernel<<>>( + d_logits, vocab, n_tokens, M, threshold_bin, d_cand_ids, fill); + + return cudaGetLastError() == cudaSuccess; +} + +size_t extract_topm_scratch_bytes(int n_tokens) { + return (size_t)n_tokens * kBins * sizeof(unsigned int) // hist + + (size_t)n_tokens * sizeof(int) // threshold_bin + + (size_t)n_tokens * sizeof(unsigned int) // above_cnt + + (size_t)n_tokens * sizeof(unsigned int); // fill +} + +} // namespace dflash::common diff --git a/server/src/common/topm_extract_cuda.h b/server/src/common/topm_extract_cuda.h new file mode 100644 index 000000000..d9245d74d --- /dev/null +++ b/server/src/common/topm_extract_cuda.h @@ -0,0 +1,22 @@ +// Fast GPU top-M candidate extractor. See the .cu for the method. +// +// d_logits : device [vocab × n_tokens] f32 (column p contiguous = position p). +// d_cand_ids: device [M × n_tokens] i32, written with position p's M candidate +// vocab ids (unordered; contains the true top-(M) approximately, +// exactly above the threshold bin). +// d_scratch : device scratch of >= extract_topm_scratch_bytes(n_tokens). +// Returns false on bad args / launch error. + +#pragma once + +#include +#include +#include + +namespace dflash::common { + +bool extract_topm_cuda(const float * d_logits, int vocab, int n_tokens, int M, + int32_t * d_cand_ids, void * d_scratch, cudaStream_t stream); +size_t extract_topm_scratch_bytes(int n_tokens); + +} // namespace dflash::common diff --git a/server/src/internal.h b/server/src/internal.h index 4ff86937a..a5001746e 100644 --- a/server/src/internal.h +++ b/server/src/internal.h @@ -556,6 +556,7 @@ struct QwenGraphInputs { bool capture_moe_router = false; // if true, expose selected expert ids for MoE layers int fa_window = 0; // sliding window for FA layers: 0 = full attention bool last_token_logits_only = false; // if true, only compute logits for last token (prefill optimization) + bool skip_lm_head = false; // if true, return hidden_states only (no full-vocab head); for the restricted greedy head ggml_tensor * parent_ids = nullptr; // [n_tokens] i32; tree mode when non-null // [n_tokens,n_head_kv] i64; non-null = step-invariant KV write via ggml_set_rows (carries kv_start). ggml_tensor * kv_write_rows = nullptr; @@ -567,6 +568,10 @@ struct QwenGraphInputs { struct QwenGraphOutputs { ggml_tensor * logits; // [vocab, n_tokens] f32 + // Final pre-LM-head hidden [n_embd, n_tokens] f32 (always set). Used by the + // candidate-restricted greedy head, which gathers only the draft top-M head + // rows instead of the full-vocab matmul. + ggml_tensor * hidden_states = nullptr; // One entry per delta-net layer (48 for qwen35-27b). Only populated when // QwenGraphInputs::capture_delta_intermediate is true. Tensors are graph // views marked as ggml_set_output() so their data persists after diff --git a/server/src/qwen35/graph_builders.cpp b/server/src/qwen35/graph_builders.cpp index b4228c3d3..86bfd5514 100644 --- a/server/src/qwen35/graph_builders.cpp +++ b/server/src/qwen35/graph_builders.cpp @@ -271,7 +271,8 @@ bool build_target_step( int kq_stride_pad, bool capture_moe_router, bool kvflash_mask, - bool capture_qk) { + bool capture_qk, + bool restricted_head) { step_graph_free(sg); // Persistent thread_local arena: rebuilt step graphs land at identical @@ -351,18 +352,29 @@ bool build_target_step( gi.last_token_logits_only = last_token_logits_only; gi.kv_write_rows = sg.kv_write_rows; gi.q_capture = capture_qk; + gi.skip_lm_head = restricted_head; QwenGraphOutputs go = build_qwen35_graph(sg.ctx, sg.gf, w, cache, gi); - if (!go.logits) return false; - sg.logits = go.logits; sg.delta_captures = std::move(go.delta_captures); sg.moe_selected = std::move(go.moe_selected); - ggml_set_output(sg.logits); - sg.argmax_tokens = ggml_argmax(sg.ctx, sg.logits); - ggml_set_name(sg.argmax_tokens, "chain_verify_argmax"); - ggml_set_output(sg.argmax_tokens); - ggml_build_forward_expand(sg.gf, sg.argmax_tokens); + if (restricted_head) { + // Candidate-restricted greedy head: expose the pre-head hidden; the + // caller runs the top-M extractor + fused Q6_K head off-graph. No + // full-vocab logits / argmax in this graph. + if (!go.hidden_states) return false; + sg.hidden_states = go.hidden_states; + ggml_set_output(sg.hidden_states); + ggml_build_forward_expand(sg.gf, sg.hidden_states); + } else { + if (!go.logits) return false; + sg.logits = go.logits; + ggml_set_output(sg.logits); + sg.argmax_tokens = ggml_argmax(sg.ctx, sg.logits); + ggml_set_name(sg.argmax_tokens, "chain_verify_argmax"); + ggml_set_output(sg.argmax_tokens); + ggml_build_forward_expand(sg.gf, sg.argmax_tokens); + } if (!sg.alloc) { sg.alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); @@ -380,7 +392,8 @@ bool build_target_step_tree( int kv_start, int n_tokens, int fa_window, - int kq_stride_pad) { + int kq_stride_pad, + bool restricted_head) { step_graph_free(sg); ggml_init_params ip{}; @@ -422,17 +435,27 @@ bool build_target_step_tree( gi.capture_layers = true; gi.capture_delta_intermediate = true; gi.parent_ids = sg.parent_ids; + gi.skip_lm_head = restricted_head; QwenGraphOutputs go = build_qwen35_graph(sg.ctx, sg.gf, w, cache, gi); - if (!go.logits) return false; - sg.logits = go.logits; sg.delta_captures = std::move(go.delta_captures); - ggml_set_output(sg.logits); - sg.argmax_tokens = ggml_argmax(sg.ctx, sg.logits); - ggml_set_name(sg.argmax_tokens, "tree_verify_argmax"); - ggml_set_output(sg.argmax_tokens); - ggml_build_forward_expand(sg.gf, sg.argmax_tokens); + if (restricted_head) { + // Candidate-restricted greedy head: expose pre-head hidden; caller runs + // the top-M extractor + fused Q6_K head per tree node off-graph. + if (!go.hidden_states) return false; + sg.hidden_states = go.hidden_states; + ggml_set_output(sg.hidden_states); + ggml_build_forward_expand(sg.gf, sg.hidden_states); + } else { + if (!go.logits) return false; + sg.logits = go.logits; + ggml_set_output(sg.logits); + sg.argmax_tokens = ggml_argmax(sg.ctx, sg.logits); + ggml_set_name(sg.argmax_tokens, "tree_verify_argmax"); + ggml_set_output(sg.argmax_tokens); + ggml_build_forward_expand(sg.gf, sg.argmax_tokens); + } if (!sg.alloc) { sg.alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); diff --git a/server/src/qwen35/graph_builders.h b/server/src/qwen35/graph_builders.h index 7557a1f7e..71b7fea60 100644 --- a/server/src/qwen35/graph_builders.h +++ b/server/src/qwen35/graph_builders.h @@ -93,7 +93,8 @@ bool build_target_step( int kq_stride_pad = KQ_MASK_PAD, bool capture_moe_router = false, bool kvflash_mask = false, - bool capture_qk = false); + bool capture_qk = false, + bool restricted_head = false); // skip full head, expose hidden for the restricted greedy head // Full target forward: DDTree tree-verify mode. bool build_target_step_tree( @@ -104,7 +105,8 @@ bool build_target_step_tree( int kv_start, int n_tokens, int fa_window = 0, - int kq_stride_pad = KQ_MASK_PAD); + int kq_stride_pad = KQ_MASK_PAD, + bool restricted_head = false); // skip full head, expose hidden for the restricted greedy head // LM-head projection: project draft hidden states through the target output matrix. bool build_lm_head_projection_step( diff --git a/server/src/qwen35/qwen35_target_graph.cpp b/server/src/qwen35/qwen35_target_graph.cpp index 742aaa329..dc05370d3 100644 --- a/server/src/qwen35/qwen35_target_graph.cpp +++ b/server/src/qwen35/qwen35_target_graph.cpp @@ -1293,22 +1293,26 @@ QwenGraphOutputs build_qwen35_graph( // 3. LM head — optionally only for the last token (prefill optimization: // reduces logits from [vocab, n_tokens] to [vocab, 1], saving ~233MB - // scratch at ubatch=384 and eliminating a large matmul). + // scratch at ubatch=384 and eliminating a large matmul). When + // in.skip_lm_head is set (candidate-restricted greedy head) the + // full-vocab matmul is skipped and only the pre-head hidden is returned. + QwenGraphOutputs og = std::move(og_early); + + if (in.last_token_logits_only && n_tokens > 1) { + out = ggml_view_2d(ctx, out, hidden, 1, out->nb[1], + (size_t)(n_tokens - 1) * out->nb[1]); + } + ggml_set_name(out, "result_norm"); + ggml_set_output(out); + ggml_build_forward_expand(gf, out); + og.hidden_states = out; + ggml_tensor * logits = nullptr; - if (w.output) { - if (in.last_token_logits_only && n_tokens > 1) { - out = ggml_view_2d(ctx, out, hidden, 1, out->nb[1], - (size_t)(n_tokens - 1) * out->nb[1]); - } + if (w.output && !in.skip_lm_head) { logits = ggml_mul_mat(ctx, w.output, out); ggml_set_name(logits, "logits"); ggml_build_forward_expand(gf, logits); - } else { - ggml_set_name(out, "result_norm"); - ggml_build_forward_expand(gf, out); } - - QwenGraphOutputs og = std::move(og_early); og.logits = logits; return og; } diff --git a/server/test/test_dflash.cpp b/server/test/test_dflash.cpp index da54394cc..5b8d10a94 100644 --- a/server/test/test_dflash.cpp +++ b/server/test/test_dflash.cpp @@ -270,6 +270,10 @@ using dflash::common::free_qwen35_layer_split_shards; #include "qwen35_layer_split_dflash_target.h" #include "common/dflash_spec_decode.h" #include "common/gguf_mmap.h" +#ifdef DFLASH27B_HAVE_TOPK_HEAD +#include "common/restricted_lm_head_cuda.h" +#include "common/topm_extract_cuda.h" +#endif using dflash::common::is_eos_tok; // ─── Layer-split daemon — extracted to src/qwen35/layer_split_daemon.{h,cpp} ─ @@ -3104,6 +3108,29 @@ int main(int argc, char ** argv) { std::vector massp95_sum(calib_Mgrid.size(), 0.0); std::vector massp95_bad(calib_Mgrid.size(), 0); + // Candidate-restricted greedy LM head (DFLASH_TOPK_HEAD=M). On the chain + // fast-rollback single-GPU path, replace the verify-side full-vocab head + // with: draft top-M extract → fused Q6_K gather/dot over those candidates → + // argmax. Approximate-greedy (exact iff target argmax ∈ draft top-M). + int g_topk_head_M = 0; + double tt_head_extract = 0, tt_head_kernel = 0; + long head_steps = 0; +#ifdef DFLASH27B_HAVE_TOPK_HEAD + if (const char * e = std::getenv("DFLASH_TOPK_HEAD")) g_topk_head_M = std::atoi(e); + int32_t * d_cand = nullptr, * d_cand_use = nullptr, * d_head_tok = nullptr; + void * d_topm_scr = nullptr, * d_head_keys = nullptr; + // Node count upper bound: chain uses q_len positions; tree uses budget+1 nodes. + const int topk_head_Nmax = std::max(q_len, ddtree_budget + 1); + if (g_topk_head_M > 0) { + cudaMalloc(&d_cand, (size_t)g_topk_head_M * q_len * sizeof(int32_t)); + cudaMalloc(&d_cand_use, (size_t)g_topk_head_M * topk_head_Nmax * sizeof(int32_t)); + cudaMalloc(&d_head_tok, (size_t)topk_head_Nmax * sizeof(int32_t)); + cudaMalloc(&d_head_keys, (size_t)topk_head_Nmax * sizeof(unsigned long long)); + cudaMalloc(&d_topm_scr, dflash::common::extract_topm_scratch_bytes(q_len)); + std::printf("[topk-head] restricted greedy head ON, M=%d\n", g_topk_head_M); + } +#endif + while (n_generated < n_gen) { const int need_commit_budget = n_gen - n_generated; @@ -3378,9 +3405,16 @@ int main(int argc, char ** argv) { const int N_actual = 1 + tree.n_nodes; // actual tree size const int N = ddtree_budget + 1; // fixed allocation size for gallocr reuse + // Candidate-restricted greedy head on the tree path: single-GPU, + // greedy only (temp>0 samples the bonus from full logits, which we + // skip). Per node, candidates come from the draft top-M at the + // node's depth (same alignment as the chain path). + const bool tree_rhead = (g_topk_head_M > 0) && !draft_hidden_bridge + && (g_sampler.temp == 0.0f); if (!build_target_step_tree(sg, w, cache, backend, /*kv_start=*/committed, /*n_tokens=*/N, - g_fa_window, g_kq_stride_pad)) { + g_fa_window, g_kq_stride_pad, + /*restricted_head=*/tree_rhead)) { std::fprintf(stderr, "ddtree verify build failed\n"); return 1; } T_verify_build = sync_us(); @@ -3463,11 +3497,40 @@ int main(int argc, char ** argv) { // after the root even though logits are valid. This is test-only; // server decode paths are unaffected. std::vector posterior(N_actual); +#ifdef DFLASH27B_HAVE_TOPK_HEAD + if (tree_rhead) { + // Per node: extract draft top-M, map node i (depth d) to draft + // row min(d+1, q_len-1), fused Q6_K gather/dot/argmax → posterior. + auto Th0 = sync_us(); + const int dvocab = (int)draft_sg.logits->ne[0]; + dflash::common::extract_topm_cuda( + (const float *)draft_sg.logits->data, dvocab, q_len, g_topk_head_M, + d_cand, d_topm_scr, 0); + for (int i = 0; i < N_actual; i++) { + const int depth = (i == 0) ? 0 : (int)tree.depths[i - 1]; + const int src = std::min(depth + 1, q_len - 1); + cudaMemcpyAsync(d_cand_use + (size_t)i * g_topk_head_M, + d_cand + (size_t)src * g_topk_head_M, + (size_t)g_topk_head_M * sizeof(int32_t), + cudaMemcpyDeviceToDevice, 0); + } + dflash::common::restricted_lm_head_q6k( + w.output->data, (int)w.output->ne[0], (int)w.output->ne[1], + (const float *)sg.hidden_states->data, N_actual, + d_cand_use, g_topk_head_M, d_head_tok, d_head_keys, 0); + cudaMemcpy(posterior.data(), d_head_tok, sizeof(int32_t) * N_actual, + cudaMemcpyDeviceToHost); + tt_head_kernel += std::chrono::duration(sync_us() - Th0).count(); + head_steps++; + } else +#endif + { ggml_backend_tensor_get(sg.logits, verify_logits_buf.data(), 0, sizeof(float) * (size_t)vocab * N_actual); for (int i = 0; i < N_actual; i++) { posterior[i] = argmax_f32(verify_logits_buf.data() + (size_t)i * vocab, vocab); } + } auto T_verify_logits_ddtree = sync_us(); tt_verify_logits += std::chrono::duration( T_verify_logits_ddtree - T_verify_compute).count(); @@ -3731,13 +3794,21 @@ int main(int argc, char ** argv) { if (!seq_verify) { const int verify_fa_window = g_fa_window; + // Restricted greedy head only on the single-GPU fast-rollback chain + // path (where draft logits use the target head over the shared vocab, + // and the commit path never reads sg.logits). + const bool rhead = (g_topk_head_M > 0) && !draft_hidden_bridge && fast_rollback; if (!build_target_step(sg, w, cache, backend, /*kv_start=*/committed, /*n_tokens=*/q_len, /*with_mask=*/true, /*capture=*/true, /*capture_delta_intermediate=*/fast_rollback, verify_fa_window, /*last_token_logits_only=*/false, - g_kq_stride_pad)) { + g_kq_stride_pad, + /*capture_moe_router=*/false, + /*kvflash_mask=*/false, + /*capture_qk=*/false, + /*restricted_head=*/rhead)) { std::fprintf(stderr, "verify build failed\n"); return 1; } T_verify_build = sync_us(); @@ -3783,6 +3854,37 @@ int main(int argc, char ** argv) { T_verify_compute = sync_us(); tt_verify_compute += std::chrono::duration(T_verify_compute - T_verify_set).count(); +#ifdef DFLASH27B_HAVE_TOPK_HEAD + if (rhead) { + // Candidate-restricted greedy head. Extract per-position draft + // top-M from the (target-vocab) draft logits, remap so verify + // position p uses draft row min(p+1, q_len-1) (the calibrated + // alignment), then fused Q6_K gather/dot/argmax over the M rows. + auto Th0 = sync_us(); + const int dvocab = (int)draft_sg.logits->ne[0]; + // All on the default stream, no intermediate syncs: extract → + // remap → fused head → D2H (the copy forces completion). + dflash::common::extract_topm_cuda( + (const float *)draft_sg.logits->data, dvocab, q_len, g_topk_head_M, + d_cand, d_topm_scr, 0); + for (int pp = 0; pp < q_len; pp++) { + const int src = std::min(pp + 1, q_len - 1); + cudaMemcpyAsync(d_cand_use + (size_t)pp * g_topk_head_M, + d_cand + (size_t)src * g_topk_head_M, + (size_t)g_topk_head_M * sizeof(int32_t), + cudaMemcpyDeviceToDevice, 0); + } + dflash::common::restricted_lm_head_q6k( + w.output->data, (int)w.output->ne[0], (int)w.output->ne[1], + (const float *)sg.hidden_states->data, q_len, + d_cand_use, g_topk_head_M, d_head_tok, d_head_keys, 0); + cudaMemcpy(target_tok.data(), d_head_tok, sizeof(int32_t) * q_len, + cudaMemcpyDeviceToHost); + tt_head_kernel += std::chrono::duration(sync_us() - Th0).count(); + head_steps++; + } else +#endif + { ggml_backend_tensor_get(sg.argmax_tokens, target_tok.data(), 0, sizeof(int32_t) * q_len); // Calibration: the batched path normally reads only the GPU argmax. @@ -3793,6 +3895,7 @@ int main(int argc, char ** argv) { ggml_backend_tensor_get(sg.logits, verify_logits_buf.data(), 0, sizeof(float) * (size_t)vocab * q_len); } + } } else { // Sequential verify: q_len independent single-token decodes. // Each call writes K/V at slot committed+i and advances SSM by 1. @@ -4219,6 +4322,12 @@ int main(int argc, char ** argv) { std::printf(" verify_set %.2f\n", avg_ms(tt_verify_set)); std::printf(" verify_compute %.2f\n", avg_ms(tt_verify_compute)); std::printf(" verify_logits %.2f\n", avg_ms(tt_verify_logits)); + if (head_steps > 0) { + std::printf(" [topk-head] M=%d steps=%ld extract=%.3f ms/step kernel=%.3f ms/step\n", + g_topk_head_M, head_steps, + tt_head_extract / 1000.0 / head_steps, + tt_head_kernel / 1000.0 / head_steps); + } std::printf(" accept %.2f\n", avg_ms(tt_accept)); std::printf(" restore_ssm %.2f\n", avg_ms(tt_restore)); std::printf(" replay_build %.2f\n", avg_ms(tt_replay_build)); From 0b2d57aeaa77d04bb05d8225e2d7b677d3f0e875 Mon Sep 17 00:00:00 2001 From: "Pramodith (via Claude Code)" Date: Fri, 19 Jun 2026 08:52:58 +0000 Subject: [PATCH 08/10] topm kernels --- server/CMakeLists.txt | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 5bfcb5ea7..42fd01d4c 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -396,6 +396,15 @@ elseif(DFLASH27B_GPU_BACKEND STREQUAL "cuda") target_sources(dflash_common PRIVATE src/flashprefill_select.cpp src/flashprefill.cpp) + # Candidate-restricted greedy LM head (top-M extractor + fused Q6_K head). + target_sources(dflash_common PRIVATE + src/common/topm_extract_cuda.cu + src/common/restricted_lm_head_cuda.cu) + set_source_files_properties( + src/common/topm_extract_cuda.cu + src/common/restricted_lm_head_cuda.cu + PROPERTIES LANGUAGE CUDA) + target_compile_definitions(dflash_common PUBLIC DFLASH27B_HAVE_TOPK_HEAD=1) # Multi-arch: scan all resolved arches and compile every applicable # flashprefill kernel variant. This lets a single binary run on mixed # GPUs (e.g. Volta sm_70 + Pascal sm_61) without "no kernel image" errors. From 000d4a7dec3b1a8aded06e05469dd5d409fbc596 Mon Sep 17 00:00:00 2001 From: "Pramodith (via Claude Code)" Date: Fri, 19 Jun 2026 11:01:38 +0000 Subject: [PATCH 09/10] delete models --- server/models | 1 - 1 file changed, 1 deletion(-) delete mode 120000 server/models diff --git a/server/models b/server/models deleted file mode 120000 index 77851bf17..000000000 --- a/server/models +++ /dev/null @@ -1 +0,0 @@ -/workspace/lucebox-hub/server/models \ No newline at end of file From 5855256a275583886b4a6ee8fa638678869092af Mon Sep 17 00:00:00 2001 From: pramodith Date: Tue, 14 Jul 2026 19:48:12 +0300 Subject: [PATCH 10/10] partial topk with entropy thresholding --- calib/calibrate.py | 127 +++++- calib/prep_prompts.py | 18 +- docs/topk-head-optimization.md | 198 +++++++++ server/CMakeLists.txt | 25 +- server/src/common/draft_entropy_cuda.cu | 275 ++++++++++++ server/src/common/draft_entropy_cuda.h | 30 ++ server/src/common/draft_topk_cuda.cu | 413 ------------------ server/src/common/draft_topk_cuda.h | 34 -- .../src/common/geometric_draft_topk_cuda.cu | 31 +- server/src/common/topm_extract_cuda.cu | 196 +++++++-- server/src/qwen35/graph_builders.cpp | 7 + server/test/test_dflash.cpp | 249 ++++++++++- server/test/test_draft_entropy_cuda.cpp | 142 ++++++ server/test/test_topm_extract_cuda.cpp | 177 ++++++++ 14 files changed, 1402 insertions(+), 520 deletions(-) create mode 100644 server/src/common/draft_entropy_cuda.cu create mode 100644 server/src/common/draft_entropy_cuda.h delete mode 100644 server/src/common/draft_topk_cuda.cu delete mode 100644 server/src/common/draft_topk_cuda.h create mode 100644 server/test/test_draft_entropy_cuda.cpp create mode 100644 server/test/test_topm_extract_cuda.cpp diff --git a/calib/calibrate.py b/calib/calibrate.py index add7875a3..b7d6c9491 100644 --- a/calib/calibrate.py +++ b/calib/calibrate.py @@ -147,24 +147,65 @@ def stage_run(args, bins, outdir: Path): def parse_block(text: str): """Return dict for one calib run, or None if no table present. - {positions, mean_rank, max_rank, vocab, - set: {metric: {M: pct}}, mass: {M: pct}, mass_bad: {M: (bad, total)}} + {positions, mean_rank, max_rank, vocab, mean_draft_entropy, + set: {metric: {M: pct}}, mass: {M: pct}, mass_bad: {M: (bad, total)}, + entropy_buckets: {bucket_id: {lo, hi, positions, + coverage: {M: pct}, head_match: {M: pct}}}} + + entropy_buckets conditions the idealized "greedy" coverage@M curve (and, + where DFLASH27B_HAVE_TOPK_HEAD was compiled in, the real restricted-head + kernel's match rate) on the draft's temperature-agnostic entropy for that + position -- the data needed to pick an entropy threshold T and an M for + the entropy-gated restricted head. mean_draft_entropy / entropy_buckets + are absent (None / {}) on logs captured before this field existed. """ m = re.search(r"\[topk-calib\] positions=(\d+)\s+mean_rank=([\d.]+)\s+" r"max_rank=(\d+)\s+\(vocab=(\d+)\)", text) if not m: return None + em = re.search(r"mean_draft_entropy=([\d.]+)", text) res = { "positions": int(m.group(1)), "mean_rank": float(m.group(2)), "max_rank": int(m.group(3)), "vocab": int(m.group(4)), + "mean_draft_entropy": float(em.group(1)) if em else None, "set": {k: {} for k in SET_METRICS}, "mass": {}, "mass_bad": {}, + "entropy_buckets": {}, } - cur = None + cur = None # current SET_METRICS label / "__mass__", for [topk-calib] lines + eb_cur = None # current bucket id, for [topk-calib-eb] coverage@M lines + head_cur = None # current bucket id, for [topk-calib-head] match@M lines for line in text.splitlines(): + ebh = re.search(r"\[topk-calib-eb\] entropy-bucket=(\d+) range=\[([\d.]+),(inf|[\d.]+)\) positions=(\d+)", line) + if ebh: + eb_cur = int(ebh.group(1)) + hi = None if ebh.group(3) == "inf" else float(ebh.group(3)) + res["entropy_buckets"].setdefault(eb_cur, { + "lo": float(ebh.group(2)), "hi": hi, "positions": int(ebh.group(4)), + "coverage": {}, "head_match": {}, + }) + continue + hh = re.search(r"\[topk-calib-head\] entropy-bucket=(\d+) positions=(\d+)", line) + if hh: + head_cur = int(hh.group(1)) + res["entropy_buckets"].setdefault(head_cur, { + "lo": None, "hi": None, "positions": int(hh.group(2)), + "coverage": {}, "head_match": {}, + }) + continue + if line.startswith("[topk-calib-eb]") and eb_cur is not None: + cm = re.search(r"coverage@M=(\d+)\s*:\s*([\d.]+)%", line) + if cm: + res["entropy_buckets"][eb_cur]["coverage"][int(cm.group(1))] = float(cm.group(2)) + continue + if line.startswith("[topk-calib-head]") and head_cur is not None: + hm = re.search(r"match@M=(\d+)\s*:\s*([\d.]+)%", line) + if hm: + res["entropy_buckets"][head_cur]["head_match"][int(hm.group(1))] = float(hm.group(2)) + continue sc = re.search(r"--- (.+?): set-coverage", line) if sc: label = sc.group(1) @@ -205,10 +246,33 @@ def aggregate(per_prompt: dict): mass_bad = {M: 0 for M in Ms} rank_w = 0.0 max_rank = 0 + + # Entropy-bucket tables: coverage@M (idealized, from calib_hist_eb) and + # head_match@M (real restricted-head kernel, only present when + # DFLASH27B_HAVE_TOPK_HEAD was compiled in). Bucketed metrics are weighted + # by each bucket's OWN position count per prompt, not the prompt total -- + # a true global position-weighted mean, since bucket occupancy varies a + # lot prompt to prompt. + bucket_ids = sorted({bid for b in ok.values() for bid in b.get("entropy_buckets", {})}) + head_Ms = sorted({M for b in ok.values() + for bd in b.get("entropy_buckets", {}).values() + for M in bd.get("head_match", {})}) + eb_range = {} + eb_pos = {bid: 0 for bid in bucket_ids} + eb_cov_sum = {bid: {M: 0.0 for M in Ms} for bid in bucket_ids} + eb_cov_w = {bid: {M: 0.0 for M in Ms} for bid in bucket_ids} + eb_head_sum = {bid: {M: 0.0 for M in head_Ms} for bid in bucket_ids} + eb_head_w = {bid: {M: 0.0 for M in head_Ms} for bid in bucket_ids} + entropy_sum = 0.0 + entropy_wt = 0.0 + for b in ok.values(): w = b["positions"] rank_w += b["mean_rank"] * w max_rank = max(max_rank, b["max_rank"]) + if b.get("mean_draft_entropy") is not None: + entropy_sum += b["mean_draft_entropy"] * w + entropy_wt += w for m in SET_METRICS: for M, v in b["set"][m].items(): set_sum[m][M] += v * w @@ -218,14 +282,38 @@ def aggregate(per_prompt: dict): mass_w[M] += w for M, (bad, _) in b["mass_bad"].items(): mass_bad[M] += bad + for bid, bd in b.get("entropy_buckets", {}).items(): + bw = bd["positions"] + eb_pos[bid] += bw + if bid not in eb_range and bd.get("lo") is not None: + eb_range[bid] = (bd["lo"], bd["hi"]) + for M, v in bd.get("coverage", {}).items(): + eb_cov_sum[bid][M] += v * bw + eb_cov_w[bid][M] += bw + for M, v in bd.get("head_match", {}).items(): + eb_head_sum[bid][M] += v * bw + eb_head_w[bid][M] += bw + set_agg = {m: {M: (set_sum[m][M] / set_w[m][M] if set_w[m][M] else None) for M in Ms} for m in SET_METRICS} mass_agg = {M: (mass_sum[M] / mass_w[M] if mass_w[M] else None) for M in Ms} + entropy_buckets_agg = { + bid: { + "range": eb_range.get(bid), + "positions": eb_pos[bid], + "coverage": {M: (eb_cov_sum[bid][M] / eb_cov_w[bid][M] if eb_cov_w[bid][M] else None) + for M in Ms}, + "head_match": {M: (eb_head_sum[bid][M] / eb_head_w[bid][M] if eb_head_w[bid][M] else None) + for M in head_Ms}, + } for bid in bucket_ids + } return { "prompts": len(ok), "failed": len(per_prompt) - len(ok), "total_positions": totpos, "mean_rank": rank_w / totpos, - "max_rank": max_rank, "Ms": Ms, + "mean_draft_entropy": (entropy_sum / entropy_wt if entropy_wt else None), + "max_rank": max_rank, "Ms": Ms, "head_Ms": head_Ms, "set": set_agg, "mass": mass_agg, "mass_bad": mass_bad, + "entropy_buckets": entropy_buckets_agg, "per_prompt": {n: {"positions": b["positions"], "mean_rank": b["mean_rank"], "max_rank": b["max_rank"]} for n, b in ok.items()}, } @@ -234,9 +322,11 @@ def aggregate(per_prompt: dict): def print_report(agg): Ms = agg["Ms"] print(f"\n=== top-k head calibration aggregate ===") + entropy_str = (f" mean_draft_entropy={agg['mean_draft_entropy']:.3f}" + if agg.get("mean_draft_entropy") is not None else "") print(f"prompts={agg['prompts']} (failed={agg['failed']}) " f"total_positions={agg['total_positions']} " - f"mean_rank={agg['mean_rank']:.2f} max_rank={agg['max_rank']}\n") + f"mean_rank={agg['mean_rank']:.2f} max_rank={agg['max_rank']}{entropy_str}\n") def cell(v, fmt): return f"{'-':>10}" if v is None else f"{v:10.{fmt}f}" hdr = "metric".ljust(28) + "".join(f"{M:>10}" for M in Ms) @@ -248,6 +338,33 @@ def cell(v, fmt): print("\n-- top_p=0.95 nucleus MASS coverage (mean %, + positions <99% covered) --") print("mass %".ljust(28) + "".join(cell(agg['mass'][M], 3) for M in Ms)) print("bad(<99%)".ljust(28) + "".join(f"{agg['mass_bad'][M]:10d}" for M in Ms)) + + eb = agg.get("entropy_buckets") or {} + if eb: + # This is the table that drives the (entropy threshold T, restricted- + # head M) decision: for a candidate T, sum "positions" over buckets + # whose range is entirely below T, then read off the smallest M whose + # coverage@M (or, better, head_match@M once head kernels ran) clears + # your target hit rate across that same bucket set. + print("\n-- entropy-gated restricted head: idealized coverage@M by entropy bucket (%) --") + print("entropy range".ljust(28) + "".join(f"{M:>10}" for M in Ms)) + for bid in sorted(eb): + lo, hi = eb[bid]["range"] or (None, None) + label = f"[{lo:.2f},{hi:.2f})" if hi is not None else (f"[{lo:.2f},inf)" if lo is not None else f"bucket {bid}") + row = f"{label} n={eb[bid]['positions']}"[:27].ljust(28) + \ + "".join(cell(eb[bid]["coverage"][M], 2) for M in Ms) + print(row) + head_Ms = agg.get("head_Ms") or [] + if head_Ms and any(eb[bid]["head_match"] for bid in eb): + print("\n-- entropy-gated restricted head: REAL kernel match@M by entropy bucket (%) --") + print("entropy range".ljust(28) + "".join(f"{M:>10}" for M in head_Ms)) + for bid in sorted(eb): + lo, hi = eb[bid]["range"] or (None, None) + label = f"[{lo:.2f},{hi:.2f})" if hi is not None else (f"[{lo:.2f},inf)" if lo is not None else f"bucket {bid}") + row = f"{label} n={eb[bid]['positions']}"[:27].ljust(28) + \ + "".join(cell(eb[bid]["head_match"].get(M), 2) for M in head_Ms) + print(row) + print("\n-- per-prompt --") for n, s in sorted(agg["per_prompt"].items()): print(f" {n:<16} pos={s['positions']:<5} mean_rank={s['mean_rank']:7.2f} " diff --git a/calib/prep_prompts.py b/calib/prep_prompts.py index 3ffb9332a..6bd35cb99 100644 --- a/calib/prep_prompts.py +++ b/calib/prep_prompts.py @@ -2,8 +2,9 @@ """Tokenize calibration prompts to int32 .bin files for test_dflash. Pulls prompts from the same datasets the benchmark uses (HumanEval / GSM8K / -MATH-500), applies the Qwen3.6 chat template, and writes one packed-int32 file -per prompt into the output dir. Each file is a prompt test_dflash can decode. +MATH-500) plus a natural-language-heavy instruction set (Alpaca), applies the +Qwen3.6 chat template, and writes one packed-int32 file per prompt into the +output dir. Each file is a prompt test_dflash can decode. python3 calib/prep_prompts.py --n 200 --out calib --thinking off @@ -13,11 +14,18 @@ import argparse, struct, sys from pathlib import Path +def _alpaca_prompt(x): + return x["instruction"] + (f"\n{x['input']}" if x["input"] else "") + DATASETS = [ # (label, hf_id, config, split, field-extractor) - ("he", "openai/openai_humaneval", None, "test", lambda x: x["prompt"]), - ("gsm", "openai/gsm8k", "main", "test", lambda x: f"Question: {x['question']}\nAnswer: "), - ("math", "HuggingFaceH4/MATH-500", None, "test", lambda x: f"Problem: {x['problem']}\nSolution: "), + ("he", "openai/openai_humaneval", None, "test", lambda x: x["prompt"]), + ("gsm", "openai/gsm8k", "main", "test", lambda x: f"Question: {x['question']}\nAnswer: "), + ("math", "HuggingFaceH4/MATH-500", None, "test", lambda x: f"Problem: {x['problem']}\nSolution: "), + # Natural-language-heavy general instructions (contrast with the + # code/math skew above: §5b/§5d both note "general" prompts calibrate + # worse than code/math, so this dataset exists to quantify that directly). + ("nl", "tatsu-lab/alpaca", None, "train", _alpaca_prompt), ] def main(): diff --git a/docs/topk-head-optimization.md b/docs/topk-head-optimization.md index 5058af103..24eb8028b 100644 --- a/docs/topk-head-optimization.md +++ b/docs/topk-head-optimization.md @@ -196,6 +196,196 @@ GPU=0 NGEN=256 bash calib/run_calib.sh | tee calib/out/calib.log python3 calib/aggregate.py calib/out/calib.log ``` +## 5d. Entropy-gated restricted head — calibrating a per-position M (2026-07-14) + +§5b/§5c calibrate a single fixed M for *every* position. The follow-on idea: gate on the +draft's own per-position prediction confidence instead of always using one M. Coverage@M is not +uniform across positions — §5b already shows this indirectly ("code/math high, general prompts +lower"); the natural per-position signal for "how uncertain is the draft here" is its **entropy**. + +**Direction (confirmed empirically, not just intuited):** low draft entropy (confident) → +high coverage → safe to use the restricted head; high draft entropy (uncertain) → lower +coverage → fall back to the full head. This is the *opposite* of "restrict harder when unsure," +which would compound the risk exactly where the shortlist is least trustworthy. + +**Entropy definition:** temperature-agnostic — computed on raw draft logits as if `T=1`, +independent of whatever `ddtree_temp` is configured for sampling. This keeps the gate a property +of the draft's *intrinsic* confidence rather than an artifact of a sampling knob, and means a +threshold calibrated once doesn't silently drift if `ddtree_temp` changes later. + +**What's implemented** (instrumentation only, same `DFLASH_TOPK_CALIB=1` chain path as §5, +`server/test/test_dflash.cpp`): +1. Per-position entropy over draft row `i+1` (the same row supplying the candidate set for + verify position `i` — the `pp → min(pp+1, q_len-1)` alignment from §5/§7), computed via a raw + max + `Σ exp(l−max)` + `Σ l·exp(l−max)` pass, no temperature scaling. +2. The existing idealized rank histogram (`calib_hist`) is now also split by a fixed entropy + bucket (`calib_hist_eb`, edges `{0, 0.5, 1, ..., 12}` nats — bounded by `ln(vocab)≈12.4` for + this model's 248k vocab) → printed as `[topk-calib-eb]` per-bucket `coverage@M` tables. +3. A **second, non-idealized check**: for a small M grid (`{128,256,512,1024,2048}`), the + *actual* `extract_topm_cuda` + `restricted_lm_head_q6k` kernels run and their output is + compared token-for-token against the exact full-head argmax, bucketed the same way → printed + as `[topk-calib-head]`. This measures the extractor's real approximation error (bin-boundary + ties aren't rank-ordered — see `topm_extract_cuda.cu`), which the idealized CPU-sorted rank + histogram in (2) can't see. Needed a small plumbing fix: `build_qwen35_graph` already always + computes+retains `hidden_states` (`server/src/internal.h:616`), but `graph_builders.cpp`'s + non-restricted branch wasn't copying it into `StepGraph::hidden_states` — one-line fix lets + calibration validate the real kernel off the *same* full-head verify build/compute, no second + graph build needed. +4. `calib/calibrate.py` parses and position-weights both new tables (per-bucket weighting uses + each bucket's own position count, not the prompt total, since bucket occupancy varies a lot + prompt to prompt) and prints them alongside the existing report; `--json` picks them up for + free via the existing `aggregate()` return dict. + +**Smoke-tested** first on a synthetic (non-language) token sequence to validate the plumbing +(no NaNs/crashes, bucket counts summed correctly, real-kernel match% tracked idealized +coverage% closely), then run for real (below). + +### Real calibration results (2026-07-14) + +30 real prompts (10 each HumanEval / GSM8K / MATH-500, Qwen3.6 chat template, thinking off), +**15,360 positions**, `n_gen=256`, chain path, Qwen3.6-27B Q4_K_M target + Q4_K_M 3.6 draft. +`mean_draft_entropy = 1.758` nats overall — this draft/target pair is confident most of the +time; the long tail (entropy > 5) is thin (only 217/15360 positions, 1.4%). + +Direction confirmed with real data, not just intuition: real kernel match@M=128 falls +monotonically as entropy rises, from **99.64%** in `[0.00,0.50)` down to **96.66%** in +`[4.00,5.00)` — low draft entropy really does predict high restricted-head fidelity. + +**The frontier that answers "pick (T, M)"** — cumulative over all positions with entropy below +T, real kernel `match@M` (i.e. restricted-head output == exact full-head argmax; NOT the +idealized coverage upper bound): + +| T (entropy <) | % of positions | match@128 | match@256 | match@512 | match@1024 | match@2048 | +|--------------:|----------------:|----------:|----------:|----------:|-----------:|-----------:| +| 0.5 | 30.7% | 99.64% | 99.70% | 99.75% | 99.81% | 99.85% | +| 1.0 | 40.9% | 99.36% | 99.48% | 99.59% | 99.68% | 99.73% | +| 2.0 | 58.0% | 99.07% | 99.34% | 99.45% | 99.56% | 99.66% | +| 3.0 | 74.7% | 98.95% | 99.25% | 99.36% | 99.49% | 99.62% | +| **4.0** | **90.4%** | 98.67% | 99.06% | 99.28% | **99.42%** | 99.55% | +| 5.0 | 98.6% | 98.50% | 98.98% | 99.23% | 99.37% | 99.51% | +| ∞ (no gate) | 100% | 98.47% | 98.98% | 99.23% | 99.38% | 99.52% | + +**Recommended starting operating point: T≈4.0, M=1024.** Routes 90.4% of positions to the +cheap restricted head at a 99.42% real match rate against exact full-head greedy (the +remaining 9.6%, entropy≥4, fall back to the full head — always exact there). That's a ~0.6% +per-position miss rate concentrated exactly where we chose to still pay for the cheap path — +tightening to T≈2.0 nearly doubles the miss rate's *rarity* (58% coverage, 99.56% match) if a +lower error budget is wanted, or loosening to T≈5.0 barely changes match rate (98.6% coverage +already captures nearly everything — the entropy>5 tail is too rare to matter much either way +in this sample). + +**Caveat also worth restating from §5b:** this is one draft/target pair (Q4 3.6 draft) and one +30-prompt sample skewed toward code/math (thinking off) — re-run before committing to numbers +for a different draft, a different prompt mix, or natural-language-heavy workloads, which §5b +showed have measurably worse coverage than code/math. + +Raw run: `calib/out/calib.json` (not checked in — `calib/out/` is gitignored; regenerate with +`python3 calib/calibrate.py --prep 30 --gpu 0 --ngen 256 --json calib/out/calib.json`). + +## 5e. Natural-language-heavy calibration + combined results (2026-07-14) + +The §5d caveat above ("re-run for natural-language-heavy workloads") turned into an actual run. +Added a 4th dataset to `calib/prep_prompts.py` — `tatsu-lab/alpaca` (general instructions, +labeled `nl`), extracted as `instruction + "\n" + input` — specifically because it's *not* +code/math, to isolate how much the §5b/§5d numbers were inflated by that skew. 10 prompts, +reusing the 30 cached code/math logs via `calib/calibrate.py --skip-existing` (only the new `nl` +prompts actually ran) to get a **40-prompt combined** aggregate cheaply. + +| | prompts | positions | mean_draft_entropy | greedy coverage@128 | +|---|---:|---:|---:|---:| +| Code/Math only (§5d) | 30 | 15,360 | 1.758 | 98.60% | +| **NL only (Alpaca)** | 10 | 11,445 | **3.979** | **93.11%** | +| Combined | 40 | 26,805 | 2.707 | 96.26% | + +Confirms the hypothesis with numbers, not just the "general prompts lower" hand-wave from §5b: +NL prompts run at more than **2× the mean entropy** of code/math and noticeably worse coverage. +The entropy gate is exactly the right lever for this — it's already conditioning on the signal +that separates these two regimes, so a single calibrated `(T, M)` should generalize across a +mixed workload better than a single fixed M would. + +**Combined (40-prompt) frontier** — same cumulative-real-match methodology as §5d, now over a +workload that isn't skewed to code/math: + +| T (entropy <) | % of positions | match@128 | match@1024 | +|--------------:|----------------:|----------:|-----------:| +| 1.0 | 28.6% | 99.48% | 99.74% | +| 2.0 | 41.6% | 99.18% | 99.61% | +| 3.0 | 55.2% | 99.03% | 99.52% | +| **4.0** | **69.5%** | 98.56% | **99.39%** | +| 5.0 | 82.9% | 97.66% | 99.16% | +| 6.0 | 93.9% | 96.79% | 98.94% | + +**Revised recommendation for a mixed/unknown workload: T≈4.0, M=1024** still gives ≈99.4% +match (about the same quality bar as the code/math-only pick), but now only **69.5%** of +positions clear the gate — down from 90.4% when the sample was code/math-skewed. That's the +real cost of not knowing your workload in advance: roughly a third of positions now pay for the +full head instead of a tenth. NL alone is worse still — at the same T=4.0/M=1024 point, NL-only +clears the gate on just 41.4% of positions (vs. code/math's 90.4%), confirming the two workloads +need to be calibrated (or at least sanity-checked) separately rather than assuming one transfers +to the other. + +**Still a one-draft, modest-sample-size result** — 40 prompts, one instruction dataset as the +NL proxy (Alpaca skews short/simple; longer conversational or reasoning-heavy NL prompts may +differ). Good enough to size the effect, not to freeze a production threshold. + +Raw runs: `calib/out/calib_combined.json` (all 40), `calib/out/calib_nl_only.json`, +`calib/out/calib_codemath_only.json` (none checked in, `calib/out/` gitignored). Regenerate with +`python3 calib/prep_prompts.py --n 40 --out calib --thinking off` then +`python3 calib/calibrate.py --gpu 0 --ngen 256 --skip-existing --json calib/out/calib_combined.json`. + +## 5f. Runtime entropy gate + e2e validation (2026-07-14) + +Wired the calibrated gate into the DDTree runtime restricted-head path (the tree-mode +`tree_rhead`/chain-mode `rhead` toggles already implemented pre-existing on this branch, gated +by `DFLASH_TOPK_HEAD=M`, previously *ungated* — always on whenever `M>0`). New: +`server/src/common/draft_entropy_cuda.{h,cu}`, a standalone kernel computing per-row +temperature-agnostic entropy directly from `draft_sg.logits` (device-resident already; the +kernel does one more full-vocab-wide device pass but no new full-vocab **host** copy, so it +doesn't reintroduce the D2H the fast decode path is built to avoid — only an `n_positions`-sized +result crosses to host). Gated via new env var `DFLASH_TOPK_HEAD_ENTROPY_T` (unset = old +ungated behavior, fully backward compatible); aggregated per verify step via **MAX** entropy +over that step's draft rows, not mean — a whole-graph decision should be as conservative as its +least-confident row. + +**De-risking first (zero new code):** before writing the gate, ran the *existing* ungated +fixed-M=1024 tree head against the exact lossless baseline on 4 real prompts (one per +code/gsm/math/nl label). 3/4 were byte-identical; the NL prompt diverged **74.6%** of tokens, +cascading from a single wrong pick at step 9 — one greedy miss early on sends the entire +autoregressive continuation down a different (locally plausible, but different) path. This +concretely justified building the gate rather than assuming it mattered. + +**Gate validation on the same 4 prompts** (`DFLASH_TOPK_HEAD_ENTROPY_T=4.0`): the NL prompt +now routes **0/74 steps** to the restricted head (100% fallback) and is **byte-identical** to +the lossless baseline; the code/math prompts still route 39–97% of steps to the cheap path and +stay byte-identical too. The gate does exactly what it's calibrated to do at the single-prompt +level. + +**e2e via `server/scripts/bench_llm.py`** (10 samples each, HumanEval/GSM8K/Math500, baseline +head-off vs. gated `T=4.0/M=1024`, identical prompts both runs): + +| Bench | speedup (off → gated) | score (off → gated) | +|---|---|---| +| HumanEval | 2.78x → 2.86x | (unscored) | +| GSM8K | 2.86x → 3.02x | 5/10 → **4/10** | +| Math500 | 3.24x → 3.36x | 1/10 → 1/10 | + +Speedup gain is small (+2.8–5.4%), matching §2's ~5% realistic ceiling. **Quality is not fully +preserved at this operating point**: one GSM8K sample (#8) flipped from correct to incorrect +under the gate — a real divergence got through despite T=4.0, consistent with calibration's own +~99.4% (not 100%) per-step match rate there compounding over dozens of gated decisions per +generation. The gate is a large improvement over the *ungated* fixed-M head (which produced +74%-token-divergent output on the NL case) but is still genuinely approximate, not lossless, at +this threshold. + +**Verdict: not ready to ship as-is.** The speedup is modest and the quality cost, while small, +is real and measured (not hypothetical). Before adopting: either (a) tighten `T` toward the +§5e frontier's more conservative end (T≈1–2, trading routed-cheap fraction for lower risk), or +(b) build the verify-and-escalate design from the original brainstorm (goal "b": run the +restricted head, detect likely misses — e.g. within-shortlist margin, or disagreement with the +draft's own top-1 — and recompute with the full head only for flagged steps) instead of +accepting a fixed statistical error rate. (a) is a config change with the current code; (b) +needs new plumbing and wasn't built this session. + ## 6. Next steps 1. **[done]** Full calibration — see §5b. Greedy validated; sampling needs a better draft. 2. Implement the restricted head for the **greedy** path in `build_qwen35_graph` / the verify graph @@ -209,6 +399,14 @@ python3 calib/aggregate.py calib/out/calib.log a BF16/stronger draft. If coverage improves: keep `top_k` finite for exactness (chain order makes it lossless when `draft-top-M ⊇ target-top-K_s`), and gate pure-top_p / pure-temperature to the full head (or over-cover with a tail-mass correction). +5. **[done, not shipped] Entropy-gated M — runtime gate wired + e2e-tested, see §5f.** + `DFLASH_TOPK_HEAD_ENTROPY_T` on the existing tree/chain restricted-head toggle, per-verify-step + (whole-graph) granularity, MAX-entropy aggregate. e2e via `bench_llm.py` at T=4.0/M=1024: +2.8 + to +5.4% speedup, but GSM8K score dropped 5/10→4/10 (one real divergence) — approximate, not + lossless, as expected, but the speedup/risk trade isn't clearly worth it yet at this + threshold. Next: either tighten T (§5e's frontier), or build the verify-and-escalate + alternative (detect likely misses and recompute with the full head, rather than accepting a + fixed statistical error rate) — see §5f's verdict. ## 7. Key code references - LM head matmul + argmax: `server/src/qwen35/graph_builders.cpp:360-365` (chain), `:432` (tree), diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index b8d2f9258..ad75f7343 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -450,13 +450,16 @@ elseif(DFLASH27B_GPU_BACKEND STREQUAL "cuda") # and take the GPU draft top-K path instead of the CPU fallback. Same macro # name as the HIP branch above (backend-neutral). target_compile_definitions(dflash_common PUBLIC DFLASH27B_HAVE_DRAFT_TOPK=1) - # Candidate-restricted greedy LM head (top-M extractor + fused Q6_K head). + # Candidate-restricted greedy LM head (top-M extractor + fused Q6_K head), + # plus the entropy kernel that gates it at runtime (see draft_entropy_cuda.h). target_sources(dflash_common PRIVATE src/common/topm_extract_cuda.cu - src/common/restricted_lm_head_cuda.cu) + src/common/restricted_lm_head_cuda.cu + src/common/draft_entropy_cuda.cu) set_source_files_properties( src/common/topm_extract_cuda.cu src/common/restricted_lm_head_cuda.cu + src/common/draft_entropy_cuda.cu PROPERTIES LANGUAGE CUDA) target_compile_definitions(dflash_common PUBLIC DFLASH27B_HAVE_TOPK_HEAD=1) # GPU port of the sample_logits chain. Compiled in by default; the path is @@ -721,6 +724,24 @@ if(DFLASH27B_TESTS) target_link_libraries(test_draft_topk_cuda PRIVATE dflash_common ${DFLASH27B_GGML_BACKEND_TARGET}) add_test(NAME draft_topk_cuda COMMAND test_draft_topk_cuda) endif() + # GPU draft entropy kernel vs a double-precision CPU reference. CUDA only: + # draft_entropy_cuda.cu is only compiled under DFLASH27B_HAVE_TOPK_HEAD + # (the cuda-backend branch above), no HIP variant exists. + if(DFLASH27B_GPU_BACKEND STREQUAL "cuda" AND EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_draft_entropy_cuda.cpp") + add_executable(test_draft_entropy_cuda test/test_draft_entropy_cuda.cpp) + target_include_directories(test_draft_entropy_cuda PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src) + target_link_libraries(test_draft_entropy_cuda PRIVATE dflash_common CUDA::cudart) + add_test(NAME draft_entropy_cuda COMMAND test_draft_entropy_cuda) + endif() + # GPU top-M candidate extractor vs a CPU reference of the same coarse-bin + # threshold algorithm. CUDA only: topm_extract_cuda.cu is only compiled + # under DFLASH27B_HAVE_TOPK_HEAD (the cuda-backend branch above). + if(DFLASH27B_GPU_BACKEND STREQUAL "cuda" AND EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_topm_extract_cuda.cpp") + add_executable(test_topm_extract_cuda test/test_topm_extract_cuda.cpp) + target_include_directories(test_topm_extract_cuda PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src) + target_link_libraries(test_topm_extract_cuda PRIVATE dflash_common CUDA::cudart) + add_test(NAME topm_extract_cuda COMMAND test_topm_extract_cuda) + endif() # GPU port of the sample_logits chain vs the CPU reference. CUDA only: # geometric_sampler_cuda.cu is compiled into dflash_common solely on the cuda backend. if(DFLASH27B_GPU_BACKEND STREQUAL "cuda" AND EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_gpu_sampler_cuda.cpp") diff --git a/server/src/common/draft_entropy_cuda.cu b/server/src/common/draft_entropy_cuda.cu new file mode 100644 index 000000000..fb2d73d68 --- /dev/null +++ b/server/src/common/draft_entropy_cuda.cu @@ -0,0 +1,275 @@ +// GPU temperature-agnostic per-position draft entropy. See .h for the +// contract/rationale. Structurally mirrors draft_topk_cuda.cu's split-grid +// reduction (bandwidth-bound over a huge vocab, tiny row count) minus the +// top-K bookkeeping: two accumulators (sumexp, weighted sum) instead of a +// register-resident sorted top-K, reduced the same way logsumexp already is. + +#include "draft_entropy_cuda.h" + +#include +#include +#include +#include + +namespace dflash::common { + +namespace { + +constexpr int kBlock = 256; +constexpr int kMaxSplit = 128; + +#define DFLASH_ENTROPY_FOLD(L) \ + do { \ + const float _l = (L); \ + if (_l > lmax) { lsum = lsum * __expf(lmax - _l) + 1.0f; \ + lwsum = lwsum * __expf(lmax - _l) + _l; lmax = _l; } \ + else { const float _e = __expf(_l - lmax); lsum += _e; lwsum += _e * _l; } \ + } while (0) + +// Pass 1: per-(position, split) partial (max, sumexp, weighted) over a raw +// (unscaled -- temperature-agnostic) logit chunk. VEC selects 16-byte float4 +// loads, gated the same way as draft_topk_partial's VEC path (same tensor, so +// the same alignment precondition applies): only used when every row's base +// pointer is 16-byte aligned (vocab % 4 == 0 and an aligned tensor). +template +__global__ void entropy_partial(const float * __restrict__ logits, + int vocab, int split, + float * __restrict__ part_max, + float * __restrict__ part_sum, + float * __restrict__ part_wsum) { + const int row = blockIdx.x; + const int s = blockIdx.y; + const int tid = threadIdx.x; + const float * __restrict__ li = logits + (size_t)row * vocab; + + float lmax = -FLT_MAX, lsum = 0.0f, lwsum = 0.0f; + + if (VEC) { + // Partition the float4s of the row into `split` contiguous chunks, + // mirroring draft_topk_partial exactly (see that file for why). + const int vocab4 = vocab >> 2; + const int chunk4 = (vocab4 + split - 1) / split; + const int b4 = s * chunk4; + const int e4 = min(b4 + chunk4, vocab4); + const float4 * __restrict__ li4 = reinterpret_cast(li); + for (int j4 = b4 + tid; j4 < e4; j4 += kBlock) { + const float4 f = li4[j4]; + DFLASH_ENTROPY_FOLD(f.x); + DFLASH_ENTROPY_FOLD(f.y); + DFLASH_ENTROPY_FOLD(f.z); + DFLASH_ENTROPY_FOLD(f.w); + } + // Tail elements past the last full float4 (only when vocab % 4 != 0); + // the last split owns them so no element is scanned twice. + if (s == split - 1) { + for (int j = (vocab4 << 2) + tid; j < vocab; j += kBlock) + DFLASH_ENTROPY_FOLD(li[j]); + } + } else { + const int chunk = (vocab + split - 1) / split; + const int begin = s * chunk; + const int end = min(begin + chunk, vocab); + for (int j = begin + tid; j < end; j += kBlock) + DFLASH_ENTROPY_FOLD(li[j]); + } + + __shared__ float s_max[kBlock]; + __shared__ float s_sum[kBlock]; + __shared__ float s_wsum[kBlock]; + s_max[tid] = lmax; s_sum[tid] = lsum; s_wsum[tid] = lwsum; + __syncthreads(); + + for (int stride = kBlock / 2; stride > 0; stride >>= 1) { + if (tid < stride) { + const float am = s_max[tid], as = s_sum[tid], aw = s_wsum[tid]; + const float bm = s_max[tid + stride], bs = s_sum[tid + stride], bw = s_wsum[tid + stride]; + const float m = fmaxf(am, bm); + s_sum[tid] = as * __expf(am - m) + bs * __expf(bm - m); + s_wsum[tid] = aw * __expf(am - m) + bw * __expf(bm - m); + s_max[tid] = m; + } + __syncthreads(); + } + + if (tid == 0) { + const int idx = row * split + s; + part_max[idx] = s_max[0]; + part_sum[idx] = s_sum[0]; + part_wsum[idx] = s_wsum[0]; + } +} + +#undef DFLASH_ENTROPY_FOLD + +// Combine one online-logsumexp triple into another: (m,s,w) := (m,s,w) ⊕ (om,os,ow). +__device__ __forceinline__ void combine3(float & m, float & s, float & w, + float om, float os, float ow) { + const float nm = fmaxf(m, om); + s = s * __expf(m - nm) + os * __expf(om - nm); + w = w * __expf(m - nm) + ow * __expf(om - nm); + m = nm; +} + +// Pass 2: merge `split` partials per row into entropy = log_z - weighted/sumexp. +// blockDim == pow2_ceil(split) <= kMaxSplit, i.e. at most 4 warps. Reduces +// within each warp via shuffle (no shared memory, no __syncthreads -- and for +// blockDim < 32, mask must be __activemask(), not a hardcoded full-warp mask: +// the lanes beyond blockDim.x aren't real executing threads, and shfl_sync's +// mask must match exactly the converged set or the result is undefined). Only +// when more than one warp is active (small n_positions can push split, and so +// blockDim, up to kMaxSplit=128) do the ≤4 per-warp partials need a shared-mem +// hop for the final cross-warp combine. +__global__ void entropy_combine(const float * __restrict__ part_max, + const float * __restrict__ part_sum, + const float * __restrict__ part_wsum, + int split, + float * __restrict__ out_entropy) { + const int row = blockIdx.x; + const int tid = threadIdx.x; // blockDim == pow2_ceil(split) + const int lane = tid & 31; + const int wid = tid >> 5; + const int n_warps = (blockDim.x + 31) >> 5; + + float m, s, w; + if (tid < split) { + const int idx = row * split + tid; + m = part_max[idx]; s = part_sum[idx]; w = part_wsum[idx]; + } else { + m = -FLT_MAX; s = 0.0f; w = 0.0f; + } + + // __shfl_down_sync only exchanges data within a 32-lane warp, never across + // warp boundaries. When n_warps==1, blockDim.x itself is <=32 (possibly a + // partial warp) so the reduction width is blockDim.x; when n_warps>1, + // blockDim.x is a multiple of 32 (pow2_ceil(split) with split>32), so + // every warp is fully populated and the per-warp width is a full 32. + const int warp_width = (n_warps == 1) ? (int)blockDim.x : 32; + const unsigned mask = __activemask(); + for (int off = warp_width >> 1; off > 0; off >>= 1) { + combine3(m, s, w, __shfl_down_sync(mask, m, off), + __shfl_down_sync(mask, s, off), + __shfl_down_sync(mask, w, off)); + } + // lane 0 of each warp now holds that warp's fully-combined partial. + + if (n_warps == 1) { + if (tid == 0) out_entropy[row] = m + logf(s) - w / s; + return; + } + + __shared__ float sh_m[4], sh_s[4], sh_w[4]; // kMaxSplit=128 -> at most 4 warps + if (lane == 0) { sh_m[wid] = m; sh_s[wid] = s; sh_w[wid] = w; } + __syncthreads(); + + if (tid == 0) { + float fm = sh_m[0], fs = sh_s[0], fw = sh_w[0]; + for (int i = 1; i < n_warps; i++) combine3(fm, fs, fw, sh_m[i], sh_s[i], sh_w[i]); + out_entropy[row] = fm + logf(fs) - fw / fs; + } +} + +struct Scratch { + int device = -1; + size_t cap = 0; // n_positions*split partial elements + float * d_pmax = nullptr; + float * d_psum = nullptr; + float * d_pwsum = nullptr; + float * d_out = nullptr; // n_positions + size_t out_cap = 0; +}; +Scratch g_scratch; + +void free_scratch() { + if (g_scratch.d_pmax) cudaFree(g_scratch.d_pmax); + if (g_scratch.d_psum) cudaFree(g_scratch.d_psum); + if (g_scratch.d_pwsum) cudaFree(g_scratch.d_pwsum); + if (g_scratch.d_out) cudaFree(g_scratch.d_out); + g_scratch = Scratch{}; +} + +bool ensure_scratch(int device, size_t n_parts, size_t n_positions) { + if (g_scratch.device == device && g_scratch.cap >= n_parts && g_scratch.out_cap >= n_positions) + return true; + free_scratch(); + if (cudaMalloc(&g_scratch.d_pmax, n_parts * sizeof(float)) != cudaSuccess) goto fail; + if (cudaMalloc(&g_scratch.d_psum, n_parts * sizeof(float)) != cudaSuccess) goto fail; + if (cudaMalloc(&g_scratch.d_pwsum, n_parts * sizeof(float)) != cudaSuccess) goto fail; + if (cudaMalloc(&g_scratch.d_out, n_positions * sizeof(float)) != cudaSuccess) goto fail; + g_scratch.device = device; + g_scratch.cap = n_parts; + g_scratch.out_cap = n_positions; + return true; +fail: + free_scratch(); + return false; +} + +inline int pow2_ceil(int x) { int p = 1; while (p < x) p <<= 1; return p; } + +// SM count of `device`, queried once and cached (topology is static for the +// process lifetime). Mirrors draft_topk_cuda.cu's sm_count/pick_split. +int sm_count(int device) { + static int cached_device = -1; + static int cached_sm = 0; + if (device != cached_device) { + cudaDeviceProp prop{}; + cached_sm = (cudaGetDeviceProperties(&prop, device) == cudaSuccess) + ? prop.multiProcessorCount + : 80; // fallback: the ~80-SM device this heuristic was tuned on + cached_device = device; + } + return cached_sm; +} + +// Aim for ~3 waves of blocks across the device's actual SM count (was a fixed +// 240 = 3*80, tuned on one ~80-SM GPU; see draft_topk_cuda.cu's pick_split for +// the same reasoning), but cap the chunk floor at ~2k elements so a block +// never scans too little to be worth launching. +int pick_split(int vocab, int n_positions, int device) { + const int target_blocks = 3 * sm_count(device); + int by_blocks = (target_blocks + n_positions - 1) / n_positions; + int by_chunk = vocab / 2048; + int split = by_blocks < by_chunk ? by_blocks : by_chunk; + if (split < 1) split = 1; + if (split > kMaxSplit) split = kMaxSplit; + return split; +} + +} // namespace + +bool extract_draft_entropy_cuda(const float * d_logits, int vocab, int n_positions, + float * out_entropy, cudaStream_t stream) { + if (!d_logits || vocab <= 0 || n_positions <= 0) return false; + + int dev = 0; + cudaGetDevice(&dev); + const int split = pick_split(vocab, n_positions, dev); + const size_t n_parts = (size_t)n_positions * split; + if (!ensure_scratch(dev, n_parts, (size_t)n_positions)) return false; + + // float4 loads are safe only when every row base is 16-byte aligned: the + // tensor base aligned and a vocab stride that is a multiple of 4 (same + // precondition as draft_topk_cuda.cu's use_vec, same tensor). + const bool use_vec = (vocab % 4 == 0) && + (reinterpret_cast(d_logits) % 16 == 0); + + const dim3 grid1((unsigned)n_positions, (unsigned)split); + if (use_vec) { + entropy_partial<<>>( + d_logits, vocab, split, g_scratch.d_pmax, g_scratch.d_psum, g_scratch.d_pwsum); + } else { + entropy_partial<<>>( + d_logits, vocab, split, g_scratch.d_pmax, g_scratch.d_psum, g_scratch.d_pwsum); + } + entropy_combine<<>>( + g_scratch.d_pmax, g_scratch.d_psum, g_scratch.d_pwsum, split, g_scratch.d_out); + + const cudaError_t e = cudaMemcpyAsync(out_entropy, g_scratch.d_out, + (size_t)n_positions * sizeof(float), + cudaMemcpyDeviceToHost, stream); + if (e != cudaSuccess) return false; + if (cudaStreamSynchronize(stream) != cudaSuccess) return false; + return cudaGetLastError() == cudaSuccess; +} + +} // namespace dflash::common diff --git a/server/src/common/draft_entropy_cuda.h b/server/src/common/draft_entropy_cuda.h new file mode 100644 index 000000000..68e4c6278 --- /dev/null +++ b/server/src/common/draft_entropy_cuda.h @@ -0,0 +1,30 @@ +// GPU temperature-agnostic entropy of the draft's per-position logit +// distribution. Used to gate the candidate-restricted greedy head at +// runtime: low draft entropy (confident) -> safe to use the cheap +// restricted head; high entropy (uncertain) -> fall back to the exact +// full-vocab head. See docs/topk-head-optimization.md §5d/§5e for the +// calibration this backs. +// +// Deliberately a standalone kernel (not folded into geometric_draft_topk_cuda.cu) +// so it stays independent of that kernel's temperature scaling and doesn't +// risk perturbing the DDTree tree-construction / sampling paths that already +// depend on it. Cost: one extra full-vocab-wide (bandwidth-bound) GPU pass +// over logits already resident on-device -- no additional host D2H beyond +// the tiny n_positions-sized entropy output, so it does not reintroduce the +// vocab-wide D2H the fast decode path is built to avoid. + +#pragma once + +#include +#include + +namespace dflash::common { + +// d_logits: device pointer, [vocab] contiguous per position, n_positions rows. +// out_entropy: HOST pointer, n_positions floats (nats), temperature-agnostic +// (raw logits, as if temperature=1 regardless of any sampling temperature). +// Returns false on any CUDA error (out_entropy left untouched). +bool extract_draft_entropy_cuda(const float * d_logits, int vocab, int n_positions, + float * out_entropy, cudaStream_t stream = nullptr); + +} // namespace dflash::common diff --git a/server/src/common/draft_topk_cuda.cu b/server/src/common/draft_topk_cuda.cu deleted file mode 100644 index 7adc5045a..000000000 --- a/server/src/common/draft_topk_cuda.cu +++ /dev/null @@ -1,413 +0,0 @@ -// GPU top-K + log-prob extraction for DDTree draft distributions. -// See draft_topk_cuda.h for the contract. Mirrors extract_draft_topk (ddtree.cpp). - -#include "draft_topk_cuda.h" - -#include -#include -#include -#include -#include - -namespace dflash::common { - -namespace { - -constexpr int kMaxK = 8; // ddtree_K is 8 in practice; K>kMaxK → CPU fallback -constexpr int kBlock = 256; // threads per block (power of two for the reduction) -constexpr int kMaxSplit = 128; // max vocab splits per position (combine-block cap) - -// The workload is tiny in rows (n_positions ~ 15) but huge in vocab (~152k), -// so it is purely DRAM-bandwidth bound: ~9 MB to read per call. The original -// one-block-per-position layout launched only ~15 blocks, leaving most of the -// GPU's SMs idle and hitting only a fraction of peak bandwidth. We instead -// split each position's vocab scan across `split` blocks (a 2D grid of -// n_positions × split) so the whole device stays busy, then a cheap second -// kernel merges the `split` partials per position into the final top-K. - -// Merge two sorted-descending top-K lists (av/ai, bv/bi) into dst (dv/di), -// keeping the K largest; ties resolve toward the lower id. dst may alias av/ai -// safely only via the temporary below, so callers pass a distinct dst when the -// inputs come from shared memory. -template -__device__ __forceinline__ void merge_topk(const float * av, const int * ai, - const float * bv, const int * bi, - float * dv, int * di) { - int ia = 0, ib = 0; -#pragma unroll - for (int k = 0; k < K; k++) { - bool takeA; - if (ib >= K) takeA = true; - else if (ia >= K) takeA = false; - else if (av[ia] > bv[ib]) takeA = true; - else if (av[ia] < bv[ib]) takeA = false; - else takeA = (ai[ia] <= bi[ib]); // tie → lower id - if (takeA) { dv[k] = av[ia]; di[k] = ai[ia]; ia++; } - else { dv[k] = bv[ib]; di[k] = bi[ib]; ib++; } - } -} - -// Fold one logit (raw value `LRAW`, vocab id `ID`) into a thread's running -// online-logsumexp (lmax,lsum) and its register-resident sorted-descending -// top-K (topv,topi). K is a compile-time constant so the unrolled bubble uses -// only fixed indices — the arrays stay in registers instead of spilling to -// local memory the way a data-dependent insertion index would. The new value -// enters at slot K-1 and bubbles up only past strictly-smaller entries, so on -// ties the earlier (lower-id, since ids ascend within a thread) entry wins — -// matching the CPU heap's first-wins behaviour. -#define DFLASH_TOPK_CONSUME(LRAW, ID) \ - do { \ - const float _l = (LRAW) * inv_t; \ - if (_l > lmax) { lsum = lsum * __expf(lmax - _l) + 1.0f; lmax = _l; } \ - else { lsum += __expf(_l - lmax); } \ - if (_l > topv[K - 1]) { \ - topv[K - 1] = _l; topi[K - 1] = (ID); \ - _Pragma("unroll") \ - for (int _p = K - 1; _p > 0; --_p) { \ - if (topv[_p] > topv[_p - 1]) { \ - const float _tv = topv[_p]; \ - topv[_p] = topv[_p - 1]; topv[_p - 1] = _tv; \ - const int _ti = topi[_p]; \ - topi[_p] = topi[_p - 1]; topi[_p - 1] = _ti; \ - } \ - } \ - } \ - } while (0) - -// ---- pass 1: per-(position, split) partial logsumexp + local top-K --------- -// Grid: (n_positions, split). Block s of row `row` scans its contiguous vocab -// chunk and writes one partial (max, sum, sorted top-K) to part_* at index -// row*split + s. Contiguous chunks keep the strided per-warp reads coalesced. -// VEC selects 16-byte float4 loads (one coalesced transaction per 4 logits) — -// only used when every row's base pointer is 16-byte aligned (vocab % 4 == 0 -// and an aligned tensor); otherwise the scalar path is used. -template -__global__ void draft_topk_partial(const float * __restrict__ logits, - int vocab, float inv_t, int split, - float * __restrict__ part_max, - float * __restrict__ part_sum, - float * __restrict__ part_v, - int32_t * __restrict__ part_i) { - const int row = blockIdx.x; - const int s = blockIdx.y; - const int tid = threadIdx.x; - const float * __restrict__ li = logits + (size_t)row * vocab; - - float lmax = -FLT_MAX; - float lsum = 0.0f; - float topv[K]; - int topi[K]; -#pragma unroll - for (int k = 0; k < K; k++) { topv[k] = -FLT_MAX; topi[k] = -1; } - - if (VEC) { - // Partition the float4s of the row into `split` contiguous chunks. - const int vocab4 = vocab >> 2; // # of full float4s - const int chunk4 = (vocab4 + split - 1) / split; - const int b4 = s * chunk4; - const int e4 = min(b4 + chunk4, vocab4); - const float4 * __restrict__ li4 = reinterpret_cast(li); - for (int j4 = b4 + tid; j4 < e4; j4 += kBlock) { - const float4 f = li4[j4]; - const int base = j4 << 2; - DFLASH_TOPK_CONSUME(f.x, base + 0); - DFLASH_TOPK_CONSUME(f.y, base + 1); - DFLASH_TOPK_CONSUME(f.z, base + 2); - DFLASH_TOPK_CONSUME(f.w, base + 3); - } - // Tail elements past the last full float4 (only when vocab % 4 != 0); - // the last split owns them so no id is scanned twice. - if (s == split - 1) { - for (int j = (vocab4 << 2) + tid; j < vocab; j += kBlock) - DFLASH_TOPK_CONSUME(li[j], j); - } - } else { - const int chunk = (vocab + split - 1) / split; - const int begin = s * chunk; - const int end = min(begin + chunk, vocab); - for (int j = begin + tid; j < end; j += kBlock) - DFLASH_TOPK_CONSUME(li[j], j); - } - - // ---- block reduction over kBlock threads ------------------------------ - __shared__ float s_max[kBlock]; - __shared__ float s_sum[kBlock]; - __shared__ float s_topv[kBlock * K]; - __shared__ int32_t s_topi[kBlock * K]; - - s_max[tid] = lmax; - s_sum[tid] = lsum; -#pragma unroll - for (int k = 0; k < K; k++) { - s_topv[tid * K + k] = topv[k]; - s_topi[tid * K + k] = topi[k]; - } - __syncthreads(); - - for (int stride = kBlock / 2; stride > 0; stride >>= 1) { - if (tid < stride) { - const float am = s_max[tid], as = s_sum[tid]; - const float bm = s_max[tid + stride], bs = s_sum[tid + stride]; - const float m = fmaxf(am, bm); - s_sum[tid] = as * __expf(am - m) + bs * __expf(bm - m); - s_max[tid] = m; - - float dv[K]; int di[K]; - merge_topk(&s_topv[tid * K], &s_topi[tid * K], - &s_topv[(tid + stride) * K], &s_topi[(tid + stride) * K], - dv, di); -#pragma unroll - for (int k = 0; k < K; k++) { - s_topv[tid * K + k] = dv[k]; - s_topi[tid * K + k] = di[k]; - } - } - __syncthreads(); - } - - if (tid == 0) { - const int idx = row * split + s; - part_max[idx] = s_max[0]; - part_sum[idx] = s_sum[0]; -#pragma unroll - for (int k = 0; k < K; k++) { - part_v[(size_t)idx * K + k] = s_topv[k]; - part_i[(size_t)idx * K + k] = s_topi[k]; - } - } -} - -#undef DFLASH_TOPK_CONSUME - -// ---- pass 2: merge the `split` partials per position into the final top-K -- -// Grid: n_positions blocks of blockDim = pow2_ceil(split) threads. Thread t -// loads partial t (or an identity when t >= split), then a shared-memory tree -// reduction merges all partials. log_prob[k] = top_v[k] - log_z. -template -__global__ void draft_topk_combine(const float * __restrict__ part_max, - const float * __restrict__ part_sum, - const float * __restrict__ part_v, - const int32_t * __restrict__ part_i, - int split, - float * __restrict__ out_lp, - int32_t * __restrict__ out_ids) { - const int row = blockIdx.x; - const int tid = threadIdx.x; // blockDim is pow2_ceil(split) - - __shared__ float s_max[kMaxSplit]; - __shared__ float s_sum[kMaxSplit]; - __shared__ float s_topv[kMaxSplit * K]; - __shared__ int32_t s_topi[kMaxSplit * K]; - - if (tid < split) { - const int idx = row * split + tid; - s_max[tid] = part_max[idx]; - s_sum[tid] = part_sum[idx]; -#pragma unroll - for (int k = 0; k < K; k++) { - s_topv[tid * K + k] = part_v[(size_t)idx * K + k]; - s_topi[tid * K + k] = part_i[(size_t)idx * K + k]; - } - } else { - s_max[tid] = -FLT_MAX; - s_sum[tid] = 0.0f; -#pragma unroll - for (int k = 0; k < K; k++) { - s_topv[tid * K + k] = -FLT_MAX; - s_topi[tid * K + k] = -1; - } - } - __syncthreads(); - - for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { - if (tid < stride) { - const float am = s_max[tid], as = s_sum[tid]; - const float bm = s_max[tid + stride], bs = s_sum[tid + stride]; - const float m = fmaxf(am, bm); - s_sum[tid] = as * __expf(am - m) + bs * __expf(bm - m); - s_max[tid] = m; - - float dv[K]; int di[K]; - merge_topk(&s_topv[tid * K], &s_topi[tid * K], - &s_topv[(tid + stride) * K], &s_topi[(tid + stride) * K], - dv, di); -#pragma unroll - for (int k = 0; k < K; k++) { - s_topv[tid * K + k] = dv[k]; - s_topi[tid * K + k] = di[k]; - } - } - __syncthreads(); - } - - if (tid == 0) { - const float log_z = s_max[0] + logf(s_sum[0]); -#pragma unroll - for (int k = 0; k < K; k++) { - out_lp[(size_t)row * K + k] = s_topv[k] - log_z; - out_ids[(size_t)row * K + k] = s_topi[k]; - } - } -} - -// Per-device scratch for the [n_positions × K] outputs plus the -// [n_positions × split × K] pass-1 partials, grown as needed. The decode loop -// is single-threaded, so a plain static cache is safe and avoids a -// cudaMalloc/cudaFree on every step. -struct Scratch { - int device = -1; - size_t cap = 0; // output elements (n_positions*K) - size_t part_cap = 0; // partial-list elements (n_positions*split*kMaxK) - float * d_lp = nullptr; - int32_t * d_ids = nullptr; - float * d_pmax = nullptr; // [n_positions*split] - float * d_psum = nullptr; // [n_positions*split] - float * d_pv = nullptr; // [n_positions*split*kMaxK] - int32_t * d_pi = nullptr; // [n_positions*split*kMaxK] -}; -Scratch g_scratch; - -void free_scratch() { - if (g_scratch.d_lp) cudaFree(g_scratch.d_lp); - if (g_scratch.d_ids) cudaFree(g_scratch.d_ids); - if (g_scratch.d_pmax) cudaFree(g_scratch.d_pmax); - if (g_scratch.d_psum) cudaFree(g_scratch.d_psum); - if (g_scratch.d_pv) cudaFree(g_scratch.d_pv); - if (g_scratch.d_pi) cudaFree(g_scratch.d_pi); - g_scratch = Scratch{}; -} - -// Allocate output + partial buffers. n = n_positions*K outputs; -// n_parts = n_positions*split partials (each carrying K entries). -bool ensure_scratch(int device, size_t n, size_t n_parts) { - const size_t n_part_lists = n_parts * (size_t)kMaxK; // upper bound on K - if (g_scratch.device == device && g_scratch.cap >= n && - g_scratch.part_cap >= n_part_lists) - return true; - free_scratch(); - if (cudaMalloc(&g_scratch.d_lp, n * sizeof(float)) != cudaSuccess) goto fail; - if (cudaMalloc(&g_scratch.d_ids, n * sizeof(int32_t)) != cudaSuccess) goto fail; - if (cudaMalloc(&g_scratch.d_pmax, n_parts * sizeof(float)) != cudaSuccess) goto fail; - if (cudaMalloc(&g_scratch.d_psum, n_parts * sizeof(float)) != cudaSuccess) goto fail; - if (cudaMalloc(&g_scratch.d_pv, n_part_lists * sizeof(float)) != cudaSuccess) goto fail; - if (cudaMalloc(&g_scratch.d_pi, n_part_lists * sizeof(int32_t)) != cudaSuccess) goto fail; - g_scratch.device = device; - g_scratch.cap = n; - g_scratch.part_cap = n_part_lists; - return true; -fail: - free_scratch(); - return false; -} - -inline int pow2_ceil(int x) { - int p = 1; - while (p < x) p <<= 1; - return p; -} - -// Choose how many blocks to split each position's vocab scan across. The work -// is bandwidth bound, so we want enough total blocks (n_positions × split) to -// saturate the device while keeping each chunk large enough that the strided -// reads stay coalesced and per-block overhead stays amortized. -int pick_split(int vocab, int n_positions) { - if (const char * v = std::getenv("DFLASH_TOPK_SPLIT")) { - int s = std::atoi(v); - if (s >= 1 && s <= kMaxSplit) return s; - } - // Aim for ~240 total blocks (≈3 waves on a ~80-SM device, measured sweet - // spot for vocab~150k), but cap the chunk floor at ~2k elements so a block - // never scans too little to be worth launching. - int by_blocks = (240 + n_positions - 1) / n_positions; - int by_chunk = vocab / 2048; - int split = by_blocks < by_chunk ? by_blocks : by_chunk; - if (split < 1) split = 1; - if (split > kMaxSplit) split = kMaxSplit; - return split; -} - -} // namespace - -bool extract_draft_topk_cuda(const void * d_logits, - int n_positions, int vocab, int K, - float * out_log_probs, - int32_t * out_token_ids, - float temperature) { - if (!d_logits || n_positions <= 0 || vocab <= 0 || K <= 0 || K > kMaxK) return false; - - cudaPointerAttributes attr{}; - if (cudaPointerGetAttributes(&attr, d_logits) != cudaSuccess) { - cudaGetLastError(); // clear the error so we don't poison the next CUDA call - return false; - } - if (attr.type != cudaMemoryTypeDevice) return false; - - int prev = 0; - cudaGetDevice(&prev); - const int dev = attr.device; - if (dev != prev) cudaSetDevice(dev); - - static const bool kProfile = std::getenv("DFLASH_TOPK_PROFILE") != nullptr; - bool ok = false; - const int split = pick_split(vocab, n_positions); - const size_t n = (size_t)n_positions * K; - const size_t n_parts = (size_t)n_positions * split; - if (ensure_scratch(dev, n, n_parts)) { - const float inv_t = 1.0f / fmaxf(1e-3f, temperature); - cudaEvent_t e_k0, e_k1, e_c1; - if (kProfile) { cudaEventCreate(&e_k0); cudaEventCreate(&e_k1); cudaEventCreate(&e_c1); cudaEventRecord(e_k0); } - - const dim3 grid1(n_positions, split); - const int comb_block = pow2_ceil(split); - const float * lp_in = static_cast(d_logits); - // float4 loads are safe only when every row base is 16-byte aligned: - // the tensor base aligned and a vocab stride that is a multiple of 4. - const bool use_vec = (vocab % 4 == 0) && - (reinterpret_cast(lp_in) % 16 == 0); - // K (and the vectorization flag) are compile-time template parameters - // so the per-thread/per-partial top-K stays register-resident; dispatch - // the runtime K to its instantiation. K>kMaxK is already rejected above. -#define DFLASH_TOPK_LAUNCH(KV, VEC) \ - draft_topk_partial<<>>( \ - lp_in, vocab, inv_t, split, \ - g_scratch.d_pmax, g_scratch.d_psum, g_scratch.d_pv, g_scratch.d_pi); \ - draft_topk_combine<<>>( \ - g_scratch.d_pmax, g_scratch.d_psum, g_scratch.d_pv, g_scratch.d_pi, \ - split, g_scratch.d_lp, g_scratch.d_ids); -#define DFLASH_TOPK_CASE(KV) \ - case KV: \ - if (use_vec) { DFLASH_TOPK_LAUNCH(KV, true) } \ - else { DFLASH_TOPK_LAUNCH(KV, false) } \ - break; - switch (K) { - DFLASH_TOPK_CASE(1) DFLASH_TOPK_CASE(2) DFLASH_TOPK_CASE(3) DFLASH_TOPK_CASE(4) - DFLASH_TOPK_CASE(5) DFLASH_TOPK_CASE(6) DFLASH_TOPK_CASE(7) DFLASH_TOPK_CASE(8) - default: break; - } -#undef DFLASH_TOPK_CASE -#undef DFLASH_TOPK_LAUNCH - - if (kProfile) cudaEventRecord(e_k1); - if (cudaGetLastError() == cudaSuccess && cudaDeviceSynchronize() == cudaSuccess) { - const cudaError_t e1 = cudaMemcpy(out_log_probs, g_scratch.d_lp, - n * sizeof(float), cudaMemcpyDeviceToHost); - const cudaError_t e2 = cudaMemcpy(out_token_ids, g_scratch.d_ids, - n * sizeof(int32_t), cudaMemcpyDeviceToHost); - ok = (e1 == cudaSuccess && e2 == cudaSuccess); - } - if (kProfile) { - cudaEventRecord(e_c1); cudaEventSynchronize(e_c1); - float k_ms = 0, c_ms = 0; - cudaEventElapsedTime(&k_ms, e_k0, e_k1); - cudaEventElapsedTime(&c_ms, e_k1, e_c1); - std::fprintf(stderr, "[topk] kernels=%.3f ms sync+copy=%.3f ms (n_pos=%d vocab=%d split=%d)\n", - k_ms, c_ms, n_positions, vocab, split); - cudaEventDestroy(e_k0); cudaEventDestroy(e_k1); cudaEventDestroy(e_c1); - } - } - if (!ok) cudaGetLastError(); - if (dev != prev) cudaSetDevice(prev); - return ok; -} - -} // namespace dflash::common diff --git a/server/src/common/draft_topk_cuda.h b/server/src/common/draft_topk_cuda.h deleted file mode 100644 index 36ec0c8fe..000000000 --- a/server/src/common/draft_topk_cuda.h +++ /dev/null @@ -1,34 +0,0 @@ -// GPU top-K + log-prob extraction for DDTree draft distributions. -// -// Drop-in device-side replacement for extract_draft_topk (ddtree.cpp): instead -// of D2H-ing the full [vocab × n_positions] draft logits and running an OpenMP -// heap top-K + online-logsumexp on the CPU, this runs the whole thing on the -// GPU directly on the logits tensor's device buffer and copies back only the -// [n_positions × K] results. -// -// Semantics match extract_draft_topk exactly: -// scaled = logit * (1 / max(1e-3, temperature)) -// log_z = logsumexp_j(scaled_j) (per position) -// out_log_probs = top-K scaled logits (desc), minus log_z -// out_token_ids = matching vocab ids; ties broken toward the lower id -// -// Returns false (caller must fall back to the CPU path) when CUDA is -// unavailable, the pointer is not device memory, K is out of range, or any -// CUDA call fails. Only compiled into CUDA builds; see CMakeLists.txt. - -#pragma once - -#include - -namespace dflash::common { - -// d_logits: device pointer to row-major [n_positions][vocab] f32 logits (the -// position stride is `vocab` floats — pass an offset pointer to skip -// leading positions). out_* are HOST buffers of size n_positions*K. -bool extract_draft_topk_cuda(const void * d_logits, - int n_positions, int vocab, int K, - float * out_log_probs, - int32_t * out_token_ids, - float temperature); - -} // namespace dflash::common diff --git a/server/src/common/geometric_draft_topk_cuda.cu b/server/src/common/geometric_draft_topk_cuda.cu index 71086c98a..84d3242ef 100644 --- a/server/src/common/geometric_draft_topk_cuda.cu +++ b/server/src/common/geometric_draft_topk_cuda.cu @@ -304,19 +304,38 @@ inline int pow2_ceil(int x) { return p; } +// SM count of `device`, queried once and cached (topology is static for the +// process lifetime — matches the existing Scratch-style single-active-device +// assumption below). Backs pick_split's device-aware block-count target. +int sm_count(int device) { + static int cached_device = -1; + static int cached_sm = 0; + if (device != cached_device) { + cudaDeviceProp prop{}; + cached_sm = (cudaGetDeviceProperties(&prop, device) == cudaSuccess) + ? prop.multiProcessorCount + : 80; // fallback: the ~80-SM device this heuristic was tuned on + cached_device = device; + } + return cached_sm; +} + // Choose how many blocks to split each position's vocab scan across. The work // is bandwidth bound, so we want enough total blocks (n_positions × split) to // saturate the device while keeping each chunk large enough that the strided // reads stay coalesced and per-block overhead stays amortized. -int pick_split(int vocab, int n_positions) { +int pick_split(int vocab, int n_positions, int device) { if (const char * v = std::getenv("DFLASH_TOPK_SPLIT")) { int s = std::atoi(v); if (s >= 1 && s <= kMaxSplit) return s; } - // Aim for ~240 total blocks (≈3 waves on a ~80-SM device, measured sweet - // spot for vocab~150k), but cap the chunk floor at ~2k elements so a block - // never scans too little to be worth launching. - int by_blocks = (240 + n_positions - 1) / n_positions; + // Aim for ~3 waves of blocks across the device's actual SM count (was a + // fixed 240 = 3*80, tuned on one ~80-SM GPU; scaling by the real SM count + // keeps the same sweet spot on smaller/larger GPUs), but cap the chunk + // floor at ~2k elements so a block never scans too little to be worth + // launching. + const int target_blocks = 3 * sm_count(device); + int by_blocks = (target_blocks + n_positions - 1) / n_positions; int by_chunk = vocab / 2048; int split = by_blocks < by_chunk ? by_blocks : by_chunk; if (split < 1) split = 1; @@ -347,7 +366,7 @@ bool geometric_extract_draft_topk_cuda(const void * d_logits, static const bool kProfile = std::getenv("DFLASH_TOPK_PROFILE") != nullptr; bool ok = false; - const int split = pick_split(vocab, n_positions); + const int split = pick_split(vocab, n_positions, dev); const size_t n = (size_t)n_positions * K; const size_t n_parts = (size_t)n_positions * split; if (ensure_scratch(dev, n, n_parts)) { diff --git a/server/src/common/topm_extract_cuda.cu b/server/src/common/topm_extract_cuda.cu index f1ade0498..5cca7bc14 100644 --- a/server/src/common/topm_extract_cuda.cu +++ b/server/src/common/topm_extract_cuda.cu @@ -9,7 +9,18 @@ // whose cumulative top-down count first reaches M, then gather. Bins strictly // above the threshold are always kept; the threshold bin fills the remainder // (its intra-bin order is irrelevant — the bin spans <0.05% of the value -// range, so coverage of the true top-M is preserved). Reads the logits twice. +// range, so coverage of the true top-M is preserved). +// +// Reads the logits twice, not three times: pass 1 (histogram) must finish +// before the threshold bin is known, so it can't be avoided or fused with the +// gather. But the two gather passes (elements strictly above the threshold +// bin, and elements equal to it) previously each did their own full-vocab scan +// — a third full read of the same data recomputing the same order_key per +// element. They're fused into one pass below: each element is classified once +// against the threshold and routed to one of two disjoint output ranges via +// two independent atomic counters (the "above" range is exactly above_cnt[p] +// elements by construction; the "eq" range fills the rest, offset by +// above_cnt[p]) — same result, one fewer full-vocab DRAM pass. #include "topm_extract_cuda.h" @@ -33,16 +44,79 @@ __device__ __forceinline__ uint32_t order_key(float f) { return (u & 0x80000000u) ? ~u : (u | 0x80000000u); } +// Per-row alignment peel: how many leading scalar elements (0-3) `col` needs +// before it reaches a 16-byte (float4) boundary. Computed from the row's +// actual runtime address, so it's correct for any vocab/base-pointer +// alignment -- no global "vocab % 4 == 0 && base % 16 == 0" precondition +// needed; each row peels its own head (and, symmetrically, its own <=3 +// element tail past the aligned body) so the vectorized body is always +// provably aligned regardless of vocab's divisibility by 4. +__device__ __forceinline__ int align_head(const float * col, int vocab) { + const uintptr_t addr = reinterpret_cast(col); + const int head = (int)(((16 - (addr & 15u)) & 15u) >> 2); + return head < vocab ? head : vocab; // degenerate: row shorter than the peel +} + +// SM count of `device`, queried once and cached (topology is static for the +// process lifetime). Mirrors draft_topk_cuda.cu's sm_count. +int sm_count(int device) { + static int cached_device = -1; + static int cached_sm = 0; + if (device != cached_device) { + cudaDeviceProp prop{}; + cached_sm = (cudaGetDeviceProperties(&prop, device) == cudaSuccess) + ? prop.multiProcessorCount + : 80; // fallback: the ~80-SM device this heuristic was tuned on + cached_device = device; + } + return cached_sm; +} + +// Blocks per position for the vocab scan: aim for ~3 waves of blocks across +// the device's actual SM count (was a fixed 64, tuned on one ~80-SM GPU — +// same reasoning as draft_topk_cuda.cu's pick_split, just without that file's +// vocab/2048 chunk floor: this kernel's grid-stride loop spans all `nsplit` +// blocks jointly rather than assigning each split a contiguous chunk, so +// there's no "chunk too small" failure mode to guard against, only atomic +// contention on the histogram at extreme nsplit — hence the cap below). +int pick_nsplit(int n_tokens, int device) { + constexpr int kMaxNsplit = 256; + const int target_blocks = 3 * sm_count(device); + int nsplit = (target_blocks + n_tokens - 1) / n_tokens; + if (nsplit < 1) nsplit = 1; + if (nsplit > kMaxNsplit) nsplit = kMaxNsplit; + return nsplit; +} + // Pass 1: per-position histogram of the top kBits of order_key over the vocab. +// Always vectorized (float4): each row peels its own <=3-element head to an +// aligned float4 boundary and <=3-element tail past the aligned body (see +// align_head), handled by one designated thread per position so they aren't +// double-counted by every block. __global__ void histogram_kernel(const float * __restrict__ logits, int vocab, int n_tokens, unsigned int * __restrict__ hist) { // [n_tokens][kBins] const int p = blockIdx.y; unsigned int * h = hist + (size_t)p * kBins; const float * __restrict__ col = logits + (size_t)p * vocab; - for (int v = blockIdx.x * blockDim.x + threadIdx.x; v < vocab; - v += gridDim.x * blockDim.x) { - atomicAdd(&h[order_key(col[v]) >> (32 - kBits)], 1u); + + const int head = align_head(col, vocab); + const int nvec = (vocab - head) >> 2; + const float4 * __restrict__ col4 = reinterpret_cast(col + head); + + for (int v4 = blockIdx.x * blockDim.x + threadIdx.x; v4 < nvec; + v4 += gridDim.x * blockDim.x) { + const float4 f = col4[v4]; + atomicAdd(&h[order_key(f.x) >> (32 - kBits)], 1u); + atomicAdd(&h[order_key(f.y) >> (32 - kBits)], 1u); + atomicAdd(&h[order_key(f.z) >> (32 - kBits)], 1u); + atomicAdd(&h[order_key(f.w) >> (32 - kBits)], 1u); + } + if (blockIdx.x == 0 && threadIdx.x == 0) { + for (int v = 0; v < head; v++) + atomicAdd(&h[order_key(col[v]) >> (32 - kBits)], 1u); + for (int v = head + (nvec << 2); v < vocab; v++) + atomicAdd(&h[order_key(col[v]) >> (32 - kBits)], 1u); } } @@ -51,6 +125,10 @@ __global__ void histogram_kernel(const float * __restrict__ logits, // per position: each thread owns a contiguous chunk of bins, computes its chunk // total, thread 0 builds the per-chunk "count above" prefix, the single thread // whose chunk straddles the M-crossing rescans just its chunk. O(kBins/T + T). +// chunk (= kBins/kThrThreads = 256, fixed by the two constants above) is always +// a multiple of 4, and h+lo is always 16-byte aligned (hist is cudaMalloc'd, +// kBins and chunk are both multiples of 4 uints) — uint4 loads are safe +// unconditionally here, no runtime alignment check needed. __global__ void threshold_kernel(const unsigned int * __restrict__ hist, int n_tokens, int M, int * __restrict__ threshold_bin, @@ -65,7 +143,12 @@ __global__ void threshold_kernel(const unsigned int * __restrict__ hist, __shared__ unsigned int part[kThrThreads]; // chunk totals __shared__ unsigned int above[kThrThreads]; // count strictly above chunk t unsigned int s = 0; - for (int b = lo; b < hi; b++) s += h[b]; + const uint4 * __restrict__ h4 = reinterpret_cast(h + lo); +#pragma unroll + for (int c4 = 0; c4 < chunk / 4; c4++) { + const uint4 v4 = h4[c4]; + s += v4.x + v4.y + v4.z + v4.w; + } part[t] = s; __syncthreads(); @@ -88,43 +171,59 @@ __global__ void threshold_kernel(const unsigned int * __restrict__ hist, } } -// Pass 3a: emit all ids in bins strictly above the threshold bin (guaranteed -// in the top-M). Uses a per-position fill counter. -__global__ void gather_above_kernel(const float * __restrict__ logits, - int vocab, int n_tokens, int M, - const int * __restrict__ threshold_bin, - int32_t * __restrict__ cand_ids, - unsigned int * __restrict__ fill) { +// Fused gather pass: classify each vocab id against the threshold bin and +// route it to one of two disjoint output ranges in a single full-vocab scan +// (replaces the previous gather_above + gather_fill pair, which each did +// their own full scan). "above" elements (key > tb) land in slots +// [0, above_cnt[p]) via their own atomic counter; "eq" elements (key == tb) +// fill the remainder [above_cnt[p], M) via a second counter offset by +// above_cnt[p]. The two ranges never collide because membership is disjoint +// and each counter only ever allocates within its own range — intra-range +// order doesn't matter (see file header). Always vectorized, same per-row +// head/tail peel as histogram_kernel (see align_head); bounded writes (never +// an early `return`) keep every thread scanning its whole strided range so no +// in-range element is skipped. +__global__ void gather_kernel(const float * __restrict__ logits, + int vocab, int n_tokens, int M, + const int * __restrict__ threshold_bin, + const unsigned int * __restrict__ above_cnt, + int32_t * __restrict__ cand_ids, + unsigned int * __restrict__ fill_above, + unsigned int * __restrict__ fill_eq) { const int p = blockIdx.y; const int tb = threshold_bin[p]; + const unsigned int ac = above_cnt[p]; const float * __restrict__ col = logits + (size_t)p * vocab; int32_t * out = cand_ids + (size_t)p * M; - for (int v = blockIdx.x * blockDim.x + threadIdx.x; v < vocab; - v += gridDim.x * blockDim.x) { - if ((int)(order_key(col[v]) >> (32 - kBits)) > tb) { - const unsigned int slot = atomicAdd(&fill[p], 1u); - if (slot < (unsigned)M) out[slot] = v; - } - } -} -// Pass 3b: fill the remaining [above_cnt, M) slots from the threshold bin. -__global__ void gather_fill_kernel(const float * __restrict__ logits, - int vocab, int n_tokens, int M, - const int * __restrict__ threshold_bin, - int32_t * __restrict__ cand_ids, - unsigned int * __restrict__ fill) { - const int p = blockIdx.y; - const int tb = threshold_bin[p]; - const float * __restrict__ col = logits + (size_t)p * vocab; - int32_t * out = cand_ids + (size_t)p * M; - for (int v = blockIdx.x * blockDim.x + threadIdx.x; v < vocab; - v += gridDim.x * blockDim.x) { - if ((int)(order_key(col[v]) >> (32 - kBits)) == tb) { - const unsigned int slot = atomicAdd(&fill[p], 1u); + auto emit = [&](int v, unsigned int key) { + if ((int)key > tb) { + const unsigned int slot = atomicAdd(&fill_above[p], 1u); if (slot < (unsigned)M) out[slot] = v; - else return; // this position is full + } else if ((int)key == tb) { + const unsigned int slot = atomicAdd(&fill_eq[p], 1u); + if (ac + slot < (unsigned)M) out[ac + slot] = v; } + }; + + const int head = align_head(col, vocab); + const int nvec = (vocab - head) >> 2; + const float4 * __restrict__ col4 = reinterpret_cast(col + head); + + for (int v4 = blockIdx.x * blockDim.x + threadIdx.x; v4 < nvec; + v4 += gridDim.x * blockDim.x) { + const float4 f = col4[v4]; + const int base = head + (v4 << 2); + emit(base + 0, order_key(f.x) >> (32 - kBits)); + emit(base + 1, order_key(f.y) >> (32 - kBits)); + emit(base + 2, order_key(f.z) >> (32 - kBits)); + emit(base + 3, order_key(f.w) >> (32 - kBits)); + } + if (blockIdx.x == 0 && threadIdx.x == 0) { + for (int v = 0; v < head; v++) + emit(v, order_key(col[v]) >> (32 - kBits)); + for (int v = head + (nvec << 2); v < vocab; v++) + emit(v, order_key(col[v]) >> (32 - kBits)); } } @@ -135,29 +234,31 @@ bool extract_topm_cuda(const float * d_logits, int vocab, int n_tokens, int M, if (vocab <= 0 || n_tokens <= 0 || M <= 0 || M > vocab) return false; // scratch layout: hist[n_tokens*kBins] u32 | threshold_bin[n_tokens] i32 | - // above_cnt[n_tokens] u32 | fill[n_tokens] u32 + // above_cnt[n_tokens] u32 | fill_above[n_tokens] u32 | + // fill_eq[n_tokens] u32 auto * base = static_cast(d_scratch); auto * hist = reinterpret_cast(base); auto * threshold_bin = reinterpret_cast(hist + (size_t)n_tokens * kBins); auto * above_cnt = reinterpret_cast(threshold_bin + n_tokens); - auto * fill = above_cnt + n_tokens; + auto * fill_above = above_cnt + n_tokens; + auto * fill_eq = fill_above + n_tokens; if (cudaMemsetAsync(hist, 0, (size_t)n_tokens * kBins * sizeof(unsigned int), stream) != cudaSuccess) return false; - cudaMemsetAsync(fill, 0, (size_t)n_tokens * sizeof(unsigned int), stream); + cudaMemsetAsync(fill_above, 0, (size_t)n_tokens * sizeof(unsigned int), stream); + cudaMemsetAsync(fill_eq, 0, (size_t)n_tokens * sizeof(unsigned int), stream); - const int nsplit = 64; // blocks per position for the vocab scan + int dev = 0; + cudaGetDevice(&dev); + const int nsplit = pick_nsplit(n_tokens, dev); dim3 grid(nsplit, n_tokens); histogram_kernel<<>>(d_logits, vocab, n_tokens, hist); + threshold_kernel<<>>( hist, n_tokens, M, threshold_bin, above_cnt); - gather_above_kernel<<>>( - d_logits, vocab, n_tokens, M, threshold_bin, d_cand_ids, fill); - // Reset fill to above_cnt so 3b continues filling from the right slot. - cudaMemcpyAsync(fill, above_cnt, (size_t)n_tokens * sizeof(unsigned int), - cudaMemcpyDeviceToDevice, stream); - gather_fill_kernel<<>>( - d_logits, vocab, n_tokens, M, threshold_bin, d_cand_ids, fill); + + gather_kernel<<>>( + d_logits, vocab, n_tokens, M, threshold_bin, above_cnt, d_cand_ids, fill_above, fill_eq); return cudaGetLastError() == cudaSuccess; } @@ -166,7 +267,8 @@ size_t extract_topm_scratch_bytes(int n_tokens) { return (size_t)n_tokens * kBins * sizeof(unsigned int) // hist + (size_t)n_tokens * sizeof(int) // threshold_bin + (size_t)n_tokens * sizeof(unsigned int) // above_cnt - + (size_t)n_tokens * sizeof(unsigned int); // fill + + (size_t)n_tokens * sizeof(unsigned int) // fill_above + + (size_t)n_tokens * sizeof(unsigned int); // fill_eq } } // namespace dflash::common diff --git a/server/src/qwen35/graph_builders.cpp b/server/src/qwen35/graph_builders.cpp index 02ef60881..0ad266876 100644 --- a/server/src/qwen35/graph_builders.cpp +++ b/server/src/qwen35/graph_builders.cpp @@ -370,6 +370,10 @@ bool build_target_step( } else { if (!go.logits) return false; sg.logits = go.logits; + // hidden_states is always computed by build_qwen35_graph (already an + // output tensor kept alive), so expose it here too at zero extra cost: + // calibration reads it to run the restricted head off the same build. + sg.hidden_states = go.hidden_states; ggml_set_output(sg.logits); sg.argmax_tokens = ggml_argmax(sg.ctx, sg.logits); ggml_set_name(sg.argmax_tokens, "chain_verify_argmax"); @@ -451,6 +455,9 @@ bool build_target_step_tree( } else { if (!go.logits) return false; sg.logits = go.logits; + // Same rationale as build_target_step: hidden_states is already + // computed/kept alive, expose it for off-graph calibration use. + sg.hidden_states = go.hidden_states; ggml_set_output(sg.logits); sg.argmax_tokens = ggml_argmax(sg.ctx, sg.logits); ggml_set_name(sg.argmax_tokens, "tree_verify_argmax"); diff --git a/server/test/test_dflash.cpp b/server/test/test_dflash.cpp index c9554f1cb..b071904d7 100644 --- a/server/test/test_dflash.cpp +++ b/server/test/test_dflash.cpp @@ -62,6 +62,7 @@ to_fp32_cuda_t ggml_get_to_fp32_cuda(ggml_type type); #include #include #include +#include #ifdef _WIN32 #define setenv(name, value, overwrite) _putenv_s(name, value) @@ -273,6 +274,7 @@ using dflash::common::free_qwen35_layer_split_shards; #ifdef DFLASH27B_HAVE_TOPK_HEAD #include "common/restricted_lm_head_cuda.h" #include "common/topm_extract_cuda.h" +#include "common/draft_entropy_cuda.h" #endif #include "common/geometric_draft_topk_cuda.h" using dflash::common::is_eos_tok; @@ -3123,6 +3125,24 @@ int main(int argc, char ** argv) { std::vector massp95_sum(calib_Mgrid.size(), 0.0); std::vector massp95_bad(calib_Mgrid.size(), 0); + // Entropy-gated restricted head: bucket each verify position by the + // temperature-agnostic (raw-logit, T=1) entropy of the SAME draft row + // (i+1) that supplies its candidate set, so the coverage@M curve above + // can be conditioned on "how confident was the draft here". Buckets are + // fixed edges in nats; entropy is bounded by ln(vocab) (~11.93 @152k), so + // the last edge is comfortably an upper bound, not a hard cap. + static const std::vector calib_entropy_edges = { + 0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0}; + auto entropy_bucket_of = [](double h) -> int { + int b = 0; + while (b + 1 < (int)calib_entropy_edges.size() && h >= calib_entropy_edges[b + 1]) b++; + return b; + }; + const int calib_n_ebuckets = (int)calib_entropy_edges.size(); + std::vector> calib_hist_eb(calib_n_ebuckets); // [bucket][rank] -> count + std::vector calib_eb_count(calib_n_ebuckets, 0); + double calib_entropy_sum = 0.0; // overall mean, for the summary header + // Candidate-restricted greedy LM head (DFLASH_TOPK_HEAD=M). On the chain // fast-rollback single-GPU path, replace the verify-side full-vocab head // with: draft top-M extract → fused Q6_K gather/dot over those candidates → @@ -3130,8 +3150,49 @@ int main(int argc, char ** argv) { int g_topk_head_M = 0; double tt_head_extract = 0, tt_head_kernel = 0; long head_steps = 0; + // Entropy gate (docs/topk-head-optimization.md §5d/§5e): per verify step, + // gate the whole-graph restricted/full-head choice on the draft's own + // temperature-agnostic entropy for that step. Default +inf preserves the + // old ungated behavior (restricted head always on whenever M>0) so this + // is opt-in and backward compatible. Aggregated via MAX over the step's + // draft rows, not mean: a whole-graph decision should be as conservative + // as its least-confident row, not smoothed out by confident ones. + // + // Declared unconditionally (not just under DFLASH27B_HAVE_TOPK_HEAD) + // because tree_rhead/rhead reference head_entropy_gate_ok() outside any + // #ifdef -- g_topk_head_M is always 0 without the macro, so the `&&` + // chain never actually calls into CUDA code there, but it must still + // compile (e.g. the HIP build, which doesn't have DFLASH27B_HAVE_TOPK_HEAD). + float g_topk_head_entropy_T = std::numeric_limits::infinity(); + std::vector head_entropy_buf; + long head_gate_pass = 0, head_gate_total = 0; + auto head_entropy_gate_ok = [&]() -> bool { +#ifdef DFLASH27B_HAVE_TOPK_HEAD + if (!std::isfinite(g_topk_head_entropy_T)) return true; + if (!dflash::common::extract_draft_entropy_cuda( + (const float *)draft_sg.logits->data, (int)draft_sg.logits->ne[0], q_len, + head_entropy_buf.data(), 0)) + return false; // conservative: fail closed to the full head on a kernel error + float max_h = -std::numeric_limits::infinity(); + for (int i = 0; i < q_len; i++) max_h = std::max(max_h, head_entropy_buf[i]); + head_gate_total++; + const bool ok = max_h <= g_topk_head_entropy_T; + if (ok) head_gate_pass++; + return ok; +#else + return true; +#endif + }; + // Tracks whether THIS iteration's verify graph was actually built + // restricted (rhead/tree_rhead), for whatever reason (M==0, bridged + // draft, sequential verify, or the entropy gate above). The calibration + // fidelity check below needs this to know when sg.hidden_states/argmax + // reflect the full head vs. the runtime-approximate one. + bool g_step_used_restricted_head = false; #ifdef DFLASH27B_HAVE_TOPK_HEAD if (const char * e = std::getenv("DFLASH_TOPK_HEAD")) g_topk_head_M = std::atoi(e); + if (const char * e = std::getenv("DFLASH_TOPK_HEAD_ENTROPY_T")) + g_topk_head_entropy_T = std::atof(e); int32_t * d_cand = nullptr, * d_cand_use = nullptr, * d_head_tok = nullptr; void * d_topm_scr = nullptr, * d_head_keys = nullptr; // Node count upper bound: chain uses q_len positions; tree uses budget+1 nodes. @@ -3142,7 +3203,35 @@ int main(int argc, char ** argv) { cudaMalloc(&d_head_tok, (size_t)topk_head_Nmax * sizeof(int32_t)); cudaMalloc(&d_head_keys, (size_t)topk_head_Nmax * sizeof(unsigned long long)); cudaMalloc(&d_topm_scr, dflash::common::extract_topm_scratch_bytes(q_len)); - std::printf("[topk-head] restricted greedy head ON, M=%d\n", g_topk_head_M); + head_entropy_buf.resize(q_len); + if (std::isfinite(g_topk_head_entropy_T)) { + std::printf("[topk-head] restricted greedy head ON, M=%d, entropy gate T=%.3f\n", + g_topk_head_M, g_topk_head_entropy_T); + } else { + std::printf("[topk-head] restricted greedy head ON, M=%d, entropy gate OFF\n", g_topk_head_M); + } + } + + // Real-kernel fidelity check (DFLASH_TOPK_CALIB=1, independent of the + // g_topk_head_M runtime toggle above): for each entropy bucket, run the + // ACTUAL extract_topm_cuda + restricted_lm_head_q6k kernels at a small M + // grid and compare their output against the exact full-head argmax. This + // measures the extractor's own approximation error, which the idealized + // CPU-sorted rank histogram (calib_hist_eb) cannot see. + static const std::vector calib_head_Mgrid = {128, 256, 512, 1024, 2048}; + const int calib_head_Mmax = calib_head_Mgrid.back(); + int32_t * d_cal_cand = nullptr, * d_cal_cand_use = nullptr, * d_cal_head_tok = nullptr; + void * d_cal_topm_scr = nullptr, * d_cal_head_keys = nullptr; + // [bucket][Midx] -> match count / total count, against the full-head ground truth. + std::vector> head_match_eb, head_total_eb; + if (g_topk_calib) { + cudaMalloc(&d_cal_cand, (size_t)calib_head_Mmax * q_len * sizeof(int32_t)); + cudaMalloc(&d_cal_cand_use, (size_t)calib_head_Mmax * q_len * sizeof(int32_t)); + cudaMalloc(&d_cal_head_tok, (size_t)q_len * sizeof(int32_t)); + cudaMalloc(&d_cal_head_keys, (size_t)q_len * sizeof(unsigned long long)); + cudaMalloc(&d_cal_topm_scr, dflash::common::extract_topm_scratch_bytes(q_len)); + head_match_eb.assign(calib_n_ebuckets, std::vector(calib_head_Mgrid.size(), 0)); + head_total_eb.assign(calib_n_ebuckets, std::vector(calib_head_Mgrid.size(), 0)); } #endif @@ -3442,9 +3531,13 @@ int main(int argc, char ** argv) { // Candidate-restricted greedy head on the tree path: single-GPU, // greedy only (temp>0 samples the bonus from full logits, which we // skip). Per node, candidates come from the draft top-M at the - // node's depth (same alignment as the chain path). + // node's depth (same alignment as the chain path). Entropy gate: + // fall back to the full head for this whole step when the draft's + // own confidence is low (see head_entropy_gate_ok above; a no-op + // when DFLASH_TOPK_HEAD_ENTROPY_T is unset). const bool tree_rhead = (g_topk_head_M > 0) && !draft_hidden_bridge - && (g_sampler.temp == 0.0f); + && (g_sampler.temp == 0.0f) && head_entropy_gate_ok(); + g_step_used_restricted_head = tree_rhead; if (!build_target_step_tree(sg, w, cache, backend, /*kv_start=*/committed, /*n_tokens=*/N, g_fa_window, g_kq_stride_pad, @@ -3863,8 +3956,11 @@ int main(int argc, char ** argv) { const int verify_fa_window = g_fa_window; // Restricted greedy head only on the single-GPU fast-rollback chain // path (where draft logits use the target head over the shared vocab, - // and the commit path never reads sg.logits). - const bool rhead = (g_topk_head_M > 0) && !draft_hidden_bridge && fast_rollback; + // and the commit path never reads sg.logits). Entropy gate: see + // head_entropy_gate_ok above (no-op when the gate env var is unset). + const bool rhead = (g_topk_head_M > 0) && !draft_hidden_bridge && fast_rollback + && head_entropy_gate_ok(); + g_step_used_restricted_head = rhead; if (!build_target_step(sg, w, cache, backend, /*kv_start=*/committed, /*n_tokens=*/q_len, /*with_mask=*/true, /*capture=*/true, @@ -3968,7 +4064,10 @@ int main(int argc, char ** argv) { // Each call writes K/V at slot committed+i and advances SSM by 1. // After the loop, target cache state is identical to the batched // path's state (both end at committed+q_len). Restore/replay below - // still apply correctly. + // still apply correctly. Always full head (build_target_step + // defaults restricted_head=false here) -- reset the flag so a + // stale true from a prior chain/tree iteration doesn't leak in. + g_step_used_restricted_head = false; std::vector single_embed(hidden); int32_t p4_single[4]; for (int i = 0; i < q_len; i++) { @@ -4026,6 +4125,8 @@ int main(int argc, char ** argv) { h[r]++; }; static std::vector> nucleus; // (draft_rank, prob) + static std::vector pos_ebucket; // per-position entropy bucket, for the fidelity check below + pos_ebucket.assign(std::max(0, q_len - 1), 0); for (int i = 0; i < q_len - 1; i++) { const float * drow = draft_logits_buf.data() + (size_t)(i + 1) * vocab; const float * trow = verify_logits_buf.data() + (size_t)i * vocab; @@ -4051,6 +4152,27 @@ int main(int argc, char ** argv) { calib_total++; calib_rank_sum += grank; if (grank > calib_max_rank) calib_max_rank = grank; + // ── temperature-agnostic (raw-logit, T=1) draft entropy over the + // SAME row (i+1) that supplied the candidate set above. Not + // scaled by ddtree_temp: the gate is meant to reflect the draft's + // intrinsic confidence, independent of whatever sampling + // temperature happens to be configured. ── + double dmaxz = drow[0]; + for (int v = 1; v < vocab; v++) if (drow[v] > dmaxz) dmaxz = drow[v]; + double dsumexp = 0.0, dweighted = 0.0; + for (int v = 0; v < vocab; v++) { + const double e = std::exp((double)drow[v] - dmaxz); + dsumexp += e; + dweighted += e * (double)drow[v]; + } + const double dlogZ = dmaxz + std::log(dsumexp); + const double entropy = dlogZ - dweighted / dsumexp; // nats + const int ebucket = entropy_bucket_of(entropy); + pos_ebucket[i] = ebucket; + bump(calib_hist_eb[ebucket], grank); + calib_eb_count[ebucket]++; + calib_entropy_sum += entropy; + // ── target softmax(temp=1) over full vocab + sorted top-MCAP ── double maxz = trow[0]; for (int v = 1; v < vocab; v++) if (trow[v] > maxz) maxz = trow[v]; @@ -4096,6 +4218,57 @@ int main(int argc, char ** argv) { if (frac < 0.99) massp95_bad[mi]++; } } + +#ifdef DFLASH27B_HAVE_TOPK_HEAD + // Real-kernel fidelity check: run the ACTUAL extract_topm_cuda + + // restricted_lm_head_q6k kernels at a small M grid and compare + // against the exact full-head argmax, bucketed by the same + // per-position entropy computed above. This measures the + // extractor's real approximation error, which the idealized + // CPU-sorted rank histogram above cannot see. Only valid when this + // step's graph was built with the full (non-restricted) head, so + // sg.hidden_states covers all q_len positions and target_tok is + // the exact ground truth (not the runtime-approximate rhead output + // compared against itself). seq_verify rebuilds sg per-token, so by + // this point sg.hidden_states only reflects the last single-token + // step -- skip there. Uses g_step_used_restricted_head (set at the + // actual rhead/tree_rhead call site) rather than recomputing the + // condition here, so it stays correct now that rhead/tree_rhead + // also depend on the entropy gate's runtime result. + const bool full_head_valid = !seq_verify && !g_step_used_restricted_head; + if (full_head_valid && q_len > 1) { + const int n_fcheck = q_len - 1; // same bound as the rank loop above + const int dvocab = (int)draft_sg.logits->ne[0]; + static std::vector head_tok_host; + head_tok_host.resize(n_fcheck); + for (size_t mi = 0; mi < calib_head_Mgrid.size(); mi++) { + const int M = calib_head_Mgrid[mi]; + dflash::common::extract_topm_cuda( + (const float *)draft_sg.logits->data, dvocab, q_len, M, + d_cal_cand, d_cal_topm_scr, 0); + // Same pp -> min(pp+1, q_len-1) alignment as the runtime + // rhead path; pp only reaches q_len-2 here so the clamp + // never actually fires. + for (int i = 0; i < n_fcheck; i++) { + cudaMemcpyAsync(d_cal_cand_use + (size_t)i * M, + d_cal_cand + (size_t)(i + 1) * M, + (size_t)M * sizeof(int32_t), + cudaMemcpyDeviceToDevice, 0); + } + dflash::common::restricted_lm_head_q6k( + w.output->data, (int)w.output->ne[0], (int)w.output->ne[1], + (const float *)sg.hidden_states->data, n_fcheck, + d_cal_cand_use, M, d_cal_head_tok, d_cal_head_keys, 0); + cudaMemcpy(head_tok_host.data(), d_cal_head_tok, + sizeof(int32_t) * n_fcheck, cudaMemcpyDeviceToHost); + for (int i = 0; i < n_fcheck; i++) { + const int b = pos_ebucket[i]; + head_total_eb[b][mi]++; + if (head_tok_host[i] == target_tok[i]) head_match_eb[b][mi]++; + } + } + } +#endif } std::printf("[step %d] committed=%d last_tok=%d\n", n_draft_steps, committed, last_tok); @@ -4341,10 +4514,17 @@ int main(int argc, char ** argv) { } if (g_topk_calib && calib_total > 0) { - std::printf("\n[topk-calib] positions=%lld mean_rank=%.3f max_rank=%d (vocab=%d)\n", + // NOTE: mean_draft_entropy is appended AFTER (vocab=%d) rather than + // inlined earlier in the line -- calib/calibrate.py's regex for this + // header anchors "max_rank=(\d+)\s+\(vocab=(\d+)\)" back-to-back, and + // re.search doesn't care about trailing text, so this keeps the + // existing parser working unmodified while still exposing the field + // for the (separately added) mean_draft_entropy=([\d.]+) pattern. + std::printf("\n[topk-calib] positions=%lld mean_rank=%.3f max_rank=%d (vocab=%d) mean_draft_entropy=%.4f\n", (long long)calib_total, (double)calib_rank_sum / (double)calib_total, - calib_max_rank, vocab); + calib_max_rank, vocab, + calib_entropy_sum / (double)calib_total); // Set-coverage: fraction of positions whose ENTIRE set is within draft top-M. auto cov = [&](const char * label, const std::vector & h) { std::printf("[topk-calib] --- %s: set-coverage (entire set in draft top-M) ---\n", label); @@ -4370,6 +4550,53 @@ int main(int argc, char ** argv) { 100.0 * massp95_sum[mi] / (double)calib_total, (long long)massp95_bad[mi], (long long)calib_total); } + + // Entropy-conditioned greedy coverage: same idealized rank histogram + // as calib_hist above, split by the draft's temperature-agnostic + // entropy bucket for this position. Answers "for positions with + // entropy < T, what M gives 99.9% coverage" by summing buckets below + // T offline (see calib/calibrate.py::aggregate). + for (int b = 0; b < calib_n_ebuckets; b++) { + if (calib_eb_count[b] == 0) continue; + const double lo = calib_entropy_edges[b]; + const bool open_ended = (b + 1 == calib_n_ebuckets); + std::printf("[topk-calib-eb] entropy-bucket=%d range=[%.2f,%s) positions=%lld\n", + b, lo, open_ended ? "inf" : std::to_string(calib_entropy_edges[b + 1]).c_str(), + (long long)calib_eb_count[b]); + for (int k : calib_Mgrid) { + int64_t c = 0; + const auto & h = calib_hist_eb[b]; + for (int r = 0; r < k && r < (int)h.size(); r++) c += h[r]; + std::printf("[topk-calib-eb] coverage@M=%-6d : %8.4f%%\n", + k, 100.0 * (double)c / (double)calib_eb_count[b]); + } + } + +#ifdef DFLASH27B_HAVE_TOPK_HEAD + // Real-kernel fidelity: match rate of the ACTUAL restricted-head + // kernels against the exact full-head argmax, by entropy bucket and M. + // This is the number that should drive the (T, M) choice -- unlike the + // coverage@M table above, it is not an idealized upper bound. + bool any_head_data = false; + for (int b = 0; b < calib_n_ebuckets; b++) + for (size_t mi = 0; mi < calib_head_Mgrid.size(); mi++) + if (head_total_eb[b][mi] > 0) any_head_data = true; + if (any_head_data) { + for (int b = 0; b < calib_n_ebuckets; b++) { + int64_t btot = 0; + for (size_t mi = 0; mi < calib_head_Mgrid.size(); mi++) btot = std::max(btot, head_total_eb[b][mi]); + if (btot == 0) continue; + std::printf("[topk-calib-head] entropy-bucket=%d positions=%lld\n", b, (long long)btot); + for (size_t mi = 0; mi < calib_head_Mgrid.size(); mi++) { + const int64_t tot = head_total_eb[b][mi]; + if (tot == 0) continue; + std::printf("[topk-calib-head] match@M=%-6d : %8.4f%%\n", + calib_head_Mgrid[mi], + 100.0 * (double)head_match_eb[b][mi] / (double)tot); + } + } + } +#endif } auto t_gen1 = std::chrono::steady_clock::now(); @@ -4395,6 +4622,12 @@ int main(int argc, char ** argv) { tt_head_extract / 1000.0 / head_steps, tt_head_kernel / 1000.0 / head_steps); } + if (head_gate_total > 0) { + std::printf(" [topk-head-gate] T=%.3f restricted=%ld/%ld steps (%.1f%%) full-head fallback=%ld steps\n", + g_topk_head_entropy_T, head_gate_pass, head_gate_total, + 100.0 * (double)head_gate_pass / (double)head_gate_total, + head_gate_total - head_gate_pass); + } std::printf(" accept %.2f\n", avg_ms(tt_accept)); std::printf(" restore_ssm %.2f\n", avg_ms(tt_restore)); std::printf(" replay_build %.2f\n", avg_ms(tt_replay_build)); diff --git a/server/test/test_draft_entropy_cuda.cpp b/server/test/test_draft_entropy_cuda.cpp new file mode 100644 index 000000000..13850a954 --- /dev/null +++ b/server/test/test_draft_entropy_cuda.cpp @@ -0,0 +1,142 @@ +// Correctness test for dflash::common::extract_draft_entropy_cuda (GPU) vs a +// double-precision CPU reference. +// +// Deliberately stresses both branches of entropy_combine's warp-level +// reduction (src/common/draft_entropy_cuda.cu): the single-warp fast path +// (n_warps==1, including a genuinely partial warp with blockDim.x < 32) and +// the multi-warp cross-warp-combine path (n_warps>1, needs split>32, which +// pick_split only reaches for a small n_positions against a large vocab — +// see the case table below for exactly how each shape is chosen to land on +// one side or the other regardless of the test GPU's actual SM count). +// +// Build: registered in server/CMakeLists.txt under DFLASH27B_TESTS (CUDA +// backend only -- draft_entropy_cuda.cu is CUDA-only, unlike its topk sibling +// which also compiles for HIP). Run: ./test_draft_entropy_cuda (0 = pass). + +#include "../src/common/draft_entropy_cuda.h" + +#include + +#include +#include +#include +#include + +using dflash::common::extract_draft_entropy_cuda; + +namespace { + +// float4/16-byte-aligned VEC path requires vocab % 4 == 0 (see the .cu); all +// cases below use such vocabs so both use_vec branches get exercised across +// the different (n, vocab) shapes. +constexpr float kTol = 8e-3f; + +struct Case { + int n; // n_positions + int vocab; +}; + +// Double-precision CPU reference: temperature-agnostic (raw-logit) entropy in +// nats, H = log_z - sum(p_i * logit_i) = log_z - weighted/sumexp. +double cpu_entropy(const float * row, int vocab) { + double m = -1e300; + for (int j = 0; j < vocab; j++) m = std::max(m, (double)row[j]); + double sumexp = 0.0, weighted = 0.0; + for (int j = 0; j < vocab; j++) { + const double e = std::exp((double)row[j] - m); + sumexp += e; + weighted += e * (double)row[j]; + } + const double log_z = m + std::log(sumexp); + return log_z - weighted / sumexp; +} + +bool run_case(const Case & c, unsigned seed) { + const size_t n_logits = (size_t)c.n * c.vocab; + std::vector h_logits(n_logits); + std::mt19937 rng(seed); + std::normal_distribution dist(0.f, 4.f); + for (auto & x : h_logits) x = dist(rng); + + std::vector cpu_h(c.n); + for (int r = 0; r < c.n; r++) + cpu_h[r] = cpu_entropy(h_logits.data() + (size_t)r * c.vocab, c.vocab); + + float * d_logits = nullptr; + cudaError_t err = cudaMalloc(&d_logits, n_logits * sizeof(float)); + if (err != cudaSuccess) { + printf(" cudaMalloc failed: %s\n", cudaGetErrorString(err)); + return false; + } + cudaMemcpy(d_logits, h_logits.data(), n_logits * sizeof(float), cudaMemcpyHostToDevice); + + std::vector gpu_h(c.n, 0.f); + bool ok = extract_draft_entropy_cuda(d_logits, c.vocab, c.n, gpu_h.data()); + cudaFree(d_logits); + + if (!ok) { + printf(" FAIL: extract_draft_entropy_cuda returned false\n"); + return false; + } + + double max_err = 0.0; + int worst = -1; + for (int r = 0; r < c.n; r++) { + const double err = std::fabs((double)gpu_h[r] - cpu_h[r]); + if (err > max_err) { max_err = err; worst = r; } + } + + const bool pass = max_err <= kTol; + printf(" [%s] n=%-4d vocab=%-7d max_err=%.3e (row=%d gpu=%.5f cpu=%.5f)\n", + pass ? "PASS" : "FAIL", c.n, c.vocab, max_err, worst, + worst >= 0 ? gpu_h[worst] : 0.f, worst >= 0 ? cpu_h[worst] : 0.0); + return pass; +} + +} // namespace + +int main() { + int dev_count = 0; + if (cudaGetDeviceCount(&dev_count) != cudaSuccess || dev_count == 0) { + printf("SKIP: no CUDA device available\n"); + return 0; + } + + const Case cases[] = { + // n=1, large vocab: by_chunk=vocab/2048=74 and by_blocks=3*SMs/1 both + // exceed 32 on any real discrete GPU -> split>32 -> pow2_ceil>=64 -> + // entropy_combine's multi-warp (n_warps>1) cross-warp-combine path. + {1, 151936}, + // n~100, same vocab: by_blocks=3*SMs/100 is small (<=32 on anything + // short of an absurdly large GPU) -> split likely single-digit -> + // pow2_ceil<32 -> entropy_combine's single-warp path with a + // genuinely partial warp (blockDim.x < 32, exercises __activemask() + // over fewer than 32 real lanes). + {100, 151936}, + // n~15 (realistic DDTree draft-batch size): split lands in the + // teens-to-32 range depending on SM count -> exercises blockDim.x + // near/at exactly 32 (a full, non-partial single warp). + {15, 151936}, + // Small/edge shapes: vocab barely above the VEC path's float4 tile, + // and a vocab not a multiple of 4 (exercises the scalar non-VEC path + // and, for the VEC-eligible cases, the tail epilogue). + {7, 1024}, + {32, 4096}, + {3, 260}, // vocab % 4 == 0, small + {5, 257}, // vocab % 4 != 0 -> scalar path + {1, 4}, // minimal vocab + }; + + int failures = 0, idx = 0; + for (const Case & c : cases) { + if (!run_case(c, /*seed=*/4242u + idx)) failures++; + idx++; + } + + if (failures) { + printf("\nFAILED: %d/%d cases\n", failures, idx); + return 1; + } + printf("\nALL PASS: %d/%d cases\n", idx, idx); + return 0; +} diff --git a/server/test/test_topm_extract_cuda.cpp b/server/test/test_topm_extract_cuda.cpp new file mode 100644 index 000000000..74ff02443 --- /dev/null +++ b/server/test/test_topm_extract_cuda.cpp @@ -0,0 +1,177 @@ +// Correctness test for dflash::common::extract_topm_cuda (GPU) vs a CPU +// reference of the same coarse-bin-threshold algorithm. +// +// extract_topm_cuda's contract is NOT "exact top-M by value" -- it partitions +// by the top 16 bits of an order-preserving key, so it's only exact up to +// that coarsening (see the .cu file header). The invariant this test checks +// is the algorithm's actual contract: every element whose coarse bin is +// strictly above the threshold bin must appear in the output (no winner +// dropped), no element whose bin is strictly below the threshold must appear +// (no non-candidate let in), the output has exactly M distinct valid ids, and +// the above/eq split sizes match a CPU-computed reference threshold exactly. +// +// Deliberately includes several vocab sizes NOT divisible by 4 -- the whole +// point of the per-row alignment-peel rewrite (align_head in the .cu) was to +// drop the old "vocab % 4 == 0" precondition, so this is the coverage that +// didn't exist before that change at all. +// +// Build: registered in server/CMakeLists.txt under DFLASH27B_TESTS, CUDA +// backend only (topm_extract_cuda.cu is CUDA-only, gated by +// DFLASH27B_HAVE_TOPK_HEAD). Run: ./test_topm_extract_cuda (0 = pass). + +#include "../src/common/topm_extract_cuda.h" + +#include + +#include +#include +#include +#include +#include +#include + +using dflash::common::extract_topm_cuda; +using dflash::common::extract_topm_scratch_bytes; + +namespace { + +constexpr int kBits = 16; // must match topm_extract_cuda.cu's kBits + +uint32_t order_key(float f) { + uint32_t u; + std::memcpy(&u, &f, sizeof(u)); + return (u & 0x80000000u) ? ~u : (u | 0x80000000u); +} + +int coarse_bin(float f) { return (int)(order_key(f) >> (32 - kBits)); } + +// CPU reference: same coarse-bin histogram + top-down threshold search the +// kernel uses. Returns (threshold_bin, above_cnt) for one row. +std::pair cpu_threshold(const float * row, int vocab, int M) { + std::vector hist(1 << kBits, 0); + for (int v = 0; v < vocab; v++) hist[coarse_bin(row[v])]++; + int64_t cum = 0; + for (int b = (1 << kBits) - 1; b >= 0; b--) { + if (cum + hist[b] >= (int64_t)M) return {b, (int)cum}; + cum += hist[b]; + } + return {0, (int)cum}; // unreachable when M <= vocab +} + +struct Case { + int n; + int vocab; + int M; +}; + +bool run_case(const Case & c, unsigned seed) { + const size_t n_logits = (size_t)c.n * c.vocab; + std::vector h_logits(n_logits); + std::mt19937 rng(seed); + std::normal_distribution dist(0.f, 4.f); + for (auto & x : h_logits) x = dist(rng); + + float * d_logits = nullptr; + int32_t * d_cand = nullptr; + void * d_scratch = nullptr; + cudaMalloc(&d_logits, n_logits * sizeof(float)); + cudaMalloc(&d_cand, (size_t)c.n * c.M * sizeof(int32_t)); + cudaMalloc(&d_scratch, extract_topm_scratch_bytes(c.n)); + cudaMemcpy(d_logits, h_logits.data(), n_logits * sizeof(float), cudaMemcpyHostToDevice); + + const bool ok = extract_topm_cuda(d_logits, c.vocab, c.n, c.M, d_cand, d_scratch, nullptr); + cudaDeviceSynchronize(); + + std::vector h_cand((size_t)c.n * c.M); + if (ok) cudaMemcpy(h_cand.data(), d_cand, h_cand.size() * sizeof(int32_t), cudaMemcpyDeviceToHost); + cudaFree(d_logits); + cudaFree(d_cand); + cudaFree(d_scratch); + + if (!ok) { + printf(" FAIL: extract_topm_cuda returned false\n"); + return false; + } + + int bad_rows = 0; + for (int r = 0; r < c.n; r++) { + const float * row = h_logits.data() + (size_t)r * c.vocab; + const auto [thr, above_cnt] = cpu_threshold(row, c.vocab, c.M); + const int32_t * out = h_cand.data() + (size_t)r * c.M; + + std::unordered_set seen; + int n_above = 0, n_eq = 0, dup = 0, out_of_range = 0, below_thr = 0; + for (int i = 0; i < c.M; i++) { + const int32_t id = out[i]; + if (id < 0 || id >= c.vocab) { out_of_range++; continue; } + if (!seen.insert(id).second) { dup++; continue; } + const int bin = coarse_bin(row[id]); + if (bin > thr) n_above++; + else if (bin == thr) n_eq++; + else below_thr++; + } + // Completeness: every element with bin > thr must be in the output. + int total_above_in_data = 0; + for (int v = 0; v < c.vocab; v++) + if (coarse_bin(row[v]) > thr) total_above_in_data++; + + const bool row_ok = (dup == 0) && (out_of_range == 0) && (below_thr == 0) && + (n_above == above_cnt) && (n_above == total_above_in_data) && + (n_above + n_eq == c.M); + if (!row_ok) { + bad_rows++; + if (bad_rows <= 4) { + printf(" row=%d BAD: thr=%d above_cnt(cpu)=%d n_above(gpu)=%d " + "n_eq(gpu)=%d dup=%d oor=%d below_thr=%d total_above_in_data=%d\n", + r, thr, above_cnt, n_above, n_eq, dup, out_of_range, below_thr, + total_above_in_data); + } + } + } + + const bool pass = (bad_rows == 0); + printf(" [%s] n=%-4d vocab=%-7d M=%-5d bad_rows=%d/%d\n", + pass ? "PASS" : "FAIL", c.n, c.vocab, c.M, bad_rows, c.n); + return pass; +} + +} // namespace + +int main() { + int dev_count = 0; + if (cudaGetDeviceCount(&dev_count) != cudaSuccess || dev_count == 0) { + printf("SKIP: no CUDA device available\n"); + return 0; + } + + const Case cases[] = { + // Realistic decode shape, vocab % 4 == 0 (the old use_vec-eligible case). + {15, 151936, 1024}, + {1, 151936, 1024}, + // vocab % 4 != 0 -- previously always the scalar fallback; now exercises + // the same vectorized path via per-row alignment peeling. + {15, 151937, 1024}, + {15, 151933, 512}, + {7, 1023, 128}, + {3, 257, 64}, + // Small/edge shapes: M close to vocab, vocab smaller than the peel width. + {4, 130, 128}, + {2, 5, 4}, + {1, 4, 1}, + {1, 3, 1}, // vocab < float4 width entirely + {32, 4096, 2048}, + }; + + int failures = 0, idx = 0; + for (const Case & c : cases) { + if (!run_case(c, /*seed=*/9001u + idx)) failures++; + idx++; + } + + if (failures) { + printf("\nFAILED: %d/%d cases\n", failures, idx); + return 1; + } + printf("\nALL PASS: %d/%d cases\n", idx, idx); + return 0; +}