Training a DFlash drafter on a B200 — the real grind behind block-diffusion speculative decoding

In the last post I said my way of learning a new technique is to build it, and I ended on DFlash — the block-diffusion speculative-decoding drafter — with a promise to bring it up on the latest SGLang. This is the follow-through: a few days ago I actually trained a DFlash drafter for a 72B target on a B200, end to end, on the latest SGLang. This is the grind.

The timing is funny, because right as I was doing it NVIDIA published “Boost inference performance up to 15× on Blackwell using DFlash speculative decoding”. DFlash went from a four-day-old LMSYS release to NVIDIA-blessed on Blackwell in about a week. That’s the pace of this field — and exactly why I’d rather be training the thing than reading about it.

Table of contents

DFlash is now the Blackwell default — with one catch

NVIDIA’s numbers are the headline: up to 15× throughput for gpt-oss-120b on eight DGX B300, ~2× interactivity over EAGLE-3 for Llama 3.1 8B, 5.8× for Gemma 4 31B and 5.1× for Qwen3 8B single-GPU. The mechanism is the one I covered last time: a block-diffusion drafter predicts a whole block of tokens in one forward, conditioned on the target’s hidden states, with those features injected into the draft model’s KV across layers — and the target still verifies, so it’s lossless.

The NVIDIA post says the migration is easy: “migrating from EAGLE to DFlash only requires updating the speculative decoding algorithm” — and a matching checkpoint. That last clause is the entire catch, and it’s the reason this post exists. KV injection ties a drafter to its specific target’s hidden states, so a DFlash drafter is not reusable across targets the way a generic draft model might be. For a 72B target, no matching checkpoint exists. You can’t download it. You have to train it. So I did.

The setup: 2×B200, the latest SGLang, SpecForge

Reason from memory first. A 72B in bf16 is ~145 GB — it won’t fit on one card with an optimizer and KV, so the target must be tensor-parallel sharded across at least two GPUs. Once it fits, the lever for speed is VRAM headroom → bigger batch (a small batch under-utilizes the 72B forward, which dominates every step). A 2×B200 pair (192 GB each) leaves ~120 GB free at tp=2 — enough for batch-4 (lands ~166/192 GB), which puts two cards at roughly an 8×H200’s throughput.

The trainer is SpecForge, which ships a full DFlash path — no recipe to reverse-engineer:

• an OnlineDFlashModel with a block-wise CE loss, a flex-attention block mask, anchor sampling, and loss decay over block positions;
• a qwen3-style DFlashDraftModel (with q/k-norm);
• and sglang / HF target wrappers to provide the frozen target’s hidden states.

The target runs on the SGLang tp=2 backend with enable_return_hidden_states (the HF backend tries to load the 72B onto one device and OOMs); hidden states are captured at 5 auto-selected layers, ~[1, 20, 39, 58, 77] of the 80. My drafter config: 5 layers, hidden 8192, block_size 16, ~2.6 B params (~5.2 GB bf16), trained from scratch (there’s no prior DFlash drafter to warm-start from) at LR 6e-4 + warmup — √4-scaled to 1.2e-3 for the batch-4 run. The data is PG19 passages self-distilled through the 72B — i.e. the target’s own greedy outputs, so the draft learns to imitate this target on this distribution.

One small-but-fatal detail: the mask token (151666) must be a reserved id inside the 152064 vocab. Let the tokenizer add a fresh <|MASK|> and it lands at id 152064 — one past the end of the reused target embedding → out-of-range → crash.

The version-drift tax (latest SGLang)

Here’s the first thing that ate a day. SpecForge’s SGLang target backend was written against an older SGLang; I needed the latest serving version — the one with Blackwell sm_100 kernels and the DFlash / Spec V2 path. That gap was seven sequential breakages, each surfacing only after I fixed the last:

1. a renamed module (a scheduler mixin that became a free function);
2. ServerArgs rejecting a kwarg the backend still emitted → filter kwargs against the live constructor signature;
3. a removed kwarg in init_model_parallel_group;
4. a world-info helper that grew from a 3-tuple to a 4-tuple return;
5. the request object’s fill_ids replaced by full_untruncated_fill_ids + fill_len, with a new assertion;
6. ForwardBatch.init_new taking the batch directly (an indirection deleted);
7. the new overlap scheduler deferring input_ids materialization — so I had to materialize it by hand after prepare_for_extend.

This is the tax of pointing a research trainer at a fast-moving inference engine. It’s terrain, not misfortune — budget for it, and capture all seven as a single patch so you only walk it once. (B200 prereq, while we’re here: the SGLang image needs Blackwell sm_100 kernels, and flex_attention on Blackwell needs torch ≥ 2.6, else fall back to --attention-backend sdpa.)

The traps that ate GPU-hours

Then the training traps — none of which throw an error, all of which waste a run:

• The max-length truncation trap. My data tokenizes to ~5.9k tokens (median seq 5860, max 5920). DFlash’s supervision lives at the end of each sequence, so any max_length below 6656 truncates the answer, empties the loss mask, and the dataset filter silently drops 100% of the data — training “runs” and learns nothing. I verified token lengths before trusting a single run; 6656 is mandatory, never a memory dial.
• OOM-hangs, not crashes. One rank OOMs and the others spin on the NCCL collective: 100% util, low power, no new logs. Since max-length can’t drop, the levers are --num-anchors (256 → 128 halves the draft forward), --sglang-mem-fraction-static (must stay above the ~83 GB/GPU model+KV footprint), and batch. DFlash helps here — it only trains num_anchors blocks, not the whole 6656-token sequence, so the drafter side is light and the sharded 72B dominates.
• The annealed-LR warm-start trap. 1.49 was undertrained (below), so I continued — and a naïve --resume would have learned nothing: the first run’s cosine LR had annealed to ≈ 0, and resuming reloads that scheduler. The fix is to load the weights but build a fresh warmup+cosine schedule, and verify the warm-start by the initial loss — a correct resume picks up at ~5.2 (the trained level), a broken one at ~11.9 (random init). Then don’t panic at the dip: re-warming to full LR deliberately kicks the weights out of their old optimum, so accuracy drops before it breaks past the previous peak.
• You can’t measure accept while training. The acceptance test needs the 72B target on a free GPU — but tp=2 training occupies both B200s. So measuring is stop → test → resume (and the resume re-warms). I monitored training accuracy as a non-disruptive proxy and ran the real accept test only after stopping.

The result: accept 1.49 → 1.71

Two epochs, ~2 h on 2×B200: a working drafter at accept 1.49 — real, but undertrained (the loss and accuracy were still climbing at the cutoff). A warm-start continue (with the fresh-LR fix) took it to 1.71, a genuine +14%, measured the same way:

2-epoch drafter:     1.654 / 1.391 / 1.461  →  mean 1.494
warm-start continue: 1.985 / 1.561 / 1.625  →  mean 1.705

And the tell that says where to go next: training accuracy plateaued at ~0.18–0.20 while the loss kept making small new lows. That pattern means I hit the ceiling of the training set, not the ceiling of the epoch count. More hours won’t move it; more and better-matched data will.

I’ll be honest about where 1.71 sits. It’s a real step up, but it’s still short of a mature EAGLE3 field head (~2–2.5), short of z-lab’s mature DFlash (~3–4), and a long way from NVIDIA’s Blackwell numbers. DFlash’s cheaper draft helps, but it doesn’t close an accept-length gap on its own. The draft being cheap matters less than the draft being right — and “right” is a data problem.

What I learned

1. Data is the ceiling. The EAGLE3 task-mismatch and the DFlash data-plateau are the same lesson twice: a drafter only learns to predict the target on the distribution you trained it on. Self-distill on the real one.
2. Version-drift is terrain. A research trainer against a moving engine is a chain of breakages — seven here. Expect it; capture the fixes as one patch.
3. Verify warm-starts by the initial loss, not by faith. ~5.2 means the weights loaded; ~11.9 means you’re training from scratch and don’t know it.
4. Acceptance length is the only metric that maps to speed, and your first checkpoints (step 250/500/1000) are the cheapest signal you’ll get — a dead run shows up in 15 minutes if you’re watching.
5. The pipeline is the durable asset. I now have a repeatable DFlash trainer for a 72B target on the latest SGLang. From here the drafter’s quality is a data problem — a far better place to stand than where I started.

Next steps

• Data-expansion retrain. Grow the self-distilled set well past 6k examples, keep it matched to the serving distribution, and push accept toward 2.5+. The pipeline is proven; this is now the main lever.
• Serve it on Spec V2 and measure in the real regime. Drop the checkpoint in and serve with --speculative-algorithm DFLASH --speculative-draft-model-path /draft --speculative-dflash-block-size 16 on the latest SGLang (the overlap scheduler is the default), then measure end-to-end τ and tok/s at the deployment TP/GPU/quantization — not a cheaper proxy that will lie to me.
• Chase the ceiling. NVIDIA’s Blackwell post shows how high it goes (up to 15× on the right model); my gap to it is now a data-and-iteration problem on a working pipeline, not an architecture mystery.

NVIDIA wrote the blog about DFlash on Blackwell. I trained the drafter on the B200. That’s the difference this whole blog and roadmap is built around — reading tells you the idea; building tells you the truth, and the truth this time was that the architecture was the easy part and the data is the work.