Pramodith/partial topk

calib/.gitignore

@@ -0,0 +1,2 @@
+*.bin
+out/

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()

calib/prep_prompts.py

@@ -0,0 +1,67 @@
+#!/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) 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
+
+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
+
+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: "),
+    # 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():
+    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("<i", int(t)))
+            written += 1
+    print(f"wrote {written} prompt bins to {out.resolve()}", file=sys.stderr)
+
+if __name__ == "__main__":
+    main()

calib/run_calib.sh

@@ -0,0 +1,23 @@
+#!/bin/bash
+# Top-k/top-p LM-head calibration over all prompt bins in calib/.
+# Chain verify path (no --ddtree) so target_tok[i] aligns with draft row i+1.
+# Each test_dflash run prints its own coverage tables; aggregate with
+# calib/aggregate.py over the captured log.
+#
+#   GPU=0 NGEN=256 bash calib/run_calib.sh | tee calib/out/calib.log
+set -u
+BIN=${BIN:-server/build/test_dflash}
+TARGET=${TARGET:-server/models/Qwen3.6-27B-Q4_K_M.gguf}
+DRAFT=${DRAFT:-server/models/draft/dflash-draft-3.6-q4_k_m.gguf}
+NGEN=${NGEN:-256}
+GPU=${GPU:-0}
+MAXCTX=${MAXCTX:-1024}
+mkdir -p calib/out
+for p in calib/*.bin; do
+  name=$(basename "$p" .bin)
+  echo "######## $name ########"
+  CUDA_VISIBLE_DEVICES=$GPU DFLASH_TOPK_CALIB=1 \
+    "$BIN" "$TARGET" "$DRAFT" "$p" "$NGEN" "calib/out/${name}_out.bin" \
+    --fast-rollback --max-ctx="$MAXCTX" 2>&1 \
+    | grep -E "topk-calib|out of memory|failed|CUDA error"
+done

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`.