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

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 12, 19:45 UTC

Revert "[Redo] #26368 (#28771)" (#29121)

Signed-off-by: Jialin Ouyang <Jialin.Ouyang@gmail.com>

Jialin Ouyang committed Nov 20, 2025 at 21:27 UTC 30b9c6774396c168bdf019b488f2cfc133b09b35

16 files changed +99 -127

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_priority_scheduler_random.py

+2 -4

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

3 import random

4 import uuid

6 -import numpy as np

6 import pytest

8 from vllm.config import VllmConfig

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

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

100 ]

101 sampled_token_ids = [

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

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

103 + for num_tokens in num_output_tokens

104 ]

105

106 return ModelRunnerOutput(

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

196 num_blocks: int,

197 ):

198 random.seed(42)

199 - np.random.seed(42)

200 -

199 seen_request_prompt_length = dict[str, int]()

200 seen_request_ids = set[str]()

201 seen_mm_hashes = set[str]()

tests/v1/core/test_scheduler.py

+38 -50

@@ -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={},

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

2624 model_output = ModelRunnerOutput(

2625 req_ids=[request1.request_id],

2626 req_id_to_index={request1.request_id: 0},

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

2627 + sampled_token_ids=[[100]],

2628 # spec_token_ids=None,

2629 logprobs=None,

2630 prompt_logprobs_dict={},

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

2830 MODEL_RUNNER_OUTPUT = ModelRunnerOutput(

2831 req_ids=req_ids,

2832 req_id_to_index=req_to_index,

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

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

2834 logprobs=None,

2835 prompt_logprobs_dict={},

2836 pooler_output=[],

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

2943 model_output = ModelRunnerOutput(

2944 req_ids=[request_low.request_id],

2945 req_id_to_index={request_low.request_id: 0},

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

2946 + sampled_token_ids=[[100]],

2947 # spec_token_ids=None,

2948 logprobs=None,

2949 prompt_logprobs_dict={},

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

2994 model_output = ModelRunnerOutput(

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

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

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

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

2998 # spec_token_ids=None,

2999 logprobs=None,

3000 prompt_logprobs_dict={},

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

3029 model_output = ModelRunnerOutput(

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

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

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

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

3033 # spec_token_ids=None,

3034 logprobs=None,

3035 prompt_logprobs_dict={},

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

3215 model_output = ModelRunnerOutput(

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

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

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

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

3219 # spec_token_ids=None,

3220 logprobs=None,

3221 prompt_logprobs_dict={},

tests/v1/kv_connector/unit/test_nixl_connector.py

+3 -4

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

11 from collections import defaultdict

12 from unittest.mock import patch

14 -import numpy as np

14 import pytest

15 import ray

16 import torch

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

826 output = ModelRunnerOutput(

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

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

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

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

830 logprobs=None,

831 prompt_logprobs_dict={},

832 pooler_output=[None],

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

907 output = ModelRunnerOutput(

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

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

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

910 + sampled_token_ids=[[123]],

911 logprobs=None,

912 prompt_logprobs_dict={},

913 pooler_output=[None],

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

965 model_output = ModelRunnerOutput(

966 req_ids=["req_0"],

967 req_id_to_index={"req_0": 0},

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

968 + sampled_token_ids=[[123]],

969 logprobs=None,

970 prompt_logprobs_dict={},

971 pooler_output=[None],

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/utils/gc_utils.py

+8 -5

@@ -53,6 +53,7 @@ class GCDebugger:

53 self.config = config

54 # Start time in micro second of this GC cycle

55 self.start_time_ns: int = time.monotonic_ns()

56 + self.num_objects: int = 0

57 # If config.top_objects is positive,

58 # compute top collected objects by object types

59 self.gc_top_collected_objects: str = ""

@@ -68,19 +69,21 @@ class GCDebugger:

69 # Before GC started, record GC start time

70 # and top collected objects

71 self.start_time_ns = time.monotonic_ns()

71 - if (top_objects := self.config.top_objects) > 0:

72 - self.gc_top_collected_objects = _compute_top_gc_collected_objects(

73 - gc.get_objects(generation), top_objects

74 - )

72 + objects = gc.get_objects(generation)

73 + self.num_objects = len(objects)

74 + self.gc_top_collected_objects = _compute_top_gc_collected_objects(

75 + objects, self.config.top_objects

76 + )

77 elif phase == "stop":

78 # After GC finished, Record GC elapsed time and

79 # optionally top collected objects

80 elpased_ms = (time.monotonic_ns() - self.start_time_ns) / 1e6

81 logger.info(

82 "GC took %.3fms to complete. "

81 - "Collected %s objects in GC generation %d.%s",

83 + "Collected %s objects (out of %d) in GC generation %d.%s",

84 elpased_ms,

85 str(info.get("collected", "?")),

86 + self.num_objects,

87 generation,

88 (

89 f" Top collected objects: \n{self.gc_top_collected_objects}"

vllm/v1/core/sched/scheduler.py

+2 -2

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

1013 continue

1014

1015 req_index = model_runner_output.req_id_to_index[req_id]

1016 - generated_token_ids: list[int] = (

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

1016 + generated_token_ids = (

1017 + sampled_token_ids[req_index] if sampled_token_ids else []

1018 )

1019

1020 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[np.ndarray]

161 + sampled_token_ids: list[list[int]]

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]] = [np.array([0]) for _ in req_ids]

223 + sampled_token_ids: list[list[int]] = [[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

+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

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

496

497 def prepare_next_token_ids_cpu(

498 self,

499 - sampled_token_ids: list[np.ndarray],

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

500 requests: dict[str, CachedRequestState],

501 gpu_input_batch: InputBatch,

502 num_scheduled_tokens: dict[str, int],

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

511 req_ids = gpu_input_batch.req_ids

512 next_token_ids: list[int] = []

513 for i, token_ids in enumerate(sampled_token_ids):

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

514 + if token_ids:

515 # Common case.

516 next_token_id = token_ids[-1]

517 else:

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

522 seq_len = req_state.num_computed_tokens + num_scheduled_tokens[req_id]

523 next_token_id = req_state.get_token_id(seq_len)

524 next_token_ids.append(next_token_id)

525 - return torch.tensor(

525 + next_token_ids = torch.tensor(

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

527 )

528 + return next_token_ids

529

530 def prepare_next_token_ids_padded(

531 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

@@ -221,16 +221,14 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput):

221 del self._sampled_token_ids

222 max_gen_len = self.sampled_token_ids_cpu.shape[-1]

223 if max_gen_len == 1:

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

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

226 - ]

224 + valid_sampled_token_ids = self.sampled_token_ids_cpu.tolist()

225 else:

226 valid_sampled_token_ids = RejectionSampler.parse_output(

227 self.sampled_token_ids_cpu,

228 self.vocab_size,

229 )

230 for i in self._invalid_req_indices:

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

231 + valid_sampled_token_ids[i].clear()

232

233 output = self._model_runner_output

234 output.sampled_token_ids = valid_sampled_token_ids

@@ -2466,7 +2464,7 @@ class GPUModelRunner(

2464 ) -> tuple[

2465 dict[str, int],

2466 LogprobsLists | None,

2469 - list[np.ndarray],

2467 + list[list[int]],

2468 dict[str, LogprobsTensors | None],

2469 list[str],

2470 dict[str, int],

@@ -2492,7 +2490,6 @@ class GPUModelRunner(

2490 num_sampled_tokens = sampler_output.sampled_token_ids.shape[0]

2491 sampled_token_ids = sampler_output.sampled_token_ids

2492 invalid_req_indices = []

2495 - valid_sampled_token_ids: list[np.ndarray]

2493 if not self.use_async_scheduling:

2494 # Get the valid generated tokens.

2495 max_gen_len = sampled_token_ids.shape[-1]

@@ -2507,7 +2504,7 @@ class GPUModelRunner(

2504 )

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

2506 for i in discard_sampled_tokens_req_indices:

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

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

2508 else:

2509 valid_sampled_token_ids = []

2510 invalid_req_indices = discard_sampled_tokens_req_indices.tolist()

@@ -2537,24 +2534,19 @@ class GPUModelRunner(

2534 [0] if spec_decode_metadata and logprobs_tensors else None

2535 )

2536 for req_idx in range(num_sampled_tokens):

2540 - sampled_ids: np.ndarray | None

2537 if self.use_async_scheduling:

2542 - sampled_ids = (

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

2544 - )

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

2539 else:

2540 sampled_ids = valid_sampled_token_ids[req_idx]

2541

2548 - num_sampled_ids: int = (

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

2550 - )

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

2543

2544 if cu_num_accepted_tokens is not None:

2545 cu_num_accepted_tokens.append(

2546 cu_num_accepted_tokens[-1] + num_sampled_ids

2547 )

2548

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

2549 + if not sampled_ids:

2550 continue

2551

2552 start_idx = self.input_batch.num_tokens_no_spec[req_idx]

@@ -2938,9 +2930,7 @@ class GPUModelRunner(

2930

2931 self.input_batch.prev_sampled_token_ids = None

2932

2941 - def propose_draft_token_ids(

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

2943 - ) -> None:

2933 + def propose_draft_token_ids(sampled_token_ids):

2934 assert spec_decode_common_attn_metadata is not None

2935 with record_function_or_nullcontext("gpu_model_runner: draft"):

2936 self._draft_token_ids = self.propose_draft_token_ids(

@@ -3113,14 +3103,14 @@ class GPUModelRunner(

3103 def propose_draft_token_ids(

3104 self,

3105 scheduler_output: "SchedulerOutput",

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

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

3107 sampling_metadata: SamplingMetadata,

3108 hidden_states: torch.Tensor,

3109 sample_hidden_states: torch.Tensor,

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

3111 spec_decode_metadata: SpecDecodeMetadata | None,

3112 common_attn_metadata: CommonAttentionMetadata,

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

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

3114 num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens

3115 spec_config = self.speculative_config

3116 assert spec_config is not None

@@ -3154,7 +3144,7 @@ class GPUModelRunner(

3144 for num_draft, tokens in zip(

3145 spec_decode_metadata.num_draft_tokens, sampled_token_ids

3146 ):

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

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

3148 offset += num_draft + 1

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

3150 hidden_states = sample_hidden_states[indices]

@@ -5150,7 +5140,7 @@ class GPUModelRunner(

5140

5141 return kv_cache_spec

5142

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

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

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

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

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

@@ -5163,4 +5153,4 @@ class GPUModelRunner(

5153 pinned.copy_(sampled_token_ids, non_blocking=True)

5154 self.transfer_event.record()

5155 self.transfer_event.synchronize()

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

5156 + return pinned.tolist()

vllm/v1/worker/tpu_model_runner.py

+3 -5

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

1262

1263 max_gen_len = selected_token_ids.shape[-1]

1264 if max_gen_len == 1:

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

1266 - row for row in selected_token_ids.numpy()

1267 - ]

1265 + valid_sampled_token_ids = selected_token_ids.tolist()

1266

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

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

1269 # the "else" case here

1270 for i in discard_sampled_tokens_req_indices:

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

1271 + valid_sampled_token_ids[i].clear()

1272

1273 # Append sampled tokens

1274 for i, req_state, seq_len in request_seq_lens:

@@ -1283,7 +1281,7 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1281 valid_mask = selected_token_ids != INVALID_TOKEN_ID

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

1283 valid_sampled_token_ids = [

1286 - seq.numpy() for seq in selected_token_ids[valid_mask].split(gen_lens)

1284 + seq.tolist() for seq in selected_token_ids[valid_mask].split(gen_lens)

1285 ]

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

1287 for i, req_state, seq_len in request_seq_lens: