Train a DFlash Drafter for Speculative Decoding

A guide for training a DFlash speculative-decoding drafter to accelerate LLM inference with NeMo AutoModel.

What is DFlash?

DFlash is a block-parallel drafter. Where an autoregressive drafter proposes one token at a time, DFlash proposes an entire block of block_size tokens in a single non-causal forward pass: the block’s first position holds the real anchor token and the rest are MASK, and the draft “denoises” the whole block at once conditioned on hidden states captured from the frozen target at target_layer_ids. The draft shares and freezes the target’s embed_tokens and lm_head, training only the draft stack, its feature projection, and its norms.

It follows the same scaffolding as the EAGLE and DSpark recipes: online target hidden-state capture, gradient accumulation, and consolidated-safetensors checkpointing.

Variants

All four share the DFlash backbone, anchor sampling, and block attention mask; only the head and the objective differ.

Variant	What it adds	Recipe	Example config
DFlash	the block-parallel baseline	`TrainDFlashRecipe`	`qwen3_dflash.yaml`
DFlash 2	a two-tap in-block convolution around every sublayer, plus a pairwise path selector that walks one coherent path through each position’s top-k candidates	`TrainDFlash2Recipe`	`qwen3_dflash2.yaml`
Domino	a serial GRU correction head that refines each block position on the previous ones	`TrainDominoRecipe`	`qwen3_domino.yaml`
JetSpec	causal in-block attention plus forward-KL distillation against the target distribution	`TrainJetSpecRecipe`	`jetspec/`

DFlash 2 keeps drafting in one pass — the convolution and the selector add parameters inside the stack rather than an autoregressive correction — so it has its own draft class whose checkpoint layout matches the published DFlash 2 drafters.

Objective

Only block positions 1..block_size-1 are supervised (position 0 is the clean anchor). Their cross-entropy is weighted by a position decay w_k = exp(-(k-1)/loss_decay_gamma), so early positions — the ones a verifier reaches first — dominate the loss.

Paper loss_decay_gamma values track the block size: 7 for block_size 16, 5 for 10, and 4 for 8.

Set loss_type: variable_prefix to train the D2SD VP-Drafter objective (arXiv:2606.04446) instead: each block draws a visible-prefix length from a truncated geometric prior (prefix_weight_base, default 0.9), only the masked suffix is supervised, and the decay restarts at the prefix boundary. This matches the regime a drafter sees when it re-drafts behind a partially accepted block.

DFlash 2 adds one term on top of the block CE, weighted by selector_loss_weight: a cross-entropy over the position’s top-k candidates, scored against the ground-truth predecessor. Positions whose true token missed the candidate list carry no selector signal and are reported as candidate_recall.

Data

Use a chat dataset of OpenAI-format messages rows. As with the other speculative recipes, use Open-PerfectBlend prompts with the responses regenerated by the target model (training is teacher-forced; regenerate before training to avoid a train/inference distribution mismatch). Point recipe_args.train_data_path at the regenerated JSONL or a Hugging Face dataset id.

mask_token_id has no default and must name a reserved, rarely-used token. Do not reuse pad: it is frequently aliased to eos, which conflates the mask signal with real content and quietly erodes acceptance. The inference runtime must fill block slots with the same id.

Run it

Example configs live under examples/speculative/dflash/.

$ torchrun --standalone --nproc_per_node=2 \
>   -m nemo_automodel.recipes.llm.train_dflash \
>   -c examples/speculative/dflash/qwen3_dflash.yaml

Swap the module and config to train a variant — for example -m nemo_automodel.recipes.llm.train_dflash2 -c examples/speculative/dflash/qwen3_dflash2.yaml.

Multi-GPU wraps the draft in DDP and replicates the frozen target. Set distributed.tp_size to shard a large target across ranks (qwen3_dflash_tp.yaml); distributed.cp_size enables context parallelism, which cannot be combined with sequence packing or with tp_size > 1.

The published drafters diverge from their target in two ways the recipe reproduces: a sliding window over the context (draft_sliding_window, which also keeps the flex BlockMask sparse on long sequences) and a different attention shape (draft_num_attention_heads and friends). See qwen3_8_27b_dflash2.yaml.

Training logs train/loss, train/accuracy, and train/accept_len. Setting recipe_args.val_data_path adds globally reduced val_loss, val_accuracy, and val_accept_len; adding a top-level wandb: block uploads all of them.

Key config fields

Field	Meaning
`target_model_name_or_path`	frozen target (e.g. `Qwen/Qwen3-8B`)
`draft_num_hidden_layers`	draft stack depth (paper: 5)
`block_size`	tokens drafted in parallel per block (paper: 16)
`num_anchors`	blocks sampled per sequence per step
`loss_decay_gamma`	block-position decay; `null` disables it
`loss_type`	`dflash` (fixed anchor) or `variable_prefix`
`target_layer_ids`	target feature layers fed to the draft (defaults to an even spread)
`mask_token_id`	reserved token id filling non-anchor block positions (required)
`attention_backend`	`flex_attention` (GPU main path) or `sdpa` (portable fallback)
`packed_sequence_size`	enables sequence packing when greater than 0
`draft_sliding_window`	bounds how far back a block reads the context; unset attends over the whole prefix
`draft_num_attention_heads` / `draft_num_key_value_heads` / `draft_head_dim`	size the draft’s attention independently of the target’s
`conv_kernel_size` / `conv_group_size`	DFlash 2 convolution taps and channels per dynamic correction
`selector_rank` / `selector_top_k` / `selector_loss_weight`	DFlash 2 path-selector width, candidates scored per position, and objective weight

Supported targets: Qwen3 and Qwen3.5 (dense and MoE), including the multimodal *ForConditionalGeneration variants such as Qwen/Qwen3.8-27B. DFlash and JetSpec drafts serve on vLLM through its dflash method.