feat(speculative): support sequence packing in the DeepSeek MLA EAGLE-3 draft

[khazic]

What

The DeepSeek MLA EAGLE-3 draft (DeepseekV3Eagle3DraftModel) shares the EAGLE-3 trainer (Eagle3TrainerModule), which already threads position_ids / seq_lens / doc_remaining for sequence packing, but the draft's forward rejected seq_lens with NotImplementedError. This wires packing through the eager MLA path exactly like the sibling dense Llama draft (LlamaEagle3DraftModel).

When seq_lens is present, the draft builds the shared [B, 1, T, T] block-causal additive mask (build_block_causal_additive_mask) instead of the plain causal + padding mask, so Block-1 attention is document-level block-causal. The TTT diagonal block is position-wise and stays document-safe; cross-document supervision is masked by the trainer's doc_remaining gate (unchanged). The MLA attention is eager-only, so packing reuses the eager mask path (no FlashAttention varlen kernel). The trainer, recipe, and target wrapper are untouched: enabling packing for a DeepSeek target now works end to end (colocated backend, cp_size=1).

Robustness

Packing inputs now fail loud instead of silently corrupting supervision:

• build_block_causal_additive_mask validates that seq_lens is 2D and each row sums to seq_length, mirroring the check the FlashAttention varlen path already performs in _seq_lens_to_cu_seqlens. A row that does not sum to T would otherwise misbucket the boundary cumsum (isolating the final token into a phantom document). This guard covers every eager packing caller: the MLA draft, the dense Llama draft, and the target wrapper.
• DeepseekV3Eagle3DraftModel.forward raises when seq_lens is passed without per-document position_ids (the global arange fallback would give every document after the first a wrong RoPE phase), matching the target-side guard.

Correctness verification

CPU unit tests (27 pass across the two touched test files), including:

• Golden parity: packing N documents into one row produces the same loss and every gradient (tight fp32 tolerance) as N padded single-document rows.
• Document isolation: perturbing only document B's inputs leaves document A's draft output bit-identical.
• Single-doc equivalence: a full-width single-document packed row matches the plain causal forward.
• Fail-loud guards: non-2D / non-summing seq_lens and missing position_ids raise.

GPU smoke (1x A100-80GB, bf16) with a tiny HF DeepSeek-V3 target wrapped in the real HFEagle3TargetModel, draft config derived from the target config exactly as the recipe does: the packed and unpacked paths both train (finite, decreasing loss), and the target's aux hidden states are document-isolated under packing (perturbing document B leaves document A's target aux states bit-identical), confirming the DeepSeek target honors the block-causal mask.

GPU smoke log

=== UNPACKED baseline ===
[unpacked] step 0 loss 4.8782 valid 111 acc 0.000
[unpacked] step 1 loss 4.8733 valid 134 acc 0.007
[unpacked] step 2 loss 4.8649 valid 106 acc 0.066
[unpacked] step 3 loss 4.8657 valid 214 acc 0.149
[unpacked] step 4 loss 4.8619 valid 168 acc 0.053
[unpacked] step 5 loss 4.8694 valid 172 acc 0.000
[unpacked] OK: 4.878 -> 4.869
=== PACKED (seq_lens/position_ids/doc_remaining) ===
[packed] step 0 loss 4.8625 valid 71 acc 0.028
[packed] step 1 loss 4.8580 valid 45 acc 0.089
[packed] step 2 loss 4.8699 valid 59 acc 0.000
[packed] step 3 loss 4.8614 valid 36 acc 0.000
[packed] step 4 loss 4.8493 valid 54 acc 0.000
[packed] step 5 loss 4.8562 valid 101 acc 0.000
[packed] OK: 4.862 -> 4.856
=== TARGET-SIDE document isolation ===
[target-isolation] max|dDocA|=0.00e+00  max|dDocB|=0.276
[target-isolation] OK: target honors the [B,1,T,T] block-causal mask (no cross-doc leak)
ALL SMOKE PASSED

Scope

DeepSeek MLA draft only, per the incremental plan. Extending packing to DFlash / DSpark / EAGLE-1/2 and resolving packing with context parallelism are separate follow-ups (each draft builds attention differently).

[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.

[HuiyingLi]

LGTM