Add configurable GQA and MLA sliding-window draft layers

docs/advanced_features/customization.md

@@ -85,6 +85,50 @@ 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.
 
+## MLA draft attention
+
+DFlash-family draft models can select Multi-head Latent Attention (MLA) in the
+draft JSON without introducing a model-specific architecture. Keep the existing
+`DFlashDraftModel`, `DominoDraftModel`, or `DSparkDraftModel` architecture and
+set `dflash_config.attention_mode` to `"mla"`:
+
+```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_rope_interleaved": true,
+    "mla_use_output_gate": false
+  }
+}
+```
+
+`q_lora_rank` may be `null` to use a direct query projection. The KV rank and
+all head dimensions must be positive except `qk_nope_head_dim`, which may be
+zero; the rotary dimension must be even. Interleaved partial RoPE is the
+default MLA convention and can be disabled with
+`dflash_config.mla_rope_interleaved: false`. The output gate is optional and
+disabled by default.
+
+MLA changes only the draft attention parameterization. The DFlash-family
+target-context injection, per-layer full/sliding masks, objectives, capture
+contract, and `eager`, `sdpa`, or `flex_attention` training backend selection
+remain unchanged. Omitting `attention_mode` preserves the existing GQA path;
+explicit `"mha"` continues to select equal query/KV head counts for DSpark.
+For a mixed projection stack, replace the scalar with an `attention_modes`
+list containing one `"gqa"`, `"mha"`, or `"mla"` entry per draft layer. This
+list composes with `layer_types`, so both SWA-GQA and SWA-MLA layers can appear
+in the same architecture; see the [training guide](../basic_usage/training.md)
+for the paired config contract.
+
 ## Draft architectures
 
 Draft classes register through `@register_draft`. The key defaults to the

docs/basic_usage/training.md

@@ -179,6 +179,41 @@ mixed layout must be edited explicitly in the draft JSON.
 
 The `eager`, `sdpa`, and `flex_attention` backends support both layouts.
 
+The context layout and projection layout are independent. Use
+`dflash_config.attention_mode` as a uniform `"gqa"`, `"mha"`, or `"mla"`
+shorthand, or use `dflash_config.attention_modes` to select one mode per draft
+layer. The per-layer list must match `num_hidden_layers`; do not set the scalar
+and list forms together.
+
+```json
+{
+  "num_hidden_layers": 4,
+  "layer_types": [
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "full_attention"
+  ],
+  "use_sliding_window": true,
+  "sliding_window": 2048,
+  "dflash_config": {
+    "attention_modes": ["gqa", "mla", "gqa", "mla"]
+  }
+}
+```
+
+The two lists are paired by layer index:
+
+| Context layout | Projection mode | Result |
+|---|---|---|
+| `sliding_attention` | `gqa` | SWA-GQA |
+| `sliding_attention` | `mla` | SWA-MLA |
+| `full_attention` | `gqa` or `mha` | full-context standard attention |
+| `full_attention` | `mla` | full-context MLA |
+
+Any configuration containing an MLA layer must also provide the standard MLA
+dimension fields described in [MLA draft attention](../advanced_features/customization.md#mla-draft-attention).
+
 Domino and DSpark need their projector/head metadata, so they require an
 explicit draft config (or a pretrained warm-start source that contains
 `config.json`). The old Domino parser exposed an optional config flag, but its