Skip to content

Speculative decoding on Qwen3.5: a fused draft step and an fp8 draft head - #1102

Draft
merceod wants to merge 59 commits into
users/atindra/pr-24-t3-specdecfrom
users/atindra/pr-36-t6-fused-draft
Draft

merceod wants to merge 59 commits into
users/atindra/pr-24-t3-specdecfrom
users/atindra/pr-36-t6-fused-draft

Conversation

@merceod

@merceod merceod commented Nov 1, 2026

Copy link
Copy Markdown
Collaborator

Stacked on #1090 (pr-24-t3-specdec). Picks the commits of #1091, #1096, #1099, #1101 first; they drop out of this diff as those merge.

Why

#1090's draft passes ran over the verify's (rows, k+1) tokens and threw k of them away after attention, and the step's bookkeeping was 88 small torch kernels plus 148 copies and fills. A 9B k=2 step at c=1 spent 886 us per drafted token on drafting (one GPU trace, shared node). After that, the 248k x 4096 vocabulary head is most of a draft pass.

What changed

Stacked on #1090, with #1091 (bringing #1085), #1096, #1099 (bringing #1092) and #1101 (bringing #1100) picked first: #1099 for 7915843, the fp8 op the draft head uses, #1096 for the compile_dynamic capture configs, #1101 because this PR's decode configs give a seeded start 22 compile-pool jobs instead of 1, which exposes the hang #1101 fixes.

  • One query a row (Qwen3_5MTPHead.draft, SpecDecodeResource.attend_draft, attend_rows): pass 0 projects all k+1 tokens and attends one query a row, later passes run one token a row.
  • mstar::qkv_norm_rope_cache: q/k norm, rope and the MTP layer's KV write in one kernel. mstar::spec_verify_inputs and mstar::spec_accept replace the bookkeeping with two Triton kernels.
  • record_draft runs sample_unfiltered at temperature 0 in the graph, into the next step's token register. GDN verify rows run on FlyDSL's chain kernel. decode_rows=True compiles a verify bucket as decode by its rows, and the MI325X decode GEMV configs gain rows 3, 6, 12 and the fc.
  • draft_head: fp8: an fp8 copy of the head with one scale for all of it, AITER wvSplitKQ at 1-2 rows, _scaled_mm from 3. 415 to 213 us at one row, +1 GB of VRAM on 9B. draft_head_variants.py measures acceptance per head variant.
  • CUDA path unchanged: decode_rows defaults to False and only the ROCm inductor_overrides read it, the rest sits behind num_draft_tokens, which no CUDA deployment declares. rust/ unchanged. Docs updated.

Verification

c plain mtp2 (#1090) mtp2 fused mtp3 fused fused k=2 / plain fused k=2 / mtp2
1 224 331 408 440 1.82x 1.23x
4 766 1155 1331 1424 1.74x 1.15x
16 2586 3529 3542 3966 1.37x 1.00x
64 7139 7880 8560 7857 1.20x 1.09x
  • k=3 fused is 1.96x plain at c=1. The c=16 fused k=2 cell is noisy: 4 of 6 runs stalled, range 2948-3933.
  • Shared node, three launches per arm, clean ms/token (best repeat's median request), 9B chat: k=2 2.33 against plain's 4.36 = 1.87x (bf16 head 1.74x), k=3 2.17 = 2.01x (bf16 head 1.86x, so k=3 needs the fp8 head). long-1k k=2: 1.59x at c=1, 1.23x at c=4, wall 1.01x at c=64.
  • Greedy correctness PASS with 0 non-tie divergences at k=2 and k=3 on 9B and 0.8B, fp8 and bf16 heads, and on every window I arm. Sampled k=2 against plain, 1600 samples over 8 prompts: p<0.01 in 0 of 32 tests on 9B and 0.8B. Head acceptance within 1 % of bf16 on chat, long and 12-language sets; tokens a step 2.44 / 2.85 (Speculative decoding as an engine resource, with Qwen3.5's MTP head as the draft #1090 2.47 / 2.85).
  • No host sync: the whole step, every draft pass's attention included, runs under set_sync_debug_mode("error"), and the fp8 head at 1, 2 and 4 rows. Three defects found reading the diff (argmax buffers sized by max_requests, a crash on an ungated head, a context past the rope table) are fixed with tests in d1e608d.
  • Tests, from the branch and integration trees, not this cut branch: ruff clean, GPU 1355 passed and 9 skipped (test/rocm plus test_kernel_ops_parity), CPU Core 10692 tests and 34 failing at 6ce8b1e, the base's set.

Not in this PR

FLUX.2-klein, Z-Image-Turbo, LTX-2.5 and Wan2.2 ship with attention_backend: auto,
the engine's ragged attention resource with the platform's default kernel
(FlashInfer FA3 on an H100, AITER's asm flash attention on MI300X/MI325X); sdpa
stays for parity runs and `flashinfer` means what it did. Wan2.2 declares a self-
and a cross-attention resource and its step (one segment per guidance branch), and
its DiT takes the engine's callables instead of calling SDPA itself.

The AITER ragged manager no longer stages every layer's output through one shared
buffer and clones it out: the kernel allocates per call, which under capture is a
graph-pool block per layer, so the per-layer copy (32 x 37 MB a step on klein-9B)
is gone. A ROCm test scores the path at the served DiT shapes against fp64 SDPA.
A bundle is a directory per (GPU product, torch, HIP or CUDA, Triton) key holding one uncompressed tar per cache (Inductor, Triton, MIOpen) and a manifest. Packing appends only new or changed files under a lock; unpacking extracts the bundle for this stack into the cache dirs a server here resolves and leaves a marker that each worker reads at startup to log the key and whether the bundle matches. mstar serve --compile-bundle unpacks before starting.
The runner hints dynamo that each static input's bucket-varying dim is dynamic (maybe_mark_dynamic, so a forward that needs a fixed size still compiles per bucket) and compiles under the platform's overrides for any count. Size 1 keeps a graph of its own. Each bucket is still captured separately. Meant for encoders over fixed-size windows, whose per-bucket Inductor compiles are most of a cold start.
MIOpen keeps each database in one file (find and perf DBs as key=value text, kernels in SQLite), so packing caches warmed separately replaced the bundle's records with the last host's. A changed database is now merged with the bundle's last copy, this host's record winning on the same key, and appended only when the merge adds something. MIOpen files are compared by content: two warm-ups can leave databases of the same size and mtime.
amd-flashinfer 0.6.18+amd.1's generated AITER prefill config names __hip_bfloat16 before anything declares it, because the JIT's -DUSE_ROCM=1 turns off the hip_bf16.h include in torch's BFloat16.h. Every bf16 model on FlashInfer's auto backend dies at setup with a ninja log. Give the JIT the header through FLASHINFER_EXTRA_CUDAFLAGS on ROCm; 0.6.18+amd.2 includes it itself.
… pool's workers

The pool is built under a config patch at the top of warmup, since the setting is part of a graph's cache key and an env default would recompile everything a golden cache holds.
A submodule can compile at load, and klein does, so the pool built at the top of warmup came too late: the first compile had already built it under inductor's defaults, quiesce included. The call moves to the top of EngineManager.build, so no model has to ask for it.
…d head, refuse a context past the rope table

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant