feat(speculative): support sequence packing in the DSpark Qwen3 draft

[khazic]

Stacks on #3004 (the DFlash packing PR) for the document-aware DFlash masks that DSpark reuses. Review/merge #3004 first; this diff is the DSpark-specific changes on top.

What

Extends sequence packing to the DSpark draft. DSpark is anchor/block-based like DFlash and reuses the same dflash masks plus the shared dspark/common.py helpers, so packing composes the same way. Scope: the online text-only Qwen3 path.

• common.py: sample_anchor_positions / build_anchor_candidate_mask require the anchor's first target to stay in its document; build_eval_mask truncates each block at its anchor's document boundary (step <= doc_remaining[anchor] before the cumprod, a partial in-document block matching the sequence-end truncation); create_position_ids gathers per-document context positions; add a context_doc_ids helper. All new params default to None, so the other drafts' calls are unchanged.
• draft_qwen3.py: forward accepts position_ids / seq_lens / doc_remaining, builds the document-aware dflash mask (ctx_doc_id / anchor_doc_id) and per-document positions, and gates supervision by document.
• core.py / target.py: thread the packing metadata (only when packing is active, so the unpacked path stays signature-compatible with every draft); the target runs block-causal (or FA varlen at batch size 1) so context hidden states do not leak across documents, and fails loud if the target forward does not accept position_ids.
• train_dspark.py: read packed_sequence_size into the LLM dataloaders, splat the metadata through _forward_batch, and fail fast on unsupported configs (multimodal, cached-target, non-Qwen3 draft, cp_size>1, FlashAttention target with micro_batch_size>1).

The Markov head is position-agnostic (token-id lookup) and needs no change. The other 4 DSpark drafts + the VLM and cached paths are gated off (documented follow-ups).

Correctness verification

CPU unit tests (14 pass; the packing helpers run on the SDPA/eager path), including:

• doc-id construction from seq_lens; anchor candidate requires the first target in-document; eval_mask truncates a block at the document boundary; per-document draft positions;
• single-document packed == unpacked trainer equivalence (seeded); a two-document packed training step with gradients; target-side document isolation of the captured hidden states; the unpacked path does not forward packing kwargs to a legacy-signature draft (so the non-Qwen3 drafts keep working); recipe helpers and fail-fast gates.

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

Training log

step 24 | loss 3.2376 | ce 5.1290 | tv 1.7975 | conf 1.1070
step 27 | loss 2.5047 | ce 5.2626 | tv 1.7849 | conf 0.3720
step 30 | loss 2.4739 | ce 5.7357 | tv 1.6559 | conf 0.4100
step 33 | loss 2.4052 | ce 5.1364 | tv 1.7021 | conf 0.3597
step 35 | loss 2.3672 | ce 4.8426 | tv 1.7004 | conf 0.3525
Finished epoch 1/1 completed_steps=35

The three-term loss falls with packing on, exercising the full path (packed dataloader -> target block-causal -> document-aware flex block mask -> anchor / position / eval-mask handling -> Markov head).

Scope

DSpark Qwen3 online text-only path only. The other 4 DSpark drafts, the VLM path, the offline-cache path, 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.