diff --git a/modelopt/torch/export/plugins/hf_spec_export.py b/modelopt/torch/export/plugins/hf_spec_export.py index fd5bb7ef086..caa93db3634 100644 --- a/modelopt/torch/export/plugins/hf_spec_export.py +++ b/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) diff --git a/modelopt/torch/speculative/config.py b/modelopt/torch/speculative/config.py index a4e6de138cb..4df57d3035f 100644 --- a/modelopt/torch/speculative/config.py +++ b/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 diff --git a/modelopt/torch/speculative/dflash/dflash_model.py b/modelopt/torch/speculative/dflash/dflash_model.py index ffffc4739d7..24f2143bf84 100644 --- a/modelopt/torch/speculative/dflash/dflash_model.py +++ b/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 diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py index 65b4314a076..7679ff0020b 100644 --- a/modelopt/torch/speculative/plugins/hf_dflash.py +++ b/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) diff --git a/modelopt/torch/speculative/plugins/hf_domino.py b/modelopt/torch/speculative/plugins/hf_domino.py index 655f2232c15..45d4e1be6c7 100644 --- a/modelopt/torch/speculative/plugins/hf_domino.py +++ b/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. diff --git a/modelopt/torch/speculative/plugins/hf_dspark.py b/modelopt/torch/speculative/plugins/hf_dspark.py index c38b86d63fa..4183a63f1d1 100644 --- a/modelopt/torch/speculative/plugins/hf_dspark.py +++ b/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] diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash.py b/tests/unit/torch/speculative/plugins/test_hf_dflash.py index cdaba1766e3..ef2eec2ad07 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_dflash.py +++ b/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.""" diff --git a/tests/unit/torch/speculative/plugins/test_hf_domino.py b/tests/unit/torch/speculative/plugins/test_hf_domino.py index cff275895b7..0030abe6f33 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_domino.py +++ b/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.""" diff --git a/tests/unit/torch/speculative/plugins/test_hf_dspark.py b/tests/unit/torch/speculative/plugins/test_hf_dspark.py index 4fd1f539955..6a0f884d33f 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_dspark.py +++ b/tests/unit/torch/speculative/plugins/test_hf_dspark.py @@ -203,6 +203,62 @@ def test_confidence_alpha_without_head_raises(self): ) +class TestDSparkSwa: + """DSpark honors dflash_swa_window_size (regression: the window was silently ignored).""" + + def _make_swa_model(self, window=6): + model = get_tiny_llama(num_hidden_layers=4) + config = _get_dspark_config() + config["dflash_swa_window_size"] = window + mtsp.convert(model, [("dflash", config)]) + return model + + def test_forward_passes_window_to_mask(self, monkeypatch): + """The training forward builds the draft mask with the configured window.""" + model = self._make_swa_model(window=6) + 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] + + def test_generate_applies_windowed_mask(self, monkeypatch): + """pseudo_speculative_generate builds (and runs with) the generation-time SWA mask.""" + model = self._make_swa_model(window=6) + model.eval() + torch.manual_seed(0) + input_ids = torch.randint(1, model.dflash_config.vocab_size, (1, SEQ_LEN)) + + masks = [] + orig = model._build_generate_swa_mask + + def spy(*args, **kwargs): + mask = orig(*args, **kwargs) + masks.append(mask) + return mask + + monkeypatch.setattr(model, "_build_generate_swa_mask", spy) + base_token, draft_tokens = model.pseudo_speculative_generate(input_ids, steps=3) + assert len(masks) == 1 and masks[0] is not None + assert masks[0].shape == (1, 1, BLOCK_SIZE, SEQ_LEN + BLOCK_SIZE) + assert base_token.shape == (1, 1) + assert draft_tokens.shape == (1, 3) + + class TestDSparkExporter: """Test the DSpark checkpoint export format (z-lab-compatible layout)."""