Running Weaver (DFlash-TfM) speculative decoding on Blackwell (sm_120) with NVFP4

weaver-nvfp4-blackwell-sm120.md

Running Weaver (DFlash-TfM) speculative decoding on Blackwell (sm_120) with NVFP4

A reproducible recipe for standing up Weaver ("Trees from Marginals", arXiv 2607.06763) speculative decoding on Qwen3.6-27B in NVFP4, using the trymirai/sglang fork, on an NVIDIA Blackwell GPU (compute capability sm_120, e.g. RTX PRO 6000). The paper targets B200; this writeup documents the four changes needed to make it load and run on sm_120, plus the measured tree-budget curve on that hardware.

Status and versions (read before reproducing)

As of 2026-07-22. This recipe patches a fast-moving dev fork at a fixed point in time. Several of the four fixes below are workarounds that will likely become unnecessary as code lands upstream, so before reproducing, check whether each is still needed:

• Has the Weaver fork merged upstream? DFLASH_TFM lived only in trymirai/sglang at this writing. If a stock SGLang release now ships --speculative-algorithm DFLASH_TFM, use it and skip the fork and all four patches.
• Is PR sgl-project/sglang#27906 already in your SGLang? It merged 2026-07-06, so a current release very likely includes it and Fix 1 becomes a no-op.
• Has the lm_head-dequant fix landed? If the DFlash/Weaver samplers already handle a quantized lm_head (in the fork or upstream), skip Fix 2.
• Does fa4 support sm_120 yet? The forced-flashinfer swap (Fix 3) exists only because fa4 asserted on Blackwell. If that is fixed, the paper's fa4/trtllm_mha backends may work and give numbers directly comparable to the paper.
• Newer checkpoint revisions? The pins below are what was tested; a newer revision of any of the three artifacts can change results.

Versions tested:

• Hardware/runtime: NVIDIA Blackwell sm_120 (RTX PRO 6000), CUDA 13.x-class.
• SGLang: pinned image lmsysorg/sglang@sha256:1d8d7976fe11a8341408b92527200502e93dd69df0a63a81c57b92e70ec6fada, with the trymirai/sglang fork (main 981ed18e, 2026-07-09) bind-mounted over it, plus the three source patches described in Fixes 1 to 3.
• Upstream PR applied on top: sgl-project/sglang#27906 (merged 2026-07-06).
• Target: nvidia/Qwen3.6-27B-NVFP4 @ 0893e1606ff3d5f97a441f405d5fc541a6bdf404
• Drafter: z-lab/Qwen3.6-27B-DFlash @ 0919688658996800f86b895034249700e9481106
• Weaver head: trymirai/weaver @ 309ceb4b1a6c44e6a3dfaeab8db1547e904254f8 (weaver/qwen36_27b_weaver.pth)

Headline result (single RTX PRO 6000 Blackwell, Spec-Bench classic-6, greedy, batch 1, 3 restart-samples/point, one pinned vllm bench serve client, decode = 1000/TPOT averaged over the 6 categories):

Tree budget	avg decode tok/s	× no-spec baseline	notes
16	202	~3.9×	above the paper-tuned 64 here
32	224	~4.3×	sm_120 optimum
64 (paper-tuned)	195	~3.7×
128	149	~2.8×	independence penalty dominates

The curve is single-peaked at 32 on sm_120 (the paper's B200/fa4 optimum is 64): it rises from 16 to 32 (202 to 224), then falls 32 → 64 → 128 (224 → 195 → 149) as the verify cost per tree token outgrows the extra acceptance (per-token accept roughly doubles at 32 vs 64). Note that budget 16 already beats the paper-tuned 64 on this hardware. Small budgets can occasionally wedge the server (mass rejects mid-run) on this build, though a clean run at 16 completed all six categories; 32 is the safe recommendation as both the fastest and the most robust point.

What Weaver is

Weaver is a small (~57M-parameter) autoregressive head bolted onto the z-lab DFlash block-diffusion drafter. DFlash predicts the top-K token marginals for several future positions in parallel; because those marginals assume independence, naive acceptance collapses as you draft further ahead. Weaver expands the marginals into a proposal tree and restores conditional dependence between positions, then the tree is verified rollback-free against the target model. Three artifacts must load together: the target model, the DFlash drafter, and the Weaver head checkpoint.

It is not in upstream SGLang; it lives in the trymirai/sglang fork and is selected with --speculative-algorithm DFLASH_TFM.

Prerequisites

• A Blackwell GPU (sm_120). NVFP4 weights + KV fit comfortably (~21 GB); 96 GB is not required.
• Docker with the NVIDIA container toolkit. No SGLang build is needed: the fork's Python is bind-mounted over a pinned stock image via PYTHONPATH (interpreted; the novel kernels are Triton and JIT-compile on first use).
• Pinned SGLang image (public): lmsysorg/sglang@sha256:1d8d7976fe11a8341408b92527200502e93dd69df0a63a81c57b92e70ec6fada.
• Three checkpoints in your Hugging Face cache:
  • Target: nvidia/Qwen3.6-27B-NVFP4 (ModelOpt MIXED_PRECISION: FP8 attention/ linear-attention + W4A16_NVFP4 MLPs + FP8 KV; note the lm_head is quantized, embed_tokens is not).
  • Drafter: z-lab/Qwen3.6-27B-DFlash (served unquantized).
  • Weaver head: trymirai/weaver (weaver/qwen36_27b_weaver.pth, ~227 MB). Note the unsloth/Qwen3.6-27B-NVFP4 compressed-tensors checkpoint does not load on this fork ("No compressed-tensors compatible scheme was found"); use the nvidia ModelOpt checkpoint.

The four sm_120 fixes

1. NVFP4 mixed-precision loading: apply upstream PR sgl-project/sglang#27906

The fork branched shortly after this PR merged and only partly contains it; without it the ModelOpt checkpoint will not load. Apply the PR to the fork checkout:

curl -sL https://github.com/sgl-project/sglang/pull/27906.diff | git apply --reject -

Most files apply clean. The one hunk that typically rejects is in the model-config loader, which must replace a model-family whitelist for MIXED_PRECISION with a content-based check so any ModelOpt-NVFP4 checkpoint is routed to the mixed loader:

if quant_algo == "MIXED_PRECISION":
    quantized_layers = json_quant_configs.get("quantized_layers") or {}
    has_modelopt_nvfp4_layers = any(
        str(layer_info.get("quant_algo", "")).upper() in ("NVFP4", "W4A16_NVFP4")
        for layer_info in quantized_layers.values()
        if isinstance(layer_info, dict)
    )
    if has_modelopt_nvfp4_layers:
        return {"quant_method": "modelopt_mixed", "quant_algo": quant_algo}
    return {"quant_method": "w4afp8", "quant_algo": quant_algo}

Without this, the pre-PR auto-detect picks w4afp8, which force-FP8s the unquantized 48-wide Gated-Delta-Net in_proj_ba and dies on a block_n=128 divisibility assert.

2. Dequantize the quantized lm_head (DFlash-family samplers)

The DFlash and Weaver samplers do a raw torch.matmul against the target lm_head. The ModelOpt checkpoint packs lm_head (W4A16), so the matmul crashes with a shape mismatch (mat1 and mat2 shapes cannot be multiplied (Nx5120 and 496640x320)).

Fix: dequantize the head once through the layer's own quantized GEMM (quant_method.apply(identity) == Wᵀ, exact for W4A16 since activations stay bf16), cache the resulting bf16 [vocab, hidden] tensor (~2.5 GB), and use it at every matmul site. Guard with the fork's should_apply_lm_head_quant_method helper so unquantized heads are untouched. Three call sites need it:

• speculative/dflash_tfm.py → _topk_from_lm_head and _weaver_residual_lm_head
• speculative/dflash_worker_v2.py → _greedy_sample_from_vocab_parallel_head (the eager path; the fused CUDA-graph sampler already falls back to eager for quantized heads)

3. Attention backends: flashinfer on all paths

The paper's B200 config uses trtllm_mha (decode) and fa4 (draft). fa4 hard-asserts on sm_120 ("Paged KV not supported on SM 12.0"). Switch all three attention backends to flashinfer, which has an explicit sm120a code path:

--decode-attention-backend flashinfer
--prefill-attention-backend flashinfer
--speculative-draft-attention-backend flashinfer

Gated-Delta-Net verification runs on the Triton GDN kernel automatically. Because the attention backend differs from the paper's, absolute tok/s here is not directly comparable to the published B200 numbers.

4. Do not force --quantization

Post-PR, the checkpoint auto-detects to the mixed loader. Passing --quantization modelopt_fp4 is rejected, and the pre-PR path is what dies on the GDN assert above. Leave quantization to auto-detection.

Serve command

Paths below are placeholders: HF_CACHE = your Hugging Face cache directory, WORKTREE = the patched fork checkout, and the Weaver .pth is referenced through a mounted artifacts directory. The fork is bind-mounted over the pinned image.

docker run -d --name weaver --gpus all --ipc=host --shm-size 32g --network host \
  -v "$HF_CACHE":"$HF_CACHE" -e HF_HUB_CACHE="$HF_CACHE" -e HF_HUB_OFFLINE=1 \
  -v "$HF_CACHE/.compilecache/triton":/root/.triton/cache \
  -v "$HF_CACHE/.compilecache/flashinfer":/root/.cache/flashinfer -e FLASHINFER_WORKSPACE_BASE=/root \
  -v "$WORKTREE":/sgl-workspace/sglang -v "$WEAVER_ARTIFACTS":/artifacts \
  -w /sgl-workspace/sglang -e PYTHONPATH=/sgl-workspace/sglang/python \
  lmsysorg/sglang@sha256:1d8d7976fe11a8341408b92527200502e93dd69df0a63a81c57b92e70ec6fada \
  python3 -m sglang.launch_server \
    --model-path nvidia/Qwen3.6-27B-NVFP4 \
    --served-model-name Qwen3.6-27B-NVFP4 \
    --dtype bfloat16 --trust-remote-code \
    --tp-size 1 --max-running-requests 1 --cuda-graph-max-bs 32 \
    --mem-fraction-static 0.75 --page-size 64 --disable-radix-cache \
    --decode-attention-backend flashinfer \
    --prefill-attention-backend flashinfer \
    --speculative-draft-attention-backend flashinfer \
    --speculative-algorithm DFLASH_TFM \
    --speculative-draft-model-path z-lab/Qwen3.6-27B-DFlash \
    --speculative-draft-model-quantization unquant \
    --speculative-dflash-tfm-path /artifacts/weaver/qwen36_27b_weaver.pth \
    --speculative-dflash-tfm-tree-budget 32 \
    --speculative-gdn-verify-kernel chunk \
    --disable-overlap-schedule --host 0.0.0.0 --port 8003

Mounting the Triton and flashinfer compile caches makes restarts fast (the JIT work is one-time). For strict reproducibility, pin the model revisions with --revision / --speculative-draft-model-revision.

Verifying it works

• Uvicorn running on ... alone is not success; a broken build can start the HTTP server and then die on the first forward pass. Always run a generation.
• The first one or two requests report 4–8 tok/s: that is one-time flashinfer/Triton JIT compilation, not a hang.
• Healthy warm signature in the scheduler log: accept len: around 4–8 (content-dependent), gen throughput 90–270 tok/s depending on workload. A low per-token accept rate (~0.1) is expected and not a problem: it is measured per drafted tree node, and most nodes in a large tree exist precisely to be pruned.

Failure modes and their causes

Symptom	Cause	Fix
Paged KV not supported on SM 12.0	fa4 draft backend on sm_120	flashinfer backends (fix 3)
... (w4afp8) does not match ... (modelopt_fp4)	forced --quantization flag	drop the flag (fix 4)
output_partition_size = 48 is not divisible by block_n = 128	pre-PR w4afp8 path force-quantizing the GDN in_proj_ba	apply PR 27906 (fix 1)
Parameter ...mlp...weight_scale not found, then a shape assert	PR 27906 missing, mixed loader can't register NVFP4 MLP params	apply PR 27906 (fix 1)
mat1 and mat2 shapes cannot be multiplied (Nx5120 and 496640x320)	quantized lm_head raw-matmul in a sampler	dequant-once patch (fix 2)
No compressed-tensors compatible scheme was found	using a compressed-tensors NVFP4 checkpoint	use the nvidia ModelOpt checkpoint

Tuning

• --speculative-dflash-tfm-tree-budget is the primary knob (the paper tunes 64; the checkpoint default is 128). Use 32 on sm_120; see the budget curve above. Larger trees cost more to verify than the extra acceptance is worth on this backend.
• --speculative-dflash-tfm-candidate-pool-size sets the top-K pool the tree draws from (checkpoint default 512).
• The underlying DFlash block size is 16 in the Weaver serve mode.

Benchmarking notes

All numbers here use a single fixed client (vllm bench serve, greedy, seed 0, --ignore-eos) so throughput is defined identically across configurations, with the target served behind the OpenAI-compatible endpoint. Three server-restart samples per point capture the ~10%-class run-to-run variance that greedy decoding shows on high-entropy prompts. Decode tok/s = 1000 / mean TPOT.