feat: add DFlash2 drafter training backend

[khazic]

Closes #61.

Summary

Adds DFlash2 drafter training support. DFlash2 (writeup, checkpoint z-lab/Qwen3.8-27B-DFlash2, code z-lab/dflash) keeps DFlash's block-diffusion drafting and target-context backbone and adds two modules on top, so this is implemented as a DFlash variant exactly the way Domino (#16, #17) and DSpark extend the same backbone: DFlash2DraftModel subclasses DFlashDraftModel and DFlash2TrainerBackend subclasses DFlashTrainerBackend, reusing the block-drafter plumbing, anchor sampling, packing and FSDP2 wrapping rather than duplicating them.

Defaults match the released checkpoint: block_size=8, two-tap conv with conv_group_size=16, selector_rank=256, selector_top_k=16.

What DFlash2 adds

Two-tap dynamic depthwise convolution wrapped around every attention and MLP sublayer, targeting the accuracy decay DFlash shows toward the end of a block:

Conv_k(x)_t = k_{t,0} * x_t + k_{t,1} * x_{t-1}

with a static per-channel kernel plus a content-adaptive correction shared by every conv_group_size channels. prepare runs before the sublayer and also emits the kernel finish applies after it, so one projection feeds both taps.

Candidate selector that keeps the drafter's top-k candidates at each block position and traces one coherent path through them:

S_t(a, b) = U_t(b) + <A(a) * H(h_t), B(b)>

U_t(b) is the drafter's own logit for candidate b, A/B are the predecessor/successor codebooks and H projects the backbone hidden state.

Two deltas over the upstream inference implementation

Both are load-bearing for training and are the only places this diverges from z-lab/dflash:

1. The convolution is applied block-locally. It is causal along the sequence axis. Upstream only ever drafts a single block, but training packs n_blocks blocks into one flat draft sequence, so a naive port would let tap 1 at a block's first position read the last position of the previous, unrelated block. GroupedDynamicCausalConv reshapes to [bsz * n_blocks, block_size, hidden] first and raises if the draft length is not a block multiple. Covered by test_conv_does_not_leak_across_block_boundaries.

2. The convolution starts as an identity passthrough. Upstream allocates base_kernel with torch.empty and relies on from_pretrained to fill it, which is fine for inference from a released checkpoint but produces uninitialized memory when training from scratch. Here tap 0 is initialized to 1, later taps to 0, and the kernel projection to zero, so a freshly built DFlash2 is numerically identical to DFlash and learns the correction from there. Covered by test_conv_is_identity_at_init.

Selector training objective

At inference the selector walks the block sequentially, feeding each choice in as the next predecessor. Training teacher-forces the predecessor to the ground-truth previous token, so all positions score in one shot. The loss is a cross-entropy over the drafter's own top-k, restricted to rows whose ground truth survived that top-k (elsewhere there is no correct choice to make), added to the DFlash CE with weight dflash2_selector_loss_weight.

test_selector_pair_scores_match_the_sequential_selector pins the two paths together: replaying the sequential trace's own predecessors through the vectorized scorer reproduces the same path.

Because the selector scores real vocabulary ids, it requires loss_mode=full_vocab and fails loud otherwise; a restricted or sampled vocabulary would make the candidate ids meaningless.

The selector's inputs are detached. It is an auxiliary head that re-ranks what the drafter already produced, so its objective must not reshape the drafter itself. torch.topk passes gradient through the selected logits and hidden_projection would otherwise push gradient back through the backbone, which would both change the drafter and make the ablation below meaningless. test_selector_loss_does_not_backprop_into_the_backbone pins this.

Acceptance length metric

The block drafters reported per-position accuracies, but those are marginals: each position is counted independently. A verifier stops at the first mismatch, so what governs the achievable speedup is the length of the correct prefix, and the marginals cannot be combined into it.

_block_acceptance_counts() is added to the shared block-drafter forward and called from the DFlash, Domino and DSpark forwards, so every block drafter emits it (Domino and DSpark carry their own copies of that forward, so they needed the call explicitly; the base trainer's aggregation was already prefix-generic for them). It is reported as mean_acceptance_length under the same convention as the existing rollout-side drafter/spec_decode/mean_acceptance_length, whose leading 1.0 is the token the target emits itself at each verification step, so the training metric predicts the served one rather than sitting a token below it on the same dashboard. The raw sum and block count are published alongside it so the ratio can be inspected.

The helper takes the drafted positions only and leaves the slicing to each caller, because the layouts genuinely differ: DFlash builds labels at arange(block_size) and masks column 0 off as the anchor, while Domino and DSpark build them at arange(block_size) + 1, so their column 0 is a real prediction. Slicing inside the helper would drop that token and, worse, stop a wrong first token from truncating the block, which is the exact failure the metric exists to catch. test_acceptance_counts_do_not_assume_an_anchor_column pins this.

DFlash additionally intersects the scoring mask with the reweighted mask, because correctness is only recorded where the reweighted mask is positive. Without that, a position zeroed by loss_decay_gamma or front_position_weight would read as a mismatch rather than as unscored and would truncate every block's prefix, pinning the metric at 1.0 while the per-position accuracies still looked healthy.

Two limits worth stating. The metric is computed from the drafter's own argmax, so for DFlash2 it does not see the candidate selector that would re-rank tokens at serve time. And a block truncated by a supervision boundary counts fully in the denominator while its numerator is capped by the supervised length, which biases the training-side number below the rollout-side one in a data-dependent way. It is a training-time signal for comparing drafters, not a prediction of served speedup.

Note on the shared DFlash file

Rather than copying DFlashTrainingModel.forward wholesale (as Domino does, ~200 lines), DFlashTrainingModel gains a documented _auxiliary_loss extension point that returns None for plain DFlash. The two conv hooks added to DFlashDecoderLayer are likewise None for every non-DFlash2 drafter, mirroring how upstream's own layer carries them. Both changes leave DFlash / Domino / DSpark / JetSpec behaviour and numerics untouched, which the regression run below confirms.

Review follow-ups

Three issues were raised in review. All three are real, and all three are confirmed against the released z-lab/Qwen3.8-27B-DFlash2 checkpoint rather than argued from the code alone.

The RoPE base was read only from rope_theta. transformers 5 moved it into a rope_parameters dict, and the released config carries only the nested spelling, so DFlashConfig fell back to its 10000.0 default where the checkpoint was trained at 1e7. Nothing failed, because modeling_dflash reads config.rope_theta and cannot tell a defaulted base from a real one. The cold-start path had the same gap on the other side, reading the base off the target config, so a target that nests it would have produced the same silent mismatch with no checkpoint involved. resolve_rope_theta now accepts both spellings, a top-level value still wins, and it is used by DFlashConfig, modeling_dflash and all four DFlash-family fallback configs, so a DFlash baseline and a DFlash2 arm cannot end up on different bases.

The resolver reads the base and discards the rest of rope_parameters, and DFlashRotaryEmbedding takes only a base, so a target using yarn / linear / llama3 scaling still gets a draft whose rotary phase diverges past the original context length. That predates this branch and applying the scaling is a larger change than this PR should carry, so the resolver warns once per distinct rope_type rather than discarding it silently.

Two selector weights were dropped on load. Upstream stores the codebooks as bare nn.Parameter tensors (candidate_selector.predecessor_codebook), while this overlay holds them in nn.Embedding modules, whose state dict spells the same tensor ...codebook.weight. Of the checkpoint's 81 tensors those two were the only mismatch, and nothing said so: the loader drops unrecognized keys with a debug log, and its required-key gate covers the DFlash backbone only. Both codebooks stayed at their random init while training carried on.

DFlashTrainerBackend gains a _CHECKPOINT_KEY_ALIASES table applied by the shared normalizer, DFlash2 fills in those two entries, and a new _validate_normalized_state hook rejects a checkpoint carrying only part of the DFlash2 modules under recognized names, so a future upstream rename cannot degrade into a silent cold start. A checkpoint carrying none of them is still accepted, since warm-starting DFlash2 from a plain DFlash backbone is a legitimate flow and the DFlash2 modules cold-start as an identity passthrough by design. The dropped-key report is promoted from debug to warning when keys were unexpected or shape-mismatched, which gives DFlash / DSpark / Domino the same visibility.

Checked against the real checkpoint's key list: 23 of 23 DFlash2 module keys match after normalization with none stray, and without the rename exactly those two codebooks are dropped.

The two engines disagreed about DFLASH2. SGLang failed loud with guidance while vLLM fell through to a bare "Unsupported speculative_algorithm". Worse, the guidance both should give (serve a trained DFlash2 checkpoint as DFLASH) was unfollowable on vLLM, whose DFlash drafter validator accepted only architectures=['DFlashDraftModel'] and so rejected DFlash2DraftModel. vLLM now carries the same DFLASH2 branch and message, accepts the DFlash2 architecture on the DFlash serve path, and _is_dflash_config classifies it so a serve-only run resolves the DFlash aux layout. DFLASH2 is dropped from the SGLang aux-hidden set: that path needs the same drafter.enable the ServerArgs override rejects DFLASH2 under, so it was unreachable.

DominoDraftModel is deliberately left out of the servable set. DFlash2's serve-as-DFlash route is the documented upstream recipe for the released checkpoint, while Domino's GRU correction head has no established DFlash runtime, so failing early stays the safer answer there.

The ablation below is unaffected by the RoPE fix. All three arms build their config through the same _build_fallback_config, so they resolved the same base before and after; the fix moves absolute numbers only for a target whose config nests the base under rope_parameters.

Testing

Contract tests (tests/integration/test_dflash2_backend_contract.py): the conv and selector modules, the block-locality invariant, identity-at-init, selector vectorized-vs-sequential agreement, the selector's contribution to the total loss, the detach boundary, the full_vocab guard, factory/aux-layer registration, config routing including the nested dflash_config block upstream checkpoints use, and the two acceptance-length semantics tests. The review follow-ups add the RoPE-base resolution on both the checkpoint and the cold-start path, the upstream codebook rename against the released key spelling, the partial-module guard and the plain-DFlash warm start it must not reject, and the DFLASH2 routing on both engines. tests/integration/test_vllm_runtime_contract.py gains the DFlash2 case for the DFlash drafter validator. A scaled rope_type is pinned to warn rather than pass silently.

Regression: the full tests/integration/ suite on this branch versus the merge-base.

merge-base (333b754):  9 failed, 201 passed
branch:                9 failed, 236 passed

The same 9 failures on both, name for name, all pre-existing and unrelated to this change (test_drafter_runtime_control_contract, test_dspark_trainer_backend::test_dspark_checkpoint_preserves_source_config_and_vllm_weight_names, test_verl_npu_vllm_compat). This run matters more than usual because the acceptance-length call was added to the Domino and DSpark forwards as well, and neither picked up a new failure.

Real-data ablation

Trained on a real multi-turn conversation dataset whose assistant turns were regenerated with the target model itself, so the drafter is fitting the distribution it will actually have to predict. 768 training conversations, a disjoint held-out set, assistant tokens only, sequences truncated to 384 tokens, target-side context hidden states pre-computed once and shared byte-for-byte by every arm. Every arm is cold-started, uses the same seed, and runs the same 4000 optimizer steps; the only variables are the algorithm and the selector loss weight.

arm	conv	selector module	selector trained	params
DFLASH	no	no	n/a	276,840,704
DFLASH2	yes	yes	yes (weight 1.0)	364,101,888
DFLASH2:0	yes	yes	no (weight 0)	364,101,888

Held-out results at step 4000:

arm	next-token acc	acceptance length
DFLASH	0.2089	0.8344
DFLASH2	0.2373	1.0000
DFLASH2:0	0.2423	1.0425

contrast	isolates	acc	acceptance length
DFLASH2:0 vs DFLASH	the convolution	+0.0334	+0.2081 (+24.9%)
DFLASH2 vs DFLASH2:0	the selector objective	-0.0050	-0.0425 (-4.1%)
DFLASH2 vs DFLASH	both together	+0.0284	+0.1656 (+19.8%)

Per-position accuracy, block positions 1 to 7:

position	1	2	3	4	5	6	7
DFLASH	0.4688	0.3156	0.2179	0.1562	0.1235	0.0978	0.0826
DFLASH2:0	0.5030	0.3622	0.2712	0.1905	0.1508	0.1192	0.0995
relative	+7.3%	+14.8%	+24.5%	+22.0%	+22.1%	+21.9%	+20.5%

The convolution accounts for the gain, and its shape matches what it is designed to fix: the improvement grows from +7.3% at position 1 to a 20 to 25% plateau from position 3 onward, which is the end-of-block decay the writeup targets.

The selector objective does not contribute positively here. The conv-only arm is ahead of full DFlash2 on both metrics at every evaluation point, and the selector's held-out lift over the drafter's own unary ranking ends at -0.019.

Two things that reading should not be stretched into:

• These metrics are computed from the drafter's own argmax and never route through the selector's re-ranking, so they cannot measure the selector doing its intended job. That job is inference-time: picking a coherent path through the top-k candidates, which needs the tree drafting this repo does not have for any DFlash-family drafter. What the numbers do show is the selector's side effect on the drafter, which should ideally be zero.
• It is not exactly zero even with the detached inputs, because clip_grad_norm_ computes one global norm over all parameters. The selector's gradients raise that norm, shrink the clipping coefficient, and so shrink the backbone's effective step. Single-sample steps make this worse. That is a plausible mechanism for the -0.5pp, and it is a property of the training setup rather than of the selector design.

Full training log, all three arms

train_loss is a single-step instantaneous value at the evaluation point, not a running mean, and it is not comparable across arms: the DFlash2 arm's total loss carries an extra selector cross-entropy term that neither other arm has. Only the held-out numbers are comparable.

=== DFLASH (baseline) ===
[ab] train_samples=768 heldout_samples=128
[ab] DFLASH trainable_params=276,840,704
[ab] DFLASH step 1000  train_loss=5.2159  heldout_acc=0.1157  accepted_len=0.3701
[ab] DFLASH step 2000  train_loss=6.3636  heldout_acc=0.1873  accepted_len=0.7176
[ab] DFLASH step 3000  train_loss=3.1573  heldout_acc=0.2059  accepted_len=0.8309
[ab] DFLASH step 4000  train_loss=5.2912  heldout_acc=0.2089  accepted_len=0.8344
[ab] ===== HELD-OUT SUMMARY =====
[ab] DFLASH         overall=0.2089  accepted_len=0.8344  per_position=[0.4688, 0.3156, 0.2179, 0.1562, 0.1235, 0.0978, 0.0826]

=== DFLASH2 (convolution + selector trained) ===
[ab] train_samples=768 heldout_samples=128
[ab] DFLASH2 trainable_params=364,101,888
[ab] DFLASH2 step 1000  train_loss=8.9757  heldout_acc=0.1604  accepted_len=0.5875  selector_acc=0.3124  unary_only=0.3093  lift=+0.0031
[ab] DFLASH2 step 2000  train_loss=10.7086  heldout_acc=0.2094  accepted_len=0.8368  selector_acc=0.3726  unary_only=0.3671  lift=+0.0056
[ab] DFLASH2 step 3000  train_loss=6.1724  heldout_acc=0.2298  accepted_len=0.9499  selector_acc=0.3656  unary_only=0.3861  lift=-0.0206
[ab] DFLASH2 step 4000  train_loss=12.7286  heldout_acc=0.2373  accepted_len=1.0000  selector_acc=0.3765  unary_only=0.3953  lift=-0.0188
[ab] ===== HELD-OUT SUMMARY =====
[ab] DFLASH2        overall=0.2373  accepted_len=1.0000  per_position=[0.492, 0.3493, 0.262, 0.1972, 0.1453, 0.1176, 0.0975]

=== DFLASH2 with selector_loss_weight=0 (convolution only) ===
[ab] train_samples=768 heldout_samples=128
[ab] DFLASH2:0 trainable_params=364,101,888
[ab] DFLASH2:0 step 1000  train_loss=4.2779  heldout_acc=0.1692  accepted_len=0.6324
[ab] DFLASH2:0 step 2000  train_loss=5.7315  heldout_acc=0.2144  accepted_len=0.8928
[ab] DFLASH2:0 step 3000  train_loss=2.8353  heldout_acc=0.2374  accepted_len=0.9988
[ab] DFLASH2:0 step 4000  train_loss=5.2135  heldout_acc=0.2423  accepted_len=1.0425
[ab] ===== HELD-OUT SUMMARY =====
[ab] DFLASH2:0      overall=0.2423  accepted_len=1.0425  per_position=[0.503, 0.3622, 0.2712, 0.1905, 0.1508, 0.1192, 0.0995]

What this ablation does not establish

• No error bars. One run per arm, no seed repeats. The +3.3pp convolution gain is consistent across all four evaluation points and across both metrics, but the -0.5pp selector effect is small enough that a single run cannot separate it from noise.
• Unequal parameter counts in the conv contrast. DFLASH lacks both modules, so DFLASH2:0 vs DFLASH mixes the convolution's design with 87M extra parameters. Only DFLASH2 vs DFLASH2:0 is a clean contrast, and there the counts are identical.
• Small scale. 768 conversations, 384 tokens, batch size 1, roughly 5 epochs. That is on the order of 0.1% of a real drafter training run, and single-sample steps make the gradients noisy.
• Acceptance length here is an offline proxy. It treats the dataset's ground-truth continuation as what the target would emit. Real speculative decoding accepts against the target's own distribution under rejection sampling, so this is correlated with, not equal to, served acceptance.
• Cold start, not a continuation of the released checkpoint. There is no matching DFlash checkpoint to start the baseline from, and the released DFlash2 checkpoint is built for a different target size, so its hidden dimensions do not fit. Absolute numbers are therefore low and should only be read across arms, never as DFlash2's achievable draft quality.

GPU training smoke

tests/special_standalone/dflash2_gpu_smoke.py drives the real training path end to end on a real target with a cold-started draft. It replays one fixed batch, so its terminal acc 1.00 is memorization and carries no generalization signal; it exists to prove the path trains and that gradients reach both the convolutions and the selector, and the ablation above is what speaks to draft quality.

The smoke emulates the production dtype setup deliberately: FSDP MixedPrecision(param_dtype=bf16) gives the forward bf16 parameters, and the surrounding torch.amp.autocast keeps cross_entropy in fp32. Emulating only one of the two fails, in a different place each time.

Out of scope

Inference-side tree drafting and serving, same as for DFlash / Domino / JetSpec in this repo. This PR is training-side only, and without tree drafting the selector cannot be evaluated at the job it was designed for.

Online co-training is out of scope for the same reason Domino's is: DFLASH2 is not an engine-level algorithm, so both engines reject the string and point at DFLASH. Train the drafter offline, then serve the checkpoint as DFLASH to use it as a frozen rollout drafter.

[tpx818]

官方脚本from pretrain时有个key mapping，这里是否也需要处理下

[khazic]

ok

[khazic]

Confirmed. The released checkpoint stores both selector codebooks as bare Parameters (candidate_selector.predecessor_codebook); this overlay holds them in nn.Embedding, so they arrived under names the model does not have and were dropped with only a debug log. Of its 81 tensors those two were the only mismatch. Fixed in 3492314: a _CHECKPOINT_KEY_ALIASES table on the base backend, plus a _validate_normalized_state hook that rejects a partial DFlash2 module set. Against the real key list, 23 of 23 module keys now match with none stray.

[tpx818]

dflash config的默认rope_theta与dflash2是否一致，dflash2好像默认是读取rope_parameters

[khazic]

You are right. The released config has no top-level rope_theta, only rope_parameters.rope_theta = 1e7, so DFlashConfig fell back to its 10000.0 default. The cold-start path had the same gap on the target side. Fixed in dcf452e: resolve_rope_theta takes both spellings and is used by DFlashConfig, modeling_dflash and all four DFlash-family fallback configs. A scaled rope_type now warns instead of being discarded silently (27c680e).

[tpx818]

co-train模式下，sgl与vllm是否支持dflash2类别？

[khazic]

Neither serves DFLASH2 as an engine method: it is a DFlash variant, like Domino. SGLang already failed loud, vLLM fell through to a generic "unsupported" string, and the advice both give (serve the trained checkpoint as DFLASH) was blocked on vLLM by an architectures allowlist. Fixed in 2ee2c99. DFlash2 is training-side only here, so online co-train is out of scope; the PR description now says so.