@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / ac86bff8cb

Fork 0

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 13, 04:30 UTC

Revert "[Core] Performance: Use list[np.ndarray] instead of list[list… (#28773)

Nick Hill committed Nov 14, 2025 at 20:24 UTC ac86bff8cb53939117a6a460af1a6c3fea829a56

12 files changed +76 -102

tests/v1/core/test_async_scheduler.py

+1 -2

@@ -2,7 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from collections import deque

5 -import numpy as np

5 import pytest

7 from vllm.v1.core.sched.output import SchedulerOutput

@@ -22,7 +21,7 @@ def _make_model_runner_output(

21 return ModelRunnerOutput(

22 req_ids=req_ids,

23 req_id_to_index={req_id: i for i, req_id in enumerate(req_ids)},

25 - sampled_token_ids=[np.array([i]) for i in range(len(req_ids))],

24 + sampled_token_ids=[[i] for i in range(len(req_ids))],

25 logprobs=None,

26 prompt_logprobs_dict={},

27 pooler_output=[],

tests/v1/core/test_scheduler.py

+32 -44

@@ -3,7 +3,6 @@

3 import dataclasses

4 from unittest.mock import Mock

6 -import numpy as np

6 import pytest

7 import torch

@@ -170,7 +169,7 @@ def test_schedule_partial_requests():

169 req_id_to_index=req_to_index,

170 # Only the first request has a sampled token id because

171 # the rest requests are still being prefilled.

173 - sampled_token_ids=[np.array([0]), np.array([]), np.array([])],

172 + sampled_token_ids=[[0], [], []],

173 logprobs=None,

174 prompt_logprobs_dict={},

175 pooler_output=[],

@@ -217,7 +216,7 @@ def test_no_mm_input_chunking():

216 model_runner_output = ModelRunnerOutput(

217 req_ids=[request.request_id for request in requests],

218 req_id_to_index=req_to_index,

220 - sampled_token_ids=[np.array([]) for _ in range(len(requests))],

219 + sampled_token_ids=[[] for _ in range(len(requests))],

220 logprobs=None,

221 prompt_logprobs_dict={},

222 pooler_output=[],

@@ -277,7 +276,7 @@ def test_schedule_concurrent_partial_requests(enable_prefix_caching: bool):

276 model_runner_output = ModelRunnerOutput(

277 req_ids=[request.request_id for request in requests],

278 req_id_to_index=req_to_index,

280 - sampled_token_ids=[np.array([]) for _ in range(len(requests))],

279 + sampled_token_ids=[[] for _ in range(len(requests))],

280 logprobs=None,

281 prompt_logprobs_dict={},

282 pooler_output=[],

@@ -301,8 +300,7 @@ def test_schedule_concurrent_partial_requests(enable_prefix_caching: bool):

300 model_runner_output = ModelRunnerOutput(

301 req_ids=[request.request_id for request in requests],

302 req_id_to_index=req_to_index,

304 - sampled_token_ids=[np.array([0]), np.array([0])]

305 - + [np.array([]) for _ in range(len(requests) - 2)],

303 + sampled_token_ids=[[0], [0]] + [[] for _ in range(len(requests) - 2)],

304 logprobs=None,

305 prompt_logprobs_dict={},

306 pooler_output=[],

@@ -349,8 +347,8 @@ def test_stop_via_update_from_output():

347 req_ids=[req.request_id for req in requests],

348 req_id_to_index={req.request_id: i for i, req in enumerate(requests)},

349 sampled_token_ids=[

352 - np.array([EOS_TOKEN_ID]),

353 - np.array([10, 11]),

350 + [EOS_TOKEN_ID],

351 + [10, 11],

352 ], # First request hits EOS, second continues

353 logprobs=None,

354 prompt_logprobs_dict={},

@@ -394,10 +392,7 @@ def test_stop_via_update_from_output():

392 model_output = ModelRunnerOutput(

393 req_ids=[req.request_id for req in requests],

394 req_id_to_index={req.request_id: i for i, req in enumerate(requests)},

397 - sampled_token_ids=[

398 - np.array([10, 42, 12]),

399 - np.array([13, 14]),

400 - ], # First request hits stop token

395 + sampled_token_ids=[[10, 42, 12], [13, 14]], # First request hits stop token

396 logprobs=None,

397 prompt_logprobs_dict={},

398 pooler_output=[],

@@ -441,10 +436,7 @@ def test_stop_via_update_from_output():

436 model_output = ModelRunnerOutput(

437 req_ids=[req.request_id for req in requests],

438 req_id_to_index={req.request_id: i for i, req in enumerate(requests)},

444 - sampled_token_ids=[

445 - np.array([10, 11, 12]),

446 - np.array([13]),

447 - ], # First request exceeds max_tokens

439 + sampled_token_ids=[[10, 11, 12], [13]], # First request exceeds max_tokens

440 logprobs=None,

441 prompt_logprobs_dict={},

442 pooler_output=[],

@@ -483,7 +475,7 @@ def test_stop_via_update_from_output():

475 model_output = ModelRunnerOutput(

476 req_ids=[requests[0].request_id],

477 req_id_to_index={requests[0].request_id: 0},

486 - sampled_token_ids=[np.array([EOS_TOKEN_ID, 10, 11])],

478 + sampled_token_ids=[[EOS_TOKEN_ID, 10, 11]],

479 logprobs=None,

480 prompt_logprobs_dict={},

481 pooler_output=[],

@@ -624,7 +616,7 @@ def test_schedule_concurrent_batches(

616 model_runner_output = ModelRunnerOutput(

617 req_ids=[requests[0].request_id],

618 req_id_to_index={requests[0].request_id: 0},

627 - sampled_token_ids=[np.array([0])],

619 + sampled_token_ids=[[0]],

620 logprobs=None,

621 prompt_logprobs_dict={},

622 pooler_output=[],

@@ -641,7 +633,7 @@ def test_schedule_concurrent_batches(

633 model_runner_output = ModelRunnerOutput(

634 req_ids=[requests[1].request_id],

635 req_id_to_index={requests[1].request_id: 0},

644 - sampled_token_ids=[np.array([0])],

636 + sampled_token_ids=[[0]],

637 logprobs=None,

638 prompt_logprobs_dict={},

639 pooler_output=[],

@@ -678,7 +670,7 @@ def test_preempt_during_execution():

670 model_runner_output0 = ModelRunnerOutput(

671 req_ids=[requests[0].request_id],

672 req_id_to_index={requests[0].request_id: 0},

681 - sampled_token_ids=[np.array([0])],

673 + sampled_token_ids=[[0]],

674 logprobs=None,

675 prompt_logprobs_dict={},

676 pooler_output=[],

@@ -695,7 +687,7 @@ def test_preempt_during_execution():

687 model_runner_output1 = ModelRunnerOutput(

688 req_ids=[requests[1].request_id],

689 req_id_to_index={requests[1].request_id: 0},

698 - sampled_token_ids=[np.array([42])],

690 + sampled_token_ids=[[42]],

691 logprobs=None,

692 prompt_logprobs_dict={},

693 pooler_output=[],

@@ -712,18 +704,14 @@ def test_preempt_during_execution():

704 @pytest.mark.parametrize(

705 "spec_tokens,output_tokens,expected",

706 [

715 - ([[1, 2, 3]], [np.array([1, 2, 3, 4])], (1, 3, 3, [1, 1, 1])), # perfect match

716 - ([[1, 2, 3]], [np.array([1, 5])], (1, 3, 1, [1, 0, 0])), # early mismatch

717 - (

718 - [[1, 2], [3]],

719 - [np.array([1, 2, 5]), np.array([3, 4])],

720 - (2, 3, 3, [2, 1]),

721 - ), # multiple sequences

722 - ([[1]], [np.array([1, 2])], (1, 1, 1, [1])), # single token sequence

723 - ([[]], [np.array([5])], (0, 0, 0, [0])), # empty sequence

707 + ([[1, 2, 3]], [[1, 2, 3, 4]], (1, 3, 3, [1, 1, 1])), # perfect match

708 + ([[1, 2, 3]], [[1, 5]], (1, 3, 1, [1, 0, 0])), # early mismatch

709 + ([[1, 2], [3]], [[1, 2, 5], [3, 4]], (2, 3, 3, [2, 1])), # multiple sequences

710 + ([[1]], [[1, 2]], (1, 1, 1, [1])), # single token sequence

711 + ([[]], [[5]], (0, 0, 0, [0])), # empty sequence

712 (

713 [[1, 2, 3], [4, 5, 6]],

726 - [np.array([1, 2, 7]), np.array([4, 8])],

714 + [[1, 2, 7], [4, 8]],

715 (2, 6, 3, [2, 1, 0]),

716 ), # multiple mismatches

717 ],

@@ -757,7 +745,7 @@ def test_schedule_spec_decoding_stats(spec_tokens, output_tokens, expected):

745 model_runner_output = ModelRunnerOutput(

746 req_ids=req_ids,

747 req_id_to_index=req_to_index,

760 - sampled_token_ids=[np.array([0]) for _ in range(len(requests))],

748 + sampled_token_ids=[[0] for _ in range(len(requests))],

749 logprobs=None,

750 prompt_logprobs_dict={},

751 pooler_output=[],

@@ -984,7 +972,7 @@ def test_kv_connector_basic(is_async: bool):

972 MODEL_RUNNER_OUTPUT = ModelRunnerOutput(

973 req_ids=req_ids,

974 req_id_to_index=req_to_index,

987 - sampled_token_ids=[np.array([1000])] * len(req_ids),

975 + sampled_token_ids=[[1000]] * len(req_ids),

976 logprobs=None,

977 prompt_logprobs_dict={},

978 pooler_output=[],

@@ -1037,7 +1025,7 @@ def test_kv_connector_basic(is_async: bool):

1025 MODEL_RUNNER_OUTPUT = ModelRunnerOutput(

1026 req_ids=req_ids,

1027 req_id_to_index=req_to_index,

1040 - sampled_token_ids=[np.array([1000])] * len(req_ids),

1028 + sampled_token_ids=[[1000]] * len(req_ids),

1029 logprobs=None,

1030 prompt_logprobs_dict={},

1031 pooler_output=[],

@@ -1100,7 +1088,7 @@ def test_external_prefix_cache_metrics():

1088 MODEL_RUNNER_OUTPUT = ModelRunnerOutput(

1089 req_ids=[r.request_id for r in requests],

1090 req_id_to_index={r.request_id: i for i, r in enumerate(requests)},

1103 - sampled_token_ids=[np.array([1000])] * NUM_REQUESTS,

1091 + sampled_token_ids=[[1000]] * NUM_REQUESTS,

1092 logprobs=None,

1093 prompt_logprobs_dict={},

1094 pooler_output=[],

@@ -1166,7 +1154,7 @@ def test_kv_connector_unable_to_allocate(use_ec_connector, ec_role):

1154 MODEL_RUNNER_OUTPUT = ModelRunnerOutput(

1155 req_ids=req_ids,

1156 req_id_to_index=req_to_index,

1169 - sampled_token_ids=[np.array([1000])] * len(req_ids),

1157 + sampled_token_ids=[[1000]] * len(req_ids),

1158 logprobs=None,

1159 prompt_logprobs_dict={},

1160 pooler_output=[],

@@ -1251,7 +1239,7 @@ def test_kv_connector_handles_preemption(use_ec_connector, ec_role):

1239 MODEL_RUNNER_OUTPUT = ModelRunnerOutput(

1240 req_ids=req_ids,

1241 req_id_to_index=req_to_index,

1254 - sampled_token_ids=[np.array([1000])] * len(req_ids),

1242 + sampled_token_ids=[[1000]] * len(req_ids),

1243 logprobs=None,

1244 prompt_logprobs_dict={},

1245 pooler_output=[],

@@ -1344,7 +1332,7 @@ def make_output(scheduler: Scheduler):

1332 return ModelRunnerOutput(

1333 req_ids=[req.request_id for req in scheduler.running],

1334 req_id_to_index={req.request_id: i for i, req in enumerate(scheduler.running)},

1347 - sampled_token_ids=[np.array([1000])] * len(scheduler.running),

1335 + sampled_token_ids=[[1000]] * len(scheduler.running),

1336 logprobs=None,

1337 prompt_logprobs_dict={},

1338 pooler_output=[],

@@ -1761,7 +1749,7 @@ def test_priority_scheduling_preemption():

1749 req_id_to_index={

1750 req.request_id: i for i, req in enumerate(low_priority_requests)

1751 },

1764 - sampled_token_ids=[np.array([100]) for _ in low_priority_requests],

1752 + sampled_token_ids=[[100] for _ in low_priority_requests],

1753 logprobs=None,

1754 prompt_logprobs_dict={},

1755 pooler_output=[],

@@ -1830,7 +1818,7 @@ def test_priority_scheduling_no_preemption_when_space_available():

1818 req_id_to_index={

1819 req.request_id: i for i, req in enumerate(low_priority_requests)

1820 },

1833 - sampled_token_ids=[np.array([100]) for _ in low_priority_requests],

1821 + sampled_token_ids=[[100] for _ in low_priority_requests],

1822 logprobs=None,

1823 prompt_logprobs_dict={},

1824 pooler_output=[],

@@ -2076,7 +2064,7 @@ def test_priority_scheduling_heap_property():

2064 model_output = ModelRunnerOutput(

2065 req_ids=[req.req_id],

2066 req_id_to_index={req.req_id: 0},

2079 - sampled_token_ids=[np.array([100])],

2067 + sampled_token_ids=[[100]],

2068 logprobs=None,

2069 prompt_logprobs_dict={},

2070 pooler_output=[],

@@ -2162,7 +2150,7 @@ def test_priority_scheduling_preemption_and_resumption_when_out_of_kv(

2150 model_output = ModelRunnerOutput(

2151 req_ids=[request_low.request_id],

2152 req_id_to_index={request_low.request_id: 0},

2165 - sampled_token_ids=[np.array([100])],

2153 + sampled_token_ids=[[100]],

2154 # spec_token_ids=None,

2155 logprobs=None,

2156 prompt_logprobs_dict={},

@@ -2193,7 +2181,7 @@ def test_priority_scheduling_preemption_and_resumption_when_out_of_kv(

2181 model_output = ModelRunnerOutput(

2182 req_ids=[req.request_id for req in requests],

2183 req_id_to_index={req.request_id: i for i, req in enumerate(requests)},

2196 - sampled_token_ids=[np.array([100]) for _ in requests],

2184 + sampled_token_ids=[[100] for _ in requests],

2185 # spec_token_ids=None,

2186 logprobs=None,

2187 prompt_logprobs_dict={},

@@ -2219,7 +2207,7 @@ def test_priority_scheduling_preemption_and_resumption_when_out_of_kv(

2207 model_output = ModelRunnerOutput(

2208 req_ids=[req.request_id for req in requests],

2209 req_id_to_index={req.request_id: i for i, req in enumerate(requests)},

2222 - sampled_token_ids=[np.array([]), np.array([100])],

2210 + sampled_token_ids=[[], [100]],

2211 # spec_token_ids=None,

2212 logprobs=None,

2213 prompt_logprobs_dict={},

tests/v1/kv_connector/unit/utils.py

+1 -2

@@ -7,7 +7,6 @@ from dataclasses import dataclass

7 from itertools import chain, count

8 from typing import Any

10 -import numpy as np

10 import torch

12 from vllm import SamplingParams

@@ -229,7 +228,7 @@ def create_model_runner_output(

228

229 # Make sampled tokens.

230 sampled_token = EOS_TOKEN_ID if use_eos else token_id

232 - sampled_token_ids = [np.array([sampled_token]) for _ in req_ids]

231 + sampled_token_ids = [[sampled_token] for _ in req_ids]

232

233 kv_connector_output = (

234 None

tests/v1/spec_decode/test_eagle.py

+1 -4

@@ -3,7 +3,6 @@

4 from unittest import mock

6 -import numpy as np

6 import pytest

7 import torch

@@ -113,9 +112,7 @@ def test_prepare_next_token_ids():

112 sampled_token_ids_tensor = torch.tensor(

113 sampled_token_ids, dtype=torch.int32, device=device

114 )

116 - sampled_token_ids_cpu = [

117 - np.array([i for i in seq if i != -1]) for seq in sampled_token_ids

118 - ]

115 + sampled_token_ids_cpu = [[i for i in seq if i != -1] for seq in sampled_token_ids]

116

117 expected_next_token_ids_cpu = [1, 4, 30, 40]

118 expected_next_token_ids_tensor = torch.tensor(

tests/v1/spec_decode/test_ngram.py

+9 -9

@@ -77,7 +77,7 @@ def test_ngram_proposer():

77 # No match.

78 token_ids_cpu = np.array([[1, 2, 3, 4, 5]])

79 result = get_ngram_proposer(min_n=2, max_n=2, k=2).propose(

80 - sampled_token_ids=[np.array([0])],

80 + sampled_token_ids=[[0]],

81 req_ids=["0"],

82 num_tokens_no_spec=np.array([len(c) for c in token_ids_cpu]),

83 token_ids_cpu=token_ids_cpu,

@@ -88,7 +88,7 @@ def test_ngram_proposer():

88 # No match for 4-gram.

89 token_ids_cpu = np.array([[1, 2, 3, 4, 1, 2, 3]])

90 result = get_ngram_proposer(min_n=4, max_n=4, k=2).propose(

91 - sampled_token_ids=[np.array([0])],

91 + sampled_token_ids=[[0]],

92 req_ids=["0"],

93 num_tokens_no_spec=np.array([len(c) for c in token_ids_cpu]),

94 token_ids_cpu=token_ids_cpu,

@@ -99,7 +99,7 @@ def test_ngram_proposer():

99 # No match for 4-gram but match for 3-gram.

100 token_ids_cpu = np.array([[1, 2, 3, 4, 1, 2, 3]])

101 result = get_ngram_proposer(min_n=3, max_n=4, k=2).propose(

102 - sampled_token_ids=[np.array([0])],

102 + sampled_token_ids=[[0]],

103 req_ids=["0"],

104 num_tokens_no_spec=np.array([len(c) for c in token_ids_cpu]),

105 token_ids_cpu=token_ids_cpu,

@@ -111,7 +111,7 @@ def test_ngram_proposer():

111 # In this case, the proposer should return the 4-gram match.

112 token_ids_cpu = np.array([[2, 3, 4, 5, 1, 2, 3, 4, 1, 2, 3, 4]])

113 result = get_ngram_proposer(min_n=3, max_n=4, k=2).propose(

114 - sampled_token_ids=[np.array([0])],

114 + sampled_token_ids=[[0]],

115 req_ids=["0"],

116 num_tokens_no_spec=np.array([len(c) for c in token_ids_cpu]),

117 token_ids_cpu=token_ids_cpu,

@@ -122,7 +122,7 @@ def test_ngram_proposer():

122 # Match for 2-gram and 3-gram, but not 4-gram.

123 token_ids_cpu = np.array([[3, 4, 5, 2, 3, 4, 1, 2, 3, 4]])

124 result = get_ngram_proposer(min_n=2, max_n=4, k=2).propose(

125 - sampled_token_ids=[np.array([0])],

125 + sampled_token_ids=[[0]],

126 req_ids=["0"],

127 num_tokens_no_spec=np.array([len(c) for c in token_ids_cpu]),

128 token_ids_cpu=token_ids_cpu,

@@ -133,7 +133,7 @@ def test_ngram_proposer():

133 # Multiple 3-gram matched, but always pick the first one.

134 token_ids_cpu = np.array([[1, 2, 3, 100, 1, 2, 3, 200, 1, 2, 3, 300, 1, 2, 3]])

135 result = get_ngram_proposer(min_n=3, max_n=3, k=2).propose(

136 - sampled_token_ids=[np.array([0])],

136 + sampled_token_ids=[[0]],

137 req_ids=["0"],

138 num_tokens_no_spec=np.array([len(c) for c in token_ids_cpu]),

139 token_ids_cpu=token_ids_cpu,

@@ -144,7 +144,7 @@ def test_ngram_proposer():

144 # check empty input

145 token_ids_cpu = np.array([[]])

146 result = get_ngram_proposer(min_n=2, max_n=2, k=2).propose(

147 - sampled_token_ids=[np.array([0])],

147 + sampled_token_ids=[[0]],

148 req_ids=["0"],

149 num_tokens_no_spec=np.array([len(c) for c in token_ids_cpu]),

150 token_ids_cpu=token_ids_cpu,

@@ -157,7 +157,7 @@ def test_ngram_proposer():

157 # second request has 3 tokens and no match. Padded with -1 for max len 5

158 token_ids_cpu = np.array([[1, 2, 3, 1, 2], [4, 5, 6, -1, -1]])

159 result = get_ngram_proposer(min_n=2, max_n=2, k=2).propose(

160 - sampled_token_ids=[np.array([0]), np.array([1])],

160 + sampled_token_ids=[[0], [1]],

161 req_ids=["0", "1"],

162 num_tokens_no_spec=np.array([5, 3]),

163 token_ids_cpu=token_ids_cpu,

@@ -181,7 +181,7 @@ def test_ngram_proposer():

181 input_2[:3] = [4, 5, 6]

182 token_ids_cpu = np.array([input_1, input_2])

183 result = ngram_proposer.propose(

184 - sampled_token_ids=[np.array([0]), np.array([1])],

184 + sampled_token_ids=[[0], [1]],

185 req_ids=["0", "1"],

186 num_tokens_no_spec=np.array([len(input_1), 3]),

187 token_ids_cpu=token_ids_cpu,

vllm/v1/core/sched/scheduler.py

+2 -2

@@ -1010,8 +1010,8 @@ class Scheduler(SchedulerInterface):

1010 continue

1011

1012 req_index = model_runner_output.req_id_to_index[req_id]

1013 - generated_token_ids: list[int] = (

1014 - sampled_token_ids[req_index].tolist() if sampled_token_ids else []

1013 + generated_token_ids = (

1014 + sampled_token_ids[req_index] if sampled_token_ids else []

1015 )

1016

1017 scheduled_spec_token_ids = (

vllm/v1/outputs.py

+1 -1

@@ -158,7 +158,7 @@ class ModelRunnerOutput:

158 # num_generated_tokens is the number of tokens

159 # generated in the current step. It can be different for

160 # each request due to speculative/jump decoding.

161 - sampled_token_ids: list[np.ndarray]

161 + sampled_token_ids: list[list[int]]

162

163 # [num_reqs, max_num_logprobs + 1]

164 # [num_reqs, max_num_logprobs + 1]

vllm/v1/sample/rejection_sampler.py

+5 -3

@@ -3,7 +3,6 @@

4 from dataclasses import replace

6 -import numpy as np

6 import torch

7 import torch.nn as nn

@@ -205,7 +204,7 @@ class RejectionSampler(nn.Module):

204 def parse_output(

205 output_token_ids: torch.Tensor,

206 vocab_size: int,

208 - ) -> list[np.ndarray]:

207 + ) -> list[list[int]]:

208 """Parse the output of the rejection sampler.

209 Args:

210 output_token_ids: The sampled token IDs in shape

@@ -221,7 +220,10 @@ class RejectionSampler(nn.Module):

220 valid_mask = (output_token_ids_np != PLACEHOLDER_TOKEN_ID) & (

221 output_token_ids_np < vocab_size

222 )

224 - return [row[valid_mask[i]] for i, row in enumerate(output_token_ids_np)]

223 + outputs = [

224 + row[valid_mask[i]].tolist() for i, row in enumerate(output_token_ids_np)

225 + ]

226 + return outputs

227

228 def apply_logits_processors(

229 self,

vllm/v1/spec_decode/eagle.py

+4 -3

@@ -484,7 +484,7 @@ class EagleProposer:

484

485 def prepare_next_token_ids_cpu(

486 self,

487 - sampled_token_ids: list[np.ndarray],

487 + sampled_token_ids: list[list[int]],

488 requests: dict[str, CachedRequestState],

489 gpu_input_batch: InputBatch,

490 num_scheduled_tokens: dict[str, int],

@@ -499,7 +499,7 @@ class EagleProposer:

499 req_ids = gpu_input_batch.req_ids

500 next_token_ids: list[int] = []

501 for i, token_ids in enumerate(sampled_token_ids):

502 - if token_ids.shape[0] > 0:

502 + if token_ids:

503 # Common case.

504 next_token_id = token_ids[-1]

505 else:

@@ -510,9 +510,10 @@ class EagleProposer:

510 seq_len = req_state.num_computed_tokens + num_scheduled_tokens[req_id]

511 next_token_id = req_state.get_token_id(seq_len)

512 next_token_ids.append(next_token_id)

513 - return torch.tensor(

513 + next_token_ids = torch.tensor(

514 next_token_ids, dtype=torch.int32, device=self.input_ids.device

515 )

516 + return next_token_ids

517

518 def prepare_next_token_ids_padded(

519 self,

vllm/v1/spec_decode/ngram_proposer.py

+3 -3

@@ -54,7 +54,7 @@ class NgramProposer:

54 # Trigger Numba JIT compilation for N-gram proposer.

55 # This usually takes less than 1 second.

56 self.propose(

57 - [np.array([])] * 1024,

57 + [[]] * 1024,

58 [""] * 1024,

59 np.zeros(1024, dtype=np.int32),

60 np.zeros((1024, self.max_model_len), dtype=np.int32),

@@ -131,7 +131,7 @@ class NgramProposer:

131

132 def propose(

133 self,

134 - sampled_token_ids: list[np.ndarray],

134 + sampled_token_ids: list[list[int]],

135 req_ids: list[str],

136 num_tokens_no_spec: np.ndarray,

137 token_ids_cpu: np.ndarray,

@@ -140,7 +140,7 @@ class NgramProposer:

140 # find which requests need ngram proposals

141 valid_ngram_requests = []

142 for i, sampled_ids in enumerate(sampled_token_ids):

143 - num_sampled_ids = sampled_ids.shape[0]

143 + num_sampled_ids = len(sampled_ids)

144 if not num_sampled_ids:

145 # Skip speculative decoding.

146 continue

vllm/v1/spec_decode/suffix_decoding.py

+4 -6

@@ -1,7 +1,5 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -import numpy as np

4 -

3 from vllm.config import VllmConfig

4 from vllm.v1.worker.gpu_input_batch import InputBatch

@@ -34,16 +32,16 @@ class SuffixDecodingProposer:

32 def propose(

33 self,

34 input_batch: InputBatch,

37 - sampled_token_ids: list[np.ndarray],

35 + sampled_token_ids: list[list[int]],

36 ) -> list[list[int]]:

37 """

38 Propose speculative tokens for each request in the input batch. Suffix Decoding

39 will speculate a dynamic number of tokens for each request every decoding step,

40 so each entry in the returned list may have different lengths.

41 """

44 - draft_token_ids: list[np.ndarray] = []

42 + draft_token_ids: list[list[int]] = []

43 for i, sampled_ids in enumerate(sampled_token_ids):

46 - if sampled_ids.shape[0] == 0:

44 + if not sampled_ids:

45 # Skip speculative decoding for partial prefills.

46 draft_token_ids.append([])

47 continue

@@ -72,7 +70,7 @@ class SuffixDecodingProposer:

70 self.suffix_cache.start_request(req_id, prompt_token_ids)

72 # Append the newly sampled ids to the suffix cache for this request.

75 - self.suffix_cache.add_active_response(req_id, sampled_ids.tolist())

73 + self.suffix_cache.add_active_response(req_id, sampled_ids)

75 # Suffix decoding only uses the most recent tokens up to max_tree_depth, so

76 # we extract the pattern from the end of the input.

vllm/v1/worker/gpu_model_runner.py

+13 -23

@@ -216,11 +216,9 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput):

216 del self._logprobs_tensors

217 del self._sampled_token_ids

218

219 - valid_sampled_token_ids: list[np.ndarray] = [

220 - row for row in self.sampled_token_ids_cpu.numpy()

221 - ]

219 + valid_sampled_token_ids = self.sampled_token_ids_cpu.tolist()

220 for i in self._invalid_req_indices:

223 - valid_sampled_token_ids[i] = np.array([])

221 + valid_sampled_token_ids[i].clear()

222

223 output = self._model_runner_output

224 output.sampled_token_ids = valid_sampled_token_ids

@@ -2341,7 +2339,7 @@ class GPUModelRunner(

2339 ) -> tuple[

2340 dict[str, int],

2341 LogprobsLists | None,

2344 - list[np.ndarray],

2342 + list[list[int]],

2343 dict[str, LogprobsTensors | None],

2344 list[str],

2345 dict[str, int],

@@ -2367,7 +2365,6 @@ class GPUModelRunner(

2365 num_sampled_tokens = sampler_output.sampled_token_ids.shape[0]

2366 sampled_token_ids = sampler_output.sampled_token_ids

2367 invalid_req_indices = []

2370 - valid_sampled_token_ids: list[np.ndarray]

2368 if not self.use_async_scheduling:

2369 # Get the valid generated tokens.

2370 max_gen_len = sampled_token_ids.shape[-1]

@@ -2382,7 +2379,7 @@ class GPUModelRunner(

2379 )

2380 # Mask out the sampled tokens that should not be sampled.

2381 for i in discard_sampled_tokens_req_indices:

2385 - valid_sampled_token_ids[int(i)] = np.array([])

2382 + valid_sampled_token_ids[int(i)].clear()

2383 else:

2384 valid_sampled_token_ids = []

2385 invalid_req_indices = discard_sampled_tokens_req_indices.tolist()

@@ -2410,24 +2407,19 @@ class GPUModelRunner(

2407 [0] if spec_decode_metadata and logprobs_tensors else None

2408 )

2409 for req_idx in range(num_sampled_tokens):

2413 - sampled_ids: np.ndarray | None

2410 if self.use_async_scheduling:

2415 - sampled_ids = (

2416 - np.array([-1]) if req_idx not in invalid_req_indices_set else None

2417 - )

2411 + sampled_ids = [-1] if req_idx not in invalid_req_indices_set else None

2412 else:

2413 sampled_ids = valid_sampled_token_ids[req_idx]

2414

2421 - num_sampled_ids: int = (

2422 - sampled_ids.shape[0] if sampled_ids is not None else 0

2423 - )

2415 + num_sampled_ids: int = len(sampled_ids) if sampled_ids else 0

2416

2417 if cu_num_accepted_tokens is not None:

2418 cu_num_accepted_tokens.append(

2419 cu_num_accepted_tokens[-1] + num_sampled_ids

2420 )

2421

2430 - if sampled_ids is None or num_sampled_ids == 0:

2422 + if not sampled_ids:

2423 continue

2424

2425 start_idx = self.input_batch.num_tokens_no_spec[req_idx]

@@ -2769,9 +2761,7 @@ class GPUModelRunner(

2761 with record_function_or_nullcontext("gpu_model_runner: sample"):

2762 sampler_output = self._sample(logits, spec_decode_metadata)

2763

2772 - def propose_draft_token_ids(

2773 - sampled_token_ids: torch.Tensor | list[np.ndarray],

2774 - ) -> None:

2764 + def propose_draft_token_ids(sampled_token_ids):

2765 assert spec_decode_common_attn_metadata is not None

2766 with record_function_or_nullcontext("gpu_model_runner: draft"):

2767 self._draft_token_ids = self.propose_draft_token_ids(

@@ -2893,14 +2883,14 @@ class GPUModelRunner(

2883 def propose_draft_token_ids(

2884 self,

2885 scheduler_output: "SchedulerOutput",

2896 - sampled_token_ids: torch.Tensor | list[np.ndarray],

2886 + sampled_token_ids: torch.Tensor | list[list[int]],

2887 sampling_metadata: SamplingMetadata,

2888 hidden_states: torch.Tensor,

2889 sample_hidden_states: torch.Tensor,

2890 aux_hidden_states: list[torch.Tensor] | None,

2891 spec_decode_metadata: SpecDecodeMetadata | None,

2892 common_attn_metadata: CommonAttentionMetadata,

2903 - ) -> torch.Tensor | list[list[int]]:

2893 + ) -> list[list[int]] | torch.Tensor:

2894 num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens

2895 if self.speculative_config.method == "ngram":

2896 assert isinstance(sampled_token_ids, list)

@@ -2932,7 +2922,7 @@ class GPUModelRunner(

2922 for num_draft, tokens in zip(

2923 spec_decode_metadata.num_draft_tokens, sampled_token_ids

2924 ):

2935 - indices.append(offset + tokens.shape[0] - 1)

2925 + indices.append(offset + len(tokens) - 1)

2926 offset += num_draft + 1

2927 indices = torch.tensor(indices, device=self.device)

2928 hidden_states = sample_hidden_states[indices]

@@ -4872,7 +4862,7 @@ class GPUModelRunner(

4862

4863 return kv_cache_spec

4864

4875 - def _to_list(self, sampled_token_ids: torch.Tensor) -> list[np.ndarray]:

4865 + def _to_list(self, sampled_token_ids: torch.Tensor) -> list[list[int]]:

4866 # This is a short term mitigation for issue mentioned in

4867 # https://github.com/vllm-project/vllm/issues/22754.

4868 # `tolist` would trigger a cuda wise stream sync, which

@@ -4885,4 +4875,4 @@ class GPUModelRunner(

4875 pinned.copy_(sampled_token_ids, non_blocking=True)

4876 self.transfer_event.record()

4877 self.transfer_event.synchronize()

4888 - return [row for row in pinned.numpy()]

4878 + return pinned.tolist()