feat(speculative): support sequence packing in the DFlash draft

[khazic]

What

Extends sequence packing to the DFlash draft. The DFlash recipe builds its dataloader with build_eagle3_dataloader but never passed packed_sequence_size, so packing was unreachable. DFlash's Q layout is draft blocks only (the context is KV, never queried), and each block's only real content is the context prefix kv_idx < anchor, so packing composes cleanly:

• Mask (dflash_mask.py): create_dflash_sdpa_mask and create_dflash_block_mask take optional ctx_doc_id / anchor_doc_id; a block's context prefix is restricted to its anchor's document (both the dense SDPA mask and the flex BlockMask mask_mod, now factored into build_dflash_mask_mod).
• Core (core.py): forward accepts position_ids / seq_lens / doc_remaining. Anchor sampling keeps whole blocks inside one document (doc_remaining >= block_size-1, the per-document analogue of the existing anchor <= seq_len-block_size bound), the context position ids reset per document, and the draft block positions are gathered from the per-document context positions so RoPE stays document-local.
• Target (target.py): generate_batch runs the frozen target with the block-causal mask (or FlashAttention varlen at batch size 1) so the captured context hidden states do not leak across documents, and carries the packing metadata to the trainer.
• Recipe (train_dflash.py): reads packed_sequence_size into both dataloaders, splats the packing metadata into generate_batch and the trainer, fails fast on incompatible configs (cp_size>1, FlashAttention target with micro_batch_size>1), and rejects packing on the JetSpec/Domino subclasses (they override the trainer forward without threading the metadata).

The packed dataset is reused unchanged. A packed document shorter than block_size yields no anchors (the whole-block-in-document bound is required because _build_block_targets gathers labels by absolute offset and does not encode boundaries); this is documented.

Correctness verification

CPU unit tests (26 pass across the packing + core + target files), including:

• doc-id construction from seq_lens; SDPA mask document restriction; the flex mask_mod materialised on CPU via create_mask and asserted equal to the SDPA mask (the flex backend is the trainer default, so its doc-gating branch is covered);
• anchor sampling keeps every block in one document; per-document draft positions;
• single-document packed == unpacked forward equivalence (loss + valid-token count); a two-document packed training step with gradients; target-side document isolation of the captured hidden states; recipe helpers and fail-fast gates.

GPU recipe smoke (1x A100-80GB, bf16), the full DFlash recipe with a real Qwen3-4B target, packing enabled, and the flex_attention backend (the compiled BlockMask path):

Training log

epoch=0 step=13 loss=4.9688 acc=0.0521
epoch=0 step=16 loss=4.4062 acc=0.0490
epoch=0 step=20 loss=4.0625 acc=0.1094
epoch=0 step=24 loss=3.7344 acc=0.0948
epoch=0 step=28 loss=3.5156 acc=0.1198
epoch=0 step=31 loss=3.3906 acc=0.2021
epoch=0 step=35 loss=3.1250 acc=0.2313
Finished epoch 1/1 completed_steps=35 skipped_short_micro_batches=0

Loss falls 5.1 -> 3.13 and block accuracy rises to ~0.23 with packing on, exercising the full path (packed dataloader -> target block-causal -> document-aware flex block mask -> anchor/position handling).

Scope

DFlash draft only (the DeepSeek MLA EAGLE-3 and EAGLE-1/2 drafts are separate PRs). The JetSpec/Domino variants and packing with context parallelism remain follow-ups.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.