JetSpec tree drafting on DFlash — lossless; width=1 clean drop-in

jetspec/.gitignore

@@ -0,0 +1,10 @@
+runs/
+profiles/
+logs/
+*.pid
+*.nsys-rep
+*.trace.json.gz
+_ref/
+_oldattempt/
+_hf_configs/
+baseline_linear.json

jetspec/README.md

@@ -0,0 +1,43 @@
+# JetSpec tree drafting on SGLang DFlash
+
+Extends SGLang's linear DFlash speculative decoding with **parallel tree drafting** (JetSpec): the DFlash draft head's per-depth block logits feed a token **tree** (top2gap construction), verified under a tree-causal mask. `tree_width=1` is the unchanged linear path; `width>1` is the new tree path.
+
+## What this delivers
+
+- **Lossless tree spec-decode** (token-exact vs linear DFlash), dense Qwen3-8B + MoE Qwen3.6-35B-A3B, cuda-graphed, EAGLE-parity direct commit (no reverify).
+- **`top2gap` construction** — per-depth fanout from the top-1/top-2 logprob gap (branch where the drafter is uncertain, chain where confident). Reproduces the paper's accept length at a fraction of the budget.
+- **`width=1` is a byte-identical, near-zero-overhead drop-in** for linear DFlash.
+- A proven FA4-exact paged-tree verify kernel (opt-in, default-off) + a full break-even / profiling analysis.
+
+## Full speedup matrix (Qwen3-8B, B300, FA4 + page_size 16, greedy)
+
+| Config | GSM8K tok/s | vs DFlash | MATH-500 tok/s | vs DFlash | accept (G / M) | lossless |
+|---|---:|---:|---:|---:|---:|:--:|
+| **DFlash w1 (native linear)** | **1170.5** | 1.00× | **1522.1** | 1.00× | 5.85 / 7.62 | baseline |
+| linear JetSpec (width=1) | 1152.4 | **0.985×** | 1488.6 | **0.978×** | 5.85 / 7.62 | ✓ exact |
+| compact tree (top2gap w8/b16) | 1002.2 | 0.856× | 1167.5 | 0.767× | 6.45 / 7.90 | ✓ exact |
+| paged tree (w8/b16) | 900.2 | 0.769× | 1060.7 | 0.697× | 6.50 / 8.02 | ✗ (default-off) |
+
+- vs **AR-greedy** (the paper's baseline): every spec config is ~4–9× faster — that comparison is trivially won and not the interesting bar.
+- The interesting bar is **SGLang's linear DFlash**, which is already FA4 + cuda-graph + fused-kernel optimized.
+
+## Why the tree could not improve further (the honest result)
+
+- **Break-even rule:** tree beats linear iff `accept_ratio > step_cost_ratio`. The tree verify costs **~1.29× (compact) / 1.44× (paged)** per step at equal node count, so it needs that much more acceptance to win.
+- **The accept gain is too small:** tree gives **+0.6 (GSM8K) / +0.28 (MATH)** accept over linear; break-even needs **+1.5 / +2.4**. The draft head's tree just doesn't concentrate enough probability to clear the verify overhead.
+- **Construction was maxed** (top2gap is the reference's sweep-winning algorithm) and **the paged kernel is a net loss** (it removes the K/V gather but a different memory layout adds attention time *and* breaks bit-exactness vs the compact verify — fp non-associativity flips greedy near-ties).
+- **The paper's own tree is slower than SGLang's linear** (paper 984/1150 tok/s vs our linear 1170/1522), so "beat SGLang linear DFlash" is a strictly harder bar than the paper's published speedup (which is vs AR-greedy). Conclusion: against an already-SOTA linear baseline, tree drafting trades verify cost for accept length and nets out **below linear** on throughput.
+
+## The good news: width=1 is a clean drop-in
+
+- **Linear JetSpec (width=1) is within ~1.5–2% of native DFlash** (1152 vs 1170, 1489 vs 1522) and **token-exact** — i.e. within measurement noise.
+- That means the tree code path adds **negligible per-step overhead when not used**, so the feature is safe to land: it never regresses the linear DFlash path, and `width>1` is available for workloads where verify is cheaper or the draft head concentrates better (e.g. weaker/AR baselines, or future paged-engine work).
+
+## Reproduce
+
+```bash
+python jetspec/bench_paper_sglang.py   # GSM8K / MATH-500, configurable width/budget
+bash jetspec/run_dflash_gate_bench.sh  # losslessness gate vs fresh linear oracle
+```
+
+Detailed run-by-run notes: `jetspec/notes/bench_results.md`; status: `jetspec/STATUS.md`.

jetspec/_oldattempt/NOTES.md

@@ -0,0 +1,43 @@
+# Old attempt (realtree) — handoff artifacts
+
+The previous Codex run removed the branch-reverify band-aid and implemented the REAL target
+tree verify, but stopped (interrupted) before it was fully lossless. It got VERY close and
+isolated the bug. Use this to continue — do NOT start over.
+
+## State of the code (already in the working tree, compiles cleanly)
+WIP edits (also backed up as `jetspec/_oldattempt/realtree_wip.patch`) touch:
+- `python/sglang/srt/speculative/dflash_worker_v2.py` (325-line rework: real tree verify, band-aid removed)
+- `python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py` (route DFlash tree through the topk>1 verify path)
+- `python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py`, `base_runner.py`, `eager_runner.py`, `model_runner.py` (cuda-graph wiring for tree decode)
+- `python/sglang/srt/speculative/{dflash_info.py,dflash_tree_utils.py,spec_info.py}`, test file
+
+## THE KEY FINDINGS (save you hours)
+1. **`jetspec/baseline_linear.json` is STALE** — it was captured under a different launch
+   mode/config (flashinfer attn / graph state). The CURRENT width=1 linear DFlash (no-graph)
+   does NOT match it. So DO NOT gate losslessness against `baseline_linear.json`.
+   **Use `jetspec/runs/8b_linear_current_nograph_31216.json` as the current-code 8B linear
+   oracle** (recaptured from the exact current code/launch). Several earlier "mismatches"
+   (incl. all of prompt 0) were FALSE ALARMS from the stale baseline.
+2. **Against the fresh oracle, the real remaining bug is OVER-ACCEPTANCE in the tree-greedy
+   accept path**, reproducing on 8B w7/b64 at **prompts 3, 5, 9 only** (first diffs ~@51/@36/@9).
+   "Over-acceptance" = the tree accept commits a branch token the target's own greedy argmax
+   would NOT have produced. Prompt 5's first diff is a simple duplicate token — easiest to trace
+   the bad accept chunk. This is an accept-selection / node-gating bug, NOT a draft-tree, KV, or
+   mask-construction bug (cap=0 root-only verification was internally consistent).
+
+## What was tried (don't repeat)
+- cachefix / tritonattn / nograph variants — all still failed because they were compared vs the
+  STALE baseline; the real comparator is the fresh oracle above.
+- cap=0 (root-only tree) — confirmed the divergence is in accept logic, not deeper tree rows.
+
+## Failing-run artifacts
+`jetspec/runs/8b_tree_w7_b64_realtree*.json`, `..._debug_prompt0_*.json` — contain per-prompt
+`losslessness.mismatches` with first_diff positions.
+
+## The fix to make
+Match the accept semantics of the reference tree verify in
+`jetspec/_ref/JetSpec/jetspec/core/tree_attention_kernel.py` and the verify stack in
+`jetspec/_ref/JetSpec/jetspec/inference_engine/compiled_verify_stack.py` /
+`paged_tree_attn.py`. The over-acceptance is almost certainly in how the accepted path is
+selected from the target's per-node logits (greedy: accept child c at node n iff
+target_argmax(logits[n]) == draft_token(c); stop at first node whose best child != target argmax).