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, and an optional output gate
• 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 — so the repo keeps one convention for attention parameterization across draft families.

Structure

• 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 (gqa/mha) and Qwen3DFlashMLAAttention implement only projection construction and rotated q/k/v
• 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
• 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
• structural MLA dims use the standard top-level HF fields; behavioral knobs (mla_rope_interleaved, mla_use_output_gate) live in dflash_config only, defaulting the rotation convention from the standard rope_interleave field

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.

Validation

• pre-commit: all hooks passed
• tests/test_modeling/test_dflash_mla.py: 10 passed, 5 subtests passed (CPU; the CUDA FlexAttention case is skipped without a GPU) — includes eager↔SDPA parity with grads, asymmetric-K/V-dim caching, an MLA spec_generate decode smoke test, and a DSpark MLA case that exercises the head-count-policy bypass
• full tests/test_modeling: 56 passed, 4 skipped, 20 subtests passed — existing GQA eager/sliding/flex-backend and Domino/DSpark tests confirm the scaffold extraction is behavior-preserving
• tests/test_runtime/test_model_loading.py + test_disaggregated_model_loading.py: 17 passed
• 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 now routes through the shared compile_friendly_flex_attention dispatch

🤖 Generated with Claude Code