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

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 14, 05:30 UTC

[V0 Deprecation] Remove multi-step scheduling (#22138)

Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu> Signed-off-by: Woosuk Kwon <woosuk@thinkingmachines.ai>

Woosuk Kwon committed Aug 12, 2025 at 20:18 UTC 71683ca6f6764f35abe23d612a0d7dbd33babe32

37 files changed +57 -3465

.buildkite/nightly-benchmarks/tests/genai-perf-tests.json

-1

@@ -12,7 +12,6 @@

12 "vllm_server_parameters": {

13 "disable_log_stats": "",

14 "gpu_memory_utilization": 0.9,

15 - "num_scheduler_steps": 10,

15 "max_num_seqs": 512,

16 "dtype": "bfloat16"

17 },

.buildkite/nightly-benchmarks/tests/nightly-tests.json

-6

@@ -36,7 +36,6 @@

36 "vllm_server_parameters": {

37 "disable_log_stats": "",

38 "gpu_memory_utilization": 0.9,

39 - "num_scheduler_steps": 10,

39 "max_num_seqs": 512,

40 "dtype": "bfloat16"

41 },

@@ -90,7 +89,6 @@

89 "vllm_server_parameters": {

90 "disable_log_stats": "",

91 "gpu_memory_utilization": 0.9,

93 - "num_scheduler_steps": 10,

92 "max_num_seqs": 512,

93 "dtype": "bfloat16"

94 },

@@ -144,7 +142,6 @@

142 "vllm_server_parameters": {

143 "disable_log_stats": "",

144 "gpu_memory_utilization": 0.9,

147 - "num_scheduler_steps": 10,

145 "max_num_seqs": 512,

146 "dtype": "bfloat16"

147 },

@@ -195,7 +192,6 @@

192 "vllm_server_parameters": {

193 "disable_log_stats": "",

194 "gpu_memory_utilization": 0.9,

198 - "num_scheduler_steps": 10,

195 "max_num_seqs": 512,

196 "dtype": "bfloat16"

197 },

@@ -248,7 +244,6 @@

244 "vllm_server_parameters": {

245 "disable_log_stats": "",

246 "gpu_memory_utilization": 0.9,

251 - "num_scheduler_steps": 10,

247 "max_num_seqs": 512,

248 "dtype": "bfloat16"

249 },

@@ -301,7 +296,6 @@

296 "vllm_server_parameters": {

297 "disable_log_stats": "",

298 "gpu_memory_utilization": 0.9,

304 - "num_scheduler_steps": 10,

299 "max_num_seqs": 512,

300 "dtype": "bfloat16"

301 },

.buildkite/test-pipeline.yaml

-22

@@ -67,7 +67,6 @@ steps:

67 - python3 standalone_tests/lazy_imports.py

68 - pytest -v -s mq_llm_engine # MQLLMEngine

69 - pytest -v -s async_engine # AsyncLLMEngine

70 - - NUM_SCHEDULER_STEPS=4 pytest -v -s async_engine/test_async_llm_engine.py

70 - pytest -v -s test_inputs.py

71 - pytest -v -s test_outputs.py

72 - pytest -v -s multimodal

@@ -773,27 +772,6 @@ steps:

772 - pytest -v -s models/test_oot_registration.py # it needs a clean process

773 - pytest -v -s plugins/lora_resolvers # unit tests for in-tree lora resolver plugins

774

776 -- label: Multi-step Tests (4 GPUs) # 36min

777 - mirror_hardwares: [amdexperimental]

778 - working_dir: "/vllm-workspace/tests"

779 - num_gpus: 4

780 - source_file_dependencies:

781 - - vllm/model_executor/layers/sampler.py

782 - - vllm/sequence.py

783 - - vllm/worker/worker_base.py

784 - - vllm/worker/worker.py

785 - - vllm/worker/multi_step_worker.py

786 - - vllm/worker/model_runner_base.py

787 - - vllm/worker/model_runner.py

788 - - vllm/worker/multi_step_model_runner.py

789 - - vllm/engine

790 - - tests/multi_step

791 - commands:

792 - # this test is quite flaky

793 - # TODO: investigate and fix.

794 - # - pytest -v -s multi_step/test_correctness_async_llm.py

795 - - pytest -v -s multi_step/test_correctness_llm.py

796 -

775 - label: Pipeline Parallelism Test # 45min

776 mirror_hardwares: [amdexperimental]

777 working_dir: "/vllm-workspace/tests"

.github/CODEOWNERS

-1

@@ -36,7 +36,6 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson

36 /tests/entrypoints @DarkLight1337 @robertgshaw2-redhat @simon-mo @aarnphm

37 /tests/kernels @tlrmchlsmth @WoosukKwon @yewentao256

38 /tests/models @DarkLight1337 @ywang96

39 -/tests/multi_step @alexm-redhat @comaniac

39 /tests/multimodal @DarkLight1337 @ywang96

40 /tests/prefix_caching @comaniac @KuntaiDu

41 /tests/quantization @mgoin @robertgshaw2-redhat @yewentao256

tests/async_engine/test_async_llm_engine.py deleted

-409

@@ -1,409 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

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

3 -

4 -import asyncio

5 -import os

6 -import uuid

7 -from asyncio import CancelledError

8 -from copy import copy

9 -from dataclasses import dataclass, field

10 -from typing import Any, Optional

11 -

12 -import pytest

13 -import pytest_asyncio

14 -import torch

15 -

16 -from vllm import SamplingParams

17 -from vllm.config import ParallelConfig

18 -from vllm.distributed import cleanup_dist_env_and_memory

19 -from vllm.engine.async_llm_engine import AsyncEngineArgs, AsyncLLMEngine

20 -from vllm.outputs import RequestOutput as RealRequestOutput

21 -from vllm.sampling_params import RequestOutputKind

22 -

23 -from ..utils import wait_for_gpu_memory_to_clear

24 -

25 -

26 -@dataclass

27 -class RequestOutput:

28 - request_id: int

29 - finished: bool = False

30 -

31 -

32 -@dataclass

33 -class MockModelConfig:

34 - use_async_output_proc = True

35 - media_io_kwargs: dict[str, dict[str, Any]] = field(default_factory=dict)

36 -

37 -

38 -class MockEngine:

39 -

40 - def __init__(self):

41 - self.step_calls = 0

42 - self.add_request_calls = 0

43 - self.abort_request_calls = 0

44 - self.request_id = None

45 - # Ugly, remove dependency when possible

46 - self.parallel_config = ParallelConfig()

47 - self.model_config = MockModelConfig()

48 -

49 - async def step_async(self, virtual_engine):

50 - # PP size is 1, ignore virtual engine

51 - self.step_calls += 1

52 - return [RequestOutput(

53 - request_id=self.request_id)] if self.request_id else []

54 -

55 - async def process_model_inputs_async(self, *args, **kwargs):

56 - pass

57 -

58 - async def stop_remote_worker_execution_loop_async(self):

59 - pass

60 -

61 - def generate(self, request_id):

62 - self.request_id = request_id

63 -

64 - def stop_generating(self):

65 - self.request_id = None

66 -

67 - def add_request(self, **kwargs):

68 - del kwargs # Unused

69 - self.add_request_calls += 1

70 - print(f'Request calls: {self.add_request_calls}')

71 -

72 - async def add_request_async(self, **kwargs):

73 - self.add_request_calls += 1

74 - return

75 -

76 - def abort_request(self, request_id):

77 - del request_id # Unused

78 - self.abort_request_calls += 1

79 -

80 - def has_unfinished_requests(self):

81 - return self.request_id is not None

82 -

83 - def has_unfinished_requests_for_virtual_engine(self, virtual_engine):

84 - return self.request_id is not None

85 -

86 -

87 -class MockAsyncLLMEngine(AsyncLLMEngine):

88 - _engine_class = MockEngine

89 -

90 -

91 -@pytest.mark.asyncio

92 -async def test_new_requests_event():

93 - params = SamplingParams()

94 -

95 - engine = MockAsyncLLMEngine()

96 - engine.start_background_loop()

97 - await asyncio.sleep(0.01)

98 - assert engine.engine.step_calls == 0

99 -

100 - await engine.add_request("1", "", params)

101 - await asyncio.sleep(0.01)

102 - assert engine.engine.add_request_calls == 1

103 - assert engine.engine.step_calls == 1

104 -

105 - await engine.add_request("2", "", params)

106 - engine.engine.generate("2")

107 - await asyncio.sleep(0)

108 - await asyncio.sleep(0)

109 - await asyncio.sleep(0)

110 - assert engine.engine.add_request_calls == 2

111 - assert engine.engine.step_calls >= 2

112 - await asyncio.sleep(0.001)

113 - assert engine.engine.step_calls >= 3

114 - engine.engine.stop_generating()

115 - await asyncio.sleep(0.001)

116 - old_step_calls = engine.engine.step_calls

117 - await asyncio.sleep(0.001)

118 - assert engine.engine.step_calls == old_step_calls

119 -

120 - await engine.add_request("3", "", params)

121 - await asyncio.sleep(0.01)

122 - assert engine.engine.add_request_calls == 3

123 - assert engine.engine.step_calls == old_step_calls + 1

124 - await asyncio.sleep(0.01)

125 - assert engine.engine.add_request_calls == 3

126 - assert engine.engine.step_calls == old_step_calls + 1

127 -

128 - engine = MockAsyncLLMEngine()

129 - assert engine.get_model_config() is not None

130 - assert engine.get_tokenizer() is not None

131 - assert engine.get_decoding_config() is not None

132 -

133 -

134 -def start_engine():

135 - wait_for_gpu_memory_to_clear(

136 - devices=list(range(torch.cuda.device_count())),

137 - threshold_bytes=2 * 2**30,

138 - timeout_s=60,

139 - )

140 -

141 - num_scheduler_steps = int(os.getenv("NUM_SCHEDULER_STEPS", "1"))

142 - print(f"Starting engine with num_scheduler_steps={num_scheduler_steps}")

143 -

144 - return AsyncLLMEngine.from_engine_args(

145 - AsyncEngineArgs(model="facebook/opt-125m",

146 - enforce_eager=True,

147 - num_scheduler_steps=num_scheduler_steps))

148 -

149 -

150 -def uid() -> str:

151 - return str(uuid.uuid4())

152 -

153 -

154 -@pytest_asyncio.fixture(scope="module")

155 -async def async_engine():

156 - # We cannot use monkeypatch since this is a module

157 - # scoped fixture and monkeypatch is function scoped.

158 - previous_value = os.getenv("VLLM_USE_V1", None)

159 - os.environ["VLLM_USE_V1"] = "0"

160 - engine = await asyncio.get_event_loop().run_in_executor(executor=None,

161 - func=start_engine)

162 - try:

163 - yield engine

164 - finally:

165 - engine.shutdown_background_loop()

166 - del engine

167 - await asyncio.sleep(0.1)

168 - cleanup_dist_env_and_memory()

169 -

170 - if previous_value:

171 - os.environ["VLLM_USE_V1"] = previous_value

172 - else:

173 - del os.environ["VLLM_USE_V1"]

174 -

175 -

176 -@pytest.fixture()

177 -def should_do_global_cleanup_after_test(request) -> bool:

178 - # So we can share the async engine fixture between these tests

179 - return False

180 -

181 -

182 -@pytest.mark.asyncio(scope="module")

183 -@pytest.mark.parametrize("stop", [None, ["a stop string"]])

184 -async def test_asyncio_run(async_engine, stop):

185 -

186 - scheduler_config = await async_engine.get_scheduler_config()

187 - num_scheduler_steps = scheduler_config.num_scheduler_steps

188 -

189 - async def run(prompt: str):

190 - sampling_params = SamplingParams(

191 - temperature=0,

192 - max_tokens=32,

193 - min_tokens=32,

194 - stop=stop,

195 - )

196 -

197 - output_count = 0

198 - final_output = None

199 - async for output in async_engine.generate(prompt,

200 - sampling_params,

201 - request_id=uid()):

202 - output_count += 1

203 - final_output = output

204 - return final_output, output_count

205 -

206 - results = await asyncio.gather(

207 - run("test0"),

208 - run("test0"),

209 - )

210 - assert len(results) == 2

211 - first, second = results

212 -

213 - # remove nondeterministic fields for comparison

214 - first[0].metrics = None

215 - second[0].metrics = None

216 - first[0].request_id = None

217 - second[0].request_id = None

218 -

219 - assert str(first) == str(second)

220 -

221 - output_count = results[0][1]

222 - if num_scheduler_steps == 1:

223 - assert output_count == 32

224 - else:

225 - assert 1 < output_count < 32

226 -

227 -

228 -@pytest.mark.asyncio(scope="module")

229 -@pytest.mark.parametrize("stop", [None, ["a stop string"]])

230 -async def test_output_kinds(async_engine, stop):

231 - """Test that output_kind works as expected and that

232 - results are equivalent across different kinds."""

233 -

234 - scheduler_config = await async_engine.get_scheduler_config()

235 - num_scheduler_steps = scheduler_config.num_scheduler_steps

236 -

237 - sampling_params = SamplingParams(

238 - temperature=0,

239 - max_tokens=32,

240 - min_tokens=32,

241 - stop=stop,

242 - )

243 -

244 - async def run(prompt: str, kind: RequestOutputKind):

245 - params = copy(sampling_params)

246 - params.output_kind = kind

247 -

248 - output_count = 0

249 - final_output = None

250 - async for output in async_engine.generate(prompt,

251 - params,

252 - request_id=uid()):

253 - output_count += 1

254 - final_output = output

255 -

256 - assert final_output is not None

257 - assert final_output.finished

258 -

259 - return (final_output.prompt_token_ids,

260 - final_output.outputs[0].token_ids,

261 - final_output.outputs[0].text, output_count)

262 -

263 - async def run_deltas(prompt: str):

264 - params = copy(sampling_params)

265 - params.output_kind = RequestOutputKind.DELTA

266 -

267 - prompt_tokens = None

268 - output_tokens: list[int] = []

269 - output_text = ""

270 - output_count = 0

271 - final_output = None

272 - async for output in async_engine.generate(prompt,

273 - params,

274 - request_id=uid()):

275 - token_ids = output.outputs[0].token_ids

276 - text = output.outputs[0].text

277 - final_output = output

278 -

279 - # Ensure we get prompt ids iff we haven't yet received output tokens

280 - if output_tokens:

281 - assert 1 <= len(token_ids) <= num_scheduler_steps

282 - assert stop or text

283 - assert not output.prompt_token_ids

284 - else:

285 - assert output.prompt_token_ids

286 - prompt_tokens = output.prompt_token_ids

287 -

288 - output_tokens.extend(token_ids)

289 - output_text += text

290 -

291 - output_count += 1

292 -

293 - assert final_output is not None

294 - assert final_output.finished

295 -

296 - return prompt_tokens, output_tokens, output_text, output_count

297 -

298 - results = await asyncio.gather(

299 - run("common input prompt", RequestOutputKind.CUMULATIVE),

300 - run("common input prompt", RequestOutputKind.FINAL_ONLY),

301 - run_deltas("common input prompt"))

302 -

303 - # Make sure outputs are the same

304 - prompt_set = set(tuple(prompt_ids) for prompt_ids, _, _, _ in results)

305 - assert len(prompt_set) == 1

306 -

307 - text_set = set(text for _, _, text, _ in results)

308 - assert len(text_set) == 1

309 -

310 - tokens_set = set(tuple(ids) for _, ids, _, _ in results)

311 - assert len(tokens_set) == 1

312 -

313 - cumulative, final, deltas = results

314 -

315 - # output message counts

316 - assert cumulative[3] == deltas[3]

317 -

318 - if num_scheduler_steps == 1:

319 - assert cumulative[3] == 32

320 - else:

321 - assert 1 < cumulative[3] < 32

322 -

323 - assert final[3] == 1

324 -

325 -

326 -@pytest.mark.asyncio(scope="module")

327 -@pytest.mark.parametrize("stop", [None, ["a stop string"]])

328 -async def test_cancellation(async_engine, stop):

329 - scheduler_config = await async_engine.get_scheduler_config()

330 - num_scheduler_steps = scheduler_config.num_scheduler_steps

331 -

332 - sampling_params = SamplingParams(

333 - temperature=0,

334 - min_tokens=13,

335 - max_tokens=13,

336 - stop=stop,

337 - )

338 -

339 - stop_at = 5 if num_scheduler_steps == 1 else 1

340 -

341 - request_id = uid()

342 -

343 - i = 0

344 - with pytest.raises(CancelledError):

345 - async for output in async_engine.generate("test2",

346 - sampling_params,

347 - request_id=request_id):

348 - assert not output.finished

349 - i += 1

350 - if i == stop_at:

351 - await async_engine.abort(request_id)

352 -

353 - assert i == stop_at

354 -

355 -

356 -@pytest.mark.asyncio(scope="module")

357 -@pytest.mark.parametrize("stop", [None, ["a stop string"]])

358 -async def test_delayed_generator(async_engine, stop):

359 - scheduler_config = await async_engine.get_scheduler_config()

360 -

361 - if scheduler_config.num_scheduler_steps != 1:

362 - pytest.skip("no need to test this one with multistep")

363 -

364 - sampling_params = SamplingParams(

365 - temperature=0,

366 - min_tokens=10,

367 - max_tokens=10,

368 - stop=stop,

369 - )

370 -

371 - stream = async_engine.generate("test3", sampling_params, request_id=uid())

372 - i = 0

373 - final_output: Optional[RealRequestOutput] = None

374 - async for output in stream:

375 - final_output = output

376 - if i == 0:

377 - # wait for generation to complete before consuming

378 - # the remaining messages

379 - await asyncio.sleep(1)

380 - if i < 9:

381 - assert not output.finished

382 - i += 1

383 -

384 - assert i == 10

385 - assert final_output is not None

386 - assert len(final_output.outputs[0].token_ids) == 10

387 - assert final_output.finished

388 -

389 -

390 -@pytest.mark.asyncio(scope="module")

391 -async def test_invalid_argument(async_engine):

392 - scheduler_config = await async_engine.get_scheduler_config()

393 -

394 - if scheduler_config.num_scheduler_steps != 1:

395 - pytest.skip("no need to test this one with multistep")

396 -

397 - sampling_params = SamplingParams(

398 - temperature=0,

399 - min_tokens=10,

400 - max_tokens=10,

401 - )

402 -

403 - # Targeting specific DP rank only supported in v1 multi-instance DP

404 - with pytest.raises(ValueError):

405 - async for _ in async_engine.generate("test",

406 - sampling_params,

407 - request_id=uid(),

408 - data_parallel_rank=0):

409 - pass

tests/config/test_config.yaml

-1

@@ -2,4 +2,3 @@ port: 12312

2 served_model_name: mymodel

3 tensor_parallel_size: 2

4 trust_remote_code: true

5 -multi_step_stream_outputs: false

tests/config/test_config_with_model.yaml

-1

@@ -4,4 +4,3 @@ port: 12312

4 served_model_name: mymodel

5 tensor_parallel_size: 2

6 trust_remote_code: true

7 -multi_step_stream_outputs: false

tests/core/test_chunked_prefill_scheduler.py

+3 -7

@@ -644,11 +644,9 @@ def test_chunked_prefill_preempt():

644 assert out.num_batched_tokens == max_num_batched_tokens

645

646

647 -@pytest.mark.parametrize("num_scheduler_steps", [1, 5])

648 -def test_chunked_prefill_spec_prefill(num_scheduler_steps):

647 +def test_chunked_prefill_spec_prefill():

648 """Verify that the num_lookahead_slots is set appropriately for an all"""

650 - """prefill batch depending on whether multi-step scheduling is enabled"""

651 - """or not"""

649 + """prefill batch."""

650 block_size = 4

651 max_seqs = 30

652 max_model_len = 200

@@ -661,7 +659,6 @@ def test_chunked_prefill_spec_prefill(num_scheduler_steps):

659 max_model_len,

660 enable_chunked_prefill=True,

661 num_lookahead_slots=num_lookahead_slots,

664 - num_scheduler_steps=num_scheduler_steps,

662 )

663 cache_config = CacheConfig(block_size, 1.0, 1, "auto")

664 cache_config.num_cpu_blocks = 16

@@ -679,8 +676,7 @@ def test_chunked_prefill_spec_prefill(num_scheduler_steps):

676 assert out.num_prefill_groups == 1

677 assert out.num_batched_tokens == max_num_batched_tokens

678 print(out.num_lookahead_slots)

682 - assert out.num_lookahead_slots == (0 if (num_scheduler_steps == 1) else

683 - num_lookahead_slots)

679 + assert out.num_lookahead_slots == 0

680

681

682 def test_chunked_prefill_max_seqs():

tests/core/test_num_computed_tokens_update.py

+4 -20

@@ -6,7 +6,6 @@ import pytest

6 from tests.conftest import VllmRunner

7 from tests.core.utils import create_dummy_prompt

8 from vllm.engine.llm_engine import LLMEngine

9 -from vllm.platforms import current_platform

9 from vllm.sequence import SequenceGroup

11 MODEL = "JackFram/llama-160m"

@@ -17,32 +16,19 @@ def add_seq_group_to_engine(engine: LLMEngine, seq_group: SequenceGroup):

16 scheduler.add_seq_group(seq_group)

20 -@pytest.mark.parametrize("num_scheduler_steps", [1, 8])

19 @pytest.mark.parametrize("enable_chunked_prefill", [False, True])

20 @pytest.mark.parametrize("enforce_eager", [False, True])

23 -def test_num_computed_tokens_update(num_scheduler_steps: int,

24 - enable_chunked_prefill: bool,

21 +def test_num_computed_tokens_update(enable_chunked_prefill: bool,

22 enforce_eager: bool):

27 - is_multi_step = num_scheduler_steps > 1

28 - is_multi_step_chunked_prefill = is_multi_step and enable_chunked_prefill

29 -

30 - if is_multi_step_chunked_prefill and current_platform.is_rocm():

31 - pytest.skip("Multi-step with Chunked-Prefill does not support "

32 - "rocm_flash_attn backend")

33 -

24 # Make a vllm engine

25 runner = VllmRunner(model_name=MODEL,

26 gpu_memory_utilization=0.7,

37 - num_scheduler_steps=num_scheduler_steps,

27 enable_chunked_prefill=enable_chunked_prefill,

28 enforce_eager=enforce_eager)

29 engine: LLMEngine = runner.llm.llm_engine

42 - # In multi-step + chunked-prefill there is no separate single prompt step.

43 - # What is scheduled will run for num_scheduler_steps always.

44 - num_prompt_steps = num_scheduler_steps \

45 - if is_multi_step_chunked_prefill else 1

31 + num_prompt_steps = 1

33 num_output_tokens_list = [4, 8, 12, 15, 16, 17]

@@ -73,10 +59,8 @@ def test_num_computed_tokens_update(num_scheduler_steps: int,

59 # Test correctness of num_computed_tokens after the decode steps

60 assert seq.data.get_num_computed_tokens(

61 ) == prompt_num_computed_tokens + decode_step_counter

76 - for _ in range(num_scheduler_steps):

77 - # decode step

78 - engine.step()

79 - decode_step_counter += 1

62 + engine.step()

63 + decode_step_counter += 1

65 # Test correctness of num_computed_tokens after the sequence finish.

66 assert seq.data.get_num_computed_tokens(

tests/engine/test_multi_step_output_processor.py deleted

-274

@@ -1,274 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

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

3 -

4 -import random

5 -from unittest.mock import MagicMock

6 -

7 -import pytest

8 -from transformers import PreTrainedTokenizer

9 -

10 -from vllm.core.scheduler import Scheduler

11 -from vllm.engine.output_processor.multi_step import MultiStepOutputProcessor

12 -from vllm.engine.output_processor.stop_checker import StopChecker

13 -from vllm.sampling_params import SamplingParams

14 -from vllm.sequence import (CompletionSequenceGroupOutput, Logprob,

15 - SequenceOutput, SequenceStatus)

16 -from vllm.transformers_utils.detokenizer import Detokenizer

17 -from vllm.utils import Counter

18 -

19 -from ..core.utils import create_seq_group

20 -

21 -

22 -@pytest.mark.parametrize("seq_output_len", [128])

23 -@pytest.mark.parametrize("num_new_tokens", [1, 12])

24 -@pytest.mark.skip_global_cleanup

25 -def test_appends_token_ids(num_new_tokens: int, seq_output_len: int):

26 - """Verify multi-step decoding appends token ids correctly.

27 -

28 - We append token ids and verify all the token ids were appended correctly.

29 - Note that ignore_eos=True.

30 - """

31 - detokenizer = MagicMock(spec=Detokenizer)

32 - scheduler = MagicMock(spec=Scheduler)

33 - stop_checker = MagicMock(spec=StopChecker)

34 - seq_counter = Counter()

35 -

36 - output_processor = MultiStepOutputProcessor(

37 - detokenizer=detokenizer,

38 - scheduler=[scheduler],

39 - seq_counter=seq_counter,

40 - get_tokenizer_for_seq=lambda _: mock_tokenizer(),

41 - stop_checker=stop_checker,

42 - )

43 -

44 - seq_group = create_seq_group(

45 - seq_prompt_len=1024,

46 - seq_output_lens=[seq_output_len],

47 - sampling_params=SamplingParams(max_tokens=seq_output_len +

48 - num_new_tokens,

49 - ignore_eos=True),

50 - )

51 -

52 - seq = seq_group.get_seqs()[0]

53 - seq.status = SequenceStatus.RUNNING

54 -

55 - new_token_ids = list(range(num_new_tokens))

56 -

57 - outputs = [

58 - CompletionSequenceGroupOutput(

59 - samples=[

60 - SequenceOutput(

61 - parent_seq_id=seq.seq_id,

62 - output_token=output_token,

63 - logprobs={output_token: Logprob(0.0)},

64 - )

65 - ],

66 - prompt_logprobs=None,

67 - ) for output_token in new_token_ids

68 - ]

69 -

70 - assert seq.get_token_ids()[-len(new_token_ids):] != new_token_ids

71 - output_processor.process_outputs(seq_group, outputs)

72 - assert seq.get_token_ids()[-len(new_token_ids):] == new_token_ids

73 -

74 -

75 -@pytest.mark.parametrize("seq_prompt_len", [1024])

76 -@pytest.mark.parametrize("seq_output_len", [128])

77 -@pytest.mark.parametrize("num_new_tokens", [5, 6, 7, 8])

78 -@pytest.mark.parametrize("max_tokens", [128 + 3])

79 -@pytest.mark.skip_global_cleanup

80 -def test_respects_max_tokens(num_new_tokens: int, seq_prompt_len: int,

81 - seq_output_len: int, max_tokens: int):

82 - """Verify tokens after max_tokens are dropped and not appended to the

83 - sequence.

84 - """

85 - detokenizer = MagicMock(spec=Detokenizer)

86 - scheduler = MagicMock(spec=Scheduler)

87 - stop_checker = MagicMock(spec=StopChecker)

88 - seq_counter = Counter()

89 -

90 - output_processor = MultiStepOutputProcessor(

91 - detokenizer=detokenizer,

92 - scheduler=[scheduler],

93 - seq_counter=seq_counter,

94 - get_tokenizer_for_seq=lambda _: mock_tokenizer(),

95 - stop_checker=stop_checker,

96 - )

97 -

98 - seq_group = create_seq_group(

99 - seq_prompt_len=seq_prompt_len,

100 - seq_output_lens=[seq_output_len],

101 - sampling_params=SamplingParams(max_tokens=max_tokens, ),

102 - )

103 -

104 - seq = seq_group.get_seqs()[0]

105 - seq.status = SequenceStatus.RUNNING

106 -

107 - new_token_ids = list(range(num_new_tokens))

108 -

109 - outputs = [

110 - CompletionSequenceGroupOutput(

111 - samples=[

112 - SequenceOutput(

113 - parent_seq_id=seq.seq_id,

114 - output_token=output_token,

115 - logprobs={output_token: Logprob(0.0)},

116 - )

117 - ],

118 - prompt_logprobs=None,

119 - ) for output_token in new_token_ids

120 - ]

121 -

122 - assert seq.get_len() == seq_prompt_len + seq_output_len

123 - output_processor.process_outputs(seq_group, outputs)

124 -

125 - # Expect the processed sequence to not go over max tokens in len.

126 - assert seq.get_len() == seq_prompt_len + max_tokens

127 -

128 - # Expect the correct tokens were appended.

129 - expected_appended_tokens = new_token_ids[:max_tokens - seq_output_len]

130 - assert seq.get_token_ids(

131 - )[-len(expected_appended_tokens):] == expected_appended_tokens

132 -

133 -

134 -@pytest.mark.parametrize("seq_prompt_len", [1024])

135 -@pytest.mark.parametrize("seq_output_len", [128])

136 -@pytest.mark.parametrize("num_new_tokens", [12])

137 -@pytest.mark.parametrize("seed", list(range(6)))

138 -@pytest.mark.skip_global_cleanup

139 -def test_respects_eos_token_id(num_new_tokens: int, seq_prompt_len: int,

140 - seq_output_len: int, seed: int):

141 - """Verify the eos token id is included in the sequence, but subsequent

142 - tokens are dropped (not appended to sequence).

143 - """

144 - random.seed(seed)

145 - detokenizer = MagicMock(spec=Detokenizer)

146 - scheduler = MagicMock(spec=Scheduler)

147 - stop_checker = MagicMock(spec=StopChecker)

148 - seq_counter = Counter()

149 -

150 - eos_token_id = 100

151 -

152 - output_processor = MultiStepOutputProcessor(

153 - detokenizer=detokenizer,

154 - scheduler=[scheduler],

155 - seq_counter=seq_counter,

156 - get_tokenizer_for_seq=lambda _: mock_tokenizer(eos_token_id),

157 - stop_checker=stop_checker,

158 - )

159 -

160 - seq_group = create_seq_group(

161 - seq_prompt_len=seq_prompt_len,

162 - seq_output_lens=[seq_output_len],

163 - sampling_params=SamplingParams(

164 - # Ensure enough space.

165 - max_tokens=seq_output_len + num_new_tokens, ),

166 - )

167 -

168 - seq = seq_group.get_seqs()[0]

169 - seq.status = SequenceStatus.RUNNING

170 -

171 - new_token_ids = list(range(num_new_tokens))

172 - assert eos_token_id not in new_token_ids

173 - eos_index = random.randint(0, len(new_token_ids) - 1)

174 - new_token_ids[eos_index] = eos_token_id

175 -

176 - outputs = [

177 - CompletionSequenceGroupOutput(

178 - samples=[

179 - SequenceOutput(

180 - parent_seq_id=seq.seq_id,

181 - output_token=output_token,

182 - logprobs={output_token: Logprob(0.0)},

183 - )

184 - ],

185 - prompt_logprobs=None,

186 - ) for output_token in new_token_ids

187 - ]

188 -

189 - assert seq.get_len() == seq_prompt_len + seq_output_len

190 - output_processor.process_outputs(seq_group, outputs)

191 -

192 - # Expect the processed sequence to not go beyond provided eos.

193 - assert seq.get_len() == seq_prompt_len + seq_output_len + (eos_index + 1)

194 -

195 - # Expect the correct tokens were appended.

196 - expected_appended_tokens = new_token_ids[:eos_index + 1]

197 - assert seq.get_token_ids(

198 - )[-len(expected_appended_tokens):] == expected_appended_tokens

199 -

200 -

201 -@pytest.mark.parametrize("seq_prompt_len", [1024])

202 -@pytest.mark.parametrize("seq_output_len", [128])

203 -@pytest.mark.parametrize("num_new_tokens", [12])

204 -@pytest.mark.parametrize("seed", list(range(6)))

205 -@pytest.mark.skip_global_cleanup

206 -def test_ignores_eos_token_id(num_new_tokens: int, seq_prompt_len: int,

207 - seq_output_len: int, seed: int):

208 - """When sampling parameters dictate that we should ignore the eos token id,

209 - ensure all token ids are appended even if the eos token id is emitted.

210 - """

211 - random.seed(seed)

212 - detokenizer = MagicMock(spec=Detokenizer)

213 - scheduler = MagicMock(spec=Scheduler)

214 - stop_checker = MagicMock(spec=StopChecker)

215 - seq_counter = Counter()

216 -

217 - eos_token_id = 100

218 -

219 - output_processor = MultiStepOutputProcessor(

220 - detokenizer=detokenizer,

221 - scheduler=[scheduler],

222 - seq_counter=seq_counter,

223 - get_tokenizer_for_seq=lambda _: mock_tokenizer(eos_token_id),

224 - stop_checker=stop_checker,

225 - )

226 -

227 - seq_group = create_seq_group(

228 - seq_prompt_len=seq_prompt_len,

229 - seq_output_lens=[seq_output_len],

230 - sampling_params=SamplingParams(

231 - # Ensure enough space.

232 - max_tokens=seq_output_len + num_new_tokens,

233 - ignore_eos=True,

234 - ),

235 - )

236 -

237 - seq = seq_group.get_seqs()[0]

238 - seq.status = SequenceStatus.RUNNING

239 -

240 - new_token_ids = list(range(num_new_tokens))

241 - assert eos_token_id not in new_token_ids

242 - eos_index = random.randint(0, len(new_token_ids) - 1)

243 - new_token_ids[eos_index] = eos_token_id

244 -

245 - outputs = [

246 - CompletionSequenceGroupOutput(

247 - samples=[

248 - SequenceOutput(

249 - parent_seq_id=seq.seq_id,

250 - output_token=output_token,

251 - logprobs={output_token: Logprob(0.0)},

252 - )

253 - ],

254 - prompt_logprobs=None,

255 - ) for output_token in new_token_ids

256 - ]

257 -

258 - assert seq.get_len() == seq_prompt_len + seq_output_len

259 - output_processor.process_outputs(seq_group, outputs)

260 -

261 - # Expect the processed sequence to go beyond eos.

262 - assert seq.get_len() == seq_prompt_len + seq_output_len + num_new_tokens

263 -

264 - # Expect the correct tokens were appended.

265 - expected_appended_tokens = new_token_ids[:seq_output_len + num_new_tokens -

266 - seq_output_len]

267 - assert seq.get_token_ids(

268 - )[-len(expected_appended_tokens):] == expected_appended_tokens

269 -

270 -

271 -def mock_tokenizer(eos_token_id=1000):

272 - tokenizer = MagicMock(spec=PreTrainedTokenizer)

273 - tokenizer.eos_token_id = eos_token_id

274 - return tokenizer

tests/entrypoints/openai/correctness/test_lmeval.py

-3

@@ -26,15 +26,12 @@ DEFAULT_ARGS = ["--max-model-len", "4096"]

26 MORE_ARGS_LIST = [

27 [], # Default

28 ["--enable-chunked-prefill"], # Chunked

29 - ["--num-scheduler-steps", "8"], # MS

30 - ["--num-scheduler-steps", "8", "--multi-step-stream-outputs"] # MS+Stream

29 ]

30 MAX_WAIT_SECONDS = None

32 if current_platform.is_tpu():

33 MORE_ARGS_LIST = [

34 [], # Default

37 - # ["--num-scheduler-steps", "8"], # Multi-step << currently fails

35 ]

36 MAX_WAIT_SECONDS = 600

tests/metrics/test_metrics.py

-39

@@ -94,45 +94,6 @@ def test_metric_counter_generation_tokens(

94 f"metric: {metric_count!r}")

97 -@pytest.mark.parametrize("model", MODELS)

98 -@pytest.mark.parametrize("max_tokens", [128, 129])

99 -@pytest.mark.parametrize("disable_async_output_proc", [True, False])

100 -def test_metric_counter_generation_tokens_multi_step(

101 - vllm_runner,

102 - example_prompts,

103 - model: str,

104 - max_tokens: int,

105 - disable_async_output_proc: bool,

106 -) -> None:

107 - num_scheduler_steps = 8

108 - with vllm_runner(

109 - model,

110 - disable_log_stats=False,

111 - gpu_memory_utilization=0.4,

112 - num_scheduler_steps=num_scheduler_steps,

113 - disable_async_output_proc=disable_async_output_proc,

114 - ) as vllm_model:

115 - vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens)

116 - tokenizer = vllm_model.llm.get_tokenizer()

117 - stat_logger = vllm_model.llm.llm_engine.stat_loggers['prometheus']

118 - metric_count = stat_logger.metrics.counter_generation_tokens.labels(

119 - **stat_logger.labels)._value.get()

120 - vllm_generation_count = 0

121 - for i in range(len(example_prompts)):

122 - vllm_output_ids, vllm_output_str = vllm_outputs[i]

123 - prompt_ids = tokenizer.encode(example_prompts[i])

124 - # vllm_output_ids contains both prompt tokens and generation tokens.

125 - # We're interested only in the count of the generation tokens.

126 - vllm_generation_count += len(vllm_output_ids) - len(prompt_ids)

127 -

128 - # The multi-step scheduling will continue to execute forward even when

129 - # encountering EOS, leading to slightly imprecise metrics.

130 - assert abs(vllm_generation_count - metric_count) <\

131 - len(example_prompts) * num_scheduler_steps, \

132 - (f"generation token count: {vllm_generation_count!r}\n"

133 - f"metric: {metric_count!r}")

134 -

135 -

97 @pytest.mark.parametrize("model", MODELS)

98 @pytest.mark.parametrize("dtype", ["float"])

99 @pytest.mark.parametrize(

tests/models/language/generation/test_hybrid.py

-26

@@ -331,32 +331,6 @@ def test_state_cleanup(

331 "could be related to finished_requests_ids")

332

333

334 -@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])

335 -@pytest.mark.parametrize("max_tokens", [64])

336 -def test_multistep_correctness(

337 - vllm_runner,

338 - example_prompts,

339 - model: str,

340 - max_tokens: int,

341 -) -> None:

342 - with vllm_runner(model, num_scheduler_steps=8,

343 - max_num_seqs=2) as vllm_model:

344 - vllm_outputs_multistep = vllm_model.generate_greedy(

345 - example_prompts, max_tokens)

346 -

347 - with vllm_runner(model, num_scheduler_steps=1,

348 - max_num_seqs=2) as vllm_model:

349 - vllm_outputs_single_step = vllm_model.generate_greedy(

350 - example_prompts, max_tokens)

351 -

352 - check_outputs_equal(

353 - outputs_0_lst=vllm_outputs_multistep,

354 - outputs_1_lst=vllm_outputs_single_step,

355 - name_0="vllm_outputs_multistep",

356 - name_1="vllm_outputs_single_step",

357 - )

358 -

359 -

334 @multi_gpu_test(num_gpus=2)

335 @pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])

336 @pytest.mark.parametrize("max_tokens", [64])

tests/multi_step/test_correctness_async_llm.py deleted

-232

@@ -1,232 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

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

3 -

4 -# Test the AsyncLLMEngine with multi-step-decoding

5 -from typing import Optional

6 -

7 -import pytest

8 -

9 -from vllm.utils import STR_BACKEND_ENV_VAR

10 -

11 -from ..models.utils import check_logprobs_close

12 -from ..utils import (completions_with_server_args, get_client_text_generations,

13 - get_client_text_logprob_generations)

14 -

15 -MODELS = [

16 - "JackFram/llama-160m",

17 -]

18 -NUM_SCHEDULER_STEPS = [8] # Multi-step decoding steps

19 -NUM_PROMPTS = [10]

20 -

21 -DEFAULT_SERVER_ARGS: list[str] = [

22 - "--distributed-executor-backend",

23 - "ray",

24 - "--gpu-memory-utilization",

25 - "0.85",

26 - "--swap-space",

27 - "16",

28 -]

29 -

30 -

31 -@pytest.mark.parametrize("model", MODELS)

32 -@pytest.mark.parametrize(("tp_size, pp_size"), [

33 - (1, 1),

34 - (2, 2),

35 -])

36 -@pytest.mark.parametrize("eager_mode", [False, True])

37 -@pytest.mark.parametrize("num_scheduler_steps", NUM_SCHEDULER_STEPS)

38 -@pytest.mark.parametrize("num_prompts", NUM_PROMPTS)

39 -@pytest.mark.parametrize("num_logprobs", [5])

40 -@pytest.mark.parametrize("is_async", [True])

41 -@pytest.mark.parametrize("attention_backend", ["FLASHINFER", "FLASH_ATTN"])

42 -@pytest.mark.parametrize("enable_chunked_prefill", [True, False])

43 -@pytest.mark.asyncio

44 -async def test_multi_step(

45 - example_prompts,

46 - model: str,

47 - tp_size: int,

48 - pp_size: int,

49 - eager_mode: int,

50 - num_scheduler_steps: int,

51 - num_prompts: int,

52 - is_async: bool,

53 - num_logprobs: Optional[int],

54 - attention_backend: str,

55 - enable_chunked_prefill: bool,

56 - monkeypatch: pytest.MonkeyPatch,

57 -) -> None:

58 - """Test vLLM engine with multi-step scheduling in an OpenAI-protocol

59 - client/server environment.

60 -

61 - Set up an engine with single-step scheduling as a ground-truth reference.

62 -

63 - Send a completions API request to both engines with the same prompts.

64 -

65 - Validate:

66 - * Generated tokens match

67 - * Generated logprobs are all very close

68 -

69 - Args:

70 - example_prompts: test fixture providing example prompts

71 - model: model under test (same for single- and multi-step engines)

72 - tp_size: degree of tensor-parallelism

73 - pp_size: degree of pipeline-parallelism

74 - eager_mode

75 - num_scheduler_steps: for multi-step scheduling, GPU-side steps per

76 - GPU -> CPU output transfer

77 - num_prompts: number of example prompts under test

78 - num_logprobs: corresponds to the `logprobs` argument to the OpenAI

79 - completions endpoint; `None` -> no logprobs

80 - """

81 - if enable_chunked_prefill and \

82 - (pp_size > 1 or attention_backend != "FLASH_ATTN"):

83 - pytest.skip("Multi-step with Chunked-Prefill only supports"

84 - "PP=1 and FLASH_ATTN backend")

85 -

86 - with monkeypatch.context() as m:

87 - m.setenv(STR_BACKEND_ENV_VAR, attention_backend)

88 -

89 - prompts = example_prompts

90 - if len(prompts) < num_prompts:

91 - prompts = prompts * ((num_prompts // len(prompts)) + 1)

92 - prompts = prompts[:num_prompts]

93 - assert len(prompts) == num_prompts

94 -

95 - server_args = DEFAULT_SERVER_ARGS + ["--enforce-eager"]

96 - ms_server_args = DEFAULT_SERVER_ARGS + \

97 - ["--num-scheduler-steps", f"{num_scheduler_steps}"]

98 -

99 - if not is_async:

100 - ms_server_args += ["--disable-async-output-proc"]

101 -

102 - if eager_mode:

103 - ms_server_args.append("--enforce-eager")

104 -

105 - if enable_chunked_prefill:

106 - ms_server_args.append("--enable-chunked-prefill")

107 -

108 - distributed_args = [

109 - "--tensor-parallel-size",

110 - str(tp_size),

111 - "--pipeline-parallel-size",

112 - str(pp_size),

113 - ]

114 -

115 - # Spin up client/server & issue completion API requests.

116 - # Default `max_wait_seconds` is 240 but was empirically

117 - # was raised 5x to 1200 *just for this test* due to

118 - # observed timeouts in GHA CI

119 - ref_completions = await completions_with_server_args(

120 - prompts,

121 - model,

122 - server_args + distributed_args,

123 - num_logprobs,

124 - max_wait_seconds=5 * 240)

125 - test_completions = await completions_with_server_args(

126 - prompts,

127 - model,

128 - ms_server_args + distributed_args,

129 - num_logprobs,

130 - max_wait_seconds=5 * 240)

131 -

132 - # Assert multi-step scheduling produces identical tokens

133 - # to single-step scheduling.

134 - ref_generations = get_client_text_generations(ref_completions)

135 - test_generations = get_client_text_generations(test_completions)

136 - assert ref_generations == test_generations

137 -

138 - # Assert multi-step scheduling produces nearly-identical logprobs

139 - # to single-step scheduling.

140 - ref_text_logprobs = get_client_text_logprob_generations(

141 - ref_completions)

142 - test_text_logprobs = get_client_text_logprob_generations(

143 - test_completions)

144 - check_logprobs_close(

145 - outputs_0_lst=ref_text_logprobs,

146 - outputs_1_lst=test_text_logprobs,

147 - name_0="hf",

148 - name_1="vllm",

149 - )

150 -

151 -

152 -@pytest.mark.parametrize(("tp_size, pp_size"), [

153 - (1, 2),

154 -])

155 -@pytest.mark.asyncio

156 -async def test_multi_step_pp_smoke(

157 - tp_size: int,

158 - pp_size: int,

159 - monkeypatch: pytest.MonkeyPatch,

160 -) -> None:

161 - """

162 - Smoke test for the vLLM engine with multi-step scheduling in an

163 - OpenAI-protocol client/server environment.

164 -

165 - This tests compares the outputs between multi-step scheduling and

166 - single-step scheduling. Notably, this test lets the engines generate

167 - more tokens (default is 5) and test for an exact match over all the

168 - tokens.

169 -

170 - Args:

171 - tp_size: degree of tensor-parallelism

172 - pp_size: degree of pipeline-parallelism

173 - eager_mode

174 - """

175 -

176 - model = "JackFram/llama-160m"

177 - num_scheduler_steps = 8

178 - attention_backend = "FLASH_ATTN"

179 - max_num_seqs = 3

180 -

181 - with monkeypatch.context() as m:

182 - m.setenv(STR_BACKEND_ENV_VAR, attention_backend)

183 -

184 - # Prompt from the ShareGPT dataset

185 - prompts = [

186 - "in the jtbd context whats a push?", # codespell:ignore

187 - "in the jtbd context whats a push?", # codespell:ignore

188 - "in the jtbd context whats a push?", # codespell:ignore

189 - "in the jtbd context whats a push?", # codespell:ignore

190 - ]

191 - # Use varying max_tokens to introduce scheduling randomness.

192 - max_tokens = [10 * i for i in range(1, len(prompts) + 1)]

193 - assert len(prompts) == len(max_tokens)

194 -

195 - test_args = [

196 - "--tensor-parallel-size",

197 - str(tp_size), "--pipeline-parallel-size",

198 - str(pp_size), "--max-num-seqs",

199 - str(max_num_seqs)

200 - ]

201 -

202 - server_args = DEFAULT_SERVER_ARGS + test_args

203 - ms_server_args = DEFAULT_SERVER_ARGS + \

204 - ["--num-scheduler-steps", f"{num_scheduler_steps}"] + \

205 - test_args

206 -

207 - # Spin up client/server & issue completion API requests.

208 - # Default `max_wait_seconds` is 240 but was empirically

209 - # was raised 3x to 720 *just for this test* due to

210 - # observed timeouts in GHA CI

211 - ref_completions = await completions_with_server_args(

212 - prompts=prompts,

213 - model_name=model,

214 - server_cli_args=server_args,

215 - num_logprobs=None,

216 - max_wait_seconds=5 * 240,

217 - max_tokens=max_tokens)

218 -

219 - test_completions = await completions_with_server_args(

220 - prompts=prompts,

221 - model_name=model,

222 - server_cli_args=ms_server_args,

223 - num_logprobs=None,

224 - max_wait_seconds=5 * 240,

225 - max_tokens=max_tokens)

226 -

227 - # Assert multi-step scheduling produces identical tokens

228 - # to single-step scheduling.

229 - ref_generations = get_client_text_generations(ref_completions)

230 - test_generations = get_client_text_generations(test_completions)

231 -

232 - assert ref_generations == test_generations

tests/multi_step/test_correctness_llm.py deleted

-383

@@ -1,383 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

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

3 -

4 -# Test the LLMEngine with multi-step-decoding

5 -

6 -import copy

7 -from typing import Optional

8 -

9 -import pytest

10 -

11 -from vllm.platforms import current_platform

12 -from vllm.utils import STR_BACKEND_ENV_VAR

13 -

14 -from ..models.utils import check_logprobs_close, check_outputs_equal

15 -

16 -MODELS = [

17 - "JackFram/llama-160m",

18 -]

19 -NUM_SCHEDULER_STEPS = [8] # Multi-step decoding steps

20 -NUM_PROMPTS = [10]

21 -

22 -

23 -@pytest.mark.parametrize("model", MODELS)

24 -@pytest.mark.parametrize("dtype", ["half"])

25 -@pytest.mark.parametrize("tp_size", [1])

26 -@pytest.mark.parametrize("enable_chunked_prefill", [False, True])

27 -@pytest.mark.parametrize("max_tokens", [5])

28 -@pytest.mark.parametrize("enforce_eager", [True, False])

29 -@pytest.mark.parametrize("num_scheduler_steps", NUM_SCHEDULER_STEPS)

30 -@pytest.mark.parametrize("num_prompts", NUM_PROMPTS)

31 -@pytest.mark.parametrize("num_logprobs", [None, 5])

32 -@pytest.mark.parametrize("attention_backend", ["FLASH_ATTN", "FLASHINFER"])

33 -def test_multi_step_llm(

34 - hf_runner,

35 - vllm_runner,

36 - example_prompts,

37 - model: str,

38 - dtype: str,

39 - tp_size: int,

40 - enable_chunked_prefill: bool,

41 - max_tokens: int,

42 - enforce_eager: int,

43 - num_scheduler_steps: int,

44 - num_prompts: int,

45 - num_logprobs: Optional[int],

46 - attention_backend: str,

47 - monkeypatch: pytest.MonkeyPatch,

48 -) -> None:

49 - """Test vLLM engine with multi-step scheduling via sync LLM Engine.

50 -

51 - Set up a HuggingFace (HF) transformers model as a ground-truth reference.

52 -

53 - Prompt them with the same example prompts.

54 -

55 - Validate:

56 - * Generated tokens match

57 - * Generated logprobs are all very close

58 -

59 - Args:

60 - hf_runner: HF transformers model runner fixture

61 - vllm_runner: vLLM model runner fixture

62 - example_prompts: test fixture providing example prompts

63 - model: model under test (same for single- and multi-step engines)

64 - dtype: tensor datatype for engine to utilize

65 - tp_size: degree of tensor-parallelism

66 - enable_chunked_prefill: chunked-prefill on/off

67 - max_tokens: the maximum number of tokens to generate

68 - enforce_eager

69 - num_scheduler_steps: for multi-step scheduling, GPU-side steps per

70 - GPU -> CPU output transfer

71 - num_prompts: number of example prompts under test

72 - num_logprobs: corresponds to the `logprobs` argument to the OpenAI

73 - completions endpoint; `None` -> 1 logprob returned.

74 - """

75 - if current_platform.is_rocm() and \

76 - (attention_backend == "FLASHINFER" or enable_chunked_prefill):

77 - pytest.skip(

78 - "Multi-Step with FLASHINFER or Chunked-Prefill is not supported"

79 - "on ROCm")

80 -

81 - with monkeypatch.context() as m:

82 - m.setenv(STR_BACKEND_ENV_VAR, attention_backend)

83 -

84 - prompts = example_prompts

85 - if len(prompts) < num_prompts:

86 - prompts = prompts * ((num_prompts // len(prompts)) + 1)

87 - prompts = prompts[:num_prompts]

88 - assert len(prompts) == num_prompts

89 -

90 - with vllm_runner(

91 - model,

92 - dtype=dtype,

93 - enforce_eager=enforce_eager,

94 - gpu_memory_utilization=0.7,

95 - tensor_parallel_size=tp_size,

96 - enable_chunked_prefill=enable_chunked_prefill,

97 - num_scheduler_steps=num_scheduler_steps,

98 - ) as vllm_model:

99 - vllm_outputs = (vllm_model.generate_greedy(prompts, max_tokens)

100 - if num_logprobs is None else

101 - vllm_model.generate_greedy_logprobs(

102 - prompts, max_tokens, num_logprobs))

103 -

104 - with hf_runner(model, dtype=dtype) as hf_model:

105 - hf_outputs = (hf_model.generate_greedy(prompts, max_tokens)

106 - if num_logprobs is None else

107 - hf_model.generate_greedy_logprobs_limit(

108 - prompts, max_tokens, num_logprobs))

109 -

110 - if num_logprobs is None:

111 - check_outputs_equal(

112 - outputs_0_lst=hf_outputs,

113 - outputs_1_lst=vllm_outputs,

114 - name_0="hf",

115 - name_1="vllm",

116 - )

117 - else:

118 - check_logprobs_close(

119 - outputs_0_lst=hf_outputs,

120 - outputs_1_lst=vllm_outputs,

121 - name_0="hf",

122 - name_1="vllm",

123 - )

124 -

125 -

126 -@pytest.mark.parametrize("model", MODELS)

127 -@pytest.mark.parametrize("dtype", ["half"])

128 -@pytest.mark.parametrize("tp_size", [1])

129 -@pytest.mark.parametrize("max_tokens", [5])

130 -@pytest.mark.parametrize("enforce_eager", [True])

131 -@pytest.mark.parametrize("num_scheduler_steps", NUM_SCHEDULER_STEPS)

132 -@pytest.mark.parametrize("num_prompts", NUM_PROMPTS)

133 -@pytest.mark.parametrize("num_logprobs,num_prompt_logprobs", [(5, 5)])

134 -@pytest.mark.parametrize("attention_backend", ["FLASH_ATTN"])

135 -def test_multi_step_llm_w_prompt_logprobs(

136 - vllm_runner,

137 - example_prompts,

138 - model: str,

139 - dtype: str,

140 - tp_size: int,

141 - max_tokens: int,

142 - enforce_eager: int,

143 - num_scheduler_steps: int,

144 - num_prompts: int,

145 - num_logprobs: Optional[int],

146 - num_prompt_logprobs: Optional[int],

147 - attention_backend: str,

148 - monkeypatch: pytest.MonkeyPatch,

149 -) -> None:

150 - """Test prompt logprobs with multi-step scheduling via sync LLM Engine.

151 -

152 - Set up a vLLM engine instance w/ single-step scheduling as a ground-truth

153 - reference.

154 -

155 - Prompt them with the same example prompts.

156 -

157 - Validate:

158 - * All generated logprobs are all very close

159 -

160 - Args:

161 - hf_runner: HF transformers model runner fixture

162 - vllm_runner: vLLM model runner fixture

163 - example_prompts: test fixture providing example prompts

164 - model: model under test (same for single- and multi-step engines)

165 - dtype: tensor datatype for engine to utilize

166 - tp_size: degree of tensor-parallelism

167 - max_tokens: the maximum number of tokens to generate

168 - enforce_eager

169 - num_scheduler_steps: for multi-step scheduling, GPU-side steps per

170 - GPU -> CPU output transfer

171 - num_prompts: number of example prompts under test

172 - num_logprobs: corresponds to the `logprobs` argument to the OpenAI

173 - completions endpoint; `None` -> no logprobs

174 - num_prompt_logprobs: number of logprobs to return for each prompt token;

175 - note that this argument is not supported by the

176 - OpenAI completions endpoint.

177 - """

178 - with monkeypatch.context() as m:

179 - m.setenv(STR_BACKEND_ENV_VAR, attention_backend)

180 -

181 - prompts = example_prompts

182 - if len(prompts) < num_prompts:

183 - prompts = prompts * ((num_prompts // len(prompts)) + 1)

184 - prompts = prompts[:num_prompts]

185 - assert len(prompts) == num_prompts

186 -

187 - with vllm_runner(

188 - model,

189 - dtype=dtype,

190 - enforce_eager=enforce_eager,

191 - gpu_memory_utilization=0.7,

192 - tensor_parallel_size=tp_size,

193 - num_scheduler_steps=num_scheduler_steps,

194 - ) as vllm_model:

195 - vllm_outputs = vllm_model.generate_greedy_logprobs(

196 - prompts,

197 - max_tokens,

198 - num_logprobs,

199 - num_prompt_logprobs=num_prompt_logprobs)

200 -

201 - with vllm_runner(

202 - model,

203 - dtype=dtype,

204 - enforce_eager=enforce_eager,

205 - gpu_memory_utilization=0.7,

206 - tensor_parallel_size=tp_size,

207 - ) as vllm_model:

208 - single_step_vllm_outputs = vllm_model.generate_greedy_logprobs(

209 - prompts,

210 - max_tokens,

211 - num_logprobs,

212 - num_prompt_logprobs=num_prompt_logprobs)

213 -

214 - check_logprobs_close(

215 - outputs_0_lst=single_step_vllm_outputs,

216 - outputs_1_lst=vllm_outputs,

217 - name_0="hf",

218 - name_1="vllm",

219 - )

220 -

221 -

222 -@pytest.mark.parametrize("model", MODELS)

223 -@pytest.mark.parametrize("dtype", ["half"])

224 -@pytest.mark.parametrize("tp_size", [1])

225 -@pytest.mark.parametrize("max_tokens", [5])

226 -@pytest.mark.parametrize("enforce_eager", [True])

227 -@pytest.mark.parametrize("num_scheduler_steps", NUM_SCHEDULER_STEPS)

228 -@pytest.mark.parametrize("num_prompts", NUM_PROMPTS)

229 -@pytest.mark.parametrize("num_logprobs", [None, 5])

230 -@pytest.mark.parametrize("attention_backend", ["FLASH_ATTN"])

231 -@pytest.mark.skipif(

232 - current_platform.is_rocm(),

233 - reason="Multi-Step + Chunked-Prefill not supported on ROCm")

234 -def test_multi_step_llm_chunked_prefill_prefix_cache(

235 - vllm_runner,

236 - example_prompts,

237 - model: str,

238 - dtype: str,

239 - tp_size: int,

240 - max_tokens: int,

241 - enforce_eager: int,

242 - num_scheduler_steps: int,

243 - num_prompts: int,

244 - num_logprobs: Optional[int],

245 - attention_backend: str,

246 - monkeypatch: pytest.MonkeyPatch,

247 -) -> None:

248 - """Test vLLM engine with multi-step+"single-step chunked prefill"+APC.

249 -

250 - Set up contrived scenario which tests for a possible failure mode of

251 - scheduling with multi-step+"single-step chunked prefill"+APC

252 -

253 - "single-step chunked prefill" here refers to the current vLLM multi-step+

254 - chunked-prefill implementation, which requires that a prefill may only

255 - be scheduled in the same step as decodes if the prefill prompt fits in a

256 - single chunk (note that "complete" multi-step+chunked-prefill would allow

257 - a prefill to span multiple chunks & multiple steps but that is not yet

258 - the case.)

259 -

260 - "APC" is short for "automatic prefix caching".

261 -

262 - This test creates a scenario where the scheduler must decide whether/how

263 - to schedule a prefill with a prompt that exceeds the available token budget.

264 - The correct behavior for multi-step+"single-step chunked prefill"+APC is to

265 - put off scheduling the prefill until a future step.

266 -

267 - Validate that:

268 - * Multi-step kernels do not raise an exception due to incorrect scheduler

269 - behavior

270 - * Generated tokens match between

271 - multi-step+"single-step chunked prefill"+APC and

272 - single-step scheduling.

273 - * (If logprobs are enabled) check logprobs are close enough

274 -

275 - Args:

276 - vllm_runner: vLLM model runner fixture

277 - example_prompts: test fixture providing example prompts

278 - model: model under test (same for single- and multi-step engines)

279 - dtype: tensor datatype for engine to utilize

280 - tp_size: degree of tensor-parallelism

281 - max_tokens: the maximum number of tokens to generate

282 - enforce_eager

283 - num_scheduler_steps: for multi-step scheduling, GPU-side steps per

284 - GPU -> CPU output transfer

285 - num_prompts: number of example prompts under test

286 - num_logprobs: corresponds to the `logprobs` argument to the OpenAI

287 - completions endpoint; `None` -> 1 logprob returned.

288 - """

289 -

290 - # Set up contrived test for correct scheduling behavior with

291 - # multi-step+"single-step chunked prefill"+APC.

292 - #

293 - # Assume block_size=16

294 - #

295 - # Assume max_num_batched_tokens=48

296 - # => Per-step token budget=48

297 - #

298 - # 1. Scheduler schedules 0th prompt (24 tokens)

299 - # => Remaining token budget=24

300 - # 2. Scheduler attempts to schedule 1st prompt (30 tokens)

301 - # * 30 tokens exceeds 24 token remaining budget

302 - # * Correct behavior: do not schedule this prompt in this step

303 - # * Incorrect behavior: schedule prompt chunk

304 - # * `do_sample=False` for this prompt in this step

305 - # * Chunk size = (remaining tokens // block size) * block size

306 - #

307 - # The Incorrect scheduling behavior - if it occurs - will cause an exception

308 - # in the model runner resulting from `do_sample=False`.

309 - with monkeypatch.context() as m:

310 - m.setenv(STR_BACKEND_ENV_VAR, attention_backend)

311 -

312 - assert len(example_prompts) >= 2

313 - challenge_prompts = copy.deepcopy(example_prompts)

314 - challenge_prompts[0] = (

315 - 'vLLM is a high-throughput and memory-efficient '

316 - 'inference and serving engine for LLMs.\n') # 24 tok

317 - challenge_prompts[1] = (

318 - 'Briefly describe the major milestones in the '

319 - 'development of artificial intelligence from 1950 to 2020.\n'

320 - ) # 30 tok

321 -

322 - # If necessary, adjust the length of `challenge_prompts` to match

323 - # `num_prompts`

324 - if len(challenge_prompts) < num_prompts:

325 - challenge_prompts = (challenge_prompts *

326 - ((num_prompts // len(challenge_prompts)) + 1))

327 - challenge_prompts = challenge_prompts[:num_prompts]

328 - assert len(challenge_prompts) == num_prompts

329 -

330 - # Single-step scheduler baseline

331 - with vllm_runner(

332 - model,

333 - dtype=dtype,

334 - enforce_eager=enforce_eager,

335 - gpu_memory_utilization=0.7,

336 - tensor_parallel_size=tp_size,

337 - num_scheduler_steps=num_scheduler_steps,

338 - max_model_len=48,

339 - max_num_batched_tokens=48,

340 - max_num_seqs=4,

341 - block_size=16,

342 - ) as vllm_model:

343 - outputs_baseline = (

344 - vllm_model.generate_greedy(challenge_prompts, max_tokens) if

345 - num_logprobs is None else vllm_model.generate_greedy_logprobs(

346 - challenge_prompts, max_tokens, num_logprobs))

347 -

348 - # multi-step+"single-step chunked prefill"+APC

349 - with vllm_runner(

350 - model,

351 - dtype=dtype,

352 - enforce_eager=enforce_eager,

353 - gpu_memory_utilization=0.7,

354 - tensor_parallel_size=tp_size,

355 - enable_chunked_prefill=True,

356 - enable_prefix_caching=True,

357 - num_scheduler_steps=num_scheduler_steps,

358 - max_model_len=48,

359 - max_num_batched_tokens=48,

360 - max_num_seqs=4,

361 - block_size=16,

362 - ) as vllm_model:

363 - outputs_w_features = (

364 - vllm_model.generate_greedy(challenge_prompts, max_tokens) if

365 - num_logprobs is None else vllm_model.generate_greedy_logprobs(

366 - challenge_prompts, max_tokens, num_logprobs))

367 -

368 - if num_logprobs is None:

369 - # No-logprobs test

370 - check_outputs_equal(

371 - outputs_0_lst=outputs_baseline,

372 - outputs_1_lst=outputs_w_features,

373 - name_0="multi-step",

374 - name_1="multi-step+features",

375 - )

376 - else:

377 - # Yes-logprobs test

378 - check_logprobs_close(

379 - outputs_0_lst=outputs_baseline,

380 - outputs_1_lst=outputs_w_features,

381 - name_0="multi-step",

382 - name_1="multi-step+features",

383 - )

tests/samplers/test_logits_processor.py deleted

-70

@@ -1,70 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

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

3 -

4 -import pytest

5 -import torch

6 -

7 -from vllm import SamplingParams

8 -

9 -MODELS = ["distilbert/distilgpt2"]

10 -

11 -

12 -@pytest.fixture(scope="function", autouse=True)

13 -def use_v0_only(monkeypatch):

14 - """

15 - This file tests V0 internals, so set VLLM_USE_V1=0.

16 - """

17 - monkeypatch.setenv('VLLM_USE_V1', '0')

18 -

19 -

20 -@pytest.mark.parametrize("model", MODELS)

21 -@pytest.mark.parametrize("dtype", ["half"])

22 -def test_logits_processor_force_generate(

23 - vllm_runner,

24 - example_prompts,

25 - model: str,

26 - dtype: str,

27 -) -> None:

28 - with vllm_runner(model, dtype=dtype) as vllm_model:

29 - tokenizer = vllm_model.llm.get_tokenizer()

30 - repeat_times = 2

31 - enforced_answers = " vLLM"

32 - vllm_token_ids = tokenizer.encode(enforced_answers,

33 - add_special_tokens=False)

34 - max_tokens = len(vllm_token_ids) * repeat_times

35 -

36 - def pick_vllm(token_ids, logits):

37 - token_id = vllm_token_ids[len(token_ids) % len(vllm_token_ids)]

38 - logits[token_id] = torch.finfo(logits.dtype).max

39 - return logits

40 -

41 - params_with_logprobs = SamplingParams(

42 - logits_processors=[pick_vllm],

43 - prompt_logprobs=3,

44 - max_tokens=max_tokens,

45 - )

46 -

47 - # test logits_processors when prompt_logprobs is not None

48 - vllm_model.llm._add_request(

49 - example_prompts[0],

50 - params=params_with_logprobs,

51 - )

52 -

53 - # test prompt_logprobs is not None

54 - vllm_model.llm._add_request(

55 - example_prompts[1],

56 - params=SamplingParams(

57 - prompt_logprobs=3,

58 - max_tokens=max_tokens,

59 - ),

60 - )

61 -

62 - # test grouped requests

63 - vllm_model.llm._add_request(

64 - example_prompts[2],

65 - params=SamplingParams(max_tokens=max_tokens),

66 - )

67 -

68 - outputs = vllm_model.llm._run_engine(use_tqdm=False)

69 -

70 - assert outputs[0].outputs[0].text == enforced_answers * repeat_times

tests/tpu/lora/test_lora.py

-1

@@ -30,7 +30,6 @@ def use_v1_only(monkeypatch: pytest.MonkeyPatch):

31 def setup_vllm(num_loras: int, tp: int) -> vllm.LLM:

32 return vllm.LLM(model="Qwen/Qwen2.5-3B-Instruct",

33 - num_scheduler_steps=1,

33 max_model_len=256,

34 max_seq_len_to_capture=256,

35 max_num_seqs=8,

tests/utils_/test_utils.py

-2

@@ -236,7 +236,6 @@ def test_config_args(parser_with_config, cli_config_file):

236 ['serve', 'mymodel', '--config', cli_config_file])

237 assert args.tensor_parallel_size == 2

238 assert args.trust_remote_code

239 - assert not args.multi_step_stream_outputs

239

240

241 def test_config_file(parser_with_config):

@@ -828,7 +827,6 @@ def test_model_specification(parser_with_config, cli_config_file,

827 ])

828 assert args.tensor_parallel_size == 2

829 assert args.trust_remote_code is True

831 - assert args.multi_step_stream_outputs is False

830 assert args.port == 12312

831

832

tests/v1/test_oracle.py

-6

@@ -58,12 +58,6 @@ def test_unsupported_configs(monkeypatch):

58 disable_async_output_proc=True,

59 ).create_engine_config()

61 - with pytest.raises(NotImplementedError):

62 - AsyncEngineArgs(

63 - model=MODEL,

64 - num_scheduler_steps=5,

65 - ).create_engine_config()

66 -

61 with pytest.raises(NotImplementedError):

62 AsyncEngineArgs(

63 model=MODEL,

tests/worker/test_model_input.py

-79

@@ -11,7 +11,6 @@ from vllm.attention.backends.utils import CommonAttentionState

11 from vllm.model_executor import SamplingMetadata

12 from vllm.model_executor.pooling_metadata import PoolingMetadata

13 from vllm.worker.model_runner import ModelInputForGPUWithSamplingMetadata

14 -from vllm.worker.multi_step_model_runner import StatefulModelInput

14 from vllm.worker.pooling_model_runner import (

15 ModelInputForGPUWithPoolingMetadata)

@@ -166,81 +165,3 @@ def test_embedding_model_runner_input():

165 None) == getattr(attn_metadata, field.name, None)

166 # Pooling metadata is not broadcast.

167 assert received_model_input.pooling_metadata is None

169 -

170 -

171 -def test_multi_step_model_runner_input():

172 - sampling_metadata = SamplingMetadata(

173 - ["seq_group"],

174 - "selected_token_indices",

175 - "categorized_sample_indices",

176 - "num_prompts",

177 - )

178 - attn_metadata = AttentionMetadata(

179 - num_prefills=1,

180 - num_prefill_tokens=2,

181 - num_decode_tokens=3,

182 - slot_mapping=torch.zeros(1),

183 - multi_modal_placeholder_index_maps=None,

184 - enable_kv_scales_calculation=True,

185 - )

186 - frozen_model_input = ModelInputForGPUWithSamplingMetadata(

187 - input_tokens=torch.ones(10),

188 - input_positions=torch.ones(10),

189 - sampling_metadata=sampling_metadata,

190 - attn_metadata=attn_metadata)

191 -

192 - model_input = StatefulModelInput(

193 - frozen_model_input=frozen_model_input,

194 - is_last_step=True,

195 - is_first_multi_step=False,

196 - current_step=4,

197 - last_sampled_token_ids=torch.ones((10, 1)),

198 - is_multi_step=True,

199 - num_queries=8,

200 - num_seqs=5,

201 - cached_outputs=[],

202 - )

203 -

204 - assert isinstance(model_input, StatefulModelInput)

205 -

206 - # Test round trip serialization.

207 - tensor_dict = model_input.as_broadcastable_tensor_dict()

208 - attn_backend = MockAttentionBackend()

209 - received_model_input = (StatefulModelInput.from_broadcasted_tensor_dict(

210 - tensor_dict, attn_backend=attn_backend))

211 -

212 - received_frozen_input = received_model_input.frozen_model_input

213 -

214 - # Check that received copy has correct values.

215 - assert isinstance(received_model_input, StatefulModelInput)

216 - assert received_frozen_input.input_tokens is not None

217 - assert (received_frozen_input.input_tokens ==

218 - frozen_model_input.input_tokens).all()

219 - assert received_frozen_input.input_positions is not None

220 - assert (received_frozen_input.input_positions ==

221 - frozen_model_input.input_positions).all()

222 - assert received_frozen_input.multi_modal_kwargs is None

223 - assert (frozen_model_input.multi_modal_kwargs ==

224 - frozen_model_input.multi_modal_kwargs)

225 - assert received_frozen_input.lora_requests is None

226 - assert (received_frozen_input.lora_requests ==

227 - frozen_model_input.lora_requests)

228 - assert received_frozen_input.lora_mapping is None

229 - assert (

230 - received_frozen_input.lora_mapping == frozen_model_input.lora_mapping)

231 - for field in dataclasses.fields(AttentionMetadata):

232 - assert getattr(received_frozen_input.attn_metadata, field.name,

233 - None) == getattr(attn_metadata, field.name, None)

234 - # For sampling metadata, only selected_token_indices is copied.

235 - assert (received_frozen_input.sampling_metadata.selected_token_indices ==

236 - sampling_metadata.selected_token_indices)

237 - assert received_frozen_input.sampling_metadata.seq_groups is None

238 -

239 - # check non frozen fields

240 - assert received_model_input.is_last_step == model_input.is_last_step

241 - assert (received_model_input.is_first_multi_step ==

242 - model_input.is_first_multi_step)

243 - assert received_model_input.current_step == model_input.current_step

244 - assert (received_model_input.last_sampled_token_ids ==

245 - model_input.last_sampled_token_ids).all()

246 - assert received_model_input.is_multi_step == model_input.is_multi_step

vllm/config/__init__.py

-2

@@ -3779,8 +3779,6 @@ class VllmConfig:

3779 f"observability_config={self.observability_config!r}, "

3780 f"seed={self.model_config.seed}, "

3781 f"served_model_name={self.model_config.served_model_name}, "

3782 - f"num_scheduler_steps={self.scheduler_config.num_scheduler_steps}, "

3783 - f"multi_step_stream_outputs={self.scheduler_config.multi_step_stream_outputs}, " # noqa

3782 f"enable_prefix_caching={self.cache_config.enable_prefix_caching}, "

3783 f"chunked_prefill_enabled={self.scheduler_config.chunked_prefill_enabled}, " # noqa

3784 f"use_async_output_proc={self.model_config.use_async_output_proc}, "

vllm/core/scheduler.py

+3 -89

@@ -929,8 +929,7 @@ class Scheduler:

929 )

930

931 def _get_prompt_limit(self, seq_group: SequenceGroup) -> int:

932 - if (self.scheduler_config.chunked_prefill_enabled

933 - and not self.scheduler_config.is_multi_step):

932 + if self.scheduler_config.chunked_prefill_enabled:

933 prompt_limit = self.scheduler_config.max_model_len

934 else:

935 prompt_limit = min(

@@ -1114,9 +1113,6 @@ class Scheduler:

1113 continue

1114

1115 num_lookahead_slots: int = 0

1117 - if self.scheduler_config.is_multi_step and enable_chunking:

1118 - num_lookahead_slots = self._get_num_lookahead_slots(

1119 - True, enable_chunking)

1116

1117 # If the sequence group cannot be allocated, stop.

1118 can_allocate = self.block_manager.can_allocate(

@@ -1195,24 +1191,6 @@ class Scheduler:

1191 partial_prefill_metadata.maybe_increment_partial_prefills(

1192 seq_group)

1193

1198 - if enable_chunking and self.scheduler_config.is_multi_step:

1199 - blocks_to_copy: List[Tuple[int, int]] = []

1200 - # init_multi_step_from_lookahead_slots happens in append_slots

1201 - self._append_slots(seq_group, blocks_to_copy, enable_chunking)

1202 - # This assert will trip when a copy-on-write happens. This is

1203 - # not a concern as the very first sequence-group block

1204 - # allocation happens above. Still, we have the assert to

1205 - # catch any edge-cases.

1206 - assert not blocks_to_copy

1207 - else:

1208 - seq_group.init_multi_step_from_lookahead_slots(

1209 - num_lookahead_slots,

1210 - num_scheduler_steps=self.scheduler_config.

1211 - num_scheduler_steps,

1212 - is_multi_step=self.scheduler_config.is_multi_step,

1213 - enable_chunking=enable_chunking,

1214 - )

1215 -

1194 seq_groups.append(

1195 ScheduledSequenceGroup(seq_group=seq_group,

1196 token_chunk_size=num_new_tokens))

@@ -1453,14 +1431,6 @@ class Scheduler:

1431 num_prefill_groups = (len(prefills.seq_groups) +

1432 len(swapped_in.prefill_seq_groups) +

1433 len(running_scheduled.prefill_seq_groups))

1456 - # If all prompts, then we set num_lookahead_slots to 0

1457 - # this allows us to go through the `no_spec` path in

1458 - # `spec_decode_worker.py`

1459 - all_prefills = len(scheduled_seq_groups) == num_prefill_groups

1460 - num_lookahead_slots = (0 if

1461 - (all_prefills

1462 - and not self.scheduler_config.is_multi_step)

1463 - else running_scheduled.num_lookahead_slots)

1434 return SchedulerOutputs(

1435 scheduled_seq_groups=scheduled_seq_groups,

1436 num_prefill_groups=num_prefill_groups,

@@ -1472,7 +1442,7 @@ class Scheduler:

1442 swapped_in.blocks_to_copy,

1443 ignored_seq_groups=prefills.ignored_seq_groups +

1444 swapped_in.infeasible_seq_groups,

1475 - num_lookahead_slots=num_lookahead_slots,

1445 + num_lookahead_slots=0,

1446 running_queue_size=len(self.running),

1447 preempted=(len(running_scheduled.preempted) +

1448 len(running_scheduled.swapped_out)),

@@ -1516,11 +1486,6 @@ class Scheduler:

1486 num_lookahead_slots = self._get_num_lookahead_slots(

1487 is_prefill, enable_chunking)

1488

1519 - if is_prefill and num_lookahead_slots > 0:

1520 - # Appending prefill slots only happens multi-step and

1521 - # chunked-prefill are enabled together.

1522 - assert self.scheduler_config.is_multi_step and enable_chunking

1523 -

1489 return self.block_manager.can_append_slots(

1490 seq_group=seq_group, num_lookahead_slots=num_lookahead_slots)

1491

@@ -1776,19 +1741,7 @@ class Scheduler:

1741 num_lookahead_slots: int = self._get_num_lookahead_slots(

1742 is_prefill, enable_chunking)

1743

1779 - seq_group.init_multi_step_from_lookahead_slots(

1780 - num_lookahead_slots,

1781 - num_scheduler_steps=self.scheduler_config.num_scheduler_steps,

1782 - is_multi_step=self.scheduler_config.is_multi_step,

1783 - enable_chunking=enable_chunking,

1784 - )

1785 -

1744 seq_status: Optional[SequenceStatus] = SequenceStatus.RUNNING

1787 - if self.scheduler_config.is_multi_step and enable_chunking:

1788 - # In multi-step chunked-prefill any sequence type can have

1789 - # slots appended.

1790 - seq_status = None

1791 -

1745 for seq in seq_group.get_seqs(status=seq_status):

1746 cows = self.block_manager.append_slots(seq, num_lookahead_slots)

1747 if len(cows) > 0:

@@ -1904,29 +1857,8 @@ class Scheduler:

1857 """The number of slots to allocate per sequence per step, beyond known

1858 token ids. Speculative decoding uses these slots to store KV activations

1859 of tokens which may or may not be accepted.

1907 -

1908 - Speculative decoding does not yet support prefill, so we do not perform

1909 - lookahead allocation for prefill.

1910 -

1911 - When chunking is enabled with multi-step, we allocate lookahead slots

1912 - for the prefills for when the prefills turn into decodes in the first

1913 - step.

1860 """

1915 - if is_prefill:

1916 - if self.scheduler_config.is_multi_step and enable_chunking:

1917 - # num_lookahead_slots was introduced in the context of decodes,

1918 - # in Speculative Decoding.

1919 - # When the num_scheduler_steps is 8, say, then the

1920 - # num_lookahead_slots is 7. Meaning, we are doing a 1-step of

1921 - # decode anyways and we wish to do 7 more.

1922 - #

1923 - # "lookaheads" for prefills, is introduced in support for

1924 - # Chunked-Prefill in Multi-Step.

1925 - return self.scheduler_config.num_lookahead_slots + 1

1926 - else:

1927 - return 0

1928 -

1929 - return self.scheduler_config.num_lookahead_slots

1861 + return 0

1862

1863 def _get_num_new_uncached_and_cached_tokens(

1864 self,

@@ -2068,24 +2000,6 @@ class Scheduler:

2000 The number of new tokens to schedule after chunking.

2001 """

2002 remaining_token_budget = budget.remaining_token_budget()

2071 - if scheduler_config.is_multi_step:

2072 - # The current multi-step + chunked prefill capability does

2073 - # not actually support chunking prompts.

2074 - #

2075 - # Therefore, `num_new_tokens` is computed in the same fashion

2076 - # for both multi-step+chunked-prefill &

2077 - # multi-step+chunked-prefill+APC

2078 - #

2079 - # Prompts with more tokens than the current remaining budget

2080 - # are postponed to future scheduler steps

2081 - if num_new_tokens > prompt_limit:

2082 - # If the seq_group is in prompt-stage, pass the

2083 - # num_new_tokens as-is so the caller can ignore

2084 - # the sequence.

2085 - return num_new_tokens

2086 -

2087 - return 0 if num_new_tokens > \

2088 - remaining_token_budget else num_new_tokens

2003

2004 # Get the number of tokens to allocate to this prefill slot

2005 prefill_slot_budget = (

vllm/engine/arg_utils.py

+7 -36

@@ -362,8 +362,6 @@ class EngineArgs:

362 lora_dtype: Optional[Union[str, torch.dtype]] = LoRAConfig.lora_dtype

363 lora_extra_vocab_size: int = LoRAConfig.lora_extra_vocab_size

364

365 - num_scheduler_steps: int = SchedulerConfig.num_scheduler_steps

366 - multi_step_stream_outputs: bool = SchedulerConfig.multi_step_stream_outputs

365 ray_workers_use_nsight: bool = ParallelConfig.ray_workers_use_nsight

366 num_gpu_blocks_override: Optional[

367 int] = CacheConfig.num_gpu_blocks_override

@@ -799,11 +797,8 @@ class EngineArgs:

797 **scheduler_kwargs["delay_factor"])

798 scheduler_group.add_argument("--preemption-mode",

799 **scheduler_kwargs["preemption_mode"])

802 - scheduler_group.add_argument("--num-scheduler-steps",

803 - **scheduler_kwargs["num_scheduler_steps"])

804 - scheduler_group.add_argument(

805 - "--multi-step-stream-outputs",

806 - **scheduler_kwargs["multi_step_stream_outputs"])

800 + # multi-step scheduling has been removed; corresponding arguments

801 + # are no longer supported.

802 scheduler_group.add_argument("--scheduling-policy",

803 **scheduler_kwargs["policy"])

804 scheduler_group.add_argument(

@@ -1257,28 +1252,11 @@ class EngineArgs:

1252 disable_log_stats=self.disable_log_stats,

1253 )

1254

1260 - # Reminder: Please update docs/features/compatibility_matrix.md

1261 - # If the feature combo become valid

1262 - if self.num_scheduler_steps > 1:

1263 - if speculative_config is not None:

1264 - raise ValueError("Speculative decoding is not supported with "

1265 - "multi-step (--num-scheduler-steps > 1)")

1266 - if self.enable_chunked_prefill and self.pipeline_parallel_size > 1:

1267 - raise ValueError("Multi-Step Chunked-Prefill is not supported "

1268 - "for pipeline-parallel-size > 1")

1269 - if current_platform.is_cpu():

1270 - logger.warning("Multi-Step (--num-scheduler-steps > 1) is "

1271 - "currently not supported for CPUs and has been "

1272 - "disabled.")

1273 - self.num_scheduler_steps = 1

1274 -

1275 - # make sure num_lookahead_slots is set the higher value depending on

1276 - # if we are using speculative decoding or multi-step

1277 - num_lookahead_slots = max(self.num_lookahead_slots,

1278 - self.num_scheduler_steps - 1)

1279 - num_lookahead_slots = num_lookahead_slots \

1280 - if speculative_config is None \

1281 - else speculative_config.num_lookahead_slots

1255 + # make sure num_lookahead_slots is set appropriately depending on

1256 + # whether speculative decoding is enabled

1257 + num_lookahead_slots = self.num_lookahead_slots

1258 + if speculative_config is not None:

1259 + num_lookahead_slots = speculative_config.num_lookahead_slots

1260

1261 scheduler_config = SchedulerConfig(

1262 runner_type=model_config.runner_type,

@@ -1292,8 +1270,6 @@ class EngineArgs:

1270 disable_chunked_mm_input=self.disable_chunked_mm_input,

1271 is_multimodal_model=model_config.is_multimodal_model,

1272 preemption_mode=self.preemption_mode,

1295 - num_scheduler_steps=self.num_scheduler_steps,

1296 - multi_step_stream_outputs=self.multi_step_stream_outputs,

1273 send_delta_data=(envs.VLLM_USE_RAY_SPMD_WORKER

1274 and parallel_config.use_ray),

1275 policy=self.scheduling_policy,

@@ -1392,11 +1368,6 @@ class EngineArgs:

1368 recommend_to_remove=True)

1369 return False

1370

1395 - if self.num_scheduler_steps != SchedulerConfig.num_scheduler_steps:

1396 - _raise_or_fallback(feature_name="--num-scheduler-steps",

1397 - recommend_to_remove=True)

1398 - return False

1399 -

1371 if self.scheduler_delay_factor != SchedulerConfig.delay_factor:

1372 _raise_or_fallback(feature_name="--scheduler-delay-factor",

1373 recommend_to_remove=True)

vllm/engine/async_llm_engine.py

+2 -24

@@ -15,7 +15,7 @@ from vllm.config import (DecodingConfig, LoRAConfig, ModelConfig,

15 from vllm.core.scheduler import SchedulerOutputs

16 from vllm.engine.arg_utils import AsyncEngineArgs

17 from vllm.engine.async_timeout import asyncio_timeout

18 -from vllm.engine.llm_engine import LLMEngine, SchedulerOutputState

18 +from vllm.engine.llm_engine import LLMEngine

19 from vllm.engine.metrics_types import StatLoggerBase

20 from vllm.engine.protocol import EngineClient

21 from vllm.executor.executor_base import ExecutorBase

@@ -308,13 +308,6 @@ class _AsyncLLMEngine(LLMEngine):

308 if not allow_async_output_proc and len(ctx.output_queue) > 0:

309 self._process_model_outputs(ctx=ctx)

310

311 - if (self.scheduler_config.is_multi_step

312 - and scheduler_outputs.num_lookahead_slots > 0):

313 - # cache the scheduler outputs for the next iteration if we have

314 - # lookahead slots

315 - self._cache_scheduler_outputs_for_multi_step(

316 - virtual_engine, seq_group_metadata_list, scheduler_outputs,

317 - allow_async_output_proc)

311 else:

312 finished_requests_ids = list()

313

@@ -351,29 +344,14 @@ class _AsyncLLMEngine(LLMEngine):

344 outputs = await self.model_executor.execute_model_async(

345 execute_model_req)

346

354 - # we need to do this here so that last step's sampled_token_ids can

355 - # be passed to the next iteration for PP.

356 - if self.scheduler_config.is_multi_step:

357 - self._update_cached_scheduler_output(virtual_engine, outputs)

347 else:

348 if len(ctx.output_queue) > 0:

349 self._process_model_outputs(ctx=ctx)

350 outputs = []

351

363 - # Finish the current step for all the sequence groups.

364 - if self.scheduler_config.is_multi_step:

365 - for seq_group in seq_group_metadata_list:

366 - seq_group.finish_step()

367 -

352 if not self._has_remaining_steps(seq_group_metadata_list):

369 - # Clear the cache if we have finished all the steps

370 - if self.scheduler_config.is_multi_step:

371 - self.cached_scheduler_outputs[

372 - virtual_engine] = SchedulerOutputState()

373 -

353 # is_first_step_output is True only when the num_steps of all

375 - # the sequences are 1. When the num_steps > 1,

376 - # multi_step_model_runner does the first-step output append.

354 + # the sequences are 1.

355 is_first_step_output: bool = False if not seq_group_metadata_list \

356 else seq_group_metadata_list[0].state.num_steps == 1

357

vllm/engine/llm_engine.py

+16 -162

@@ -25,7 +25,6 @@ from vllm.engine.metrics_types import StatLoggerBase, Stats

25 from vllm.engine.output_processor.interfaces import (

26 SequenceGroupOutputProcessor)

27 from vllm.engine.output_processor.stop_checker import StopChecker

28 -from vllm.engine.output_processor.util import create_output_by_sequence_group

28 from vllm.entrypoints.openai.logits_processors import (

29 get_logits_processors as get_openai_logits_processors)

30 from vllm.executor.executor_base import ExecutorBase

@@ -91,7 +90,7 @@ class OutputData(NamedTuple):

91 class SchedulerContext:

94 - def __init__(self, multi_step_stream_outputs: bool = False):

93 + def __init__(self) -> None:

94 self.output_queue: Deque[OutputData] = deque()

95 self.request_outputs: List[Union[RequestOutput,

96 PoolingRequestOutput]] = []

@@ -99,8 +98,6 @@ class SchedulerContext:

98 List[SequenceGroupMetadata]] = None

99 self.scheduler_outputs: Optional[SchedulerOutputs] = None

100

102 - self.multi_step_stream_outputs: bool = multi_step_stream_outputs

103 -

101 def append_output(self, outputs: List[SamplerOutput],

102 seq_group_metadata_list: List[SequenceGroupMetadata],

103 scheduler_outputs: SchedulerOutputs, is_async: bool,

@@ -303,8 +300,7 @@ class LLMEngine:

300 ]

301

302 self.scheduler_contexts = [

306 - SchedulerContext(multi_step_stream_outputs=self.scheduler_config.

307 - multi_step_stream_outputs)

303 + SchedulerContext()

304 for _ in range(self.parallel_config.pipeline_parallel_size)

305 ]

306

@@ -683,8 +679,7 @@ class LLMEngine:

679 "Priority scheduling is not enabled.")

680

681 if isinstance(params, SamplingParams) \

686 - and params.logits_processors \

687 - and self.scheduler_config.num_scheduler_steps > 1:

682 + and params.logits_processors:

683 raise ValueError(

684 "Logits processors are not supported in multi-step decoding")

685

@@ -868,45 +863,6 @@ class LLMEngine:

863

864 return

865

871 - def _update_num_computed_tokens_for_multi_step_prefill(

872 - self, seq_group: SequenceGroup,

873 - seq_group_meta: SequenceGroupMetadata,

874 - is_first_step_output: Optional[bool]):

875 - """

876 - This function updates num_computed_tokens for prompt sequences

877 - when Multi-Step is enabled.

878 -

879 - seq_group: SequenceGroup to update the num_computed_tokens for.

880 - seq_group_meta: Metadata of the given SequenceGroup.

881 - is_first_step_output: Optional[bool] -

882 - When available, is_first_step_output indicates if the appended

883 - output token is the output of the first-step in multi-step.

884 - A value of None indicates that outputs from all steps in

885 - in multi-step are submitted in a single burst.

886 - """

887 -

888 - assert self.scheduler_config.is_multi_step

889 -

890 - if not seq_group_meta.is_prompt:

891 - # num_computed_token updates for multi-step decodes happen after

892 - # the tokens are appended to the sequence.

893 - return

894 -

895 - do_update: bool = False

896 - if self.scheduler_config.chunked_prefill_enabled:

897 - # In multi-step + chunked-prefill case, the prompt sequences

898 - # that are scheduled are fully processed in the first step.

899 - do_update = is_first_step_output is None or is_first_step_output

900 - else:

901 - # Normal multi-step decoding case. In this case prompt-sequences

902 - # are actually single-stepped. Always update in this case.

903 - assert seq_group.state.num_steps == 1

904 - do_update = True

905 -

906 - if do_update:

907 - seq_group.update_num_computed_tokens(

908 - seq_group_meta.token_chunk_size)

909 -

866 def _process_model_outputs(self,

867 ctx: SchedulerContext,

868 request_id: Optional[str] = None) -> None:

@@ -939,33 +895,8 @@ class LLMEngine:

895

896 has_multiple_outputs: bool = len(outputs) > 1

897 outputs_by_sequence_group: List[List[SequenceGroupOutput]]

942 - if has_multiple_outputs:

943 - assert self.scheduler_config.is_multi_step or \

944 - self.speculative_config

945 - # Organize outputs by [step][sequence group] instead of

946 - # [sequence group][step].

947 - if self.scheduler_config.is_multi_step:

948 - outputs_by_sequence_group = create_output_by_sequence_group(

949 - outputs, len(seq_group_metadata_list))

950 - elif self.speculative_config:

951 - # Decodes are multi-steps while prefills are not, outputting at

952 - # most 1 token. Separate them so that we can trigger chunk

953 - # processing without having to pad or copy over prompts K times

954 - # to match decodes structure (costly with prompt_logprobs).

955 - num_prefills = sum(sg.is_prompt

956 - for sg in seq_group_metadata_list)

957 - prefills, decodes = outputs[:num_prefills], outputs[

958 - num_prefills:]

959 - outputs_by_sequence_group = create_output_by_sequence_group(

960 - decodes,

961 - num_seq_groups=len(seq_group_metadata_list) - num_prefills)

962 - outputs_by_sequence_group = [p.outputs for p in prefills

963 - ] + outputs_by_sequence_group

964 - # We have outputs for multiple steps submitted in a single burst,

965 - # so invalidate is_first_step_output.

966 - is_first_step_output = None

967 - else:

968 - outputs_by_sequence_group = outputs

898 + assert not has_multiple_outputs

899 + outputs_by_sequence_group = outputs

900

901 # Determine the requests we need to operate on

902 if request_id:

@@ -1006,13 +937,8 @@ class LLMEngine:

937 output = [outputs_by_sequence_group[0][i]]

938

939 if not is_async:

1009 - if self.scheduler_config.is_multi_step:

1010 - # Updates happen only if the sequence is prefill

1011 - self._update_num_computed_tokens_for_multi_step_prefill(

1012 - seq_group, seq_group_meta, is_first_step_output)

1013 - else:

1014 - seq_group.update_num_computed_tokens(

1015 - seq_group_meta.token_chunk_size or 0)

940 + seq_group.update_num_computed_tokens(

941 + seq_group_meta.token_chunk_size or 0)

942

943 if outputs:

944 for o in outputs:

@@ -1074,15 +1000,6 @@ class LLMEngine:

1000 for scheduler in self.scheduler:

1001 scheduler.free_finished_seq_groups()

1002

1077 - # For multi-step without streaming, don't create outputs each iteration

1078 - if not is_last_step and not ctx.multi_step_stream_outputs:

1079 - # Immediately process request outputs here (if callback is given)

1080 - if (finished_now

1081 - and self.process_request_outputs_callback is not None):

1082 - self.process_request_outputs_callback(ctx.request_outputs)

1083 - ctx.request_outputs.clear()

1084 - return

1085 -

1003 # Create the outputs

1004 for i in indices:

1005 if i in skip or i in finished_before or i in finished_now:

@@ -1101,13 +1018,7 @@ class LLMEngine:

1018 if request_output:

1019 ctx.request_outputs.append(request_output)

1020

1104 - # For multi-step with streaming, create outputs each iteration

1105 - if not is_last_step and ctx.multi_step_stream_outputs:

1106 - # Immediately process request outputs here (if callback is given)

1107 - if self.process_request_outputs_callback is not None:

1108 - self.process_request_outputs_callback(ctx.request_outputs)

1109 - ctx.request_outputs.clear()

1110 - return

1021 + # Create outputs only after processing the scheduler's results

1022

1023 for seq_group in scheduler_outputs.ignored_seq_groups:

1024 params = seq_group.sampling_params

@@ -1157,16 +1068,10 @@ class LLMEngine:

1068 if seq_group.is_finished():

1069 continue

1070

1160 - if self.scheduler_config.is_multi_step:

1161 - # Updates happen only if the sequence is prefill

1162 - self._update_num_computed_tokens_for_multi_step_prefill(

1163 - seq_group, seq_group_metadata,

1164 - seq_group.state.num_steps == 1)

1165 - else:

1166 - token_chunk_size = (seq_group_metadata.token_chunk_size

1167 - if seq_group_metadata.token_chunk_size

1168 - is not None else 0)

1169 - seq_group.update_num_computed_tokens(token_chunk_size)

1071 + token_chunk_size = (seq_group_metadata.token_chunk_size

1072 + if seq_group_metadata.token_chunk_size

1073 + is not None else 0)

1074 + seq_group.update_num_computed_tokens(token_chunk_size)

1075

1076 if seq_group_metadata.do_sample:

1077 assert len(sequence_group_outputs.samples) == 1, (

@@ -1177,16 +1082,8 @@ class LLMEngine:

1082 assert len(seq_group.seqs) == 1

1083 seq = seq_group.seqs[0]

1084

1180 - if self.scheduler_config.is_multi_step:

1181 - is_prefill_append = seq.data.get_num_uncomputed_tokens(

1182 - ) == 0

1183 - seq.append_token_id(sample.output_token, sample.logprobs,

1184 - sample.output_embed)

1185 - if not is_prefill_append:

1186 - seq_group.update_num_computed_tokens(1)

1187 - else:

1188 - seq.append_token_id(sample.output_token, sample.logprobs,

1189 - sample.output_embed)

1085 + seq.append_token_id(sample.output_token, sample.logprobs,

1086 + sample.output_embed)

1087

1088 def step(self) -> List[Union[RequestOutput, PoolingRequestOutput]]:

1089 """Performs one decoding iteration and returns newly generated results.

@@ -1289,13 +1186,6 @@ class LLMEngine:

1186 if not allow_async_output_proc and len(ctx.output_queue) > 0:

1187 self._process_model_outputs(ctx=ctx)

1188

1292 - if (self.scheduler_config.is_multi_step

1293 - and scheduler_outputs.num_lookahead_slots > 0):

1294 - # cache the scheduler outputs for the next iteration if we have

1295 - # lookahead slots

1296 - self._cache_scheduler_outputs_for_multi_step(

1297 - virtual_engine, seq_group_metadata_list, scheduler_outputs,

1298 - allow_async_output_proc)

1189 else:

1190 finished_requests_ids = list()

1191

@@ -1345,10 +1235,6 @@ class LLMEngine:

1235 # Raise so the caller is notified that this request failed

1236 raise

1237

1348 - # We need to do this here so that last step's sampled_token_ids can

1349 - # be passed to the next iteration for PP.

1350 - if self.scheduler_config.is_multi_step:

1351 - self._update_cached_scheduler_output(virtual_engine, outputs)

1238 else:

1239 # Nothing scheduled => If there is pending async postprocessor,

1240 # then finish it here.

@@ -1357,19 +1243,9 @@ class LLMEngine:

1243 # No outputs in this case

1244 outputs = []

1245

1360 - # Finish the current step for all the sequence groups.

1361 - if self.scheduler_config.is_multi_step:

1362 - for seq_group in seq_group_metadata_list:

1363 - seq_group.finish_step()

1364 -

1246 if not self._has_remaining_steps(seq_group_metadata_list):

1366 - # clear the cache if we have finished all the steps.

1367 - if self.scheduler_config.is_multi_step:

1368 - self.cached_scheduler_outputs[0] = SchedulerOutputState()

1369 -

1247 # is_first_step_output is True only when the num_steps of all

1371 - # the sequences are 1. When the num_steps > 1,

1372 - # multi_step_model_runner does the first-step output append.

1248 + # the sequences are 1.

1249 is_first_step_output: bool = False if not seq_group_metadata_list \

1250 else seq_group_metadata_list[0].state.num_steps == 1

1251

@@ -1453,22 +1329,7 @@ class LLMEngine:

1329 def _has_remaining_steps(

1330 self, seq_group_metadata_list: Optional[List[SequenceGroupMetadata]]

1331 ) -> bool:

1456 - if (not self.scheduler_config.is_multi_step

1457 - or not seq_group_metadata_list):

1458 - return False

1459 -

1460 - # TODO(will) this is a sanity check for nowto make sure that all the

1461 - # seqs are on the same steps. Eventually we will want to do some sort of

1462 - # dynamic scheduling when doing multi-step decoding.

1463 - ref_remaining_steps = seq_group_metadata_list[0].state.remaining_steps

1464 - if any([

1465 - seq_group.state.remaining_steps != ref_remaining_steps

1466 - for seq_group in seq_group_metadata_list[1:]

1467 - ]):

1468 - raise AssertionError("All running sequence groups should "

1469 - "have the same remaining steps.")

1470 -

1471 - return ref_remaining_steps > 0

1332 + return False

1333

1334 def _cache_scheduler_outputs_for_multi_step(

1335 self, virtual_engine: int,

@@ -1497,13 +1358,6 @@ class LLMEngine:

1358

1359 def _get_last_sampled_token_ids(

1360 self, virtual_engine: int) -> Optional[torch.Tensor]:

1500 - cached_last_output = self.cached_scheduler_outputs[

1501 - virtual_engine].last_output

1502 - if (self.scheduler_config.is_multi_step

1503 - and self.parallel_config.pipeline_parallel_size > 1

1504 - and cached_last_output is not None

1505 - and cached_last_output.sampled_token_ids_cpu is not None):

1506 - return cached_last_output.sampled_token_ids_cpu

1361 return None

1362

1363 def add_logger(self, logger_name: str, logger: StatLoggerBase) -> None:

vllm/engine/output_processor/interfaces.py

+6 -20

@@ -36,27 +36,13 @@ class SequenceGroupOutputProcessor(ABC):

36 ):

37 """Create an output processor.

39 - This returns a single-step output processor if num_lookahead_slots is

40 - zero, else returns a multi-step output processor.

39 + Multi-step scheduling is no longer supported. Always return a

40 + single-step output processor.

41 """

42 - if scheduler_config.num_lookahead_slots == 0:

43 - # Importing here to avoid cycle.

44 - from vllm.engine.output_processor.single_step import (

45 - SingleStepOutputProcessor)

46 - return SingleStepOutputProcessor(scheduler_config, detokenizer,

47 - scheduler, seq_counter,

48 - stop_checker)

49 - else:

50 - # Importing here to avoid cycle.

51 - from vllm.engine.output_processor.multi_step import (

52 - MultiStepOutputProcessor)

53 - return MultiStepOutputProcessor(

54 - detokenizer,

55 - scheduler,

56 - seq_counter,

57 - get_tokenizer_for_seq,

58 - stop_checker,

59 - )

42 + from vllm.engine.output_processor.single_step import (

43 + SingleStepOutputProcessor)

44 + return SingleStepOutputProcessor(scheduler_config, detokenizer,

45 + scheduler, seq_counter, stop_checker)

47 @abstractmethod

48 def process_outputs(self, sequence_group: SequenceGroup,

vllm/engine/output_processor/multi_step.py deleted

-211

@@ -1,211 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

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

3 -

4 -import functools

5 -from typing import Callable, List, cast

6 -

7 -from vllm.core.scheduler import Scheduler

8 -from vllm.engine.output_processor.interfaces import (

9 - SequenceGroupOutputProcessor)

10 -from vllm.engine.output_processor.single_step import (

11 - single_step_process_prompt_logprob)

12 -from vllm.engine.output_processor.stop_checker import StopChecker

13 -from vllm.logger import init_logger

14 -from vllm.sampling_params import SamplingParams

15 -from vllm.sequence import (VLLM_INVALID_TOKEN_ID,

16 - CompletionSequenceGroupOutput, Sequence,

17 - SequenceGroup, SequenceGroupOutput, SequenceOutput,

18 - SequenceStatus)

19 -from vllm.transformers_utils.detokenizer import Detokenizer

20 -from vllm.transformers_utils.tokenizer import AnyTokenizer

21 -from vllm.utils import Counter

22 -

23 -logger = init_logger(__name__)

24 -

25 -

26 -class MultiStepOutputProcessor(SequenceGroupOutputProcessor):

27 - """SequenceGroupOutputProcessor which handles logic related to

28 - detokenization and stopping conditions. It specializes to "multi-step

29 - decoding", where vLLM's worker may generate multiple tokens per invocation.

30 - This is currently mutually exclusive with advanced sampling techniques like

31 - beam search, which motivates the separation of this logic from the single

32 - step output processor.

33 -

34 - This class is responsible for things such as correctly appending all new

35 - token ids to their sequence, detokenizing new token ids, truncating new

36 - output tokens after an eos token, and correctly handling the case where the

37 - number of new output tokens per sequence differs in a single batch.

38 - """

39 -

40 - def __init__(

41 - self,

42 - detokenizer: Detokenizer,

43 - scheduler: List[Scheduler],

44 - seq_counter: Counter,

45 - get_tokenizer_for_seq: Callable[[Sequence], AnyTokenizer],

46 - stop_checker: StopChecker,

47 - ):

48 - self.detokenizer = detokenizer

49 - self.scheduler = scheduler

50 - self.seq_counter = seq_counter

51 - self.get_tokenizer_for_seq = get_tokenizer_for_seq

52 - self.stop_checker = stop_checker

53 -

54 - def process_prompt_logprob(self, seq_group: SequenceGroup,

55 - outputs: List[SequenceGroupOutput]) -> None:

56 - """Process prompt logprobs associated with each step of a multi-step-

57 - scheduled computation.

58 -

59 - Args:

60 - seq_group: the outputs are associated with this

61 - [`SequenceGroup`][vllm.sequence.SequenceGroup]

62 - outputs: the

63 - [`SequenceGroupOutput`][vllm.sequence.SequenceGroupOutput]s

64 - for all scheduler steps

65 - """

66 - for output in outputs:

67 - # Concatenate single-step prompt logprob processing results.

68 - assert isinstance(output, CompletionSequenceGroupOutput)

69 - single_step_process_prompt_logprob(self, seq_group, output)

70 -

71 - @staticmethod

72 - @functools.lru_cache

73 - def _log_prompt_logprob_unsupported_warning_once():

74 - # Reminder: Please update docs/features/compatibility_matrix.md

75 - # If the feature combo become valid

76 - logger.warning(

77 - "Prompt logprob is not supported by multi step workers. "

78 - "(e.g., speculative decode uses multi step workers).")

79 -

80 - def process_outputs(self,

81 - sequence_group: SequenceGroup,

82 - outputs: List[SequenceGroupOutput],

83 - is_async: bool = False) -> None:

84 - """Append new tokens in the outputs to sequences in the sequence group.

85 -

86 - This only supports sequence groups of size 1. It supports greater than

87 - one new token per sequence.

88 -

89 - This applies logic like stop condition checking and detokenization.

90 - It also handles cases where there are tokens emitted after

91 - the EOS token.

92 -

93 - is_async - Indicates whether this postprocessor runs in

94 - parallel with the GPU forward pass and is processing

95 - tokens from the previous step. If this is true, then

96 - no tokens need to be appended since it is already done

97 - externally (before the next schedule() call)

98 - """

99 - # Sequences can be in RUNNING or FINISHED_ABORTED state

100 - # once scheduled, as a sequence is moved to FINISHED_ABORTED

101 - # if a client disconnects from the api server.

102 - seqs = sequence_group.get_seqs(status=SequenceStatus.RUNNING)

103 - if seqs is None:

104 - seqs = sequence_group.get_seqs(

105 - status=SequenceStatus.FINISHED_ABORTED)

106 -

107 - assert seqs, "Expected RUNNING or FINISHED_ABORTED sequences"

108 - assert len(seqs) == 1, (

109 - "Beam search not supported in multi-step decoding.")

110 - seq = seqs[0]

111 - seq_id = seq.seq_id

112 - # This method is defined in the more generic

113 - # SequenceGroupOutputProcessor, but here we assume that the outputs are

114 - # of a more specific type.

115 - assert all([

116 - isinstance(output, CompletionSequenceGroupOutput)

117 - for output in outputs

118 - ])

119 - compl_outputs = cast(List[CompletionSequenceGroupOutput], outputs)

120 - assert all([

121 - seq_id == output.samples[0].parent_seq_id

122 - for output in compl_outputs

123 - ])

124 -

125 - if is_async:

126 - # Async case: We process tokens one by one. Here, we know the token

127 - # was already appended, so we only need to do the rest of the

128 - # postprocessor: Detokenization + stopping logic

129 - self._process_decode_and_stop(seq, sequence_group.sampling_params)

130 - else:

131 - # Standard multi-step case

132 -

133 - # Since there's only one sequence per sequence group,

134 - # we can take the first sample.

135 - samples = [output.samples[0] for output in compl_outputs]

136 -

137 - # entries in sample tokens may be invalid (eg. due to spec decode

138 - # rejecting tokens).

139 - valid_samples = [

140 - sample for sample in samples

141 - if sample.output_token != VLLM_INVALID_TOKEN_ID

142 - ]

143 -

144 - # When both spec-decode and pre-fill chunking are enabled, we

145 - # don't have guaranteed samples here (e.g. all -1s).

146 - if valid_samples:

147 - self._process_seq_outputs(seq, valid_samples,

148 - sequence_group.sampling_params)

149 -

150 - def _process_decode_and_stop(self, seq: Sequence,

151 - sampling_params: SamplingParams) -> None:

152 - new_char_count = 0

153 - if sampling_params.detokenize and self.detokenizer:

154 - new_char_count = self.detokenizer.decode_sequence_inplace(

155 - seq, sampling_params)

156 -

157 - # TODO(sang): Support lora.

158 - self.stop_checker.maybe_stop_sequence(

159 - seq,

160 - new_char_count=new_char_count,

161 - sampling_params=sampling_params,

162 - )

163 -

164 - def _process_seq_outputs(self, seq: Sequence,

165 - valid_samples: List[SequenceOutput],

166 - sampling_params: SamplingParams) -> None:

167 - output_token_ids = [sample.output_token for sample in valid_samples]

168 - output_logprobs = [sample.logprobs for sample in valid_samples]

169 - output_embeds = [sample.output_embed for sample in valid_samples]

170 -

171 - # Truncate to max_tokens if necessary.

172 - remaining_tokens = sampling_params.max_tokens - (seq.get_output_len() +

173 - len(output_token_ids))

174 - if remaining_tokens < 0:

175 - output_token_ids = output_token_ids[:remaining_tokens]

176 -

177 - # Truncate any tokens after EOS. This is required as spec decode

178 - # generates a fixed number of tokens without evaluating stopping

179 - # conditions within the block. This can cause an eos token to be

180 - # unintentionally ignored.

181 - if not sampling_params.ignore_eos and self.detokenizer:

182 - eos_token_id = self.get_tokenizer_for_seq(seq).eos_token_id

183 - # Avoiding .index calls as exception throwing in the happy path

184 - # is expensive.

185 - for i in range(len(output_token_ids)):

186 - if output_token_ids[i] == eos_token_id:

187 - output_token_ids = output_token_ids[:i + 1]

188 - break

189 -

190 - is_prefill_sampled_token = seq.data.get_num_uncomputed_tokens() == 0

191 - # Incrementally append tokens to the sequence, as if we had only one new

192 - # token.

193 - for output_token_id, output_logprob, output_embed in zip(

194 - output_token_ids, output_logprobs, output_embeds):

195 - seq.append_token_id(

196 - token_id=output_token_id,

197 - logprobs=output_logprob,

198 - token_embed=output_embed,

199 - )

200 -

201 - if is_prefill_sampled_token:

202 - is_prefill_sampled_token = False

203 - else:

204 - # Update num_computed_tokens iff the sampled token is not from

205 - # a prefill step.

206 - seq.data.update_num_computed_tokens(1)

207 -

208 - self._process_decode_and_stop(seq, sampling_params)

209 -

210 - if seq.is_finished():

211 - break

vllm/platforms/cuda.py

+2 -12

@@ -118,20 +118,10 @@ class CudaPlatformBase(Platform):

118 @classmethod

119 def check_and_update_config(cls, vllm_config: "VllmConfig") -> None:

120 parallel_config = vllm_config.parallel_config

121 - scheduler_config = vllm_config.scheduler_config

121 model_config = vllm_config.model_config

122

123 if parallel_config.worker_cls == "auto":

125 - if scheduler_config.is_multi_step:

126 - if envs.VLLM_USE_V1:

127 - raise NotImplementedError(

128 - "Multi-step scheduling is not supported (and not "

129 - "needed) on vLLM V1. Please launch without "

130 - "--num-scheduler-steps.")

131 - else:

132 - parallel_config.worker_cls = \

133 - "vllm.worker.multi_step_worker.MultiStepWorker"

134 - elif vllm_config.speculative_config:

124 + if vllm_config.speculative_config:

125 if not envs.VLLM_USE_V1:

126 raise NotImplementedError(

127 "Speculative decoding is not supported on vLLM V0.")

@@ -139,7 +129,7 @@ class CudaPlatformBase(Platform):

129 else:

130 if envs.VLLM_USE_V1:

131 parallel_config.worker_cls = \

142 - "vllm.v1.worker.gpu_worker.Worker"

132 + "vllm.v1.worker.gpu_worker.Worker"

133 else:

134 parallel_config.worker_cls = "vllm.worker.worker.Worker"

135

vllm/platforms/rocm.py

+2 -12

@@ -327,18 +327,8 @@ class RocmPlatform(Platform):

327 cache_config.block_size = 16

328

329 parallel_config = vllm_config.parallel_config

330 - scheduler_config = vllm_config.scheduler_config

330 if parallel_config.worker_cls == "auto":

332 - if scheduler_config.is_multi_step:

333 - if envs.VLLM_USE_V1:

334 - raise NotImplementedError(

335 - "Multi-step scheduling is not supported (and not "

336 - "needed) on vLLM V1. Please launch without "

337 - "--num-scheduler-steps.")

338 - else:

339 - parallel_config.worker_cls = \

340 - "vllm.worker.multi_step_worker.MultiStepWorker"

341 - elif vllm_config.speculative_config:

331 + if vllm_config.speculative_config:

332 if not envs.VLLM_USE_V1:

333 raise NotImplementedError(

334 "Speculative decoding is not supported on vLLM V0.")

@@ -346,7 +336,7 @@ class RocmPlatform(Platform):

336 else:

337 if envs.VLLM_USE_V1:

338 parallel_config.worker_cls = \

349 - "vllm.v1.worker.gpu_worker.Worker"

339 + "vllm.v1.worker.gpu_worker.Worker"

340 else:

341 parallel_config.worker_cls = "vllm.worker.worker.Worker"

342

vllm/platforms/tpu.py

+1 -6

@@ -133,18 +133,13 @@ class TpuPlatform(Platform):

133 parallel_config = vllm_config.parallel_config

134 scheduler_config = vllm_config.scheduler_config

135 if parallel_config.worker_cls == "auto":

136 - if scheduler_config.is_multi_step:

137 - raise NotImplementedError(

138 - "Multi-step scheduling is not supported (and not "

139 - "needed) on vLLM V1. Please launch without "

140 - "--num-scheduler-steps.")

136 parallel_config.worker_cls = "vllm.v1.worker.tpu_worker.TPUWorker"

137

138 assert not vllm_config.speculative_config, (

139 "Speculative decoding is not yet supported for TPU backend")

140

141 if scheduler_config.is_multimodal_model and not \

147 - scheduler_config.disable_chunked_mm_input:

142 + scheduler_config.disable_chunked_mm_input:

143 logger.warning("TPU does not support running Multimodal models"\

144 " without setting `--disable_chunked_mm_input`. " \

145 "Forcing --disable_chunked_mm_input.")

vllm/sequence.py

-38

@@ -794,35 +794,6 @@ class SequenceGroup:

794 def lora_int_id(self) -> int:

795 return self.lora_request.lora_int_id if self.lora_request else 0

796

797 - def init_multi_step(self, num_steps: int) -> None:

798 - self.state.num_steps = num_steps

799 - self.state.current_step = 0

800 -

801 - def init_multi_step_from_lookahead_slots(self, num_lookahead_slots: int,

802 - num_scheduler_steps: int,

803 - is_multi_step: bool,

804 - enable_chunking: bool) -> None:

805 -

806 - if not is_multi_step:

807 - self.init_multi_step(num_steps=num_scheduler_steps)

808 - return

809 -

810 - # Multi-Step case

811 - is_prefill = self.is_prefill()

812 -

813 - # The asserts below reflect the expectations of the current system.

814 - if is_prefill and enable_chunking:

815 - assert num_lookahead_slots == num_scheduler_steps

816 - self.init_multi_step(num_steps=num_lookahead_slots)

817 - else:

818 - is_decode: bool = not is_prefill

819 - # If it is a prefill, num_lookahead_slots must be 0

820 - assert num_lookahead_slots == 0 or is_decode

821 - # If it is a decode, num_lookahead_slots + 1 must match

822 - # the scheduler steps.

823 - assert num_lookahead_slots + 1 == num_scheduler_steps or is_prefill

824 - self.init_multi_step(num_steps=num_lookahead_slots + 1)

825 -

797 def set_last_token_time(self, now: float) -> None:

798 """Sets the last token time for Request level timings."""

799 # If still in prefill phase, assertion fails.

@@ -1367,15 +1338,6 @@ class ExecuteModelRequest(

1338 # Async callback

1339 async_callback: Optional[Callable] = None

1340

1370 - @property

1371 - def is_first_multi_step(self) -> bool:

1372 - # TODO(will) make this be able to handle batches with variable number of

1373 - # steps

1374 - assert len(self.seq_group_metadata_list) > 0

1375 - first_seq_group = self.seq_group_metadata_list[0]

1376 - assert first_seq_group.state is not None

1377 - return first_seq_group.state.current_step == 0

1378 -

1341 @property

1342 def is_last_step(self) -> bool:

1343 # TODO(will) make this be able to handle batches with variable number of

vllm/worker/model_runner.py

+2 -5

@@ -508,8 +508,7 @@ class ModelInputForGPUBuilder(ModelRunnerInputBuilderBase[ModelInputForGPU]):

508 if inter_data.is_prompt:

509 context_len = seq_data.get_num_computed_tokens()

510 seq_len = min(seq_len, context_len + token_chunk_size)

511 - elif self.runner.scheduler_config.is_multi_step or \

512 - self.runner.model_config.is_encoder_decoder:

511 + elif self.runner.model_config.is_encoder_decoder:

512 context_len = seq_len - 1

513 else:

514 context_len = seq_data.get_num_computed_tokens()

@@ -778,9 +777,7 @@ class ModelInputForGPUBuilder(ModelRunnerInputBuilderBase[ModelInputForGPU]):

777 int: Returns the determined number of padding sequences. If

778 CUDA graphs is not viable, returns -1.

779 """

781 - is_mscp: bool = self.runner.scheduler_config.is_multi_step and \

782 - self.runner.scheduler_config.chunked_prefill_enabled

783 - decode_only = self.decode_only or is_mscp

780 + decode_only = self.decode_only

781 if not decode_only:

782 # Early exit so we can treat num_seqs as the batch_size below.

783 return -1

vllm/worker/multi_step_model_runner.py deleted

-908

@@ -1,908 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

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

3 -

4 -import dataclasses

5 -import functools

6 -from dataclasses import dataclass, field

7 -from typing import (TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple,

8 - Union)

9 -

10 -import torch

11 -

12 -from vllm.distributed import get_pp_group

13 -from vllm.logger import init_logger

14 -from vllm.model_executor.layers.sampler import (PromptLogprobs, SampleLogprobs,

15 - SamplerOutput,

16 - SamplingMetadata, get_logprobs,

17 - get_pythonized_sample_results)

18 -from vllm.platforms import current_platform

19 -from vllm.sequence import (CompletionSequenceGroupOutput, IntermediateTensors,

20 - Logprob, SequenceGroupMetadata, SequenceOutput)

21 -from vllm.utils import PyObjectCache, async_tensor_h2d, current_stream

22 -from vllm.worker.model_runner import (GPUModelRunnerBase,

23 - ModelInputForGPUWithSamplingMetadata)

24 -from vllm.worker.model_runner_base import (

25 - BroadcastableModelInput, _init_attn_metadata_from_tensor_dict,

26 - _init_frozen_model_input_from_tensor_dict,

27 - _init_sampling_metadata_from_tensor_dict)

28 -

29 -from ..model_executor.model_loader.tensorizer import TensorizerConfig

30 -

31 -if TYPE_CHECKING:

32 - from vllm.attention.backends.abstract import AttentionBackend

33 -

34 -logger = init_logger(__name__)

35 -

36 -MULTI_STEP_ATTENTION_BACKENDS = [

37 - "FLASH_ATTN", "ROCM_FLASH", "FLASHINFER", "NO_ATTENTION"

38 -]

39 -MULTI_STEP_CHUNKED_PREFILL_ATTENTION_BACKENDS = ["FLASH_ATTN", "FLASHINFER"]

40 -

41 -def _get_supported_attention_backends(chunked_prefill_enabled: bool) \

42 - -> List[str]:

43 - if chunked_prefill_enabled:

44 - return MULTI_STEP_CHUNKED_PREFILL_ATTENTION_BACKENDS

45 - else:

46 - return MULTI_STEP_ATTENTION_BACKENDS

47 -

48 -

49 -def seq_output_builder():

50 - return SequenceOutput(

51 - 0, 0,

52 - {0: Logprob(logprob=float('inf'), rank=None, decoded_token=None)})

53 -

54 -

55 -def completion_seq_group_output_builder():

56 - return CompletionSequenceGroupOutput([], None)

57 -

58 -

59 -# Used by pythonization to reduce python object allocations

60 -class PythonizationCache:

61 -

62 - def __init__(self):

63 - self.cached_seq_output = PyObjectCache(seq_output_builder)

64 - self.cached_completion_seq_group_output = PyObjectCache(

65 - completion_seq_group_output_builder)

66 -

67 - def reset(self):

68 - self.cached_seq_output.reset()

69 - self.cached_completion_seq_group_output.reset()

70 -

71 -

72 -@dataclass

73 -class ModelOutput:

74 - """The output of a single model forward pass.

75 -

76 - The sampler_output_ready_event is set when the tensors in

77 - sampler_output are ready (the model+sampler forward pass has

78 - completed). We use the event to synchronize the GPU->CPU transfer,

79 - which we want to only run when the data has been written to the

80 - GPU tensors. Until the event is ready, the tensors in sampler_output

81 - will have garbage data.

82 -

83 - There are two scenarios:

84 - 1. The output tensors are ready and we can pythonize them immediately.

85 - 2. The output tensors are not ready and we need to wait for the event to be

86 - ready.

87 - """

88 - sampler_output: SamplerOutput

89 - sampler_output_ready_event: torch.cuda.Event

90 - sampled_token_ids: Optional[torch.Tensor] = None

91 - pythonized: bool = False

92 - # On-device tensor containing the logprobs of each token.

93 - logprobs: Optional["torch.Tensor"] = None

94 - pythonization_cache: Optional[PythonizationCache] = None

95 -

96 - def pythonize(self, input_metadata: "StatefulModelInput",

97 - copy_stream: torch.cuda.Stream,

98 - pinned_sampled_token_buffer: torch.Tensor) -> None:

99 - """Pythonize the output. Blocking."""

100 - if not self.pythonized:

101 - self._pythonize_sampler_output(input_metadata, copy_stream,

102 - pinned_sampled_token_buffer, True)

103 - self.pythonized = True

104 -

105 - def maybe_pythonize(self, input_metadata: "StatefulModelInput",

106 - copy_stream: torch.cuda.Stream,

107 - pinned_sampled_token_buffer: torch.Tensor) -> None:

108 - """Pythonize the output if ready, else return None. Non-blocking."""

109 - if not self.pythonized:

110 - self.pythonized = self._pythonize_sampler_output(

111 - input_metadata, copy_stream, pinned_sampled_token_buffer,

112 - False)

113 -

114 - def _pythonize_sampler_output(self, input_metadata: "StatefulModelInput",

115 - copy_stream: torch.cuda.Stream,

116 - pinned_sampled_token_buffer: torch.Tensor,

117 - blocking: bool) -> bool:

118 - """

119 - If blocking is set, will block until the forward pass for the output is

120 - ready and pythonize the output. Upon completing Pythonization, erases

121 - self.logprobs (note that a non-blocking call that is performed when

122 - the sampler output is not yet ready, will not erase self.logprobs.)

123 - """

124 - assert self.sampled_token_ids is not None

125 - if not blocking and not self.sampler_output_ready_event.query():

126 - return False

127 -

128 - if blocking:

129 - self.sampler_output_ready_event.synchronize()

130 - with torch.cuda.stream(copy_stream):

131 - _pythonize_sampler_output(input_metadata, self.sampler_output,

132 - pinned_sampled_token_buffer,

133 - self.sampled_token_ids, self.logprobs,

134 - self.pythonization_cache)

135 -

136 - # Erase the logprobs GPU-side tensor.

137 - # Note that although _pythonize_sampler_output() runs in its

138 - # own CUDA stream, nonetheless _pythonize_sampler_output()

139 - # cannot return until Pythonization is complete; therefore

140 - # we know that by the time the CPU reaches this point,

141 - # `self.logprobs` is no longer needed.

142 - self.logprobs = None

143 - return True

144 -

145 -

146 -@dataclass(frozen=False)

147 -class StatefulModelInput(BroadcastableModelInput):

148 - # actual frozen model input dataclass passed to _base_model_runner

149 - frozen_model_input: Optional[ModelInputForGPUWithSamplingMetadata] = None

150 -

151 - # list of model outputs for each step, may not be all pythonized

152 - cached_outputs: List[ModelOutput] = field(default_factory=list)

153 -

154 - # used to pass sampled token ids from the last step to the current step for

155 - # TP workers. Used to append to end of outputs and used by advance_step

156 - last_sampled_token_ids: Optional[torch.Tensor] = None

157 - current_step: int = 0

158 - is_multi_step: bool = True

159 - is_last_step: bool = False

160 - is_first_multi_step: bool = False

161 - base_output_proc_callback: Optional[Callable] = None

162 - # ping-pong data structures for multi-step to wait on the previous step

163 - step_cuda_events: List[current_platform.Event] = field(

164 - default_factory=lambda: [current_platform.Event(blocking=True)] * 2)

165 - num_seqs: int = -1

166 - num_queries: int = -1

167 - num_single_step_prefills: int = 0

168 -

169 - def as_broadcastable_tensor_dict(self) -> Dict[str, Any]:

170 - assert self.frozen_model_input is not None

171 - tensor_dict = self.frozen_model_input.as_broadcastable_tensor_dict()

172 - new_tensor_dict = {

173 - 'last_sampled_token_ids': self.last_sampled_token_ids,

174 - 'current_step': self.current_step,

175 - 'is_multi_step': self.is_multi_step,

176 - 'is_last_step': self.is_last_step,

177 - 'is_first_multi_step': self.is_first_multi_step,

178 - 'num_seqs': self.num_seqs,

179 - 'num_queries': self.num_queries,

180 - 'num_single_step_prefills': self.num_single_step_prefills,

181 - }

182 - tensor_dict.update(new_tensor_dict)

183 - return tensor_dict

184 -

185 - @classmethod

186 - def from_broadcasted_tensor_dict(

187 - cls,

188 - tensor_dict: Dict[str, Any],

189 - attn_backend: Optional["AttentionBackend"] = None,

190 - ) -> "StatefulModelInput":

191 - tensor_dict = _init_sampling_metadata_from_tensor_dict(tensor_dict)

192 - if attn_backend is not None:

193 - tensor_dict = _init_attn_metadata_from_tensor_dict(

194 - attn_backend, tensor_dict)

195 - tensor_dict = _init_frozen_model_input_from_tensor_dict(

196 - ModelInputForGPUWithSamplingMetadata, tensor_dict)

197 -

198 - return cls(**tensor_dict)

199 -

200 - def record_step_event(self, current_stream: torch.cuda.Stream):

201 - # record the event for the current step so that the next step can sync

202 - # on it. We modulo by 2 to keep the events in a circular buffer and

203 - # support any attn backends that may be supported in the future. ie

204 - # Flashinfer would want two DecodeWrappers to overlap the CPU and GPU.

205 - self.step_cuda_events[self.current_step & 1] = \

206 - torch.cuda.Event(blocking=True)

207 - self.step_cuda_events[self.current_step & 1].record(current_stream)

208 -

209 - def wait_previous_step(self):

210 - # These cuda events are an explicit synchronization to ensure that

211 - # advance_step() (for other attn backends that may be supported in the

212 - # future) do not clobber any data structures that is also used by any

213 - # enqueued forwards steps. For distributed case, only a single event is

214 - # needed, but for single GPU case, since we can let the CPU run much

215 - # further ahead, two events allow us to overlap the advance_step with

216 - # the previous forward (ie using two DecodeWrappers for flashinfer

217 - # backend)

218 - self.step_cuda_events[(self.current_step + 1) & 1].wait()

219 -

220 - def add_sampler_output(self,

221 - sampler_output: SamplerOutput,

222 - sampled_token_ids: Optional[torch.Tensor] = None):

223 - self.cached_outputs.append(

224 - ModelOutput(sampler_output=sampler_output,

225 - sampler_output_ready_event=None,

226 - sampled_token_ids=sampled_token_ids,

227 - pythonized=False))

228 -

229 - def maybe_advance_sampling_metadata(self, device: str, pin_memory: bool):

230 - """

231 - sampling_metadata.selected_token_indices is constructed for the

232 - first-step in Multi-Step. However, when chunked-prefill is enabled with

233 - multi-step, the scheduled prompts are fully processed in the

234 - first-step and are processed as decodes in the rest of the steps.

235 - This function updates the sampling_metadata.selected_token_indices

236 - to account for this conversion.

237 -

238 - Example:

239 - Let 2 prompts and 2 decodes be scheduled together. Let the

240 - num-tokens to process for the 2 prompts be 5 and 8 respectively.

241 -

242 - In that case, sampling_metadata.sampled_token_indices will be,

243 - [4, 12, 13, 14] as it is constructed for the first-step in

244 - multi-step.

245 - However, the prompts turns to decodes after the first-step

246 - and the num-tokens for the previously-prompt sequences will

247 - be 1 and 1 as they are decodes now. The self.sampled_token_indices

248 - must be updated to [0,1,2,3].

249 - """

250 - assert self.current_step == 1 and self.num_single_step_prefills > 0

251 - if not get_pp_group().is_last_rank:

252 - return

253 -

254 - assert self.frozen_model_input is not None

255 - assert self.frozen_model_input.sampling_metadata is not None

256 - self.frozen_model_input.sampling_metadata.selected_token_indices = \

257 - async_tensor_h2d(list(range(self.num_queries)),

258 - dtype=torch.long,

259 - target_device=device,

260 - pin_memory=pin_memory)

261 -

262 - def maybe_advance_frozen_model_input(self, device: str, pin_memory: bool):

263 - """

264 - Advancing the datastructures of StatefulModelInput::frozen_model_input

265 - is only required when prefills are scheduled with decodes to run in

266 - multi-step. This advancement/correction is required to account for

267 - the conversion of Prefills to Decodes after the first multi-step.

268 - """

269 - if self.current_step != 1 or self.num_single_step_prefills == 0:

270 - return

271 -

272 - assert self.frozen_model_input is not None

273 - fmi = self.frozen_model_input

274 -

275 - # Truncate input_tokens

276 - assert fmi.input_tokens is not None

277 - assert fmi.input_tokens.shape[0] >= self.num_seqs

278 - fmi_new_input_tokens: torch.Tensor = fmi.input_tokens[:self.num_seqs]

279 -

280 - # Update frozen_model_input::input_positions.

281 - assert fmi.input_positions is not None

282 - assert fmi.input_positions.shape[0] >= self.num_seqs

283 - fmi_new_input_positions: torch.Tensor = fmi.input_positions[:self.

284 - num_seqs]

285 -

286 - # Assert unsupported

287 - assert fmi.lora_mapping is None

288 - assert fmi.lora_requests is not None

289 - assert len(fmi.lora_requests) == 0

290 - assert fmi.attn_metadata is not None

291 - assert fmi.multi_modal_kwargs is not None

292 - assert len(fmi.multi_modal_kwargs) == 0

293 -

294 - self.frozen_model_input = dataclasses.replace(

295 - self.frozen_model_input,

296 - input_tokens=fmi_new_input_tokens,

297 - input_positions=fmi_new_input_positions)

298 -

299 - self.maybe_advance_sampling_metadata(device, pin_memory)

300 -

301 -

302 -# MutableModelInputForGPUWithMultiStepMetadata is not subclass of

303 -# ModelInputForGPU but it wraps the actual input dataclass and adds multi-step

304 -# metadata

305 -# mypy: disable-error-code=type-var

306 -class MultiStepModelRunner(GPUModelRunnerBase[StatefulModelInput]):

307 - # mypy: enable-error-code=type-var

308 -

309 - def __init__(self, base_model_runner: GPUModelRunnerBase, *args, **kwargs):

310 -

311 - super().__init__(*args, **kwargs)

312 -

313 - # Check attention backend support.

314 - supported_attention_backends: List[str] = \

315 - _get_supported_attention_backends(

316 - self.scheduler_config.chunked_prefill_enabled)

317 - if self.attn_backend.get_name() not in supported_attention_backends:

318 - ms_config_str: str = "Multi-Step + Chunked-Prefill" \

319 - if self.scheduler_config.chunked_prefill_enabled \

320 - else "Multi-Step"

321 - raise ValueError(

322 - f"{ms_config_str} not supported for attention backend: "

323 - f"{self.attn_backend.get_name()}. Set VLLM_ATTENTION_BACKEND "

324 - f"to a value from {supported_attention_backends}.")

325 -

326 - # uses the base model runner to execute the model and wraps it with

327 - # multi-step logic

328 - self._base_model_runner: GPUModelRunnerBase = base_model_runner

329 -

330 - self.is_multi_step = self.scheduler_config.is_multi_step

331 - self.pinned_sampled_token_ids: Optional[torch.Tensor] = None

332 -

333 - # Using the PythonizationCache in Pipeline-Parallel clobbers the

334 - # SequenceOutput and CompletionSequenceGroupOutput object.

335 - # When cache-reset happens at the last step of a multi-step

336 - # execution, there may be other on-going single-step/multi-step

337 - # executions. The current caching implementation does not check

338 - # for this.

339 - self.pythonization_cache = PythonizationCache() \

340 - if self.parallel_config.pipeline_parallel_size == 1 else None

341 -

342 - @functools.cached_property

343 - def _copy_stream(self):

344 - # used to copy tensors from GPU to CPU asynchronously

345 - return torch.cuda.Stream()

346 -

347 - def make_model_input_from_broadcasted_tensor_dict(

348 - self, tensor_dict: Dict[str, Any]) -> StatefulModelInput:

349 - model_input = (StatefulModelInput.from_broadcasted_tensor_dict(

350 - tensor_dict,

351 - attn_backend=self.attn_backend,

352 - ))

353 - return model_input

354 -

355 - def prepare_model_input(

356 - self,

357 - seq_group_metadata_list: List[SequenceGroupMetadata],

358 - virtual_engine: int = 0,

359 - finished_requests_ids: Optional[List[str]] = None

360 - ) -> StatefulModelInput:

361 - frozen_model_input: ModelInputForGPUWithSamplingMetadata = \

362 - self._base_model_runner.prepare_model_input(

363 - seq_group_metadata_list,

364 - virtual_engine,

365 - finished_requests_ids)

366 -

367 - assert frozen_model_input.query_lens is not None

368 - assert frozen_model_input.seq_lens is not None

369 - assert frozen_model_input.attn_metadata is not None

370 - num_queries = len(frozen_model_input.query_lens)

371 - num_seqs = len(frozen_model_input.seq_lens)

372 - num_single_step_prefills = frozen_model_input.attn_metadata.num_prefills

373 -

374 - model_input = StatefulModelInput(

375 - frozen_model_input=frozen_model_input,

376 - num_seqs=num_seqs,

377 - num_queries=num_queries,

378 - num_single_step_prefills=num_single_step_prefills)

379 -

380 - return model_input

381 -

382 - def _async_process_outputs(self, model_input: StatefulModelInput,

383 - output_proc_callback: Callable):

384 - # Proceed with pythonization and output_proc in order.

385 - # Stop on the first one that fails to pythonize

386 - output_proc_callback()

387 -

388 - cont = True

389 - for step_num, model_output in enumerate(model_input.cached_outputs):

390 - if not model_output.pythonized:

391 - model_output.maybe_pythonize(model_input, self._copy_stream,

392 - self.pinned_sampled_token_ids)

393 - if model_output.pythonized:

394 - ctx = output_proc_callback.keywords["ctx"]

395 - ctx.append_output(

396 - outputs=[model_output.sampler_output],

397 - seq_group_metadata_list=ctx.seq_group_metadata_list,

398 - scheduler_outputs=ctx.scheduler_outputs,

399 - is_async=False,

400 - is_last_step=False,

401 - is_first_step_output=step_num == 0)

402 -

403 - output_proc_callback()

404 - else:

405 - cont = False

406 -

407 - if not cont:

408 - break

409 -

410 - def _final_process_outputs(

411 - self, model_input: StatefulModelInput,

412 - output_proc_callback: Optional[Callable]) -> List[SamplerOutput]:

413 - assert model_input.frozen_model_input is not None

414 -

415 - has_async_callback = output_proc_callback is not None

416 -

417 - outputs = []

418 - for step_num, output in enumerate(model_input.cached_outputs):

419 - is_last_step = step_num == len(model_input.cached_outputs) - 1

420 -

421 - # For non-async case:

422 - # -- We simply add the outputs

423 - # For async case:

424 - # -- Invoke callback, pythonize, add to callback queue and repeat

425 - # -- For last output, just add to callback queue

426 - if has_async_callback:

427 - assert output_proc_callback is not None

428 -

429 - # Invoke callback before pythonize (to overlap with GPU)

430 - output_proc_callback()

431 -

432 - # Pythonize

433 - if not output.pythonized:

434 - output.pythonize(model_input, self._copy_stream,

435 - self.pinned_sampled_token_ids)

436 -

437 - # For non last step, add to callback queue to chain

438 - # callbacks=>pythonize pairs (for GPU overlap)

439 - if not is_last_step:

440 - ctx = output_proc_callback.keywords[ # type: ignore

441 - "ctx"] # type: ignore

442 - ctx.append_output(

443 - outputs=[output.sampler_output],

444 - seq_group_metadata_list=ctx.

445 - seq_group_metadata_list,

446 - scheduler_outputs=ctx.scheduler_outputs,

447 - is_async=False,

448 - is_last_step=False,

449 - is_first_step_output=step_num == 0)

450 - else:

451 - outputs.append(output.sampler_output)

452 - else:

453 - output.pythonize(model_input, self._copy_stream,

454 - self.pinned_sampled_token_ids)

455 - outputs.append(output.sampler_output)

456 -

457 - return outputs

458 -

459 - @torch.inference_mode()

460 - def execute_model(

461 - self,

462 - model_input: StatefulModelInput,

463 - kv_caches: List[torch.Tensor],

464 - intermediate_tensors: Optional[IntermediateTensors] = None,

465 - num_steps: int = 1,

466 - ) -> Optional[Union[List[SamplerOutput], IntermediateTensors]]:

467 - """

468 - Execute the model for a single step and update multi-step

469 - metadata

470 - """

471 - assert num_steps == 1, "MultiStepModelRunner only supports num_steps=1"

472 - frozen_model_input = model_input.frozen_model_input

473 - assert frozen_model_input is not None

474 -

475 - # path for warm up runs

476 - if not model_input.is_multi_step:

477 - return self._base_model_runner.execute_model(

478 - frozen_model_input, None, intermediate_tensors, num_steps)

479 -

480 - # make sure we skip the sampler on the lask rank and only pythonize

481 - # if CPU is ahead.

482 - if self.is_driver_worker and get_pp_group().is_last_rank:

483 - if self.pinned_sampled_token_ids is None:

484 - self.pinned_sampled_token_ids = torch.zeros(

485 - (self.scheduler_config.max_num_seqs, 1),

486 - dtype=torch.long,

487 - device="cpu",

488 - pin_memory=True)

489 -

490 - self._base_model_runner.sampler.include_gpu_probs_tensor = True

491 - if frozen_model_input.sampling_metadata:

492 - frozen_model_input.sampling_metadata.skip_sampler_cpu_output = (

493 - True)

494 -

495 - # some pre-execute model logic for multi-step:

496 - # - if it's the first step, we need to reset the sampling tensors

497 - # - if it's not the first step, we need to advance the step using the

498 - # appended sampler output from last iteration

499 - # - also maybe pythonize if CPU is ahead of GPU

500 -

501 - stream = current_stream()

502 - if not model_input.is_first_multi_step:

503 - # Explicitly block on the previous step's forward to make sure we

504 - # don't clobber any GPU tensors still in use.

505 - # This is not needed for flashattn backend, but for other attn

506 - # backends such as flashinfer that performs extra CPU operations on

507 - # input metadata we may need to synchronize any CPU operations that

508 - # might clobber enqueued forwards. (prevents CPU from running too

509 - # far ahead if needed)

510 - model_input.wait_previous_step()

511 - model_input = self._advance_step(

512 - model_input, model_input.cached_outputs[-1].sampler_output)

513 -

514 - # frozen_model_input may have been updated

515 - frozen_model_input = model_input.frozen_model_input

516 - assert frozen_model_input is not None

517 -

518 - if model_input.base_output_proc_callback is None:

519 - assert frozen_model_input is not None

520 - model_input.base_output_proc_callback = \

521 - frozen_model_input.async_callback

522 -

523 - if frozen_model_input.async_callback is not None:

524 - assert model_input.base_output_proc_callback is not None

525 - async_callback = functools.partial(

526 - self._async_process_outputs,

527 - model_input=model_input,

528 - output_proc_callback=model_input.base_output_proc_callback)

529 -

530 - model_input.frozen_model_input = dataclasses.replace( # type: ignore

531 - model_input.frozen_model_input,

532 - async_callback=async_callback)

533 - # Update the local instance

534 - frozen_model_input = model_input.frozen_model_input

535 - assert frozen_model_input is not None

536 -

537 - # Execute the model

538 - output = self._base_model_runner.execute_model(frozen_model_input,

539 - None,

540 - intermediate_tensors,

541 - num_steps=1)

542 -

543 - # record the event for the current step so that the next step can sync

544 - model_input.record_step_event(stream)

545 -

546 - if get_pp_group().is_last_rank and self.is_driver_worker:

547 - assert isinstance(output, list)

548 - assert len(

549 - output

550 - ) == 1, "MultiStepModelRunner requires single-step base_models"

551 -

552 - # event for the pythonization so that we only pythonize if the

553 - # tensors are ready. May be able to be combined with the step event

554 - output_ready_event = torch.cuda.Event()

555 - output_ready_event.record(stream)

556 - if self.parallel_config.pipeline_parallel_size > 1:

557 - output[0].sampled_token_ids_cpu = output[

558 - 0].sampled_token_ids.cpu()

559 - model_input.cached_outputs.append(

560 - ModelOutput(output[0], output_ready_event,

561 - output[0].sampled_token_ids, False,

562 - output[0].logprobs, self.pythonization_cache))

563 -

564 - # These GPU tensors are not required by multi-step;

565 - # erase them to ensure they are not pythonized or

566 - # transferred to CPU

567 - output[0].sampled_token_ids = None

568 - output[0].sampled_token_probs = None

569 - output[0].logprobs = None

570 -

571 - # Pythonize the output if CPU is ahead and the previous step is

572 - # ready.

573 - if frozen_model_input.async_callback is None:

574 - for model_output in model_input.cached_outputs:

575 - model_output.maybe_pythonize(model_input,

576 - self._copy_stream,

577 - self.pinned_sampled_token_ids)

578 -

579 - model_input.current_step += 1

580 -

581 - if not get_pp_group().is_last_rank:

582 - # Should be IntermediateTensors

583 - assert isinstance(output, IntermediateTensors)

584 - return output

585 - if not self.is_driver_worker:

586 - return []

587 -

588 - # Pythonize the output and block if needed since it is the last step

589 - if model_input.is_last_step:

590 - outputs = self._final_process_outputs(

591 - model_input, model_input.base_output_proc_callback)

592 - if self.pythonization_cache:

593 - self.pythonization_cache.reset()

594 - return outputs

595 -

596 - # should be [SamplerOutput]

597 - return output

598 -

599 - def _update_sampling_metadata(self, sampling_metadata: SamplingMetadata,

600 - num_seqs: Optional[int], num_queries: int):

601 -

602 - assert sampling_metadata.num_prompts == 0

603 - assert len(sampling_metadata.seq_groups) == num_queries

604 - assert sampling_metadata.selected_token_indices.shape == (

605 - num_queries, )

606 - # assert sampling_metadata.categorized_sample_indices == TODO: Add if needed # noqa: E501

607 -

608 - # Verify that all sequences are decodes

609 - for i in range(num_queries):

610 - seq_group = sampling_metadata.seq_groups[i]

611 -

612 - assert seq_group.is_prompt is False # No prompt

613 - assert seq_group.prompt_logprob_indices == [] # No prompt

614 - assert seq_group.sample_indices == [i] # Simple

615 - assert seq_group.seq_len is None # Decode

616 - assert seq_group.query_len is None # Decode

617 -

618 - def _advance_step(self, model_input: StatefulModelInput,

619 - out: SamplerOutput) -> StatefulModelInput:

620 -

621 - model_input.maybe_advance_frozen_model_input(self.device,

622 - self.pin_memory)

623 - frozen_model_input = model_input.frozen_model_input

624 - assert frozen_model_input is not None

625 - assert frozen_model_input.input_tokens is not None

626 - assert frozen_model_input.input_tokens.shape[0] == model_input.num_seqs

627 - assert frozen_model_input.attn_metadata is not None

628 -

629 - sampled_token_ids = model_input.cached_outputs[-1].sampled_token_ids

630 - num_seqs = model_input.num_seqs

631 - num_queries = model_input.num_queries

632 - frozen_model_input = model_input.frozen_model_input

633 - assert frozen_model_input is not None

634 - attn_metadata = frozen_model_input.attn_metadata

635 - assert attn_metadata is not None

636 -

637 - turn_prefills_into_decodes: bool = model_input.current_step == 1 and \

638 - model_input.num_single_step_prefills != 0

639 - attn_metadata.advance_step(

640 - frozen_model_input,

641 - sampled_token_ids,

642 - self.block_size,

643 - num_seqs,

644 - num_queries,

645 - turn_prefills_into_decodes=turn_prefills_into_decodes)

646 -

647 - return model_input

648 -

649 - def load_model(self) -> None:

650 - self._base_model_runner.load_model()

651 - self.model_memory_usage = self._base_model_runner.model_memory_usage

652 -

653 - def save_sharded_state(

654 - self,

655 - path: str,

656 - pattern: Optional[str] = None,

657 - max_size: Optional[int] = None,

658 - ) -> None:

659 - return self._base_model_runner.save_sharded_state(

660 - path, pattern, max_size)

661 -

662 - def save_tensorized_model(self,

663 - tensorizer_config: TensorizerConfig) -> None:

664 - return self._base_model_runner.save_tensorized_model(tensorizer_config)

665 -

666 - def profile_run(self) -> None:

667 - return self._base_model_runner.profile_run()

668 -

669 - def remove_all_loras(self):

670 - return self._base_model_runner.remove_all_loras()

671 -

672 - def capture_model(self, kv_caches: List[List]) -> None:

673 - return self._base_model_runner.capture_model(kv_caches)

674 -

675 - @property

676 - def vocab_size(self) -> int:

677 - return self._base_model_runner.vocab_size

678 -

679 -

680 -DeferredLogprobsReturnType = Tuple[Optional[List[Optional[PromptLogprobs]]],

681 - Optional[List[SampleLogprobs]]]

682 -

683 -

684 -def deferred_pythonize_logprobs(

685 - output: SamplerOutput,

686 - sampling_metadata: SamplingMetadata,

687 - logprobs_tensor: Optional[torch.Tensor],

688 -) -> DeferredLogprobsReturnType:

689 - """Perform deferred logprob Pythonization.

690 -

691 - 1. Pythonize GPU-side sampler result tensors into CPU-side sampler result.

692 - 2. Pythonize GPU-side logprobs tensor into CPU-side logprobs lists,

693 - utilizing the Pythonized sampler result computed in step 1.

694 -

695 - These deferred computations are not required for single-step scheduling

696 - or the `profile_run()` phase of multi-step scheduling.

697 -

698 - Args:

699 - output: sampler output (under deferred Pythonization)

700 - sampling_metadata

701 -

702 - Returns:

703 - prompt_logprobs (CPU), sample_logprobs (CPU)

704 - """

705 -

706 - # - Deferred pythonization of sample result

707 - sampler_result = get_pythonized_sample_results(

708 - output.deferred_sample_results_args)

709 -

710 - # - Erase the GPU-side deferred sample_result

711 - # computation args to ensure it is never

712 - # pythonized or transferred to CPU

713 - output.deferred_sample_results_args = None

714 -

715 - # - Deferred pythonization of logprobs

716 - (

717 - prompt_logprobs,

718 - sample_logprobs,

719 - ) = get_logprobs(logprobs_tensor, sampling_metadata, sampler_result)

720 - assert len(prompt_logprobs) == len(sampling_metadata.seq_groups)

721 - assert len(sample_logprobs) == len(sampling_metadata.seq_groups)

722 -

723 - return prompt_logprobs, sample_logprobs

724 -

725 -

726 -def _pythonize_sampler_output(

727 - model_input: StatefulModelInput,

728 - output: SamplerOutput,

729 - pinned_sampled_token_buffer: torch.Tensor,

730 - sampled_token_ids: torch.Tensor,

731 - logprobs_tensor: Optional[torch.Tensor],

732 - cache: Optional[PythonizationCache],

733 -) -> None:

734 - """ This function is only called when the output tensors are ready.

735 - See [`ModelOutput`][vllm.worker.multi_step_model_runner.ModelOutput].

736 -

737 - Modifies `output.outputs` and `pinned_sampled_token_buffer` in-place,

738 - adding a Pythonized output data structure

739 - ([`CompletionSequenceGroupOutput`][vllm.sequence.CompletionSequenceGroupOutput])

740 - for each [`SequenceGroup`][vllm.sequence.SequenceGroup].

741 -

742 - Args:

743 - model_input

744 - output: sampler output

745 - pinned_sampled_token_token_buffer: CPU-side pinned memory

746 - (receives copy of

747 - GPU-side token buffer.)

748 - sampled_token_ids: GPU-side token buffer

749 - logprobs_tensor: GPU-side tensor containing

750 - logprobs computed during sampling

751 - """

752 -

753 - assert model_input.frozen_model_input is not None

754 -

755 - frozen_model_input = model_input.frozen_model_input

756 - assert frozen_model_input.sampling_metadata is not None

757 - sampling_metadata = frozen_model_input.sampling_metadata

758 - # samples generation should have been skipped

759 - assert not output.outputs

760 -

761 - pinned_buffer = pinned_sampled_token_buffer[:model_input.num_queries]

762 -

763 - # We guarantee output tensors are ready, so it is safe to

764 - # pythonize the sampler output & obtain CPU-side logprobs.

765 - #

766 - # However we should check whether logprobs pythonization may

767 - # be skipped entirely, i.e. because no logprobs were requested

768 - # or pythonization was not deferred. To that end,

769 - #

770 - # * `prompt_logprobs_are_requested_for_prefill` signals that

771 - # there are *any* prefill-phase requests which specify that

772 - # prompt logprobs should be returned.

773 - #

774 - # * `any_logprobs_are_requested` signals that there are any

775 - # requests which (1) specify that sample logprobs should be

776 - # returned, or (2) are in the prefill phase AND specify that

777 - # prompt logprobs should be returned.

778 - #

779 - # Later on, these flags cause adjustments to the pythonization

780 - # process to accommodate logprobs.

781 -

782 - seq_groups = sampling_metadata.seq_groups

783 - prompt_logprobs_are_requested_for_prefill = any([

784 - sg.sampling_params.prompt_logprobs is not None and sg.is_prompt

785 - for sg in seq_groups

786 - ])

787 - any_logprobs_are_requested = (

788 - prompt_logprobs_are_requested_for_prefill

789 - or any([sg.sampling_params.logprobs is not None for sg in seq_groups]))

790 -

791 - if prompt_logprobs_are_requested_for_prefill:

792 - # CPU GPU sync, after gathering *only* sampled tokens (since

793 - # requesting prompt logprobs leads `sampled_token_ids` to

794 - # include prompt token ids in addition to sampled token ids.)

795 - sample_idx_tensor = torch.tensor(

796 - [sdx for sg in seq_groups for sdx in sg.sample_indices])

797 - pinned_buffer = pinned_buffer.copy_(

798 - sampled_token_ids[sample_idx_tensor, :], non_blocking=False)

799 - else:

800 - # CPU GPU sync

801 - pinned_buffer = pinned_buffer.copy_(sampled_token_ids,

802 - non_blocking=False)

803 -

804 - # this will not block as the tensors are already on CPU

805 - samples_list = pinned_buffer.tolist()

806 -

807 - skip_sampler_cpu_output = (

808 - frozen_model_input.sampling_metadata.skip_sampler_cpu_output)

809 -

810 - # *Don't* skip logprobs pythonization *if*:

811 - # * Any requests require logprobs to be returned in this

812 - # iteration AND

813 - # * These requests are being scheduled in a fashion which

814 - # defers pythonization (i.e. multi-step scheduling.)

815 - do_pythonize_logprobs = (skip_sampler_cpu_output

816 - and any_logprobs_are_requested)

817 - (

818 - prompt_logprobs,

819 - sample_logprobs,

820 - ) = (deferred_pythonize_logprobs(output, sampling_metadata,

821 - logprobs_tensor)

822 - if do_pythonize_logprobs else (None, None))

823 -

824 - for sgdx, (seq_group,

825 - sample_result) in enumerate(zip(seq_groups, samples_list)):

826 - # Reminder: Please update docs/features/compatibility_matrix.md

827 - # If the feature combo become valid

828 - # (Check for Guided Decoding)

829 - if seq_group.sampling_params.logits_processors:

830 - assert len(seq_group.sampling_params.logits_processors) == 0, (

831 - "Logits Processors are not supported in multi-step decoding")

832 -

833 - if do_pythonize_logprobs:

834 - assert prompt_logprobs is not None

835 - assert sample_logprobs is not None

836 -

837 - (

838 - group_prompt_logprobs,

839 - group_sample_logprobs,

840 - ) = ( # Utilize deferred pythonization results

841 - prompt_logprobs[sgdx],

842 - sample_logprobs[sgdx],

843 - )

844 - elif any_logprobs_are_requested:

845 - (

846 - group_prompt_logprobs,

847 - group_sample_logprobs,

848 - ) = (

849 - # profile_run: use already-computed logprobs

850 - output.outputs[sgdx].prompt_logprobs,

851 - [sample.logprobs for sample in output.outputs[sgdx].samples])

852 -

853 - seq_ids = seq_group.seq_ids

854 - next_token_ids = sample_result

855 - parent_ids = [0]

856 - seq_outputs: List[SequenceOutput]

857 -

858 - if cache is not None:

859 - completion_seq_group_output: CompletionSequenceGroupOutput = \

860 - cache.cached_completion_seq_group_output.get_object()

861 - completion_seq_group_output.samples.clear()

862 - seq_outputs = completion_seq_group_output.samples

863 - else:

864 - seq_outputs = []

865 -

866 - for tdx, (parent_id,

867 - next_token_id) in enumerate(zip(parent_ids, next_token_ids)):

868 - if cache is not None:

869 - seq_output: SequenceOutput = cache.cached_seq_output.get_object(

870 - )

871 - seq_output.parent_seq_id = seq_ids[parent_id]

872 - seq_output.output_token = next_token_id

873 -

874 - if any_logprobs_are_requested:

875 - seq_output.logprobs = group_sample_logprobs[tdx]

876 - else:

877 - logprobs = next(iter(seq_output.logprobs.values()))

878 - seq_output.logprobs.clear()

879 -

880 - logprobs.logprob = float('inf')

881 - logprobs.rank = None

882 - logprobs.decoded_token = None

883 -

884 - seq_output.logprobs[next_token_id] = logprobs

885 -

886 - seq_outputs.append(seq_output)

887 -

888 - else:

889 - seq_outputs.append(

890 - SequenceOutput(seq_ids[parent_id], next_token_id,

891 - (group_sample_logprobs[tdx]

892 - if any_logprobs_are_requested else {

893 - next_token_id:

894 - Logprob(logprob=float('inf'),

895 - rank=None,

896 - decoded_token=None)

897 - })))

898 - if cache is not None:

899 - completion_seq_group_output.prompt_logprobs = \

900 - group_prompt_logprobs if any_logprobs_are_requested else None

901 - output.outputs.append(completion_seq_group_output)

902 - else:

903 - output.outputs.append(

904 - CompletionSequenceGroupOutput(

905 - seq_outputs, (group_prompt_logprobs

906 - if any_logprobs_are_requested else None)))

907 -

908 - assert len(output.outputs) > 0

vllm/worker/multi_step_neuron_model_runner.py deleted

-84

@@ -1,84 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

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

3 -

4 -from importlib.util import find_spec

5 -from typing import List, Optional

6 -

7 -import torch

8 -

9 -from vllm.config import VllmConfig

10 -from vllm.model_executor.layers.sampler import SamplerOutput

11 -from vllm.multimodal import MultiModalKwargs

12 -from vllm.sequence import IntermediateTensors

13 -from vllm.worker.neuron_model_runner import (ModelInputForNeuron,

14 - NeuronModelRunner)

15 -

16 -

17 -class MultiStepNeuronModelRunner(NeuronModelRunner):

18 - """A model runner for multi step decoding using the transformers_neuronx

19 - framework"""

20 -

21 - def __init__(

22 - self,

23 - vllm_config: VllmConfig,

24 - ):

25 - super().__init__(vllm_config)

26 - self.speculation_config = self.speculative_config

27 - from transformers_neuronx.config import GenerationConfig

28 - self.speculation_config.draft_model_config.neuron_sampling_params = (

29 - GenerationConfig(

30 - max_length=self.scheduler_config.max_model_len,

31 - do_sample=True,

32 - per_batch_line=True,

33 - top_k=[self._MAX_NEURON_SAMPLING_TOP_K] \

34 - * self.scheduler_config.max_num_seqs,

35 - top_p=[1.0] * self.scheduler_config.max_num_seqs,

36 - temperature=[1.0] * self.scheduler_config.max_num_seqs,

37 - dynamic=True,

38 - global_top_k=self._MAX_NEURON_SAMPLING_TOP_K

39 - ))

40 -

41 - def load_model(self) -> None:

42 - if find_spec("transformers_neuronx") is not None:

43 - from vllm.model_executor.model_loader.neuron import (

44 - get_neuron_eagle_speculation_model,

45 - get_neuron_speculation_model)

46 - if self.speculation_config.speculative_token_tree is not None:

47 - self.model = get_neuron_eagle_speculation_model(

48 - self.model_config,

49 - parallel_config=self.parallel_config,

50 - scheduler_config=self.scheduler_config,

51 - speculation_config=self.speculation_config)

52 - else:

53 - self.model = get_neuron_speculation_model(

54 - self.model_config,

55 - parallel_config=self.parallel_config,

56 - scheduler_config=self.scheduler_config,

57 - speculation_config=self.speculation_config)

58 - else:

59 - raise NotImplementedError(

60 - "Supports only Transformer-NeuronX based models.")

61 -

62 - @torch.inference_mode()

63 - def execute_model(

64 - self,

65 - model_input: ModelInputForNeuron,

66 - kv_caches: Optional[List[torch.Tensor]] = None,

67 - intermediate_tensors: Optional[IntermediateTensors] = None,

68 - num_steps: int = 1,

69 - ) -> Optional[List[SamplerOutput]]:

70 - logits = self.model(

71 - input_ids=model_input.input_tokens,

72 - positions=model_input.input_positions,

73 - input_block_ids=model_input.input_block_ids,

74 - **MultiModalKwargs.as_kwargs(

75 - model_input.multi_modal_kwargs or {},

76 - device=self.device,

77 - ),

78 - )

79 -

80 - output = self.model.sample(

81 - logits=logits,

82 - sampling_metadata=model_input.sampling_metadata,

83 - )

84 - return output

vllm/worker/multi_step_neuronx_distributed_model_runner.py deleted

-63

@@ -1,63 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

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

3 -from typing import List, Optional

4 -

5 -import torch

6 -

7 -from vllm.config import VllmConfig

8 -from vllm.model_executor.layers.sampler import SamplerOutput

9 -from vllm.multimodal import MultiModalKwargs

10 -from vllm.sequence import IntermediateTensors

11 -from vllm.worker.neuronx_distributed_model_runner import (

12 - NeuronxDistributedModelRunner)

13 -

14 -

15 -class MultiStepNeuronxDistributedModelRunner(NeuronxDistributedModelRunner):

16 - """A model runner for multi-step decoding using the

17 - neuronx-distributed-inference framework"""

18 -

19 - def __init__(

20 - self,

21 - vllm_config: VllmConfig,

22 - ):

23 - super().__init__(vllm_config)

24 -

25 - def load_model(self) -> None:

26 - from vllm.model_executor.model_loader.neuronx_distributed import (

27 - get_neuron_speculation_model)

28 - self.model = get_neuron_speculation_model(

29 - self.model_config,

30 - parallel_config=self.parallel_config,

31 - scheduler_config=self.scheduler_config,

32 - speculation_config=self.speculative_config)

33 -

34 - @torch.inference_mode()

35 - def execute_model(

36 - self,

37 - model_input,

38 - kv_caches: Optional[List[torch.Tensor]] = None,

39 - intermediate_tensors: Optional[IntermediateTensors] = None,

40 - num_steps: int = 1,

41 - ) -> Optional[List[SamplerOutput]]:

42 - sampling_params = torch.tensor([[

43 - seq_group.sampling_params.top_k,

44 - seq_group.sampling_params.top_p,

45 - seq_group.sampling_params.temperature,

46 - ] for seq_group in model_input.sampling_metadata.seq_groups])

47 -

48 - logits = self.model(

49 - input_ids=model_input.input_tokens,

50 - positions=model_input.input_positions,

51 - input_block_ids=model_input.input_block_ids,

52 - sampling_params=sampling_params,

53 - **MultiModalKwargs.as_kwargs(

54 - model_input.multi_modal_kwargs or {},

55 - device=self.device,

56 - ),

57 - )

58 -

59 - output = self.model.sample(

60 - logits=logits,

61 - sampling_metadata=model_input.sampling_metadata,

62 - )

63 - return output

vllm/worker/multi_step_worker.py deleted

-197

@@ -1,197 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

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

3 -

4 -import dataclasses

5 -from dataclasses import dataclass

6 -from typing import Dict, List, Optional, Tuple

7 -

8 -import torch

9 -

10 -from vllm.distributed import broadcast_tensor_dict, get_pp_group

11 -from vllm.model_executor.layers.sampler import SamplerOutput

12 -from vllm.sequence import ExecuteModelRequest

13 -from vllm.worker.model_runner_base import BroadcastableModelInput

14 -from vllm.worker.multi_step_model_runner import (MultiStepModelRunner,

15 - StatefulModelInput)

16 -from vllm.worker.worker import Worker, WorkerInput

17 -

18 -

19 -@dataclass

20 -class MultiStepState:

21 - worker_input: WorkerInput

22 - model_input: StatefulModelInput

23 -

24 -

25 -class MultiStepWorker(Worker):

26 -

27 - def __init__(self, *args, **kwargs):

28 - super().__init__(*args, **kwargs)

29 - base_model_runner = self.model_runner

30 - # for multi-step model, wrap the model runner with MultiStepModelRunner

31 - self.model_runner = MultiStepModelRunner(

32 - base_model_runner,

33 - vllm_config=base_model_runner.vllm_config,

34 - kv_cache_dtype=self.cache_config.cache_dtype,

35 - is_driver_worker=base_model_runner.is_driver_worker,

36 - )

37 -

38 - pipeline_parallel_size = self.parallel_config.pipeline_parallel_size

39 - self.multi_step_states: List[

40 - Optional[MultiStepState]] = [None] * pipeline_parallel_size

41 - self.temp_output = None

42 -

43 - def _get_driver_input_and_broadcast(

44 - self, execute_model_req: ExecuteModelRequest

45 - ) -> Tuple[BroadcastableModelInput, WorkerInput, Dict[str, torch.Tensor]]:

46 - """

47 - Get the driver input and broadcast it to other workers.

48 - """

49 - assert self.is_driver_worker

50 - virtual_engine = execute_model_req.virtual_engine

51 - is_first_multi_step = execute_model_req.is_first_multi_step

52 - if is_first_multi_step:

53 - # on first step we prepare the worker input and model input normally

54 - worker_input: WorkerInput = self.prepare_worker_input(

55 - execute_model_req=execute_model_req)

56 - model_input: StatefulModelInput = (

57 - self.model_runner.prepare_model_input(

58 - execute_model_req.seq_group_metadata_list,

59 - execute_model_req.virtual_engine,

60 - execute_model_req.finished_requests_ids))

61 -

62 - if execute_model_req.async_callback:

63 - model_input.frozen_model_input = dataclasses.replace( # type: ignore

64 - model_input.frozen_model_input,

65 - async_callback=execute_model_req.async_callback)

66 - else:

67 - # on subsequent steps we reuse the worker input and model input

68 - multi_step_state = self.multi_step_states[virtual_engine]

69 - worker_input = multi_step_state.worker_input

70 - model_input = multi_step_state.model_input

71 - frozen_model_input = model_input.frozen_model_input

72 - assert frozen_model_input is not None

73 - assert frozen_model_input.attn_metadata is not None

74 - # clear the cached metadata so that it can be recomputed on

75 - # the workers.

76 - frozen_model_input.attn_metadata._cached_prefill_metadata = None

77 - frozen_model_input.attn_metadata._cached_decode_metadata = None

78 -

79 - model_input.is_first_multi_step = is_first_multi_step

80 - model_input.is_last_step = execute_model_req.is_last_step

81 -

82 - if not is_first_multi_step:

83 - # we broadcast the last sampled token ids to all TP workers so they

84 - # can update their model input metadata in-place.

85 - self._prepare_last_sampled_token_ids_for_tp_workers(

86 - execute_model_req=execute_model_req, model_input=model_input)

87 -

88 - if self.do_metadata_broadcast:

89 - broadcast_data = worker_input.as_broadcastable_tensor_dict()

90 - broadcast_data.update(model_input.as_broadcastable_tensor_dict())

91 - broadcast_tensor_dict(broadcast_data, src=0)

92 -

93 - # Retuning empty dict here to keep this compatible with

94 - # `LocalOrDistributedWorkerBase._get_driver_input_and_broadcast`

95 - return model_input, worker_input, {}

96 -

97 - def _prepare_last_sampled_token_ids_for_tp_workers(

98 - self,

99 - execute_model_req: ExecuteModelRequest,

100 - model_input: StatefulModelInput,

101 - ) -> None:

102 - """

103 - Prepare the last sampled token ids for TP workers. If it's the last

104 - PP rank, then the last sampled token ids are already in the model_input.

105 - If it is NOT the last PP rank, then we need to get the last sampled

106 - token that is cached in the execute_model_req.

107 - """

108 - if get_pp_group().is_last_rank:

109 - assert model_input.cached_outputs[

110 - -1].sampler_output.sampled_token_ids is None

111 - assert model_input.cached_outputs[-1].sampled_token_ids is not None

112 - model_input.last_sampled_token_ids = model_input.cached_outputs[

113 - -1].sampled_token_ids

114 - # free sampled token ids from the previous step if it has been

115 - # pythonized. Cannot free the last sampled token ids because

116 - # we need it for GPU advance_step.

117 - for output in model_input.cached_outputs[:-1]:

118 - if output.pythonized:

119 - output.sampled_token_ids = None

120 - else:

121 - # otherwise we need to get the cached sampled token ids from the

122 - # execute_model_req

123 - assert execute_model_req.last_sampled_token_ids is not None

124 - model_input.last_sampled_token_ids = (

125 - execute_model_req.last_sampled_token_ids.cuda())

126 - model_input.add_sampler_output(

127 - SamplerOutput(outputs=[], sampled_token_ids=None),

128 - model_input.last_sampled_token_ids)

129 -

130 - # free sampled token ids from the previous step.

131 - # TODO(will) we could reuse the sampled token ids tensor from

132 - # the previous step instead.

133 - for output in model_input.cached_outputs[:-1]:

134 - output.sampled_token_ids = None

135 - assert model_input.cached_outputs[-1].sampled_token_ids is not None

136 -

137 - def prepare_input(

138 - self,

139 - execute_model_req: Optional[ExecuteModelRequest] = None,

140 - ) -> Optional[Tuple[StatefulModelInput, WorkerInput, Dict[str,

141 - torch.Tensor]]]:

142 - """

143 - Depending on the current state of the request and multi step worker,

144 - this method may skip the normal _prepare_model_input and

145 - _prepare_worker_input methods and instead used cached values.

146 - """

147 - if self.is_driver_worker:

148 - if execute_model_req is None:

149 - if self.do_metadata_broadcast:

150 - # This signals that there's no more requests to process for

151 - # now. All workers are running infinite loop with

152 - # broadcast_tensor_dict, and it stops the loop when the

153 - # driver broadcasts an empty input. Send an empty input to

154 - # notify all other workers to stop their execution loop.

155 - broadcast_tensor_dict({}, src=0)

156 - return None

157 -

158 - virtual_engine = execute_model_req.virtual_engine

159 - (model_input, worker_input,

160 - kwargs) = self._get_driver_input_and_broadcast(execute_model_req)

161 - assert isinstance(model_input, StatefulModelInput)

162 - if execute_model_req.is_first_multi_step:

163 - # cache the worker input and model input for the next steps

164 - self.multi_step_states[virtual_engine] = MultiStepState(

165 - worker_input=worker_input, model_input=model_input)

166 - # if TP workers

167 - else:

168 - broadcast_data = self._get_worker_input_from_broadcast()

169 - # if the driver has sent an empty input, we should stop the worker

170 - # loop

171 - if broadcast_data is None:

172 - return None

173 - model_input, worker_input, kwargs = broadcast_data

174 - assert isinstance(model_input, StatefulModelInput)

175 - virtual_engine = worker_input.virtual_engine

176 - if model_input.is_first_multi_step:

177 - pass

178 - # TODO(will) Can cache the worker input and model input for the

179 - # next steps. See below for details

180 - else:

181 - # TODO(will) possible to also cache and reuse the cached worker

182 - # input and model input. The idea is essentially the delta

183 - # optimization for model_inputs. Where the TP workers can cache

184 - # the model input states and we only broadcast the delta need

185 - # for the next step (sampled_token_ids from the previous step)

186 -

187 - assert isinstance(model_input, StatefulModelInput)

188 - # we need to update the last sampled token ids in the model

189 - # input for the workers so that they can run inplace

190 - # advance_step

191 - model_input.add_sampler_output(

192 - SamplerOutput(outputs=[], sampled_token_ids=None),

193 - model_input.last_sampled_token_ids)

194 -

195 - assert model_input is not None

196 - assert worker_input is not None

197 - return model_input, worker_input, kwargs

vllm/worker/neuron_worker.py

+9 -13

@@ -64,25 +64,21 @@ class NeuronWorker(LocalOrDistributedWorkerBase):

64 assert (self.lora_config

65 is None), ("LoRA is not supported for TransformersNeuronX "

66 "framework.")

67 - from vllm.worker.multi_step_neuron_model_runner import (

68 - MultiStepNeuronModelRunner)

67 if self.speculative_config is not None:

70 - return MultiStepNeuronModelRunner(vllm_config=vllm_config)

71 - else:

72 - return NeuronModelRunner(vllm_config=vllm_config)

68 + raise NotImplementedError(

69 + "Speculative decoding is not supported for TransformersNeuronX"

70 + )

71 + return NeuronModelRunner(vllm_config=vllm_config)

73 def get_neuronx_distributed_model_runner(self, vllm_config):

75 - from vllm.worker.multi_step_neuronx_distributed_model_runner import (

76 - MultiStepNeuronxDistributedModelRunner)

74 from vllm.worker.neuronx_distributed_model_runner import (

75 NeuronxDistributedModelRunner)

76 if self.speculative_config is not None:

80 - assert (self.lora_config

81 - is None), "LoRA is not supported for Speculative Decoding"

82 - return MultiStepNeuronxDistributedModelRunner(

83 - vllm_config=vllm_config)

84 - else:

85 - return NeuronxDistributedModelRunner(vllm_config=vllm_config)

77 + assert (self.lora_config is None), (

78 + "LoRA is not supported for Speculative Decoding")

79 + raise NotImplementedError(

80 + "Speculative decoding is not supported for NeuronxDistributed")

81 + return NeuronxDistributedModelRunner(vllm_config=vllm_config)

83 def init_device(self) -> None:

84 self.init_distributed_environment()