[Bugfix] Bound GDN speculative decode query widths (#55504)

Signed-off-by: DimensionSTP <65501090+DimensionSTP@users.noreply.github.com>
Co-authored-by: OpenAI Codex <noreply@openai.com>

tests/v1/worker/test_mamba_hybrid_model_state.py

@@ -314,6 +314,17 @@ def test_padded_rows_keep_the_sentinel(monkeypatch):
     assert metadata.num_decode_draft_tokens_cpu.tolist() == [0, 0, -1, -1]
 
 
+def test_rows_wider_than_the_speculative_state_window_keep_the_sentinel(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    state = _mamba_hybrid_state(num_speculative_tokens=7)
+    input_batch = _input_batch([12, 8, 1], None, [False, False, False])
+
+    metadata = _prepare_attn_metadata(state, monkeypatch, input_batch)
+
+    assert metadata.num_decode_draft_tokens_cpu.tolist() == [-1, 0, 0]
+
+
 def test_chunked_prefill_tail_of_two_or_three_tokens_keeps_the_sentinel(
     monkeypatch,
 ):

vllm/v1/worker/gpu/model_states/mamba_hybrid.py

@@ -288,7 +288,14 @@ class MambaHybridModelState(DefaultModelState):
             # Test request state, not num_scheduled_tokens == draft_count+1:
             # adaptive rewrites num_scheduled_tokens to an even split, so that
             # equality rarely holds and would demote every verify row to decode.
-            is_decode = ~is_prefilling_np & (input_batch.num_scheduled_tokens > 0)
+            is_decode = (
+                ~is_prefilling_np
+                & (input_batch.num_scheduled_tokens > 0)
+                & (
+                    input_batch.num_scheduled_tokens
+                    <= self.vllm_config.num_speculative_tokens + 1
+                )
+            )
             num_decode_draft_tokens_np[: input_batch.num_reqs] = np.where(
                 is_decode, num_draft_tokens_per_req, -1
             )