From 10cc2f6ae2c9ba7cc5841ece95e27d0562aef1bf Mon Sep 17 00:00:00 2001 From: DimensionSTP <65501090+DimensionSTP@users.noreply.github.com> Date: Sat, 10 Oct 2026 12:00:16 +0900 Subject: [PATCH] [Bugfix] Bound GDN speculative decode query widths (#55504) Signed-off-by: DimensionSTP <65501090+DimensionSTP@users.noreply.github.com> Co-authored-by: OpenAI Codex --- tests/v1/worker/test_mamba_hybrid_model_state.py | 11 +++++++++++ vllm/v1/worker/gpu/model_states/mamba_hybrid.py | 9 ++++++++- 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/tests/v1/worker/test_mamba_hybrid_model_state.py b/tests/v1/worker/test_mamba_hybrid_model_state.py index fa54ccfced..0dcc4725d7 100644 --- a/tests/v1/worker/test_mamba_hybrid_model_state.py +++ b/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, ): diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index c45304efa5..94c57c1968 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/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 )