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

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

[Redo] #26368 (#28771)

Signed-off-by: Jialin Ouyang <Jialin.Ouyang@gmail.com> Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk> Co-authored-by: Jialin Ouyang <Jialin.Ouyang@gmail.com>

Cyrus Leung committed Nov 15, 2025 at 14:47 UTC 98b4d389ed27f09fd185ade889a02f640a3ff0b4

15 files changed +122 -91

tests/v1/core/test_async_scheduler.py

+2 -1

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

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

3 from collections import deque

5 +import numpy as np

6 import pytest

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

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

22 return ModelRunnerOutput(

23 req_ids=req_ids,

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

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

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

26 logprobs=None,

27 prompt_logprobs_dict={},

28 pooler_output=[],

tests/v1/core/test_priority_scheduler_random.py

+4 -2

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

3 import random

4 import uuid

6 +import numpy as np

7 import pytest

9 from vllm.config import VllmConfig

@@ -99,8 +100,7 @@ def _mock_execute_model(

100 random.randint(*num_output_tokens_range) for _ in range(len(request_ids))

101 ]

102 sampled_token_ids = [

102 - [random.randint(0, 100) for _ in range(num_tokens)]

103 - for num_tokens in num_output_tokens

103 + np.random.randint(0, 100, size=num_tokens) for num_tokens in num_output_tokens

104 ]

105

106 return ModelRunnerOutput(

@@ -196,6 +196,8 @@ def test_priority_scheduling_blast(

196 num_blocks: int,

197 ):

198 random.seed(42)

199 + np.random.seed(42)

200 +

201 seen_request_prompt_length = dict[str, int]()

202 seen_request_ids = set[str]()

203 seen_mm_hashes = set[str]()

tests/v1/core/test_scheduler.py

+50 -38

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

3 import dataclasses

4 from unittest.mock import Mock

6 +import numpy as np

7 import pytest

8 import torch

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

170 req_id_to_index=req_to_index,

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

172 # the rest requests are still being prefilled.

172 - sampled_token_ids=[[0], [], []],

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

174 logprobs=None,

175 prompt_logprobs_dict={},

176 pooler_output=[],

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

217 model_runner_output = ModelRunnerOutput(

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

219 req_id_to_index=req_to_index,

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

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

221 logprobs=None,

222 prompt_logprobs_dict={},

223 pooler_output=[],

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

277 model_runner_output = ModelRunnerOutput(

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

279 req_id_to_index=req_to_index,

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

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

281 logprobs=None,

282 prompt_logprobs_dict={},

283 pooler_output=[],

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

301 model_runner_output = ModelRunnerOutput(

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

303 req_id_to_index=req_to_index,

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

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

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

306 logprobs=None,

307 prompt_logprobs_dict={},

308 pooler_output=[],

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

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

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

351 sampled_token_ids=[

350 - [EOS_TOKEN_ID],

351 - [10, 11],

352 + np.array([EOS_TOKEN_ID]),

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

354 ], # First request hits EOS, second continues

355 logprobs=None,

356 prompt_logprobs_dict={},

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

394 model_output = ModelRunnerOutput(

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

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

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

397 + sampled_token_ids=[

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

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

400 + ], # First request hits stop token

401 logprobs=None,

402 prompt_logprobs_dict={},

403 pooler_output=[],

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

441 model_output = ModelRunnerOutput(

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

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

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

444 + sampled_token_ids=[

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

446 + np.array([13]),

447 + ], # First request exceeds max_tokens

448 logprobs=None,

449 prompt_logprobs_dict={},

450 pooler_output=[],

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

483 model_output = ModelRunnerOutput(

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

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

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

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

487 logprobs=None,

488 prompt_logprobs_dict={},

489 pooler_output=[],

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

624 model_runner_output = ModelRunnerOutput(

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

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

619 - sampled_token_ids=[[0]],

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

628 logprobs=None,

629 prompt_logprobs_dict={},

630 pooler_output=[],

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

641 model_runner_output = ModelRunnerOutput(

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

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

636 - sampled_token_ids=[[0]],

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

645 logprobs=None,

646 prompt_logprobs_dict={},

647 pooler_output=[],

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

678 model_runner_output0 = ModelRunnerOutput(

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

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

673 - sampled_token_ids=[[0]],

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

682 logprobs=None,

683 prompt_logprobs_dict={},

684 pooler_output=[],

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

695 model_runner_output1 = ModelRunnerOutput(

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

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

690 - sampled_token_ids=[[42]],

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

699 logprobs=None,

700 prompt_logprobs_dict={},

701 pooler_output=[],

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

712 @pytest.mark.parametrize(

713 "spec_tokens,output_tokens,expected",

714 [

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

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

724 (

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

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

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

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

728 ), # multiple mismatches

729 ],

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

757 model_runner_output = ModelRunnerOutput(

758 req_ids=req_ids,

759 req_id_to_index=req_to_index,

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

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

761 logprobs=None,

762 prompt_logprobs_dict={},

763 pooler_output=[],

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

984 MODEL_RUNNER_OUTPUT = ModelRunnerOutput(

985 req_ids=req_ids,

986 req_id_to_index=req_to_index,

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

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

988 logprobs=None,

989 prompt_logprobs_dict={},

990 pooler_output=[],

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

1037 MODEL_RUNNER_OUTPUT = ModelRunnerOutput(

1038 req_ids=req_ids,

1039 req_id_to_index=req_to_index,

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

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

1041 logprobs=None,

1042 prompt_logprobs_dict={},

1043 pooler_output=[],

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

1100 MODEL_RUNNER_OUTPUT = ModelRunnerOutput(

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

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

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

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

1104 logprobs=None,

1105 prompt_logprobs_dict={},

1106 pooler_output=[],

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

1166 MODEL_RUNNER_OUTPUT = ModelRunnerOutput(

1167 req_ids=req_ids,

1168 req_id_to_index=req_to_index,

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

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

1170 logprobs=None,

1171 prompt_logprobs_dict={},

1172 pooler_output=[],

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

1251 MODEL_RUNNER_OUTPUT = ModelRunnerOutput(

1252 req_ids=req_ids,

1253 req_id_to_index=req_to_index,

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

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

1255 logprobs=None,

1256 prompt_logprobs_dict={},

1257 pooler_output=[],

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

1344 return ModelRunnerOutput(

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

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

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

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

1348 logprobs=None,

1349 prompt_logprobs_dict={},

1350 pooler_output=[],

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

1761 req_id_to_index={

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

1763 },

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

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

1765 logprobs=None,

1766 prompt_logprobs_dict={},

1767 pooler_output=[],

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

1830 req_id_to_index={

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

1832 },

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

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

1834 logprobs=None,

1835 prompt_logprobs_dict={},

1836 pooler_output=[],

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

2076 model_output = ModelRunnerOutput(

2077 req_ids=[req.req_id],

2078 req_id_to_index={req.req_id: 0},

2067 - sampled_token_ids=[[100]],

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

2080 logprobs=None,

2081 prompt_logprobs_dict={},

2082 pooler_output=[],

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

2162 model_output = ModelRunnerOutput(

2163 req_ids=[request_low.request_id],

2164 req_id_to_index={request_low.request_id: 0},

2153 - sampled_token_ids=[[100]],

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

2166 # spec_token_ids=None,

2167 logprobs=None,

2168 prompt_logprobs_dict={},

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

2193 model_output = ModelRunnerOutput(

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

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

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

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

2197 # spec_token_ids=None,

2198 logprobs=None,

2199 prompt_logprobs_dict={},

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

2219 model_output = ModelRunnerOutput(

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

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

2210 - sampled_token_ids=[[], [100]],

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

2223 # spec_token_ids=None,

2224 logprobs=None,

2225 prompt_logprobs_dict={},

@@ -2624,7 +2636,7 @@ def test_ec_connector_with_partial_cache_hit_multi_round(use_kv_connector):

2636 model_output = ModelRunnerOutput(

2637 req_ids=[request1.request_id],

2638 req_id_to_index={request1.request_id: 0},

2627 - sampled_token_ids=[[100]],

2639 + sampled_token_ids=[np.array([100])],

2640 # spec_token_ids=None,

2641 logprobs=None,

2642 prompt_logprobs_dict={},

@@ -2830,7 +2842,7 @@ def test_ec_connector_unable_to_allocate(use_kv_connector):

2842 MODEL_RUNNER_OUTPUT = ModelRunnerOutput(

2843 req_ids=req_ids,

2844 req_id_to_index=req_to_index,

2833 - sampled_token_ids=[[1000]] * len(req_ids),

2845 + sampled_token_ids=[np.array([1000])] * len(req_ids),

2846 logprobs=None,

2847 prompt_logprobs_dict={},

2848 pooler_output=[],

@@ -2943,7 +2955,7 @@ def test_priority_scheduling_ec_connector_preemption_and_resumption(

2955 model_output = ModelRunnerOutput(

2956 req_ids=[request_low.request_id],

2957 req_id_to_index={request_low.request_id: 0},

2946 - sampled_token_ids=[[100]],

2958 + sampled_token_ids=[np.array([100])],

2959 # spec_token_ids=None,

2960 logprobs=None,

2961 prompt_logprobs_dict={},

@@ -2994,7 +3006,7 @@ def test_priority_scheduling_ec_connector_preemption_and_resumption(

3006 model_output = ModelRunnerOutput(

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

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

2997 - sampled_token_ids=[[100] for _ in requests],

3009 + sampled_token_ids=[np.array([100]) for _ in requests],

3010 # spec_token_ids=None,

3011 logprobs=None,

3012 prompt_logprobs_dict={},

@@ -3029,7 +3041,7 @@ def test_priority_scheduling_ec_connector_preemption_and_resumption(

3041 model_output = ModelRunnerOutput(

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

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

3032 - sampled_token_ids=[[100], [100, 200]],

3044 + sampled_token_ids=[np.array([100]), np.array([100, 200])],

3045 # spec_token_ids=None,

3046 logprobs=None,

3047 prompt_logprobs_dict={},

@@ -3215,7 +3227,7 @@ def test_ec_connector_allocate_encoder_tokens_with_external_load(use_kv_connecto

3227 model_output = ModelRunnerOutput(

3228 req_ids=[request1.request_id, request2.request_id],

3229 req_id_to_index={request1.request_id: 0, request2.request_id: 1},

3218 - sampled_token_ids=[[100], [121]],

3230 + sampled_token_ids=[np.array([100]), np.array([121])],

3231 # spec_token_ids=None,

3232 logprobs=None,

3233 prompt_logprobs_dict={},

tests/v1/kv_connector/unit/test_nixl_connector.py

+4 -3

@@ -11,6 +11,7 @@ import uuid

11 from collections import defaultdict

12 from unittest.mock import patch

14 +import numpy as np

15 import pytest

16 import ray

17 import torch

@@ -826,7 +827,7 @@ def test_kv_connector_stats_aggregation():

827 output = ModelRunnerOutput(

828 req_ids=[f"req_{i}"],

829 req_id_to_index={f"req_{i}": 0},

829 - sampled_token_ids=[[123]], # dummy token

830 + sampled_token_ids=[np.array([123])], # dummy token

831 logprobs=None,

832 prompt_logprobs_dict={},

833 pooler_output=[None],

@@ -907,7 +908,7 @@ def test_multi_kv_connector_stats_aggregation():

908 output = ModelRunnerOutput(

909 req_ids=[f"req_{i}"],

910 req_id_to_index={f"req_{i}": 0},

910 - sampled_token_ids=[[123]],

911 + sampled_token_ids=[np.array([123])],

912 logprobs=None,

913 prompt_logprobs_dict={},

914 pooler_output=[None],

@@ -965,7 +966,7 @@ def test_scheduler_kv_connector_stats_aggregation():

966 model_output = ModelRunnerOutput(

967 req_ids=["req_0"],

968 req_id_to_index={"req_0": 0},

968 - sampled_token_ids=[[123]],

969 + sampled_token_ids=[np.array([123])],

970 logprobs=None,

971 prompt_logprobs_dict={},

972 pooler_output=[None],

tests/v1/kv_connector/unit/utils.py

+2 -1

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

7 from itertools import chain, count

8 from typing import Any

10 +import numpy as np

11 import torch

13 from vllm import SamplingParams

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

229

230 # Make sampled tokens.

231 sampled_token = EOS_TOKEN_ID if use_eos else token_id

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

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

233

234 kv_connector_output = (

235 None

tests/v1/spec_decode/test_eagle.py

+4 -1

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

4 from unittest import mock

6 +import numpy as np

7 import pytest

8 import torch

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

113 sampled_token_ids_tensor = torch.tensor(

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

115 )

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

116 + sampled_token_ids_cpu = [

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

118 + ]

119

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

121 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=[[0]],

80 + sampled_token_ids=[np.array([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=[[0]],

91 + sampled_token_ids=[np.array([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=[[0]],

102 + sampled_token_ids=[np.array([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=[[0]],

114 + sampled_token_ids=[np.array([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=[[0]],

125 + sampled_token_ids=[np.array([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=[[0]],

136 + sampled_token_ids=[np.array([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=[[0]],

147 + sampled_token_ids=[np.array([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=[[0], [1]],

160 + sampled_token_ids=[np.array([0]), np.array([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=[[0], [1]],

184 + sampled_token_ids=[np.array([0]), np.array([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 = (

1014 - sampled_token_ids[req_index] if sampled_token_ids else []

1013 + generated_token_ids: list[int] = (

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

1015 )

1016

1017 scheduled_spec_token_ids = (

vllm/v1/outputs.py

+2 -2

@@ -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[list[int]]

161 + sampled_token_ids: list[np.ndarray]

162

163 # [num_reqs, max_num_logprobs + 1]

164 # [num_reqs, max_num_logprobs + 1]

@@ -220,7 +220,7 @@ def make_empty_encoder_model_runner_output(

220 req_id_to_index: dict[str, int] = {rid: idx for idx, rid in enumerate(req_ids)}

221

222 # No tokens generated yet ⇒ one empty list per request

223 - sampled_token_ids: list[list[int]] = [[0] for _ in req_ids]

223 + sampled_token_ids: list[list[int]] = [np.array([0]) for _ in req_ids]

224

225 # Pooler outputs are not available yet ⇒ use None placeholders

226 pooler_output: list[torch.Tensor | None] = [None for _ in req_ids]

vllm/v1/sample/rejection_sampler.py

+3 -5

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

4 from dataclasses import replace

6 +import numpy as np

7 import torch

8 import torch.nn as nn

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

205 def parse_output(

206 output_token_ids: torch.Tensor,

207 vocab_size: int,

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

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

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

210 Args:

211 output_token_ids: The sampled token IDs in shape

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

221 valid_mask = (output_token_ids_np != PLACEHOLDER_TOKEN_ID) & (

222 output_token_ids_np < vocab_size

223 )

223 - outputs = [

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

225 - ]

226 - return outputs

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

225

226 def apply_logits_processors(

227 self,

vllm/v1/spec_decode/eagle.py

+3 -4

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

484

485 def prepare_next_token_ids_cpu(

486 self,

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

487 + sampled_token_ids: list[np.ndarray],

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:

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

503 # Common case.

504 next_token_id = token_ids[-1]

505 else:

@@ -510,10 +510,9 @@ 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 - next_token_ids = torch.tensor(

513 + return torch.tensor(

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

515 )

516 - return next_token_ids

516

517 def prepare_next_token_ids_padded(

518 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 - [[]] * 1024,

57 + [np.array([])] * 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[list[int]],

134 + sampled_token_ids: list[np.ndarray],

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 = len(sampled_ids)

143 + num_sampled_ids = sampled_ids.shape[0]

144 if not num_sampled_ids:

145 # Skip speculative decoding.

146 continue

vllm/v1/spec_decode/suffix_decoding.py

+6 -4

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 +import numpy as np

4 +

5 from vllm.config import VllmConfig

6 from vllm.v1.worker.gpu_input_batch import InputBatch

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

34 def propose(

35 self,

36 input_batch: InputBatch,

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

37 + sampled_token_ids: list[np.ndarray],

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

39 """

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

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

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

43 """

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

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

45 for i, sampled_ids in enumerate(sampled_token_ids):

44 - if not sampled_ids:

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

47 # Skip speculative decoding for partial prefills.

48 draft_token_ids.append([])

49 continue

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

72 self.suffix_cache.start_request(req_id, prompt_token_ids)

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

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

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

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

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

vllm/v1/worker/gpu_model_runner.py

+23 -13

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

216 del self._logprobs_tensors

217 del self._sampled_token_ids

218

219 - valid_sampled_token_ids = self.sampled_token_ids_cpu.tolist()

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

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

221 + ]

222 for i in self._invalid_req_indices:

221 - valid_sampled_token_ids[i].clear()

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

224

225 output = self._model_runner_output

226 output.sampled_token_ids = valid_sampled_token_ids

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

2341 ) -> tuple[

2342 dict[str, int],

2343 LogprobsLists | None,

2342 - list[list[int]],

2344 + list[np.ndarray],

2345 dict[str, LogprobsTensors | None],

2346 list[str],

2347 dict[str, int],

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

2367 num_sampled_tokens = sampler_output.sampled_token_ids.shape[0]

2368 sampled_token_ids = sampler_output.sampled_token_ids

2369 invalid_req_indices = []

2370 + valid_sampled_token_ids: list[np.ndarray]

2371 if not self.use_async_scheduling:

2372 # Get the valid generated tokens.

2373 max_gen_len = sampled_token_ids.shape[-1]

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

2382 )

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

2384 for i in discard_sampled_tokens_req_indices:

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

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

2386 else:

2387 valid_sampled_token_ids = []

2388 invalid_req_indices = discard_sampled_tokens_req_indices.tolist()

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

2410 [0] if spec_decode_metadata and logprobs_tensors else None

2411 )

2412 for req_idx in range(num_sampled_tokens):

2413 + sampled_ids: np.ndarray | None

2414 if self.use_async_scheduling:

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

2415 + sampled_ids = (

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

2417 + )

2418 else:

2419 sampled_ids = valid_sampled_token_ids[req_idx]

2420

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

2421 + num_sampled_ids: int = (

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

2423 + )

2424

2425 if cu_num_accepted_tokens is not None:

2426 cu_num_accepted_tokens.append(

2427 cu_num_accepted_tokens[-1] + num_sampled_ids

2428 )

2429

2422 - if not sampled_ids:

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

2431 continue

2432

2433 start_idx = self.input_batch.num_tokens_no_spec[req_idx]

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

2769 with record_function_or_nullcontext("gpu_model_runner: sample"):

2770 sampler_output = self._sample(logits, spec_decode_metadata)

2771

2764 - def propose_draft_token_ids(sampled_token_ids):

2772 + def propose_draft_token_ids(

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

2774 + ) -> None:

2775 assert spec_decode_common_attn_metadata is not None

2776 with record_function_or_nullcontext("gpu_model_runner: draft"):

2777 self._draft_token_ids = self.propose_draft_token_ids(

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

2893 def propose_draft_token_ids(

2894 self,

2895 scheduler_output: "SchedulerOutput",

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

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

2897 sampling_metadata: SamplingMetadata,

2898 hidden_states: torch.Tensor,

2899 sample_hidden_states: torch.Tensor,

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

2901 spec_decode_metadata: SpecDecodeMetadata | None,

2902 common_attn_metadata: CommonAttentionMetadata,

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

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

2904 num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens

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

2906 assert isinstance(sampled_token_ids, list)

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

2932 for num_draft, tokens in zip(

2933 spec_decode_metadata.num_draft_tokens, sampled_token_ids

2934 ):

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

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

2936 offset += num_draft + 1

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

2938 hidden_states = sample_hidden_states[indices]

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

4872

4873 return kv_cache_spec

4874

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

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

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

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

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

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

4885 pinned.copy_(sampled_token_ids, non_blocking=True)

4886 self.transfer_event.record()

4887 self.transfer_event.synchronize()

4878 - return pinned.tolist()

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

vllm/v1/worker/tpu_model_runner.py

+5 -3

@@ -1254,13 +1254,15 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1254

1255 max_gen_len = selected_token_ids.shape[-1]

1256 if max_gen_len == 1:

1257 - valid_sampled_token_ids = selected_token_ids.tolist()

1257 + valid_sampled_token_ids: list[np.ndarray] = [

1258 + row for row in selected_token_ids.numpy()

1259 + ]

1260

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

1262 # TODO: Keep in sync with gpu_model_runner.py, in particular

1263 # the "else" case here

1264 for i in discard_sampled_tokens_req_indices:

1263 - valid_sampled_token_ids[i].clear()

1265 + valid_sampled_token_ids[i] = np.array([])

1266

1267 # Append sampled tokens

1268 for i, req_state, seq_len in request_seq_lens:

@@ -1273,7 +1275,7 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1275 valid_mask = selected_token_ids != INVALID_TOKEN_ID

1276 gen_lens = valid_mask.sum(dim=1).tolist()

1277 valid_sampled_token_ids = [

1276 - seq.tolist() for seq in selected_token_ids[valid_mask].split(gen_lens)

1278 + seq.numpy() for seq in selected_token_ids[valid_mask].split(gen_lens)

1279 ]

1280 self.input_batch.num_tokens[:num_reqs] += gen_lens

1281 for i, req_state, seq_len in request_seq_lens: