Add config-selectable MLA draft attention

[maocheng23]

Summary

• add config-selectable Multi-head Latent Attention to the DFlash-family draft backbone: dflash_config.attention_mode: "mla" keeps the existing DFlashDraftModel, DominoDraftModel, or DSparkDraftModel architecture and swaps only the attention projections; omitting the option preserves the current GQA path
• support the standard MLA projection fields (optional low-rank Q, shared compressed KV latent), partial interleaved or NeoX RoPE selected by the standard rope_interleave field, and DeepSeek YaRN mscale_all_dim logit scaling
• preserve the DFlash-family target-context injection, full/sliding masks, objectives, capture contract, and eager/SDPA/FlexAttention backend selection

MLA is wired exactly the way gqa/mha already are — family-local, selected per mode inside the shared decoder layer — and every MLA parameter is a standard top-level HF config field: the mode introduces no new behavioral knobs.

Structure (two commits)

1. refactor(dflash): extract shared attention scaffold — pure code motion, no behavior change: Qwen3DFlashAttentionBase owns everything the modes must agree on (KV-cache updates, flex/eager/sdpa backend dispatch including the SPECFORGE_FLEX_ATTENTION_BACKEND override, fully-masked-query zeroing, output projection); Qwen3DFlashAttention keeps its name, constructor, and state-dict keys and implements only _init_projections/_compute_qkv. The subclass contract (scaling, num_key_value_groups, o_proj) is documented and asserted at init.
2. Add config-selectable MLA draft attention — additive: this commit does not touch the scaffold or the GQA class again (its only deletions in dflash.py are the decoder-layer class lookup and the model-level rotary dim).

Details:

• attention_mode has a single source of truth (resolve_dflash_attention_mode + _DFLASH_ATTENTION_CLASSES); DSpark imports it instead of re-declaring the mode set, and "mha" now enforces equal query/KV head counts for every family (previously only DSpark checked this)
• the rotary embedding stays model-level: DFlashDraftModel builds it with the mode-appropriate dim (qk_rope_head_dim under MLA) and all modes consume the same position_embeddings pipeline
• MLA dimension validation runs once at model init, next to the existing resolve_dflash_attention_layout boundary, and CPU-side in apply_dflash_overrides before GPU allocation

Docs note the serving boundary explicitly: MLA is a training-side mode (train, spec_generate, export --to hf); SGLang serving of DFlash-family drafts currently implements the GQA/MHA layout only, so the export gate rejects MLA.

Validation

• pre-commit: all hooks passed
• commit 1 alone passes the full pre-existing tests/test_modeling suite unchanged (46 passed, 4 skipped, 15 subtests) — the scaffold extraction is behavior-preserving
• branch tip: tests/test_modeling 63 passed, 5 skipped, 24 subtests — test_dflash_mla.py covers eager↔SDPA parity with grads, hand-computed RoPE reference values for both conventions, asymmetric-K/V-dim caching, an MLA spec_generate decode smoke test, the DSpark MLA head-count-policy bypass, and MHA head-count enforcement
• tests/test_runtime/test_model_loading.py + test_disaggregated_model_loading.py + tests/test_scripts/test_gate_orchestration.py: 27 passed, 8 subtests
• tested with the repository-pinned Transformers 5.8.1 and PyTorch 2.11.0; the CUDA FlexAttention forward/backward case needs a GPU re-run on this revision since the flex path routes through the shared compile_friendly_flex_attention dispatch

🤖 Generated with Claude Code

[maocheng23]

Force-pushed a rework (also rebased on current main, no conflicts):

• Scope cut — no custom knobs. Removed the two speculative dflash_config options. mla_use_output_gate is gone: standard DeepSeek MLA has no attention output gate and nothing consumes it; dropping it also removes the only behavioral hook the scaffold inserted into the shared forward. mla_rope_interleaved is gone: the standard HF rope_interleave field on the draft config already selects the rotation convention, so a second surface for the same choice was redundant. Every MLA parameter is now a standard top-level HF field.
• History split for reviewability. Commit 1 is the pure scaffold extraction (no behavior change — the pre-existing test_modeling suite passes on it unchanged); commit 2 adds MLA and never touches the scaffold or GQA class again.
• Tightened validation. attention_mode: "mha" now requires equal query/KV head counts for every family (previously only DSpark enforced it), and the scaffold's subclass contract is documented and asserted at init.