[Bugfix][MRV2] Treat padded prompt tails as spec-decode rows for hybrid models (#58434)
Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Nick Hill committed
Sep 24, 2026 at 18:29 UTC
7dbd0a8d22d20eb938994ad7a9158e02f30c9ce6
2 files changed
+88
-1
tests/v1/worker/test_mamba_hybrid_model_state.py
+80
@@ -4,9 +4,16 @@
4
from types import SimpleNamespace
5
from unittest.mock import Mock
6
7
+import numpy as np
8
import pytest
9
import torch
10
11
+from tests.v1.attention.test_gdn_metadata_builder import (
12
+ BLOCK_SIZE,
13
+ DEVICE,
14
+ _create_gdn_builder,
15
+)
16
+from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata
17
from vllm.config.compilation import CUDAGraphMode
18
from vllm.platforms import current_platform
19
from vllm.v1.attention.backends.recoverssm_metadata import (
@@ -59,6 +66,79 @@ def test_prepare_attn_forwards_positions(monkeypatch: pytest.MonkeyPatch) -> Non
66
assert build_attn_metadata.call_args.kwargs["positions"] is positions
67
68
69
+def test_padded_prompt_tail_builds_as_spec_decode(
70
+ monkeypatch: pytest.MonkeyPatch,
71
+) -> None:
72
+ """A one-token prompt tail over prior state, padded with K placeholder
73
+ drafts (e.g. a P/D decode-node arrival), must reach the GDN builder as a
74
+ spec-decode row. Built as a prefill, the placeholder tokens are folded into
75
+ the recurrent state and can't be rolled back.
76
+ """
77
+ k = 3
78
+ state = object.__new__(MambaHybridModelState)
79
+ state.vllm_config = SimpleNamespace(num_speculative_tokens=k)
80
+ state.model_config = SimpleNamespace(max_model_len=8192)
81
+ state._align_mode = False
82
+ state.recoverssm = None
83
+ state.num_accepted_tokens_gpu = torch.ones(4, dtype=torch.int32)
84
+
85
+ # A verify decode, a padded prompt tail (128 of 129 prompt tokens already
86
+ # computed), and a fresh prompt chunk of the same length.
87
+ query_lens = [k + 1] * 3
88
+ seq_lens = [50, 128 + k + 1, k + 1]
89
+ is_prefilling = [False, True, True]
90
+ query_start_loc = np.array([0, 4, 8, 12], dtype=np.int32)
91
+ input_batch = SimpleNamespace(
92
+ num_reqs=3,
93
+ num_tokens=12,
94
+ num_reqs_after_padding=3,
95
+ num_tokens_after_padding=12,
96
+ idx_mapping=torch.arange(3),
97
+ query_start_loc_np=query_start_loc,
98
+ query_start_loc=torch.from_numpy(query_start_loc),
99
+ num_scheduled_tokens=np.array(query_lens, dtype=np.int32),
100
+ num_draft_tokens_per_req=np.array([k, k, 0], dtype=np.int32),
101
+ max_query_len=None,
102
+ seq_lens_cpu_upper_bound=torch.tensor(seq_lens, dtype=torch.int32),
103
+ seq_lens=torch.tensor(seq_lens, dtype=torch.int32),
104
+ is_prefilling_np=np.array(is_prefilling),
105
+ prefill_len_np=np.array([40, 129, 100], dtype=np.int32),
106
+ num_computed_prefill_tokens_np=np.array([40, 128, 0], dtype=np.int32),
107
+ dcp_local_seq_lens=None,
108
+ positions=torch.zeros(12, dtype=torch.int64),
109
+ prompt_lens=None,
110
+ )
111
+ build_attn_metadata = Mock(return_value={})
112
+ monkeypatch.setattr(mamba_hybrid, "build_attn_metadata", build_attn_metadata)
113
+ state.prepare_attn(
114
+ input_batch=input_batch,
115
+ cudagraph_mode=CUDAGraphMode.NONE,
116
+ block_tables=(),
117
+ slot_mappings=torch.empty(0, dtype=torch.int64),
118
+ attn_groups=[],
119
+ kv_cache_config=Mock(),
120
+ )
121
+ mamba_metadata = build_attn_metadata.call_args.kwargs[
122
+ "model_specific_attn_metadata"
123
+ ]
124
+
125
+ builder = _create_gdn_builder(num_speculative_tokens=k)
126
+ common = create_common_attn_metadata(
127
+ BatchSpec(seq_lens=seq_lens, query_lens=query_lens), BLOCK_SIZE, DEVICE
128
+ ).replace(is_prefilling=torch.tensor(is_prefilling))
129
+ meta = builder.build(
130
+ common_prefix_len=0,
131
+ common_attn_metadata=common,
132
+ num_accepted_tokens=mamba_metadata.num_accepted_tokens,
133
+ num_decode_draft_tokens_cpu=mamba_metadata.num_decode_draft_tokens_cpu,
134
+ )
135
+
136
+ # Only the fresh prompt chunk needs the prefill kernels.
137
+ assert meta.num_spec_decodes == 2
138
+ assert meta.num_prefills == 1
139
+ assert meta.num_prefill_tokens == k + 1
140
+
141
+
142
@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA")
143
@pytest.mark.parametrize(("num_sampled", "expected_value"), [(0, 1), (3, 3)])
144
def test_postprocess_state_scalar_with_int32_mapping(
vllm/v1/worker/gpu/model_states/mamba_hybrid.py
+8
-1
@@ -281,7 +281,14 @@ class MambaHybridModelState(DefaultModelState):
281
# Test request state, not num_scheduled_tokens == draft_count+1:
282
# adaptive rewrites num_scheduled_tokens to an even split, so that
283
# equality rarely holds and would demote every verify row to decode.
284
- is_decode = (~input_batch.is_prefilling_np) & (
284
+ # A one-token prompt tail over prior state that the scheduler padded
285
+ # with placeholder drafts is also a spec-decode row: the prefill
286
+ # kernels can't roll the placeholders back.
287
+ num_computed = input_batch.num_computed_prefill_tokens_np
288
+ is_prompt_tail = (num_computed > 0) & (
289
+ input_batch.prefill_len_np - num_computed == 1
290
+ )
291
+ is_decode = (~input_batch.is_prefilling_np | is_prompt_tail) & (
292
input_batch.num_scheduled_tokens > 0
293
)
294
spec_decode_mask = (num_draft_tokens_per_req > 0) & is_decode