[Feat]: SWA (sliding-window attention) support for DFlash drafter training

modelopt/torch/export/plugins/hf_spec_export.py

@@ -399,6 +399,21 @@ def _export_config(self):
         else:
             config["layer_types"] = ["full_attention"] * draft_config.num_hidden_layers
 
+        # Sliding-window attention: all draft layers use non-causal SWA (MiMo-style). vLLM's
+        # _resolve_layer_attention reads dflash_config.use_swa + swa_window_size; with
+        # layer_types left all "full_attention" it applies a non-causal sliding window to
+        # every draft layer (window from swa_window_size / top-level sliding_window).
+        swa_window = getattr(self.model, "dflash_swa_window_size", None)
+        if swa_window is not None:
+            config["sliding_window"] = swa_window
+            config["dflash_config"].update(
+                {
+                    "use_swa": True,
+                    "swa_window_size": swa_window,
+                    "causal": False,
+                }
+            )
+
         # Inject the export-time YaRN rope_scaling from the dflash_export_rope_scaling
         # config field (empty dict disables). Mirrors eagle's eagle_export_rope_scaling.
         export_rope_scaling = getattr(self.model, "dflash_export_rope_scaling", None)

modelopt/torch/speculative/config.py

@@ -149,6 +149,19 @@ class DFlashConfig(ModeloptBaseConfig):
         description="Whether to use torch.compile on DFlash forward/loss methods.",
     )
 
+    dflash_swa_window_size: int | None = ModeloptField(
+        default=None,
+        description=(
+            "Sliding-window attention (SWA) window size for the DFlash draft. When set, ALL "
+            "draft layers use non-causal sliding-window attention (MiMo-style): each draft "
+            "query attends only to context positions within `dflash_swa_window_size` tokens "
+            "before it, while block-internal attention stays bidirectional. None (default) "
+            "keeps full attention over all context. Must be >= dflash_block_size. Exported to "
+            "the draft config as dflash_config.use_swa/swa_window_size (+ top-level "
+            "sliding_window) so vLLM applies the same window at inference."
+        ),
+    )
+
     dflash_export_rope_scaling: dict = ModeloptField(
         default={},
         description=(
@@ -221,6 +234,14 @@ def _check_dpace_alpha(self) -> "DFlashConfig":
         # is rejected even if it only becomes active after a later objective override.
         if not 0.0 < self.dflash_dpace_alpha <= 1.0:
             raise ValueError(f"dflash_dpace_alpha must be in (0, 1], got {self.dflash_dpace_alpha}")
+        if self.dflash_swa_window_size is not None:
+            # Block-internal attention is left un-windowed (bidirectional), so the window must
+            # cover a full block; otherwise the effective inference window would differ.
+            if self.dflash_swa_window_size < self.dflash_block_size:
+                raise ValueError(
+                    f"dflash_swa_window_size ({self.dflash_swa_window_size}) must be >= "
+                    f"dflash_block_size ({self.dflash_block_size})."
+                )
         return self
 
 

modelopt/torch/speculative/dflash/dflash_model.py

@@ -48,4 +48,5 @@ def modify(self, config):
         self.dflash_num_anchors = config.dflash_num_anchors
         self.dflash_report_acc = config.dflash_report_acc
         self.dflash_use_torch_compile = config.dflash_use_torch_compile
+        self.dflash_swa_window_size = config.dflash_swa_window_size
         self.dflash_export_rope_scaling = config.dflash_export_rope_scaling

modelopt/torch/speculative/plugins/hf_dflash.py

@@ -396,9 +396,16 @@ def _build_position_ids(self, seq_len, anchor_positions, device):
         return torch.cat([ctx_pos, draft_pos], dim=1)
 
     def _build_draft_attention_mask(
-        self, seq_len, anchor_positions, block_keep_mask, n_blocks, dtype, device
+        self, seq_len, anchor_positions, block_keep_mask, n_blocks, dtype, device, window=None
     ):
-        """Build SDPA attention mask: context (causal) + draft (bidirectional within block)."""
+        """Build SDPA attention mask: context (causal) + draft (bidirectional within block).
+
+        When ``window`` is not None, all layers use non-causal sliding-window attention
+        (MiMo-style): each draft query only sees context positions within ``window`` tokens
+        before its own position. Block-internal attention stays bidirectional and is left
+        un-windowed (the config enforces ``window >= block_size``, so a full block always
+        fits inside the window and windowing it would be a no-op).
+        """
         bsz = anchor_positions.shape[0]
         block_size = self.dflash_block_size
         q_len = n_blocks * block_size
@@ -412,6 +419,12 @@ def _build_draft_attention_mask(
 
         # Context: kv < S and kv < anchor
         mask_ctx = (kv_indices < seq_len) & (kv_indices < anchor_exp)
+
+        # Sliding window on the context: keep only context kv whose real position is within
+        # `window` tokens before the query's real position (anchor + position-in-block).
+        if window is not None:
+            q_real_pos = anchor_exp + (q_indices % block_size)  # [B, 1, q_len, 1]
+            mask_ctx = mask_ctx & (kv_indices > q_real_pos - window)
         # Draft: kv >= S and same block
         is_draft = kv_indices >= seq_len
         kv_block_ids = (kv_indices - seq_len) // block_size
@@ -426,6 +439,29 @@ def _build_draft_attention_mask(
         attn_mask.masked_fill_(~final_mask, torch.finfo(dtype).min)
         return attn_mask
 
+    def _build_generate_swa_mask(self, ctx_len, bsz, dtype, device):
+        """Generation-time SWA mask [B, 1, block_size, ctx_len + block_size], or None.
+
+        Returns None with full attention (KV cache with no mask): all positions attend
+        freely to context and each other within the block. With sliding-window attention,
+        each block query only sees context within ``dflash_swa_window_size`` tokens before
+        its real position (ctx_len + position-in-block), matching training and vLLM
+        inference; block kv stays fully visible (bidirectional / un-windowed).
+        """
+        if self.dflash_swa_window_size is None:
+            return None
+        window = self.dflash_swa_window_size
+        block_size = self.dflash_block_size
+        kv_len = ctx_len + block_size
+        kv_idx = torch.arange(kv_len, device=device).view(1, 1, 1, -1)
+        q_real_pos = torch.arange(ctx_len, ctx_len + block_size, device=device).view(1, 1, -1, 1)
+        is_ctx = kv_idx < ctx_len
+        # Context kv kept iff within the window; block kv (>= ctx_len) always visible.
+        keep = (~is_ctx) | (kv_idx > q_real_pos - window)
+        attn_mask = torch.zeros(bsz, 1, block_size, kv_len, device=device, dtype=dtype)
+        attn_mask.masked_fill_(~keep, torch.finfo(dtype).min)
+        return attn_mask
+
     def _compute_loss(
         self, logits, input_ids, anchor_positions, block_keep_mask, loss_mask, base_logits=None
     ):
@@ -648,7 +684,13 @@ def forward(
         )
         full_pos = self._build_position_ids(seq_len, anchor_positions, device)
         attn_mask = self._build_draft_attention_mask(
-            seq_len, anchor_positions, block_keep_mask, n_blocks, target_hidden.dtype, device
+            seq_len,
+            anchor_positions,
+            block_keep_mask,
+            n_blocks,
+            target_hidden.dtype,
+            device,
+            window=self.dflash_swa_window_size,
         )
 
         # 5. Draft forward
@@ -776,16 +818,14 @@ def pseudo_speculative_generate(self, input_ids, steps=1):
         block_positions = torch.arange(ctx_len, ctx_len + block_size, device=device)
         pos_ids = torch.cat([ctx_positions, block_positions]).unsqueeze(0).expand(bsz, -1)
 
-        # No attention mask at inference
-        # which uses KV cache with no mask. All positions attend freely to
-        # context and each other within the block.
+        attn_mask = self._build_generate_swa_mask(ctx_len, bsz, target_hidden.dtype, device)
 
         # Draft forward
         draft_hidden = self.dflash_module(
             noise_embedding=noise_embedding,
             target_hidden=target_hidden,
             position_ids=pos_ids,
-            attention_mask=None,
+            attention_mask=attn_mask,
         )
 
         # Logits on positions 1..block_size-1 (skip anchor at position 0)

modelopt/torch/speculative/plugins/hf_domino.py

@@ -354,7 +354,13 @@ def forward(
         )
         full_pos = self._build_position_ids(seq_len, anchor_positions, device)
         attn_mask = self._build_draft_attention_mask(
-            seq_len, anchor_positions, block_keep_mask, n_blocks, target_hidden.dtype, device
+            seq_len,
+            anchor_positions,
+            block_keep_mask,
+            n_blocks,
+            target_hidden.dtype,
+            device,
+            window=self.dflash_swa_window_size,
         )
 
         # 5. Draft backbone forward.

modelopt/torch/speculative/plugins/hf_dspark.py

@@ -381,7 +381,13 @@ def forward(
         )
         full_pos = self._build_position_ids(seq_len, anchor_positions, device)
         attn_mask = self._build_draft_attention_mask(
-            seq_len, anchor_positions, block_keep_mask, n_blocks, target_hidden.dtype, device
+            seq_len,
+            anchor_positions,
+            block_keep_mask,
+            n_blocks,
+            target_hidden.dtype,
+            device,
+            window=self.dflash_swa_window_size,
         )
 
         # 5. Draft backbone forward.
@@ -459,7 +465,7 @@ def pseudo_speculative_generate(self, input_ids, steps=1):
             noise_embedding=noise_embedding,
             target_hidden=target_hidden,
             position_ids=pos_ids,
-            attention_mask=None,
+            attention_mask=self._build_generate_swa_mask(ctx_len, bsz, target_hidden.dtype, device),
         )
         backbone_logits = self._base_model_lm_head(draft_hidden)  # [B, block_size, V]
 

tests/unit/torch/speculative/plugins/test_hf_dflash.py

@@ -372,6 +372,75 @@ def test_no_sliding_window_without_config(self):
         assert attn.sliding_window is None
 
 
+class TestDFlashSwaMask:
+    """Test all-layer non-causal sliding-window attention mask (MiMo-style)."""
+
+    def test_window_masks_context_beyond_window(self):
+        """Context beyond the window (relative to each query's real position) is masked out."""
+        model = get_tiny_llama(num_hidden_layers=4)
+        config = _get_dflash_config(block_size=4)
+        window = 6
+        config["dflash_swa_window_size"] = window
+        mtsp.convert(model, [("dflash", config)])
+
+        seq_len = 16
+        block_size = 4
+        # One block anchored at position 10 → query real positions [10, 11, 12, 13].
+        anchor_positions = torch.tensor([[10]], dtype=torch.long)
+        block_keep_mask = torch.tensor([[True]])
+        dtype = torch.float32
+        device = torch.device("cpu")
+
+        mask = model._build_draft_attention_mask(
+            seq_len, anchor_positions, block_keep_mask, 1, dtype, device, window=window
+        )
+        neg = torch.finfo(dtype).min
+        attend = mask > neg / 2  # True where a position is attended (additive mask == 0)
+
+        # Context kv are positions [0, seq_len). For query k (real pos 10 + k) only context
+        # positions in (10 + k - window, 10) are visible.
+        for k in range(block_size):
+            q_real = 10 + k
+            for c in range(seq_len):
+                visible = attend[0, 0, k, c].item()
+                if c < 10:  # context strictly before the anchor
+                    assert visible == (c > q_real - window), (
+                        f"query k={k} (pos {q_real}), context c={c}: "
+                        f"expected visible={c > q_real - window}, got {visible}"
+                    )
+
+    def test_window_is_subset_of_full(self):
+        """The windowed mask attends to a subset of what the full-attention mask attends to."""
+        model = get_tiny_llama(num_hidden_layers=4)
+        config = _get_dflash_config(block_size=4)
+        config["dflash_swa_window_size"] = 6
+        mtsp.convert(model, [("dflash", config)])
+
+        args = (
+            16,
+            torch.tensor([[10]]),
+            torch.tensor([[True]]),
+            1,
+            torch.float32,
+            torch.device("cpu"),
+        )
+        full = model._build_draft_attention_mask(*args, window=None)
+        windowed = model._build_draft_attention_mask(*args, window=6)
+        neg = torch.finfo(torch.float32).min
+        # Everything masked by full attention must also be masked by the windowed mask.
+        assert ((full <= neg / 2) <= (windowed <= neg / 2)).all()
+        # The window strictly removes some connections (it is not a no-op here).
+        assert (windowed <= neg / 2).sum() > (full <= neg / 2).sum()
+
+    def test_window_smaller_than_block_rejected(self):
+        """A window smaller than the block size is rejected at config validation."""
+        model = get_tiny_llama(num_hidden_layers=4)
+        config = _get_dflash_config(block_size=4)
+        config["dflash_swa_window_size"] = 2  # < block_size
+        with pytest.raises(ValueError, match="dflash_swa_window_size"):
+            mtsp.convert(model, [("dflash", config)])
+
+
 class TestValidateOnline:
     """Test validate_online acceptance counting logic."""
 
@@ -513,6 +582,33 @@ def test_export_config_fields(self, tmp_path):
         assert "vocab_size" in cfg
         assert "layer_types" in cfg
         assert len(cfg["layer_types"]) == NUM_DRAFT_LAYERS
+        # Without SWA configured, no sliding-window fields are emitted.
+        assert "sliding_window" not in cfg
+        assert "use_swa" not in cfg["dflash_config"]
+
+    def test_export_swa_fields(self, tmp_path):
+        """With dflash_swa_window_size set, exported config carries vLLM's SWA fields."""
+        model = get_tiny_llama(num_hidden_layers=4)
+        config = _get_dflash_config()
+        config["dflash_swa_window_size"] = 256
+        mtsp.convert(model, [("dflash", config)])
+
+        exporter = model.get_exporter()
+        export_dir = tmp_path / "exported"
+        exporter.export(export_dir)
+
+        with open(export_dir / "config.json") as f:
+            cfg = json.load(f)
+
+        # vLLM _resolve_layer_attention reads these; all-full layer_types + use_swa=True
+        # → non-causal sliding window on every draft layer.
+        assert cfg["sliding_window"] == 256
+        assert cfg["dflash_config"]["use_swa"] is True
+        assert cfg["dflash_config"]["swa_window_size"] == 256
+        assert cfg["dflash_config"]["causal"] is False
+        # The pre-existing dflash_config keys must survive the update.
+        assert "mask_token_id" in cfg["dflash_config"]
+        assert "target_layer_ids" in cfg["dflash_config"]
 
     def test_export_tensor_count(self, tmp_path):
         """Exported model should have the right number of tensors."""

tests/unit/torch/speculative/plugins/test_hf_domino.py

@@ -157,6 +157,37 @@ def test_lambda_zero_uses_final_only(self):
         assert abs(out.loss.item() - out.domino_metrics["final_loss"]) < 1e-4
 
 
+class TestDominoSwa:
+    """Domino honors dflash_swa_window_size (regression: the window was silently ignored)."""
+
+    def test_forward_passes_window_to_mask(self, monkeypatch):
+        """The training forward builds the draft mask with the configured window."""
+        model = get_tiny_llama(num_hidden_layers=4)
+        config = _get_domino_config()
+        config["dflash_swa_window_size"] = 6
+        mtsp.convert(model, [("dflash", config)])
+        model.train()
+
+        torch.manual_seed(0)
+        input_ids = torch.randint(1, model.dflash_config.vocab_size, (2, SEQ_LEN))
+
+        windows = []
+        orig = model._build_draft_attention_mask
+
+        def spy(*args, **kwargs):
+            windows.append(kwargs.get("window"))
+            return orig(*args, **kwargs)
+
+        monkeypatch.setattr(model, "_build_draft_attention_mask", spy)
+        out = model(
+            input_ids=input_ids,
+            attention_mask=torch.ones_like(input_ids),
+            labels=input_ids.clone(),
+        )
+        assert out.loss.dim() == 0
+        assert windows == [6]
+
+
 class TestLambdaSchedule:
     """Test the lambda_base curriculum schedule."""