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 12, 06:00 UTC
[BugFix] Fix index error in ngram_proposer (#29779)
Signed-off-by: Bradley <bradley.b.pitt@gmail.com>
usberkeley committed
Dec 2, 2025 at 12:48 UTC
81fe3f82af2e30c93ddb106b7a84b7730b982e66
2 files changed
+35
-3
tests/v1/spec_decode/test_ngram.py
+33
-1
@@ -2,7 +2,11 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import numpy as np
4
5
-from vllm.config import ModelConfig, SpeculativeConfig, VllmConfig
5
+from vllm.config import (
6
+ ModelConfig,
7
+ SpeculativeConfig,
8
+ VllmConfig,
9
+)
10
from vllm.v1.spec_decode.ngram_proposer import (
11
NgramProposer,
12
_find_longest_matched_ngram_and_propose_tokens,
@@ -167,6 +171,34 @@ def test_ngram_proposer():
171
assert np.array_equal(result[0], np.array([3, 1]))
172
assert np.array_equal(result[1], np.array([]))
173
174
+ # Test non-contiguous indices: requests 0 and 2 need proposals,
175
+ # request 1 is in prefill
176
+ proposer = get_ngram_proposer(min_n=2, max_n=2, k=2)
177
+ max_model_len = 20
178
+ token_ids_cpu = np.zeros((3, max_model_len), dtype=np.int32)
179
+ token_ids_cpu[0, :5] = [1, 2, 3, 1, 2]
180
+ token_ids_cpu[1, :3] = [4, 5, 6]
181
+ token_ids_cpu[2, :5] = [7, 8, 9, 7, 8]
182
+ num_tokens_no_spec = np.array([5, 3, 5], dtype=np.int32)
183
+ sampled_token_ids = [[2], [], [8]] # Empty list for request 1 simulates prefill
184
+ result = proposer.propose(
185
+ sampled_token_ids=sampled_token_ids,
186
+ req_ids=["0", "1", "2"],
187
+ num_tokens_no_spec=num_tokens_no_spec,
188
+ token_ids_cpu=token_ids_cpu,
189
+ spec_decode_unsupported_reqs=(),
190
+ )
191
+ assert len(result) == 3
192
+ assert np.array_equal(result[0], [3, 1])
193
+ assert len(result[1]) == 0
194
+ assert np.array_equal(result[2], [9, 7])
195
+ # Verify internal arrays written to correct indices
196
+ assert proposer.valid_ngram_num_drafts[0] == 2
197
+ assert proposer.valid_ngram_num_drafts[1] == 0
198
+ assert proposer.valid_ngram_num_drafts[2] == 2
199
+ assert np.array_equal(proposer.valid_ngram_draft[0, :2], [3, 1])
200
+ assert np.array_equal(proposer.valid_ngram_draft[2, :2], [9, 7])
201
+
202
# test if 0 threads available: can happen if TP size > CPU count
203
ngram_proposer = get_ngram_proposer(min_n=2, max_n=2, k=2)
204
ngram_proposer.num_numba_thread_available = 0
vllm/v1/spec_decode/ngram_proposer.py
+2
-2
@@ -196,9 +196,9 @@ def batch_propose_numba(
196
k=k,
197
)
198
199
- valid_ngram_num_drafts[i] = drafter_output.shape[0]
199
+ valid_ngram_num_drafts[idx] = drafter_output.shape[0]
200
if len(drafter_output):
201
- valid_ngram_draft[i, : drafter_output.shape[0]] = drafter_output
201
+ valid_ngram_draft[idx, : drafter_output.shape[0]] = drafter_output
202
203
204
@jit(nopython=True)