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

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

[V1] Prompt logprobs + APC compatibility; prompt logprobs reqs cannot fill APC (#13949)

afeldman-nm committed Mar 7, 2025 at 20:48 UTC ef640440791a461a181e6d497965701462c166b3

9 files changed +292 -162

tests/v1/core/test_prefix_caching.py

+110 -2

@@ -1,6 +1,8 @@

1 # SPDX-License-Identifier: Apache-2.0

2 """Compare the with and without prefix caching."""

4 +from typing import Optional

5 +

6 import pytest

8 from vllm.multimodal.inputs import MultiModalKwargs, PlaceholderRange

@@ -15,7 +17,8 @@ from vllm.v1.core.kv_cache_utils import (BlockHashType, KVCacheBlock,

17 def make_request(request_id,

18 prompt_token_ids,

19 mm_positions=None,

18 - mm_hashes=None):

20 + mm_hashes=None,

21 + prompt_logprobs: Optional[int] = None):

22 if mm_positions is None:

23 multi_modal_inputs = None

24 else:

@@ -28,7 +31,8 @@ def make_request(request_id,

31 multi_modal_inputs=multi_modal_inputs,

32 multi_modal_hashes=mm_hashes,

33 multi_modal_placeholders=mm_positions,

31 - sampling_params=SamplingParams(max_tokens=17),

34 + sampling_params=SamplingParams(max_tokens=17,

35 + prompt_logprobs=prompt_logprobs),

36 eos_token_id=100,

37 arrival_time=0,

38 lora_request=None,

@@ -144,6 +148,110 @@ def test_prefill():

148 assert manager.block_pool.free_block_queue.free_list_tail is None

149

150

151 +def test_prefill_plp():

152 + '''Test prefill with APC and some prompt logprobs (plp) requests.

153 +

154 + 1. Schedule plp request and validate APC block allocation

155 + 2. Schedule non-plp request and validate blocks

156 + 3. Schedule plp request; no hit should occur; validate blocks

157 + '''

158 + manager = KVCacheManager(

159 + block_size=16,

160 + num_gpu_blocks=10,

161 + max_model_len=8192,

162 + sliding_window=None,

163 + enable_caching=True,

164 + num_preallocate_tokens=16,

165 + )

166 +

167 + # Complete 3 blocks (48 tokens)

168 + common_token_ids = [i for i in range(3) for _ in range(16)]

169 +

170 + # Request #0 is a prompt logprobs request

171 + # Fully cache miss

172 + # Incomplete 1 block (7 tokens)

173 + unique_token_ids = [3] * 7

174 + all_token_ids = common_token_ids + unique_token_ids

175 + req0 = make_request("0", all_token_ids, prompt_logprobs=5)

176 + computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)

177 + assert len(manager.req_to_block_hashes[req0.request_id]) == 3

178 + assert not computed_blocks

179 + assert num_computed_tokens == 0

180 + blocks = manager.allocate_slots(req0, 55, computed_blocks)

181 + assert [b.block_id for b in blocks] == [0, 1, 2, 3, 4]

182 + req0_block_hashes = [b.block_hash for b in blocks]

183 +

184 + # Check full block metadata

185 + parent_block_hash = None

186 + for block_id in (0, 1, 2):

187 + block_tokens = tuple(all_token_ids[block_id * 16:(block_id + 1) * 16])

188 + block_hash = hash_block_tokens(parent_block_hash, block_tokens)

189 + assert manager.block_pool.blocks[block_id].block_hash == block_hash

190 + assert manager.block_pool.blocks[block_id].ref_cnt == 1

191 + parent_block_hash = block_hash.hash_value

192 +

193 + # Check partial/preallocated block metadata

194 + for block_id in (3, 4):

195 + assert manager.block_pool.blocks[block_id].block_hash is None

196 + assert manager.block_pool.blocks[block_id].ref_cnt == 1

197 +

198 + # Request #1 is a non-prompt-logprobs request:

199 + # Cache hit in the common prefix when the original block is still in use.

200 + # Incomplete 1 block (5 tokens)

201 + unique_token_ids = [3] * 5

202 + req1 = make_request("1", common_token_ids + unique_token_ids)

203 + computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)

204 + assert len(manager.req_to_block_hashes[req1.request_id]) == 3

205 + assert [b.block_id for b in computed_blocks] == [0, 1, 2]

206 + assert num_computed_tokens == 3 * 16

207 + num_new_tokens = 53 - 3 * 16

208 + blocks = manager.allocate_slots(req1, num_new_tokens, computed_blocks)

209 + assert [b.block_id for b in blocks] == [5, 6]

210 + for block in computed_blocks:

211 + assert block.ref_cnt == 2

212 +

213 + # At this point, we should have 3 free blocks left.

214 + assert manager.block_pool.free_block_queue.num_free_blocks == 3

215 +

216 + manager.free(req0)

217 + manager.free(req1)

218 +

219 + # All blocks should be available.

220 + assert manager.block_pool.free_block_queue.num_free_blocks == 10

221 + # The order should be

222 + # [unallocated (7, 8, 9)]

223 + # [unique_req0 (4, 3)]

224 + # [unique_req1 (6, 5)]

225 + # [common (2, 1, 0)]

226 + assert [

227 + b.block_id

228 + for b in manager.block_pool.free_block_queue.get_all_free_blocks()

229 + ] == [7, 8, 9, 4, 3, 6, 5, 2, 1, 0]

230 +

231 + # Request #2 is a prompt-logprobs request:

232 + # NO cache hit in the common prefix; duplicates request #0 cached blocks

233 + unique_token_ids = [3] * 6

234 + req2 = make_request("2",

235 + common_token_ids + unique_token_ids,

236 + prompt_logprobs=5)

237 + computed_blocks, num_computed_tokens = manager.get_computed_blocks(req2)

238 + assert len(manager.req_to_block_hashes[req2.request_id]) == 3

239 + assert not computed_blocks

240 + assert num_computed_tokens == 0

241 + blocks = manager.allocate_slots(req2, 55, computed_blocks)

242 + block_ids = [b.block_id for b in blocks]

243 + # Duplicate cached blocks have different ids but same hashes vs request #0

244 + assert [b.block_hash for b in blocks] == req0_block_hashes

245 + assert block_ids != [0, 1, 2, 3, 4]

246 +

247 + # Request #2 block hashes are valid since request #0 hashes are.

248 + # Check block reference counts.

249 + for block_id in block_ids:

250 + assert manager.block_pool.blocks[block_id].ref_cnt == 1

251 +

252 + manager.free(req2)

253 +

254 +

255 def test_decode():

256 manager = KVCacheManager(

257 block_size=16,

tests/v1/core/test_scheduler.py

+49 -12

@@ -1,6 +1,8 @@

1 # SPDX-License-Identifier: Apache-2.0

2 from typing import Optional

4 +import pytest

5 +

6 from vllm.config import CacheConfig, ModelConfig, SchedulerConfig, VllmConfig

7 from vllm.multimodal.inputs import MultiModalKwargs, PlaceholderRange

8 from vllm.sampling_params import SamplingParams

@@ -16,7 +18,21 @@ def create_scheduler(

18 model: str = "facebook/opt-125m",

19 max_num_seqs: int = 16,

20 max_num_batched_tokens: int = 8192,

21 + enable_prefix_caching: Optional[bool] = None,

22 ) -> Scheduler:

23 + '''Create scheduler under test.

24 +

25 + Args:

26 + model: model under test

27 + max_num_seqs: max sequences to schedule

28 + max_num_batch_tokens: max num tokens to batch

29 + enable_prefix_caching: optionally force APC config

30 + (True/False) or use default

31 + (None)

32 +

33 + Returns:

34 + :class:`Scheduler` instance

35 + '''

36 scheduler_config = SchedulerConfig(

37 max_num_seqs=max_num_seqs,

38 max_num_batched_tokens=max_num_batched_tokens,

@@ -31,11 +47,16 @@ def create_scheduler(

47 dtype="float16",

48 seed=42,

49 )

50 + # Cache config, optionally force APC

51 + kwargs_cache = ({} if enable_prefix_caching is None else {

52 + 'enable_prefix_caching': enable_prefix_caching

53 + })

54 cache_config = CacheConfig(

55 block_size=16,

56 gpu_memory_utilization=0.9,

57 swap_space=0,

58 cache_dtype="auto",

59 + **kwargs_cache,

60 )

61 vllm_config = VllmConfig(

62 scheduler_config=scheduler_config,

@@ -54,16 +75,16 @@ def create_scheduler(

75 )

57 -def create_requests(

58 - num_requests: int,

59 - num_tokens: int = 10,

60 - mm_positions: Optional[list[PlaceholderRange]] = None,

61 - max_tokens: int = 16,

62 - stop_token_ids: Optional[list[int]] = None,

63 -):

78 +def create_requests(num_requests: int,

79 + num_tokens: int = 10,

80 + mm_positions: Optional[list[PlaceholderRange]] = None,

81 + max_tokens: int = 16,

82 + stop_token_ids: Optional[list[int]] = None,

83 + prompt_logprobs: Optional[int] = None):

84 sampling_params = SamplingParams(ignore_eos=False,

85 max_tokens=max_tokens,

66 - stop_token_ids=stop_token_ids)

86 + stop_token_ids=stop_token_ids,

87 + prompt_logprobs=prompt_logprobs)

88 requests = []

89 for i in range(num_requests):

90 if mm_positions is not None:

@@ -122,9 +143,18 @@ def test_get_num_unfinished_requests():

143 assert scheduler.get_num_unfinished_requests() == len(requests) - i - 1

144

145

125 -def test_schedule():

126 - scheduler = create_scheduler()

127 - requests = create_requests(num_requests=10)

146 +@pytest.mark.parametrize("enable_prefix_caching, prompt_logprobs", [

147 + (None, None),

148 + (True, 5),

149 +])

150 +def test_schedule(enable_prefix_caching: Optional[bool],

151 + prompt_logprobs: Optional[int]):

152 + '''Test scheduling.

153 + Two cases: default APC/no prompt logprobs; APC=True + prompt logprobs

154 + '''

155 + scheduler = create_scheduler(enable_prefix_caching=enable_prefix_caching)

156 + requests = create_requests(num_requests=10,

157 + prompt_logprobs=prompt_logprobs)

158 for request in requests:

159 scheduler.add_request(request)

160

@@ -427,14 +457,21 @@ def test_stop_via_update_from_output():

457 assert list(requests[0].output_token_ids) == [EOS_TOKEN_ID, 10, 11]

458

459

430 -def test_schedule_concurrent_batches():

460 +@pytest.mark.parametrize("enable_prefix_caching, prompt_logprobs", [

461 + (None, None),

462 + (True, 5),

463 +])

464 +def test_schedule_concurrent_batches(enable_prefix_caching: Optional[bool],

465 + prompt_logprobs: Optional[int]):

466 scheduler = create_scheduler(

467 max_num_batched_tokens=1024,

468 max_num_seqs=2,

469 + enable_prefix_caching=enable_prefix_caching,

470 )

471 requests = create_requests(

472 num_requests=2,

473 num_tokens=512,

474 + prompt_logprobs=prompt_logprobs,

475 )

476

477 # Schedule the first request.

tests/v1/engine/test_async_llm.py

-36

@@ -6,7 +6,6 @@ from typing import Optional

7 import pytest

9 -from tests.v1.engine.utils import PLP_APC_UNSUPPORTED_MSG

9 from vllm import SamplingParams

10 from vllm.assets.image import ImageAsset

11 from vllm.engine.arg_utils import AsyncEngineArgs

@@ -72,41 +71,6 @@ async def generate(engine: AsyncLLM,

71 return count, request_id

75 -@pytest.mark.parametrize(

76 - "output_kind", [RequestOutputKind.DELTA, RequestOutputKind.FINAL_ONLY])

77 -@pytest.mark.asyncio

78 -async def test_async_llm_refuses_prompt_logprobs_with_apc(

79 - monkeypatch, output_kind: RequestOutputKind):

80 - """Test passes if AsyncLLM raises an exception when it is configured

81 - for automatic prefix caching and it receives a request with

82 - prompt_logprobs enabled, which is incompatible."""

83 - # TODO(rickyx): Remove monkeypatch VLLM_USE_V1 setting once we have a

84 - # better way to test V1 so that in the future when we switch, we don't

85 - # have to change all the tests.

86 - monkeypatch.setenv("VLLM_USE_V1", "1")

87 - # Create AsyncLLM engine with APC

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

89 - enable_prefix_caching=True,

90 - gpu_memory_utilization=0.8,

91 - disable_log_requests=True)

92 - engine = AsyncLLM.from_engine_args(apc_engine_args)

93 - try:

94 - with pytest.raises(ValueError) as excinfo:

95 - # Issue a request with prompt logprobs enabled, which should fail

96 - await asyncio.create_task(

97 - generate(engine,

98 - "request-0",

99 - TEXT_PROMPT,

100 - output_kind,

101 - 10,

102 - prompt_logprobs=5))

103 - # Validate exception string is correct

104 - assert str(excinfo.value) == PLP_APC_UNSUPPORTED_MSG

105 - finally:

106 - # Shut down engine

107 - engine.shutdown()

108 -

109 -

74 @pytest.mark.parametrize(

75 "output_kind", [RequestOutputKind.DELTA, RequestOutputKind.FINAL_ONLY])

76 @pytest.mark.parametrize("engine_args_and_prompt",

tests/v1/engine/test_llm_engine.py

-15

@@ -5,7 +5,6 @@ from typing import Optional

6 import pytest

8 -from tests.v1.engine.utils import PLP_APC_UNSUPPORTED_MSG

8 from vllm import LLM, SamplingParams

10 MODEL = "facebook/opt-125m"

@@ -98,17 +97,3 @@ def test_parallel_sampling(vllm_model, example_prompts) -> None:

97 raise AssertionError(

98 f"{len(completion_counts)} unique completions; expected"

99 f" {n}. Repeats: {repeats}")

101 -

102 -

103 -def test_llm_engine_refuses_prompt_logprobs_with_apc(vllm_model_apc):

104 - """Test passes if LLMEngine raises an exception when it is configured

105 - for automatic prefix caching and it receives a request with

106 - prompt_logprobs enabled, which is incompatible."""

107 - model: LLM = vllm_model_apc.model

108 - with pytest.raises(ValueError) as excinfo:

109 - model.generate(

110 - "Hello, my name is",

111 - SamplingParams(temperature=0.8, top_p=0.95, prompt_logprobs=5))

112 -

113 - # Validate exception string is correct

114 - assert str(excinfo.value) == PLP_APC_UNSUPPORTED_MSG

tests/v1/engine/utils.py

-3

@@ -30,9 +30,6 @@ FULL_STRINGS = [

30 STOP_STRINGS = ["I love working on", "company by far", "brother in"]

31 PROMPT_LEN = 5

33 -PLP_APC_UNSUPPORTED_MSG = ("Prefix caching with prompt logprobs not yet "

34 - "supported on VLLM V1.")

35 -

33 random.seed(42)

tests/v1/sample/test_logprobs.py

+85 -60

@@ -1,24 +1,34 @@

1 # SPDX-License-Identifier: Apache-2.0

3 import itertools

4 +from collections.abc import Generator

6 import pytest

7 import torch

9 from tests.kernels.utils import override_backend_env_variable

10 from tests.v1.sample.utils import (

11 + BatchLogprobsComposition, BatchLogprobsSpecType,

12 assert_incr_detok_str_matches_non_incr_detok_str,

13 compute_correct_cumulative_logprob, get_test_batch)

14 from vllm import SamplingParams

14 -from ...conftest import VllmRunner

16 +from ...conftest import HfRunner, VllmRunner

18 MODEL = "meta-llama/Llama-3.2-1B-Instruct"

19 DTYPE = "half"

21 +NONE = BatchLogprobsComposition.NONE

22 +SAMPLE = BatchLogprobsComposition.SAMPLE

23 +PROMPT = BatchLogprobsComposition.PROMPT

24 +SAMPLE_PROMPT = BatchLogprobsComposition.SAMPLE_PROMPT

20 -@pytest.fixture(scope="module")

21 -def vllm_model(vllm_runner):

26 +

27 +@pytest.fixture(

28 + scope="module",

29 + # Parameterize APC

30 + params=[False, True])

31 +def vllm_model(vllm_runner, request) -> Generator[VllmRunner, None, None]:

32 with vllm_runner(

33 MODEL,

34 dtype=DTYPE,

@@ -31,22 +41,22 @@ def vllm_model(vllm_runner):

41 enforce_eager=True,

42 #TODO: enable this once we support it for

43 # prompt logprobs.

34 - enable_prefix_caching=False,

44 + enable_prefix_caching=request.param,

45 gpu_memory_utilization=0.5,

46 ) as vllm_model:

47 yield vllm_model

50 @pytest.fixture(scope="module")

41 -def hf_model(hf_runner):

51 +def hf_model(hf_runner) -> Generator[HfRunner, None, None]:

52 with hf_runner(MODEL, dtype=DTYPE) as hf_model:

53 yield hf_model

56 def _repeat_logprob_config(

57 test_prompts,

48 - logprob_prompt_logprob_list: list[tuple],

49 -) -> list[tuple]:

58 + logprob_prompt_logprob_list: BatchLogprobsSpecType,

59 +) -> BatchLogprobsSpecType:

60 """Ensure each test prompt has a logprob config.

62 A logprob config specifies the optional (i.e.

@@ -91,42 +101,17 @@ def _repeat_logprob_config(

101 return logprob_prompt_logprob_list

102

103

94 -def _test_case_get_logprobs_and_prompt_logprobs(

95 - hf_model,

96 - vllm_model,

97 - batch_logprobs_composition: str,

104 +def _run_and_validate(

105 + vllm_model: VllmRunner,

106 + test_prompts: list[str],

107 + vllm_sampling_params: SamplingParams,

108 + hf_logprobs: list[list[torch.Tensor]],

109 + hf_outputs: list[tuple[list[int], str]],

110 + logprob_prompt_logprob_list: BatchLogprobsSpecType,

111 temperature: float,

99 - example_prompts,

112 + max_tokens: int,

113 + do_apc: bool,

114 ) -> None:

101 - test_prompts = example_prompts

102 -

103 - max_tokens = 5

104 - hf_outputs = hf_model.generate_greedy(

105 - test_prompts,

106 - max_tokens=max_tokens,

107 - )

108 - hf_logprobs = hf_model.generate_greedy_logprobs(

109 - test_prompts,

110 - max_tokens=max_tokens,

111 - )

112 -

113 - # Batch has mixed sample params

114 - # (different logprobs/prompt logprobs combos)

115 - logprob_prompt_logprob_list = get_test_batch(batch_logprobs_composition)

116 -

117 - # Ensure that each test prompt has a logprob config for testing

118 - logprob_prompt_logprob_list = _repeat_logprob_config(

119 - test_prompts, logprob_prompt_logprob_list)

120 - # Generate SamplingParams

121 - vllm_sampling_params = [

122 - SamplingParams(max_tokens=max_tokens,

123 - logprobs=num_lp,

124 - prompt_logprobs=num_plp,

125 - temperature=temperature,

126 - seed=1984)

127 - for num_lp, num_plp in logprob_prompt_logprob_list

128 - ]

129 -

115 vllm_results = vllm_model.model.generate(

116 test_prompts, sampling_params=vllm_sampling_params)

117

@@ -267,14 +252,13 @@ def _test_case_get_logprobs_and_prompt_logprobs(

252 assert vllm_result.prompt_logprobs is None

253

254

270 -#@pytest.mark.skip_global_cleanup

255 @pytest.mark.parametrize("batch_logprobs_composition",

272 - ["NONE", "SAMPLE", "PROMPT", "SAMPLE_PROMPT"])

256 + [NONE, SAMPLE, PROMPT, SAMPLE_PROMPT])

257 @pytest.mark.parametrize("temperature", [0.0, 2.0])

258 def test_get_logprobs_and_prompt_logprobs(

259 hf_model,

260 vllm_model,

277 - batch_logprobs_composition: str,

261 + batch_logprobs_composition: BatchLogprobsComposition,

262 temperature: float,

263 example_prompts,

264 ) -> None:

@@ -292,25 +276,70 @@ def test_get_logprobs_and_prompt_logprobs(

276 batch_logprobs_composition controls the logprobs configurations for

277 requests in the batch under test.

278

279 + APC tests run two test iterations so that cache hits occur.

280 +

281 + To save time, only test one APC-enabled scenario

282 + (sample & prompt logprobs enabled, temperature>0.0).

283 +

284 Args:

296 - hf_model

297 - vllm_model

285 + hf_model: HuggingFace reference model fixture

286 + vllm_model: vLLM model fixture

287 batch_logprobs_composition: logprobs configuration for test batch

299 - example_prompts

300 - monkeypatch

288 + temperature: "temperature" sampling parameter

289 + example_prompts: example prompt fixture

290 """

302 - _test_case_get_logprobs_and_prompt_logprobs(

303 - hf_model=hf_model,

304 - vllm_model=vllm_model,

305 - batch_logprobs_composition=batch_logprobs_composition,

306 - temperature=temperature,

307 - example_prompts=example_prompts)

291 + do_apc = vllm_model.model.llm_engine.cache_config.enable_prefix_caching

292 + if do_apc and (temperature < 2.0

293 + or batch_logprobs_composition != SAMPLE_PROMPT):

294 + # Skip some test-cases to save time.

295 + pytest.skip()

296 + test_prompts = example_prompts

297 +

298 + max_tokens = 5

299 + hf_outputs = hf_model.generate_greedy(

300 + test_prompts,

301 + max_tokens=max_tokens,

302 + )

303 + hf_logprobs = hf_model.generate_greedy_logprobs(

304 + test_prompts,

305 + max_tokens=max_tokens,

306 + )

307 +

308 + # Batch has mixed sample params

309 + # (different logprobs/prompt logprobs combos)

310 + logprob_prompt_logprob_list = get_test_batch(batch_logprobs_composition)

311 +

312 + # Ensure that each test prompt has a logprob config for testing

313 + logprob_prompt_logprob_list = _repeat_logprob_config(

314 + test_prompts, logprob_prompt_logprob_list)

315 + # Generate SamplingParams

316 + vllm_sampling_params = [

317 + SamplingParams(max_tokens=max_tokens,

318 + logprobs=num_lp,

319 + prompt_logprobs=num_plp,

320 + temperature=temperature,

321 + seed=1984)

322 + for num_lp, num_plp in logprob_prompt_logprob_list

323 + ]

324 + for _ in range(2 if do_apc else 1):

325 + _run_and_validate(

326 + vllm_model=vllm_model,

327 + test_prompts=test_prompts,

328 + vllm_sampling_params=vllm_sampling_params,

329 + hf_logprobs=hf_logprobs,

330 + hf_outputs=hf_outputs,

331 + logprob_prompt_logprob_list=logprob_prompt_logprob_list,

332 + temperature=temperature,

333 + max_tokens=max_tokens,

334 + do_apc=do_apc)

335

336

337 def test_max_logprobs(monkeypatch):

338 """vLLM v1 engine should fail a request with `logprobs > max_logprobs`

339

340 Should also fail for `prompt_logprobs > max_logprobs`

341 +

342 + APC should not matter as this test checks basic request validation.

343

344 Args:

345 monkeypatch

@@ -330,14 +359,12 @@ def test_max_logprobs(monkeypatch):

359 runner.generate(["Hello world"], sampling_params=bad_sampling_params)

360

361

333 -def test_none_logprobs(vllm_model, example_prompts, monkeypatch):

362 +def test_none_logprobs(vllm_model, example_prompts):

363 """Engine should return `logprobs` and `prompt_logprobs` as `None`

364

365 Args:

366 vllm_model: vLLM model fixture

367 example_prompts: list of example prompts (test fixture)

339 - monkeypatch: supports editing env vars and rolling back changes

340 - after the test

368 """

369 max_tokens = 5

370

@@ -356,14 +383,12 @@ def test_none_logprobs(vllm_model, example_prompts, monkeypatch):

383 assert results_logprobs_none[i].prompt_logprobs is None

384

385

359 -def test_zero_logprobs(vllm_model, example_prompts, monkeypatch):

386 +def test_zero_logprobs(vllm_model, example_prompts):

387 """Engine should return sampled token and prompt token logprobs

388

389 Args:

390 vllm_model: vLLM model fixture

391 example_prompts: list of example prompts (test fixture)

365 - monkeypatch: supports editing env vars and rolling back changes

366 - after the test

392 """

393 max_tokens = 5

394

tests/v1/sample/utils.py

+24 -9

@@ -1,27 +1,42 @@

1 # SPDX-License-Identifier: Apache-2.0

3 import re

4 +from enum import Enum

5 +from typing import Optional

7 from vllm import CompletionOutput

8 -def get_test_batch(batch_logprobs_composition: str) -> list[tuple]:

10 +class BatchLogprobsComposition(Enum):

11 + """Types of logprobs configs to include in test batch"""

12 + NONE = 0

13 + SAMPLE = 1

14 + PROMPT = 2

15 + SAMPLE_PROMPT = 3

16 +

17 +

18 +BatchLogprobsSpecType = list[tuple[Optional[int], Optional[int]]]

19 +

20 +

21 +def get_test_batch(

22 + batch_logprobs_composition: BatchLogprobsComposition

23 +) -> BatchLogprobsSpecType:

24 """Generate logprobs configs for a batch of requests

26 A given request's logprobs configuration is (1) num_sample_logprobs and (2)

27 num_prompt_logprobs. The batch logprobs configuration is the list of request

28 logprobs configs.

15 - batch_logprobs_composition == "NONE" yields a batch with no sample or prompt

30 + batch_logprobs_composition == NONE yields a batch with no sample or prompt

31 logprobs

18 - batch_logprobs_composition == "SAMPLE" yields a batch with some requests

33 + batch_logprobs_composition == SAMPLE yields a batch with some requests

34 configured for sample logprobs only, and others configured for no logprobs

21 - batch_logprobs_composition == "PROMPT" yields a batch with some requests

36 + batch_logprobs_composition == PROMPT yields a batch with some requests

37 configured for prompt logprobs only, and others configured for no logprobs

24 - batch_logprobs_composition == "SAMPLE_PROMPT" yields a batch with some

39 + batch_logprobs_composition == SAMPLE_PROMPT yields a batch with some

40 requests configured for sample logprobs and prompt logprobs, some configured

41 for only sample logprobs or only prompt logprobs, and some configured for

42 no logprobs

@@ -34,10 +49,10 @@ def get_test_batch(batch_logprobs_composition: str) -> list[tuple]:

49 list of (Optional[num_sample_logprobs], Optional[num_prompt_logprobs])

50 tuples

51 """

37 - if batch_logprobs_composition == "NONE":

52 + if batch_logprobs_composition == BatchLogprobsComposition.NONE:

53 # No requests with sample or prompt logprobs

54 return [(None, None)]

40 - elif batch_logprobs_composition == "SAMPLE":

55 + elif batch_logprobs_composition == BatchLogprobsComposition.SAMPLE:

56 # Requests requiring sample logprobs or no logprobs

57 return [

58 (None, None),

@@ -45,7 +60,7 @@ def get_test_batch(batch_logprobs_composition: str) -> list[tuple]:

60 (5, None),

61 (3, None),

62 ]

48 - elif batch_logprobs_composition == "PROMPT":

63 + elif batch_logprobs_composition == BatchLogprobsComposition.PROMPT:

64 # Requests requiring prompt logprobs or no logprobs

65 return [

66 (None, None),

@@ -53,7 +68,7 @@ def get_test_batch(batch_logprobs_composition: str) -> list[tuple]:

68 (None, 6),

69 (None, 5),

70 ]

56 - elif batch_logprobs_composition == "SAMPLE_PROMPT":

71 + elif batch_logprobs_composition == BatchLogprobsComposition.SAMPLE_PROMPT:

72 # Requests requiring either no logprobs, just

73 # sample logprobs, just prompt logprobs, or

74 # both sample and prompt logprobs

vllm/v1/core/kv_cache_manager.py

+24 -19

@@ -105,8 +105,6 @@ class KVCacheManager:

105 # Prefix caching is disabled.

106 return [], 0

107

108 - computed_blocks = []

109 -

108 # The block hashes for the request may already be computed

109 # if the scheduler has tried to schedule the request before.

110 block_hashes = self.req_to_block_hashes[request.request_id]

@@ -114,24 +112,31 @@ class KVCacheManager:

112 block_hashes = hash_request_tokens(self.block_size, request)

113 self.req_to_block_hashes[request.request_id] = block_hashes

114

117 - for block_hash in block_hashes:

118 - # block_hashes is a chain of block hashes. If a block hash is not

119 - # in the cached_block_hash_to_id, the following block hashes are

120 - # not computed yet for sure.

121 - if cached_block := self.block_pool.get_cached_block(block_hash):

122 - computed_blocks.append(cached_block)

123 - else:

124 - break

125 -

115 self.prefix_cache_stats.requests += 1

127 - self.prefix_cache_stats.queries += len(block_hashes)

128 - self.prefix_cache_stats.hits += len(computed_blocks)

129 -

130 - # NOTE(woosuk): Since incomplete blocks are not eligible for

131 - # sharing, `num_computed_tokens` is always a multiple of

132 - # `block_size`.

133 - num_computed_tokens = len(computed_blocks) * self.block_size

134 - return computed_blocks, num_computed_tokens

116 + if request.sampling_params.prompt_logprobs is None:

117 + # Check for cache hits

118 + computed_blocks = []

119 + for block_hash in block_hashes:

120 + # block_hashes is a chain of block hashes. If a block hash

121 + # is not in the cached_block_hash_to_id, the following

122 + # block hashes are not computed yet for sure.

123 + if cached_block := self.block_pool.get_cached_block(

124 + block_hash):

125 + computed_blocks.append(cached_block)

126 + else:

127 + break

128 +

129 + self.prefix_cache_stats.queries += len(block_hashes)

130 + self.prefix_cache_stats.hits += len(computed_blocks)

131 +

132 + # NOTE(woosuk): Since incomplete blocks are not eligible for

133 + # sharing, `num_computed_tokens` is always a multiple of

134 + # `block_size`.

135 + num_computed_tokens = len(computed_blocks) * self.block_size

136 + return computed_blocks, num_computed_tokens

137 + else:

138 + # Skip cache hits for prompt logprobs

139 + return [], 0

140

141 def allocate_slots(

142 self,

vllm/v1/engine/processor.py

-6

@@ -72,12 +72,6 @@ class Processor:

72 f"Requested prompt logprobs of {params.prompt_logprobs}, "

73 f"which is greater than max allowed: {max_logprobs}")

75 - # TODO(andy): enable this in follow up by recomputing.

76 - if (params.prompt_logprobs is not None

77 - and self.cache_config.enable_prefix_caching):

78 - raise ValueError("Prefix caching with prompt logprobs not yet "

79 - "supported on VLLM V1.")

80 -

75 def _validate_sampling_params(

76 self,

77 params: SamplingParams,