Add config-selectable MLA draft attention

docs/advanced_features/customization.md

@@ -85,6 +85,47 @@ EAGLE3 offline sequence parallelism is selected with
 tracking are also config features rather than custom launchers; see the
 [training guide](../basic_usage/training.md) for their validated combinations.
 
+## DFlash-family attention modes
+
+DFlash-family draft models select their attention parameterization with
+`dflash_config.attention_mode`: `gqa` (the default), `mha`, or `mla`. The mode
+swaps only the attention projections inside the shared decoder layer; the
+`DFlashDraftModel`, `DominoDraftModel`, and `DSparkDraftModel` architectures,
+target-context injection, per-layer full/sliding masks, objectives, capture
+contract, and `eager`/`sdpa`/`flex_attention` backend selection are identical
+across modes. Multi-head Latent Attention is therefore a draft JSON change,
+not a new architecture:
+
+```json
+{
+  "architectures": ["DSparkDraftModel"],
+  "hidden_size": 4096,
+  "num_attention_heads": 32,
+  "q_lora_rank": 1536,
+  "kv_lora_rank": 512,
+  "qk_nope_head_dim": 128,
+  "qk_rope_head_dim": 64,
+  "v_head_dim": 128,
+  "dflash_config": {
+    "projector_type": "dspark",
+    "attention_mode": "mla"
+  }
+}
+```
+
+MLA dimensions and behavior use the standard top-level Hugging Face fields:
+`q_lora_rank` (`null` selects a direct query projection), `kv_lora_rank`, the
+head dims (`qk_nope_head_dim` may be zero; `qk_rope_head_dim` must be even),
+and `rope_interleave` for the rotation convention (omitted means interleaved,
+the DeepSeek default; `false` selects NeoX-style half rotation). Omitting
+`attention_mode` preserves the existing GQA path; `"mha"` requires equal
+query/KV head counts.
+
+MLA is a training-side mode: checkpoints train, evaluate through
+`spec_generate`, and export through `--to hf`. SGLang serving of DFlash-family
+drafts currently implements the GQA/MHA layout only, so plan benchmarks
+accordingly.
+
 ## Draft architectures
 
 Draft classes register through `@register_draft`. The key defaults to the

scripts/gates/normalize_dflash_export.py

@@ -81,6 +81,15 @@ def normalize_export(config_path: str, expected_block_size: int) -> Dict[str, An
             "export is not DFlash-family: "
             f"dflash_config.projector_type={projector_type!r}"
         )
+    attention_mode = method_config.get("attention_mode", "gqa")
+    if not isinstance(attention_mode, str) or attention_mode.lower() not in {
+        "gqa",
+        "mha",
+    }:
+        raise ValueError(
+            "SGLang DFlash-family serving supports only GQA/MHA exports, "
+            f"got dflash_config.attention_mode={attention_mode!r}"
+        )
 
     if projector_type == "dspark":
         _normalize_dspark(config, method_config)

specforge/algorithms/model_providers.py

@@ -479,6 +479,7 @@ def apply_dflash_overrides(cfg: Config, draft_config: Any) -> None:
     from specforge.modeling.draft.dflash import (
         build_target_layer_ids,
         resolve_dflash_attention_layout,
+        validate_dflash_attention_config,
     )
 
     requested_layers = cfg.model.draft_num_hidden_layers
@@ -502,6 +503,7 @@ def apply_dflash_overrides(cfg: Config, draft_config: Any) -> None:
         }
 
     resolve_dflash_attention_layout(draft_config)
+    validate_dflash_attention_config(draft_config)
 
 
 __all__ = [