feat(dflash2): add MLA sliding-window training support

tests/test_dflash2.py

@@ -157,6 +157,56 @@ def _batch(seed=0, all_masked=False):
 
 
 class TestDFlash2Config(unittest.TestCase):
+    def test_mla_mode_builds_k3_mla_attention(self):
+        from torchspec.models.draft.dspark import K3DSparkMLAAttention
+
+        config = DFlash2Config(
+            **_tiny_config_kwargs(
+                num_hidden_layers=3,
+                attention_mode="mla",
+                layer_types=["sliding_attention", "sliding_attention", "full_attention"],
+                sliding_window=2048,
+                is_causal=False,
+                q_lora_rank=4,
+                kv_lora_rank=4,
+                qk_nope_head_dim=2,
+                qk_rope_head_dim=2,
+                v_head_dim=2,
+                rope_parameters={"rope_type": "default", "rope_theta": 10000.0},
+            )
+        )
+
+        model = DFlash2DraftModel(config)
+
+        self.assertEqual(config.attention_mode, "mla")
+        self.assertTrue(
+            all(isinstance(layer.self_attn, K3DSparkMLAAttention) for layer in model.layers)
+        )
+
+    def test_mixed_layer_types_build_per_layer_mask_policies(self):
+        config = DFlash2Config(
+            **_tiny_config_kwargs(
+                num_hidden_layers=3,
+                layer_types=["sliding_attention", "sliding_attention", "full_attention"],
+                sliding_window=2048,
+                is_causal=False,
+            )
+        )
+        model = DFlash2Model(
+            DFlash2DraftModel(config),
+            block_size=config.block_size,
+            num_anchors=1,
+        )
+
+        self.assertEqual(
+            [model._block_mask_options_for_layer(i) for i in range(3)],
+            [
+                {"is_causal": False, "sliding_window": 2048},
+                {"is_causal": False, "sliding_window": 2048},
+                {"is_causal": False, "sliding_window": None},
+            ],
+        )
+
     def test_repository_config_dispatches_to_dflash2(self):
         config_path = ROOT / "torchspec" / "config" / "dflash2_draft_config.json"
         self.assertTrue(config_path.is_file())
@@ -372,7 +422,6 @@ def test_invalid_config_is_rejected(self):
         invalid_layer_types = (
             (["full_attention", "full_attention"], 1, "one entry"),
             (["linear_attention"], 1, "Unsupported"),
-            (["full_attention", "sliding_attention"], 2, "mixed"),
         )
         for layer_types, num_hidden_layers, error in invalid_layer_types:
             with self.subTest(layer_types=layer_types):

torchspec/models/dflash.py

@@ -267,6 +267,17 @@ def _create_noise_embed(
     def _block_mask_options(self) -> dict:
         return {}
 
+    def _block_mask_options_for_layer(self, layer_id: int) -> dict:
+        """Return the attention-mask options for one draft layer.
+
+        DFlash historically used one mask for every layer.  Keeping this
+        method separate from ``_block_mask_options`` lets model variants use
+        per-layer attention schedules without changing the public training
+        wrapper API.
+        """
+        del layer_id
+        return self._block_mask_options()
+
     def _compute_logits(
         self,
         draft_hidden: torch.Tensor,
@@ -331,21 +342,32 @@ def _draft_backbone(
 
         block_mask = None
         if device.type == "cuda":
-            mask_mod = _create_dflash_mask_mod(
-                anchor_positions=anchor_positions,
-                block_keep_mask=block_keep_mask,
-                ctx_len=seq_len,
-                block_size=self.block_size,
-                **self._block_mask_options(),
-            )
-            block_mask = compile_friendly_create_block_mask(
-                mask_mod=mask_mod,
-                B=bsz,
-                H=None,
-                Q_LEN=draft_len,
-                KV_LEN=kv_len,
-                device=device,
-            )
+            # A mixed full/SWA schedule needs a distinct BlockMask for each
+            # attention policy.  Cache equal policies so the common uniform
+            # schedule still creates exactly one mask.
+            mask_cache = {}
+            block_masks = []
+            for layer_id in range(self.draft_model.num_layers):
+                options = self._block_mask_options_for_layer(layer_id)
+                cache_key = tuple(sorted(options.items()))
+                if cache_key not in mask_cache:
+                    mask_mod = _create_dflash_mask_mod(
+                        anchor_positions=anchor_positions,
+                        block_keep_mask=block_keep_mask,
+                        ctx_len=seq_len,
+                        block_size=self.block_size,
+                        **options,
+                    )
+                    mask_cache[cache_key] = compile_friendly_create_block_mask(
+                        mask_mod=mask_mod,
+                        B=bsz,
+                        H=None,
+                        Q_LEN=draft_len,
+                        KV_LEN=kv_len,
+                        device=device,
+                    )
+                block_masks.append(mask_cache[cache_key])
+            block_mask = block_masks[0] if len(mask_cache) == 1 else block_masks
 
         # 6. Draft model forward — pass embeddings directly
         draft_hidden = self.draft_model(

torchspec/models/dflash2.py

@@ -48,28 +48,41 @@ def __init__(self, *args, selector_loss_alpha: float = 1.0, **kwargs):
         attention_types = set(layer_types or ["full_attention"])
         if not attention_types <= {"full_attention", "sliding_attention"}:
             raise ValueError(f"Unsupported DFlash2 layer types: {sorted(attention_types)}")
-        if len(attention_types) > 1:
-            raise ValueError("DFlash2 training does not support mixed full and sliding layers")
-
-        uses_sliding_window = "sliding_attention" in attention_types
         explicit_causality = getattr(config, "is_causal", None)
-        self.attention_is_causal = (
-            uses_sliding_window if explicit_causality is None else bool(explicit_causality)
-        )
         configured_window = getattr(config, "sliding_window", None)
-        self.sliding_window = (
-            int(configured_window)
-            if uses_sliding_window and configured_window is not None
-            else None
-        )
-        if self.sliding_window is not None and self.sliding_window < 1:
-            raise ValueError(f"sliding_window must be positive, got {self.sliding_window}")
+        if "sliding_attention" in attention_types:
+            if configured_window is None:
+                raise ValueError(
+                    "DFlash2 sliding_attention layers require an explicit positive sliding_window"
+                )
+            configured_window = int(configured_window)
+            if configured_window < 1:
+                raise ValueError(f"sliding_window must be positive, got {configured_window}")
+
+        normalized_layer_types = layer_types or ["full_attention"] * config.num_hidden_layers
+        self.layer_block_mask_options = []
+        for layer_type in normalized_layer_types:
+            uses_sliding_window = layer_type == "sliding_attention"
+            self.layer_block_mask_options.append(
+                {
+                    "is_causal": (
+                        uses_sliding_window
+                        if explicit_causality is None
+                        else bool(explicit_causality)
+                    ),
+                    "sliding_window": configured_window if uses_sliding_window else None,
+                }
+            )
+
+        # Retain the uniform-policy attributes for callers that inspect them.
+        self.attention_is_causal = self.layer_block_mask_options[0]["is_causal"]
+        self.sliding_window = self.layer_block_mask_options[0]["sliding_window"]
 
     def _block_mask_options(self) -> dict:
-        return {
-            "is_causal": self.attention_is_causal,
-            "sliding_window": self.sliding_window,
-        }
+        return self.layer_block_mask_options[0]
+
+    def _block_mask_options_for_layer(self, layer_id: int) -> dict:
+        return self.layer_block_mask_options[layer_id]
 
     def _compute_logits(
         self,

torchspec/models/draft/dflash.py

@@ -446,7 +446,8 @@ def forward(
             context_feature: [B, ctx_len, D] — projected context from target
             draft_position_ids: [B, draft_len]
             context_position_ids: [B, ctx_len]
-            block_mask: FlexAttention BlockMask
+            block_mask: FlexAttention BlockMask, or a per-layer sequence of
+                BlockMasks for mixed attention schedules
             noise_embedding: [B, draft_len, D] — pre-computed embeddings (from training wrapper)
 
         Returns:
@@ -457,13 +458,16 @@ def forward(
         else:
             draft_hidden = self.embed_tokens(draft_input_ids).to(context_feature.dtype)
 
-        for layer in self.layers:
+        for layer_id, layer in enumerate(self.layers):
+            layer_block_mask = block_mask
+            if isinstance(block_mask, (list, tuple)):
+                layer_block_mask = block_mask[layer_id]
             draft_hidden = layer(
                 draft_hidden=draft_hidden,
                 context_hidden=context_feature,
                 draft_position_ids=draft_position_ids,
                 context_position_ids=context_position_ids,
-                block_mask=block_mask,
+                block_mask=layer_block_mask,
             )
 
         return self.final_norm(draft_hidden)

torchspec/models/draft/dflash2.py

@@ -44,6 +44,7 @@
 )
 
 _NESTED_CONFIG_FIELDS = (
+    "attention_mode",
     "block_size",
     "conv_group_size",
     "conv_kernel_size",
@@ -101,6 +102,15 @@ def __init__(
         conv_group_size: int = 16,
         selector_rank: int = 256,
         selector_top_k: int = 16,
+        attention_mode: str = "gqa",
+        q_lora_rank: int | None = 1536,
+        kv_lora_rank: int = 512,
+        qk_nope_head_dim: int = 128,
+        qk_rope_head_dim: int = 64,
+        v_head_dim: int = 128,
+        mla_use_output_gate: bool = False,
+        rope_scaling: dict | None = None,
+        rope_parameters: dict | None = None,
         input_embedding_scale: float = 1.0,
         output_multiplier: float = 1.0,
         final_logit_softcapping: float | None = None,
@@ -119,6 +129,7 @@ def __init__(
         conv_group_size = int(nested.get("conv_group_size", conv_group_size))
         selector_rank = int(nested.get("selector_rank", selector_rank))
         selector_top_k = int(nested.get("selector_top_k", selector_top_k))
+        attention_mode = str(nested.get("attention_mode", attention_mode)).lower()
         input_embedding_scale = float(nested.get("input_embedding_scale", input_embedding_scale))
         output_multiplier = float(nested.get("output_multiplier", output_multiplier))
         final_logit_softcapping = nested.get("final_logit_softcapping", final_logit_softcapping)
@@ -181,8 +192,10 @@ def __init__(
         kwargs["mask_token_id"] = int(
             nested.get("mask_token_id", kwargs.get("mask_token_id", 151669))
         )
-        rope_parameters = kwargs.get("rope_parameters") or {}
-        kwargs.setdefault("rope_theta", rope_parameters.get("rope_theta", 10000.0))
+        rope_parameters = rope_scaling if rope_scaling is not None else rope_parameters
+        rope_parameters = dict(rope_parameters) if rope_parameters is not None else None
+        rope_parameters_for_validation = rope_parameters or {}
+        kwargs.setdefault("rope_theta", rope_parameters_for_validation.get("rope_theta", 10000.0))
         kwargs.pop("model_type", None)
 
         hidden_act = kwargs.get("hidden_act", "silu")
@@ -194,11 +207,22 @@ def __init__(
             raise ValueError("DFlash2 training requires attention_dropout=0")
         if kwargs.get("fc_norm", False):
             raise ValueError("DFlash2 training does not support fc_norm=True")
-        if kwargs.get("rope_scaling") is not None:
-            raise ValueError("DFlash2 training does not support rope_scaling")
-        unsupported_rope_keys = set(rope_parameters) - {"rope_theta", "rope_type"}
-        if rope_parameters.get("rope_type", "default") != "default" or unsupported_rope_keys:
-            raise ValueError("DFlash2 training supports only default rope_parameters")
+        if attention_mode not in {"gqa", "mla"}:
+            raise ValueError(
+                f"DFlash2 attention_mode must be 'gqa' or 'mla', got {attention_mode!r}"
+            )
+        if attention_mode == "gqa":
+            if rope_scaling is not None:
+                raise ValueError("GQA DFlash2 training does not support rope_scaling")
+            unsupported_rope_keys = set(rope_parameters_for_validation) - {
+                "rope_theta",
+                "rope_type",
+            }
+            if (
+                rope_parameters_for_validation.get("rope_type", "default") != "default"
+                or unsupported_rope_keys
+            ):
+                raise ValueError("GQA DFlash2 training supports only default rope_parameters")
 
         num_hidden_layers = int(kwargs.get("num_hidden_layers", 5))
         use_sliding_window = bool(kwargs.get("use_sliding_window", False))
@@ -226,8 +250,6 @@ def __init__(
         attention_types = set(layer_types)
         if not attention_types <= {"full_attention", "sliding_attention"}:
             raise ValueError(f"Unsupported DFlash2 layer types: {sorted(attention_types)}")
-        if len(attention_types) > 1:
-            raise ValueError("DFlash2 training does not support mixed full and sliding layers")
         if "sliding_attention" in attention_types:
             if sliding_window is None:
                 raise ValueError(
@@ -275,6 +297,31 @@ def __init__(
             )
 
         super().__init__(**kwargs)
+        self.attention_mode = attention_mode
+        self.q_lora_rank = q_lora_rank
+        self.kv_lora_rank = int(kv_lora_rank)
+        self.qk_nope_head_dim = int(qk_nope_head_dim)
+        self.qk_rope_head_dim = int(qk_rope_head_dim)
+        self.v_head_dim = int(v_head_dim)
+        self.mla_use_output_gate = bool(mla_use_output_gate)
+        self.rope_parameters = rope_parameters
+        if self.attention_mode == "mla":
+            # DeepSeekMLAAttention's shared rotary builder reads the legacy
+            # ``rope_scaling`` name.  Keep it populated for Transformers
+            # versions where it is not an alias of ``rope_parameters``.
+            self.rope_scaling = rope_parameters
+        if rope_parameters is not None:
+            nested_rope_theta = rope_parameters.get("rope_theta")
+            if nested_rope_theta is not None:
+                self.rope_theta = float(nested_rope_theta)
+            if rope_parameters.get("rope_type", rope_parameters.get("type")) == "yarn":
+                for key, default in {
+                    "beta_fast": 32.0,
+                    "beta_slow": 1.0,
+                    "mscale": 1.0,
+                    "mscale_all_dim": 0.0,
+                }.items():
+                    rope_parameters.setdefault(key, default)
         self.block_size = block_size
         self.conv_kernel_size = conv_kernel_size
         self.conv_group_size = conv_group_size
@@ -286,6 +333,7 @@ def __init__(
 
         nested.update(
             {
+                "attention_mode": self.attention_mode,
                 "block_size": self.block_size,
                 "conv_kernel_size": self.conv_kernel_size,
                 "conv_group_size": self.conv_group_size,
@@ -418,6 +466,10 @@ def score_candidates(
 
 class DFlash2DecoderLayer(DFlashDecoderLayer):
     def __init__(self, config: DFlash2Config):
+        if config.attention_mode == "mla":
+            from torchspec.models.draft.dspark import K3DSparkMLAAttention
+
+            self.attention_class = K3DSparkMLAAttention
         super().__init__(config)
         conv_args = {
             "hidden_size": config.hidden_size,