A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 18:30 UTC
[Bugfix][Async] Fix async spec decoding with hybrid models (#38556)
Signed-off-by: SandishKumarHN <sandishkumarhn@gmail.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: SandishKumarHN <sandishkumarhn@gmail.com>
Matthew Bonanni committed
Mar 31, 2026 at 11:08 UTC
757068dc65f6760fec1692ae587d8be13e100ee2
6 files changed
+177
-36
tests/v1/spec_decode/test_backup_token_async_spec.py
new
+147
@@ -0,0 +1,147 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Regression tests for the backup token fix in prepare_next_token_ids_padded.
4
+
5
+Fixes #38098: with async scheduling, seq_lens_cpu is inflated by unaccepted
6
+draft token placeholders, causing get_token_id() to return -1.
7
+"""
8
+
9
+from __future__ import annotations
10
+
11
+import numpy as np
12
+import pytest
13
+import torch
14
+
15
+
16
+class _FakeRequest:
17
+ def __init__(self, prompt_tokens: list[int], output_tokens: list[int]):
18
+ self.num_prompt_tokens = len(prompt_tokens)
19
+ self._prompt = prompt_tokens
20
+ self._output = output_tokens
21
+
22
+ @property
23
+ def num_tokens(self) -> int:
24
+ return self.num_prompt_tokens + len(self._output)
25
+
26
+ def get_token_id(self, idx: int) -> int:
27
+ if idx < self.num_prompt_tokens:
28
+ return self._prompt[idx]
29
+ out_idx = idx - self.num_prompt_tokens
30
+ if out_idx < len(self._output):
31
+ return self._output[out_idx]
32
+ return -1 # out of range
33
+
34
+
35
+class _FakeInputBatch:
36
+ def __init__(
37
+ self,
38
+ req_ids: list[str],
39
+ num_tokens_no_spec: list[int],
40
+ vocab_size: int = 32000,
41
+ ):
42
+ self.req_ids = req_ids
43
+ self.num_reqs = len(req_ids)
44
+ self.vocab_size = vocab_size
45
+ self.num_tokens_no_spec = np.array(num_tokens_no_spec, dtype=np.int64)
46
+
47
+
48
+def _make_requests(
49
+ req_ids: list[str],
50
+ prompt_lens: list[int],
51
+ output_lens: list[int],
52
+) -> dict[str, _FakeRequest]:
53
+ requests = {}
54
+ for rid, plen, olen in zip(req_ids, prompt_lens, output_lens):
55
+ requests[rid] = _FakeRequest(list(range(plen)), list(range(1000, 1000 + olen)))
56
+ return requests
57
+
58
+
59
+def _backup_buggy(
60
+ seq_lens_cpu: torch.Tensor,
61
+ requests: dict[str, _FakeRequest],
62
+ batch: _FakeInputBatch,
63
+) -> list[int]:
64
+ """Old logic: uses seq_lens_cpu directly (may be inflated)."""
65
+ n = batch.num_reqs
66
+ return [
67
+ requests[batch.req_ids[i]].get_token_id(int(seq_lens_cpu[i])) for i in range(n)
68
+ ]
69
+
70
+
71
+def _backup_fixed(
72
+ requests: dict[str, _FakeRequest],
73
+ batch: _FakeInputBatch,
74
+) -> list[int]:
75
+ """New logic: uses num_tokens_no_spec - 1 (last committed token)."""
76
+ n = batch.num_reqs
77
+ idx = (batch.num_tokens_no_spec[:n] - 1).tolist()
78
+ return [requests[batch.req_ids[i]].get_token_id(int(idx[i])) for i in range(n)]
79
+
80
+
81
+class TestBackupTokenAsyncSpec:
82
+ def test_no_inflation_fixed_returns_last_token(self):
83
+ req_ids = ["r0", "r1"]
84
+ requests = _make_requests(req_ids, [3, 3], [2, 2])
85
+ batch = _FakeInputBatch(req_ids, [5, 5])
86
+ # idx = 5-1 = 4 → output[1] = 1001
87
+ assert _backup_fixed(requests, batch) == [1001, 1001]
88
+
89
+ def test_inflation_buggy_returns_placeholder(self):
90
+ req_ids = ["r0", "r1"]
91
+ requests = _make_requests(req_ids, [3, 3], [2, 2])
92
+ batch = _FakeInputBatch(req_ids, [5, 5])
93
+ # inflated by 3 spec tokens → idx 8 is out of range
94
+ seq_lens = torch.tensor([8, 8], dtype=torch.int64)
95
+ assert _backup_buggy(seq_lens, requests, batch) == [-1, -1]
96
+
97
+ def test_inflation_fixed_returns_correct_token(self):
98
+ req_ids = ["r0", "r1"]
99
+ requests = _make_requests(req_ids, [3, 3], [2, 2])
100
+ batch = _FakeInputBatch(req_ids, [5, 5])
101
+ assert _backup_fixed(requests, batch) == [1001, 1001]
102
+
103
+ def test_mixed_inflation_per_request(self):
104
+ req_ids = ["r0", "r1", "r2"]
105
+ requests = {
106
+ "r0": _FakeRequest([0, 1], [1000, 1001, 1002]),
107
+ "r1": _FakeRequest([0, 1, 2, 3], [2000]),
108
+ "r2": _FakeRequest([0], [3000, 3001, 3002, 3003]),
109
+ }
110
+ batch = _FakeInputBatch(req_ids, [5, 5, 5])
111
+ seq_lens = torch.tensor([7, 9, 5], dtype=torch.int64)
112
+
113
+ assert _backup_buggy(seq_lens, requests, batch) == [-1, -1, -1]
114
+ assert _backup_fixed(requests, batch) == [1002, 2000, 3003]
115
+
116
+ def test_prefill_only_request(self):
117
+ """No output tokens yet — backup should be the last prompt token."""
118
+ req_ids = ["r0"]
119
+ requests = {"r0": _FakeRequest([10, 20, 30], [])}
120
+ batch = _FakeInputBatch(req_ids, [3])
121
+ # idx = 3-1 = 2 → prompt[2] = 30
122
+ assert _backup_fixed(requests, batch) == [30]
123
+
124
+ @pytest.mark.parametrize("num_spec_tokens", [1, 2, 3, 4, 5])
125
+ def test_various_spec_token_counts(self, num_spec_tokens: int):
126
+ req_ids = ["r0"]
127
+ requests = {"r0": _FakeRequest([0, 1, 2], list(range(1000, 1005)))}
128
+ batch = _FakeInputBatch(req_ids, [8])
129
+ # idx = 8-1 = 7 → output[4] = 1004
130
+ assert _backup_fixed(requests, batch) == [1004]
131
+
132
+ def test_buggy_code_was_always_off_by_one(self):
133
+ """The original code used seq_len as index, which is always one past
134
+ the end of output_token_ids even without async inflation."""
135
+ req_ids = ["r0"]
136
+ requests = {"r0": _FakeRequest([0, 1, 2], [1000, 1001])}
137
+ batch = _FakeInputBatch(req_ids, [5])
138
+
139
+ # no inflation: seq_len == num_tokens == 5 → idx 5 is out of range
140
+ seq_lens = torch.tensor([5], dtype=torch.int64)
141
+ assert _backup_buggy(seq_lens, requests, batch) == [-1]
142
+ assert _backup_fixed(requests, batch) == [1001]
143
+
144
+ # with inflation: still -1, fixed still correct
145
+ seq_lens_inf = torch.tensor([8], dtype=torch.int64)
146
+ assert _backup_buggy(seq_lens_inf, requests, batch) == [-1]
147
+ assert _backup_fixed(requests, batch) == [1001]
tests/v1/spec_decode/test_eagle.py
+4
-12
@@ -3,6 +3,7 @@
3
4
from unittest import mock
5
6
+import numpy as np
7
import pytest
8
import torch
9
@@ -127,16 +128,14 @@ def test_prepare_next_token_ids():
128
129
num_requests = 4
130
num_speculative_tokens = 4
130
- batch_spec = BatchSpec(
131
- seq_lens=[num_speculative_tokens + 1] * num_requests,
132
- query_lens=[num_speculative_tokens + 1] * num_requests,
133
- )
134
-
131
req_ids = [f"req_{i + 1}" for i in range(num_requests)]
132
mock_input_batch = mock.MagicMock(spec=InputBatch)
133
mock_input_batch.req_ids = req_ids
134
mock_input_batch.num_reqs = num_requests
135
mock_input_batch.vocab_size = 100
136
+ mock_input_batch.num_tokens_no_spec = np.array(
137
+ [num_speculative_tokens + 1] * num_requests
138
+ )
139
140
mock_num_scheduled_tokens = {req_id: 0 for req_id in req_ids}
141
mock_requests = {}
@@ -181,19 +180,12 @@ def test_prepare_next_token_ids():
180
181
assert torch.equal(next_token_ids_from_cpu, expected_next_token_ids_tensor)
182
184
- common_attn_metadata = create_common_attn_metadata(
185
- batch_spec,
186
- block_size=BLOCK_SIZE,
187
- device=device,
188
- )
189
-
183
expected_valid_sampled_tokens_count = torch.tensor(
184
[2, 5, 0, 0], dtype=torch.int32, device=device
185
)
186
187
next_token_ids_from_padded, valid_sampled_tokens_count = (
188
proposer.prepare_next_token_ids_padded(
196
- common_attn_metadata.seq_lens_cpu,
189
sampled_token_ids_tensor,
190
mock_requests,
191
mock_input_batch,
tests/v1/spec_decode/test_extract_hidden_states.py
+2
-12
@@ -3,6 +3,7 @@
3
4
from unittest import mock
5
6
+import numpy as np
7
import pytest
8
import torch
9
@@ -132,16 +133,12 @@ def test_prepare_next_token_ids_padded():
133
device = torch.device(current_platform.device_type)
134
135
num_requests = 4
135
- batch_spec = BatchSpec(
136
- seq_lens=[5] * num_requests,
137
- query_lens=[5] * num_requests,
138
- )
139
-
136
req_ids = [f"req_{i + 1}" for i in range(num_requests)]
137
mock_input_batch = mock.MagicMock(spec=InputBatch)
138
mock_input_batch.req_ids = req_ids
139
mock_input_batch.num_reqs = num_requests
140
mock_input_batch.vocab_size = 100
141
+ mock_input_batch.num_tokens_no_spec = np.array([5] * num_requests)
142
143
mock_requests = {}
144
for req_id in req_ids:
@@ -174,12 +171,6 @@ def test_prepare_next_token_ids_padded():
171
172
proposer = _create_proposer(num_speculative_tokens=1)
173
177
- common_attn_metadata = create_common_attn_metadata(
178
- batch_spec,
179
- block_size=16,
180
- device=device,
181
- )
182
-
174
# valid_sampled_tokens_count tracks if token is valid (not -1 and in vocab range)
175
# It doesn't depend on whether the request is discarded
176
expected_valid_sampled_tokens_count = torch.tensor(
@@ -187,7 +178,6 @@ def test_prepare_next_token_ids_padded():
178
)
179
180
next_token_ids, valid_sampled_tokens_count = proposer.prepare_next_token_ids_padded(
190
- common_attn_metadata.seq_lens_cpu,
181
sampled_token_ids,
182
mock_requests,
183
mock_input_batch,
vllm/v1/spec_decode/eagle.py
+1
-2
@@ -859,7 +859,6 @@ class SpecDecodeBaseProposer:
859
860
def prepare_next_token_ids_padded(
861
self,
862
- seq_lens_cpu: torch.Tensor,
862
sampled_token_ids: torch.Tensor,
863
requests: dict[str, CachedRequestState],
864
gpu_input_batch: InputBatch,
@@ -874,7 +873,7 @@ class SpecDecodeBaseProposer:
873
"""
874
# Precompute get_token_id for when there is no valid next token
875
num_reqs = gpu_input_batch.num_reqs
877
- seq_lens_list = seq_lens_cpu[:num_reqs].tolist()
876
+ seq_lens_list = (gpu_input_batch.num_tokens_no_spec[:num_reqs] - 1).tolist()
877
self.backup_next_token_ids.np[:num_reqs] = np.array(
878
[
879
requests[gpu_input_batch.req_ids[i]].get_token_id(seq_lens_list[i])
vllm/v1/spec_decode/extract_hidden_states.py
+1
-2
@@ -286,7 +286,6 @@ class ExtractHiddenStatesProposer:
286
287
def prepare_next_token_ids_padded(
288
self,
289
- seq_lens: torch.Tensor,
289
sampled_token_ids: torch.Tensor,
290
requests: dict[str, CachedRequestState],
291
gpu_input_batch: InputBatch,
@@ -303,7 +302,7 @@ class ExtractHiddenStatesProposer:
302
device = sampled_token_ids.device
303
304
# Compute backup tokens for discarded / invalid requests
306
- seq_lens_list = seq_lens[:num_reqs].tolist()
305
+ seq_lens_list = (gpu_input_batch.num_tokens_no_spec[:num_reqs] - 1).tolist()
306
backup_tokens_gpu = torch.tensor(
307
[
308
requests[gpu_input_batch.req_ids[i]].get_token_id(seq_lens_list[i])
vllm/v1/worker/gpu_model_runner.py
+22
-8
@@ -210,7 +210,7 @@ from .utils import (
210
if TYPE_CHECKING:
211
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
212
from vllm.v1.spec_decode.ngram_proposer import NgramProposer
213
- from vllm.v1.worker.encoder_cudagraph import EncoderCudaGraphManager
213
+ from vllm.v1.worker.gpu.mm.encoder_cudagraph import EncoderCudaGraphManager
214
215
logger = init_logger(__name__)
216
@@ -1939,9 +1939,24 @@ class GPUModelRunner(
1939
# _update_states_after_model_execute for hybrid models).
1940
if self.num_accepted_tokens_event is not None:
1941
self.num_accepted_tokens_event.synchronize()
1942
- self.num_accepted_tokens.np[:num_reqs] = (
1943
- self.input_batch.num_accepted_tokens_cpu[:num_reqs]
1944
- )
1942
+ # Async mode: condense() reordered indices, use prev_positions mapping
1943
+ if self.use_async_scheduling and prev_req_id_to_index:
1944
+ prev_idx = self.prev_positions.np[:num_reqs]
1945
+ new_mask = prev_idx < 0
1946
+ self.num_accepted_tokens.np[:num_reqs] = (
1947
+ self.input_batch.num_accepted_tokens_cpu[
1948
+ np.where(new_mask, 0, prev_idx)
1949
+ ]
1950
+ )
1951
+ self.num_accepted_tokens.np[:num_reqs][new_mask] = 1
1952
+ self.input_batch.num_accepted_tokens_cpu[:num_reqs] = (
1953
+ self.num_accepted_tokens.np[:num_reqs]
1954
+ )
1955
+ else:
1956
+ # Non-async mode: use values directly
1957
+ self.num_accepted_tokens.np[:num_reqs] = (
1958
+ self.input_batch.num_accepted_tokens_cpu[:num_reqs]
1959
+ )
1960
self.num_accepted_tokens.np[num_reqs:].fill(1)
1961
self.num_accepted_tokens.copy_to_gpu()
1962
else:
@@ -4220,7 +4235,6 @@ class GPUModelRunner(
4235
assert spec_decode_common_attn_metadata is not None
4236
next_token_ids, valid_sampled_tokens_count = (
4237
self.drafter.prepare_next_token_ids_padded(
4223
- self.optimistic_seq_lens_cpu,
4238
sampled_token_ids,
4239
self.requests,
4240
self.input_batch,
@@ -4587,7 +4601,6 @@ class GPUModelRunner(
4601
)
4602
next_token_ids, valid_sampled_tokens_count = (
4603
self.drafter.prepare_next_token_ids_padded(
4590
- self.optimistic_seq_lens_cpu,
4604
sampled_token_ids,
4605
self.requests,
4606
self.input_batch,
@@ -4632,7 +4645,6 @@ class GPUModelRunner(
4645
)
4646
next_token_ids, valid_sampled_tokens_count = (
4647
self.drafter.prepare_next_token_ids_padded(
4635
- self.optimistic_seq_lens_cpu,
4648
sampled_token_ids,
4649
self.requests,
4650
self.input_batch,
@@ -5990,7 +6002,9 @@ class GPUModelRunner(
6002
SupportsEncoderCudaGraph,
6003
supports_encoder_cudagraph,
6004
)
5993
- from vllm.v1.worker.encoder_cudagraph import EncoderCudaGraphManager
6005
+ from vllm.v1.worker.gpu.mm.encoder_cudagraph import (
6006
+ EncoderCudaGraphManager,
6007
+ )
6008
6009
raw_model = self.get_model()
6010
if supports_encoder_cudagraph(raw_model):