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

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 5, 04:45 UTC

[Core][Metrics][BugFix] Replace num_cached_tokens/num_external_computed_tokens with PrefillStats (#37460)

Related to `Counters can only be incremented by non-negative amounts` error with the `vllm:prompt_tokens_by_source_total` metric. Signed-off-by: Mark McLoughlin <markmc@redhat.com> Co-authored-by: Or Ozeri <or@ozery.com>

Mark McLoughlin committed Apr 14, 2026 at 09:00 UTC d3af8c18317c0dc008d42e4367fbb9045cfb7bf6

9 files changed +187 -102

tests/v1/core/test_async_scheduler.py

+18 -11

@@ -153,7 +153,6 @@ def test_prefix_caching_for_prefill_dedup():

153 same_prompt=True,

154 block_size=BLOCK_SIZE,

155 )

156 - requests_copy = requests.copy()

156

157 # Two requests with the same prompt.

158 req0 = requests.pop(0)

@@ -167,26 +166,31 @@ def test_prefix_caching_for_prefill_dedup():

166 # Make sure prefix caching de-duplicates the prompts in the same step,

167 # so all the blocks except the last are shared between the two requests.

168 assert len(sched_output.num_scheduled_tokens) == 2

170 - num_blocks = num_prompt_tokens // BLOCK_SIZE

171 - assert req0.num_cached_tokens == 0

172 - assert req1.num_cached_tokens >= num_blocks * BLOCK_SIZE

169 + assert sched_output.num_scheduled_tokens[req0.request_id] == num_prompt_tokens

170 + assert (

171 + sched_output.num_scheduled_tokens[req1.request_id]

172 + == num_prompt_tokens % BLOCK_SIZE

173 + )

174

175 sched_outputs.append(scheduler.schedule())

176 while sched_outputs:

177 + added_req = None

178 if requests:

177 - scheduler.add_request(requests.pop(0))

179 + added_req = requests.pop(0)

180 + scheduler.add_request(added_req)

181 sched_output = sched_outputs.popleft()

182 model_runner_output = _make_model_runner_output(sched_output)

183 scheduler.update_from_output(sched_output, model_runner_output)

184 sched_output = scheduler.schedule()

185 if sched_output.num_scheduled_tokens:

186 sched_outputs.append(sched_output)

187 + if added_req:

188 + assert (

189 + sched_output.num_scheduled_tokens[added_req.request_id]

190 + == num_prompt_tokens % BLOCK_SIZE

191 + )

192

185 - # Other requests scheduled after the two requests should also get

186 - # prefix cache hit.

193 assert scheduler.get_num_unfinished_requests() == 0

188 - for req in requests_copy[1:]:

189 - assert req.num_cached_tokens >= num_blocks * BLOCK_SIZE

194

195

196 def test_prefix_caching_for_multi_turn():

@@ -243,12 +247,15 @@ def test_prefix_caching_for_multi_turn():

247 # Schedule the next-turn requests.

248 for req in next_turn_requests:

249 scheduler.add_request(req)

246 - sched_outputs.append(scheduler.schedule())

250 + sched_output = scheduler.schedule()

251 + sched_outputs.append(sched_output)

252

253 # Make sure the next-turn requests get prefix cache hit by the previous

254 # requests.

255 for req in next_turn_requests:

251 - assert req.num_cached_tokens == req.num_prompt_tokens // BLOCK_SIZE * BLOCK_SIZE

256 + assert sched_output.num_scheduled_tokens[req.request_id] == (

257 + req.num_prompt_tokens % BLOCK_SIZE

258 + )

259

260

261 def test_abort_request_when_structured_output_fsm_cannot_advance():

tests/v1/engine/test_output_processor.py

+10 -4

@@ -84,6 +84,7 @@ def test_incremental_detokenization(

85 engine_core = MockEngineCore(

86 tokens_list=dummy_test_vectors.generation_tokens,

87 + prompts_list=dummy_test_vectors.prompt_tokens,

88 request_ids=[req.request_id for req in requests],

89 )

@@ -506,6 +507,7 @@ def test_logprobs_processor(

507

508 engine_core = MockEngineCore(

509 tokens_list=dummy_test_vectors.generation_tokens,

510 + prompts_list=dummy_test_vectors.prompt_tokens,

511 generated_logprobs_raw=None

512 if num_sample_logprobs is None

513 else dummy_test_vectors.generation_logprobs,

@@ -691,6 +693,7 @@ def test_stop_token(

693

694 engine_core = MockEngineCore(

695 tokens_list=[generation_tokens],

696 + prompts_list=dummy_test_vectors.prompt_tokens,

697 generated_logprobs_raw=[generation_logprobs] if do_logprobs else None,

698 prompt_logprobs_raw=None,

699 eos_token_id=sampling_params.eos_token_id,

@@ -794,6 +797,7 @@ def test_stop_string(

797

798 engine_core = MockEngineCore(

799 tokens_list=dummy_test_vectors.generation_tokens,

800 + prompts_list=dummy_test_vectors.prompt_tokens,

801 generated_logprobs_raw=dummy_test_vectors.generation_logprobs

802 if num_sample_logprobs

803 else None,

@@ -917,6 +921,7 @@ def test_iteration_stats(dummy_test_vectors):

921

922 engine_core = MockEngineCore(

923 dummy_test_vectors.generation_tokens,

924 + dummy_test_vectors.prompt_tokens,

925 request_ids=[req.request_id for req in requests],

926 )

927

@@ -927,7 +932,7 @@ def test_iteration_stats(dummy_test_vectors):

932 inactive_request = requests[num_active]

933

934 # First iteration has 2 prefills.

930 - outputs = engine_core.get_outputs()[:num_active]

935 + outputs = engine_core.get_outputs(num_active)

936 iteration_stats = IterationStats()

937 output_processor.process_outputs(outputs, engine_core_timestamp, iteration_stats)

938 total_prompt_tokens = sum(

@@ -941,7 +946,7 @@ def test_iteration_stats(dummy_test_vectors):

946 assert iteration_stats.num_generation_tokens == num_active

947

948 # Just decodes in this step.

944 - outputs = engine_core.get_outputs()[:num_active]

949 + outputs = engine_core.get_outputs(num_active)

950 iteration_stats = IterationStats()

951 output_processor.process_outputs(outputs, engine_core_timestamp, iteration_stats)

952

@@ -951,7 +956,7 @@ def test_iteration_stats(dummy_test_vectors):

956 # Add a new request - prefill and 2 decodes in this step.

957 output_processor.add_request(inactive_request, None)

958 num_active += 1

954 - outputs = engine_core.get_outputs()[:num_active]

959 + outputs = engine_core.get_outputs(num_active)

960 iteration_stats = IterationStats()

961 output_processor.process_outputs(outputs, engine_core_timestamp, iteration_stats)

962 total_prompt_tokens = len(dummy_test_vectors.prompt_tokens[num_active - 1])

@@ -960,7 +965,7 @@ def test_iteration_stats(dummy_test_vectors):

965 assert iteration_stats.num_generation_tokens == num_active

966

967 # Just decodes in this step.

963 - outputs = engine_core.get_outputs()[:num_active]

968 + outputs = engine_core.get_outputs(num_active)

969 iteration_stats = IterationStats()

970 output_processor.process_outputs(outputs, engine_core_timestamp, iteration_stats)

971

@@ -1003,6 +1008,7 @@ def test_lora_request_tracking(log_stats: bool, dummy_test_vectors):

1008

1009 engine_core = MockEngineCore(

1010 dummy_test_vectors.generation_tokens,

1011 + dummy_test_vectors.prompt_tokens,

1012 request_ids=[req.request_id for req in requests],

1013 )

1014

tests/v1/engine/utils.py

+27 -6

@@ -11,6 +11,7 @@ from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast

12 from vllm.engine.arg_utils import EngineArgs

13 from vllm.v1.engine import EngineCoreOutput, FinishReason

14 +from vllm.v1.metrics.stats import PrefillStats

15 from vllm.v1.outputs import LogprobsLists, LogprobsTensors

17 GeneralTokenizerType: TypeAlias = PreTrainedTokenizer | PreTrainedTokenizerFast

@@ -330,6 +331,7 @@ class MockEngineCore:

331 def __init__(

332 self,

333 tokens_list: list[list[int]],

334 + prompts_list: list[list[int]],

335 # For each request, for each sampled token offset,

336 # a tuple of

337 # (list of topk token ids, list of sample logprob vals, rank)

@@ -346,12 +348,13 @@ class MockEngineCore:

348 ) -> None:

349 self.num_requests = len(tokens_list)

350 self.tokens_list = tokens_list

349 - self.current_idx = 0

351 + self.prompts_list = prompts_list

352 self.generated_logprobs_raw = generated_logprobs_raw

353 self.do_logprobs = generated_logprobs_raw is not None

354 self.prompt_logprobs_raw = prompt_logprobs_raw

355 self.do_prompt_logprobs = prompt_logprobs_raw is not None

356 self.request_finished = [False for _ in range(self.num_requests)]

357 + self.request_token_idx = [0 for _ in range(self.num_requests)]

358 self.eos_token_id = eos_token_id

359 self.stop_token_ids = stop_token_ids

360 self.request_ids = (

@@ -360,14 +363,18 @@ class MockEngineCore:

363 else [f"request-{i}" for i in range(self.num_requests)]

364 )

365

363 - def get_outputs(self) -> list[EngineCoreOutput]:

366 + def get_outputs(self, num_active: int = -1) -> list[EngineCoreOutput]:

367 do_logprobs = self.do_logprobs

368 do_prompt_logprobs = self.do_prompt_logprobs

366 - token_idx = self.current_idx

369

370 outputs = []

369 - for req_idx, token_ids in enumerate(self.tokens_list):

371 + for req_idx, (token_ids, prompt_token_ids) in enumerate(

372 + zip(self.tokens_list, self.prompts_list)

373 + ):

374 + if num_active != -1 and req_idx >= num_active:

375 + break

376 if not self.request_finished[req_idx]:

377 + token_idx = self.request_token_idx[req_idx]

378 if do_logprobs:

379 assert self.generated_logprobs_raw is not None

380 (logprobs_token_ids_, logprobs_, sampled_token_ranks_) = (

@@ -381,19 +388,32 @@ class MockEngineCore:

388 else:

389 logprobs = None

390 if do_prompt_logprobs:

384 - if self.current_idx == 0:

391 + if token_idx == 0:

392 assert self.prompt_logprobs_raw is not None

393 prompt_logprobs = self.prompt_logprobs_raw[req_idx]

394 else:

395 prompt_logprobs = None

396 else:

397 prompt_logprobs = None

398 +

399 + # Add prefill_stats on first output (prefill) for this request

400 + if token_idx == 0:

401 + prefill_stats = PrefillStats()

402 + prefill_stats.set(

403 + num_prompt_tokens=len(prompt_token_ids),

404 + num_local_cached_tokens=0,

405 + num_external_cached_tokens=0,

406 + )

407 + else:

408 + prefill_stats = None

409 +

410 new_token_id = token_ids[token_idx]

411 output = EngineCoreOutput(

412 request_id=self.request_ids[req_idx],

413 new_token_ids=[new_token_id],

414 new_logprobs=logprobs,

415 new_prompt_logprobs_tensors=prompt_logprobs,

416 + prefill_stats=prefill_stats,

417 )

418 if token_idx == len(token_ids) - 1:

419 output.finish_reason = FinishReason.LENGTH

@@ -407,5 +427,6 @@ class MockEngineCore:

427 self.request_finished[req_idx] = True

428 outputs.append(output)

429

410 - self.current_idx += 1

430 + self.request_token_idx[req_idx] += 1

431 +

432 return outputs

tests/v1/metrics/test_stats.py

+58 -29

@@ -1,7 +1,12 @@

1 # SPDX-License-Identifier: Apache-2.0

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

3 from vllm.v1.engine import FinishReason

4 -from vllm.v1.metrics.stats import IterationStats, PromptTokenStats, RequestStateStats

4 +from vllm.v1.metrics.stats import (

5 + IterationStats,

6 + PrefillStats,

7 + PromptTokenStats,

8 + RequestStateStats,

9 +)

12 def test_iteration_stats_repr():

@@ -114,15 +119,18 @@ def test_prompt_token_stats_all_computed():

119 stats = PromptTokenStats()

120

121 # Case 1: No caching (All tokens computed locally)

117 - stats.update_from_output(

118 - num_cached_tokens=0,

119 - num_external_computed_tokens=0,

120 - prompt_len=1000,

122 + prefill_stats = PrefillStats()

123 + prefill_stats.set(

124 + num_prompt_tokens=1000,

125 + num_local_cached_tokens=0,

126 + num_external_cached_tokens=0,

127 )

128 + stats.update_from_output(prefill_stats)

129

130 assert stats.computed == 1000

131 assert stats.local_cache_hit == 0

132 assert stats.external_kv_transfer == 0

133 + assert stats.cached_tokens == 0

134 assert stats.total == 1000

135

136

@@ -131,15 +139,19 @@ def test_prompt_token_stats_partial_local_cache():

139 stats = PromptTokenStats()

140

141 # Case 2: Partial local cache

134 - stats.update_from_output(

135 - num_cached_tokens=300,

136 - num_external_computed_tokens=0,

137 - prompt_len=1000,

142 + prefill_stats = PrefillStats()

143 + prefill_stats.set(

144 + num_prompt_tokens=1000,

145 + num_local_cached_tokens=300,

146 + num_external_cached_tokens=0,

147 )

148 + stats.update_from_output(prefill_stats)

149

150 assert stats.computed == 700

151 assert stats.local_cache_hit == 300

152 assert stats.external_kv_transfer == 0

153 + assert stats.cached_tokens == 300

154 + assert stats.total == 1000

155

156

157 def test_prompt_token_stats_partial_external_transfer():

@@ -147,15 +159,19 @@ def test_prompt_token_stats_partial_external_transfer():

159 stats = PromptTokenStats()

160

161 # Case 3: Partial external transfer

150 - stats.update_from_output(

151 - num_cached_tokens=500,

152 - num_external_computed_tokens=500,

153 - prompt_len=1000,

162 + prefill_stats = PrefillStats()

163 + prefill_stats.set(

164 + num_prompt_tokens=1000,

165 + num_local_cached_tokens=0,

166 + num_external_cached_tokens=500,

167 )

168 + stats.update_from_output(prefill_stats)

169

170 assert stats.computed == 500

171 assert stats.local_cache_hit == 0

172 assert stats.external_kv_transfer == 500

173 + assert stats.cached_tokens == 500

174 + assert stats.total == 1000

175

176

177 def test_prompt_token_stats_mixed_sources():

@@ -163,47 +179,60 @@ def test_prompt_token_stats_mixed_sources():

179 stats = PromptTokenStats()

180

181 # Case 4: Mixed sources

166 - stats.update_from_output(

167 - num_cached_tokens=600,

168 - num_external_computed_tokens=200,

169 - prompt_len=1000,

182 + prefill_stats = PrefillStats()

183 + prefill_stats.set(

184 + num_prompt_tokens=1000,

185 + num_local_cached_tokens=400,

186 + num_external_cached_tokens=200,

187 )

188 + stats.update_from_output(prefill_stats)

189

190 assert stats.computed == 400

191 assert stats.local_cache_hit == 400

192 assert stats.external_kv_transfer == 200

193 + assert stats.cached_tokens == 600

194 + assert stats.total == 1000

195

196

197 def test_prompt_token_stats_full_local_cache_recompute():

198 """Test full local cache triggers last token recomputation.

199

180 - When all tokens are cached, the scheduler reduces num_cached_tokens by 1

181 - to force the model to recompute the last token.

200 + When all tokens are cached, the scheduler forces the model to recompute

201 + the last token (num_computed_tokens=1), with the rest from cache.

202 """

203 stats = PromptTokenStats()

204

185 - # Case 5: Full local cache (999 cached after reduction, 1 recomputed)

186 - stats.update_from_output(

187 - num_cached_tokens=999,

188 - num_external_computed_tokens=0,

189 - prompt_len=1000,

205 + # Case 5: Full local cache (999 cached, 1 recomputed)

206 + prefill_stats = PrefillStats()

207 + prefill_stats.set(

208 + num_prompt_tokens=1000,

209 + num_local_cached_tokens=999,

210 + num_external_cached_tokens=0,

211 )

212 + stats.update_from_output(prefill_stats)

213

214 assert stats.computed == 1

215 assert stats.local_cache_hit == 999

216 + assert stats.external_kv_transfer == 0

217 + assert stats.cached_tokens == 999

218 + assert stats.total == 1000

219

220

221 def test_prompt_token_stats_full_external_transfer_recompute():

222 """Test full external transfer triggers last token recomputation."""

223 stats = PromptTokenStats()

224

200 - # Case 6: Full external transfer (999 cached after reduction, 1 recomputed)

201 - stats.update_from_output(

202 - num_cached_tokens=999,

203 - num_external_computed_tokens=999,

204 - prompt_len=1000,

225 + # Case 6: Full external transfer (999 from external, 1 recomputed)

226 + prefill_stats = PrefillStats()

227 + prefill_stats.set(

228 + num_prompt_tokens=1000,

229 + num_local_cached_tokens=0,

230 + num_external_cached_tokens=999,

231 )

232 + stats.update_from_output(prefill_stats)

233

234 assert stats.computed == 1

235 assert stats.local_cache_hit == 0

236 assert stats.external_kv_transfer == 999

237 + assert stats.cached_tokens == 999

238 + assert stats.total == 1000

vllm/v1/core/sched/scheduler.py

+11 -12

@@ -629,7 +629,6 @@ class Scheduler(SchedulerInterface):

629 step_skipped_waiting.prepend_request(request)

630 continue

631

632 - request.num_external_computed_tokens = ext_tokens

632 num_external_computed_tokens = ext_tokens

633

634 connector_prefix_cache_queries = (

@@ -642,6 +641,15 @@ class Scheduler(SchedulerInterface):

641 num_new_local_computed_tokens + num_external_computed_tokens

642 )

643 assert num_computed_tokens <= request.num_tokens

644 +

645 + # Track first scheduled prefill, not post-preemption repeat prefills

646 + if request.prefill_stats is not None:

647 + assert num_computed_tokens <= request.num_prompt_tokens

648 + request.prefill_stats.set(

649 + num_prompt_tokens=request.num_prompt_tokens,

650 + num_local_cached_tokens=num_new_local_computed_tokens,

651 + num_external_cached_tokens=num_external_computed_tokens,

652 + )

653 else:

654 # KVTransfer: WAITING reqs have num_computed_tokens > 0

655 # after async KV recvs are completed.

@@ -826,9 +834,6 @@ class Scheduler(SchedulerInterface):

834 token_budget -= num_new_tokens

835 request.status = RequestStatus.RUNNING

836 request.num_computed_tokens = num_computed_tokens

829 - # Count the number of prefix cached tokens.

830 - if request.num_cached_tokens < 0:

831 - request.num_cached_tokens = num_computed_tokens

837 # Encoder-related.

838 if encoder_inputs_to_schedule:

839 scheduled_encoder_inputs[request_id] = encoder_inputs_to_schedule

@@ -1466,10 +1471,9 @@ class Scheduler(SchedulerInterface):

1471 pooling_output=pooler_output,

1472 stop_reason=request.stop_reason,

1473 events=request.take_events(),

1474 + prefill_stats=request.take_prefill_stats(),

1475 kv_transfer_params=kv_transfer_params,

1476 trace_headers=request.trace_headers,

1471 - num_cached_tokens=request.num_cached_tokens,

1472 - num_external_computed_tokens=request.num_external_computed_tokens,

1477 routed_experts=routed_experts,

1478 num_nans_in_logits=request.num_nans_in_logits,

1479 )

@@ -1496,7 +1500,6 @@ class Scheduler(SchedulerInterface):

1500 finish_reason=request.get_finished_reason(),

1501 events=request.take_events(),

1502 trace_headers=request.trace_headers,

1499 - num_cached_tokens=request.num_cached_tokens,

1503 )

1504 )

1505

@@ -2071,10 +2074,6 @@ class Scheduler(SchedulerInterface):

2074 if request.num_computed_tokens == request.num_tokens:

2075 request.num_computed_tokens = request.num_tokens - 1

2076

2074 - # Count the number of prefix cached tokens.

2075 - if request.num_cached_tokens < 0:

2076 - request.num_cached_tokens = request.num_computed_tokens

2077 -

2077 self.finished_recving_kv_req_ids.remove(request.request_id)

2078

2079 def _try_promote_blocked_waiting_request(self, request: Request) -> bool:

@@ -2221,7 +2220,7 @@ class Scheduler(SchedulerInterface):

2220 req_num_computed_tokens - request.num_computed_tokens

2221 )

2222 total_affected_tokens += num_affected_tokens

2224 - request.num_external_computed_tokens -= num_affected_tokens

2223 +

2224 # collect invalid block and all downstream dependent blocks

2225 if evict_blocks:

2226 blocks_to_evict.update(req_block_ids[idx:])

vllm/v1/engine/__init__.py

+4 -5

@@ -15,7 +15,7 @@ from vllm.lora.request import LoRARequest

15 from vllm.multimodal.inputs import MultiModalFeatureSpec

16 from vllm.pooling_params import PoolingParams

17 from vllm.sampling_params import SamplingParams

18 -from vllm.v1.metrics.stats import SchedulerStats

18 +from vllm.v1.metrics.stats import PrefillStats, SchedulerStats

19 from vllm.v1.outputs import LogprobsLists, LogprobsTensors

20 from vllm.v1.serial_utils import UtilityResult

@@ -171,10 +171,9 @@ class EngineCoreOutput(

171 kv_transfer_params: dict[str, Any] | None = None

172

173 trace_headers: Mapping[str, str] | None = None

174 - # The number of tokens with prefix cache hits (local + external).

175 - num_cached_tokens: int = 0

176 - # The number of tokens computed remotely (original count from connector).

177 - num_external_computed_tokens: int = 0

174 +

175 + prefill_stats: PrefillStats | None = None

176 +

177 routed_experts: np.ndarray | None = None

178 # The number of NaNs in logits.

179 # A value greater than 0 indicates that the output is corrupted.

vllm/v1/engine/output_processor.py

+7 -3

@@ -617,8 +617,13 @@ class OutputProcessor:

617 stop_reason = engine_core_output.stop_reason

618 kv_transfer_params = engine_core_output.kv_transfer_params

619 routed_experts = engine_core_output.routed_experts

620 - req_state.num_cached_tokens = engine_core_output.num_cached_tokens

621 - req_state.is_prefilling = False

620 +

621 + if req_state.is_prefilling:

622 + if engine_core_output.prefill_stats is not None:

623 + req_state.num_cached_tokens = (

624 + engine_core_output.prefill_stats.num_cached_tokens

625 + )

626 + req_state.is_prefilling = False

627

628 if pooling_output is None:

629 assert req_state.detokenizer is not None

@@ -776,7 +781,6 @@ class OutputProcessor:

781 engine_core_output,

782 engine_core_timestamp,

783 req_state.is_prefilling,

779 - req_state.prompt_len,

784 req_state.stats,

785 self.lora_states,

786 req_state.lora_name,

vllm/v1/metrics/stats.py

+43 -27

@@ -239,6 +239,40 @@ class FinishedRequestStats:

239 num_cached_tokens: int = 0

240

241

242 +@dataclass

243 +class PrefillStats:

244 + """Breakdown of a scheduled prefill computation.

245 +

246 + Fields:

247 + num_prompt_tokens: Total number of tokens to be prefilled.

248 + num_computed_tokens: Tokens to be prefilled locally (actual compute work).

249 + num_cached_tokens: Tokens to be prefilled without actual compute work.

250 + num_local_cached_tokens: Tokens to be prefilled from local prefix cache.

251 + num_external_cached_tokens: Tokens to be prefilled from external KV transfer.

252 + """

253 +

254 + num_prompt_tokens: int = 0

255 + num_computed_tokens: int = 0

256 + num_cached_tokens: int = 0

257 + num_local_cached_tokens: int = 0

258 + num_external_cached_tokens: int = 0

259 +

260 + def set(

261 + self,

262 + num_prompt_tokens: int,

263 + num_local_cached_tokens: int,

264 + num_external_cached_tokens: int,

265 + ):

266 + num_cached_tokens = num_local_cached_tokens + num_external_cached_tokens

267 + assert num_cached_tokens <= num_prompt_tokens

268 +

269 + self.num_prompt_tokens = num_prompt_tokens

270 + self.num_computed_tokens = num_prompt_tokens - num_cached_tokens

271 + self.num_cached_tokens = num_cached_tokens

272 + self.num_local_cached_tokens = num_local_cached_tokens

273 + self.num_external_cached_tokens = num_external_cached_tokens

274 +

275 +

276 @dataclass

277 class PromptTokenStats:

278 """Breakdown of prompt tokens by source.

@@ -267,28 +301,14 @@ class PromptTokenStats:

301 cached_tokens: int = 0

302 total: int = 0

303

270 - def update_from_output(

271 - self,

272 - num_cached_tokens: int,

273 - num_external_computed_tokens: int,

274 - prompt_len: int,

275 - ) -> None:

304 + def update_from_output(self, prefill_stats: PrefillStats) -> None:

305 """Update stats from a prefill output."""

277 - self.computed += prompt_len - num_cached_tokens

278 - self.external_kv_transfer += num_external_computed_tokens

279 - # FIXME(yifan): local_cache_hit can go negative after preemption.

280 - # num_cached_tokens is a one-time snapshot from first scheduling and

281 - # is never reset on preemption, while num_external_computed_tokens is

282 - # overwritten on re-scheduling. If CPU offload finds more tokens on

283 - # the second pass than the original total, the subtraction underflows.

284 - # A fundamental fix is to track the first-time num_external_computed_tokens

285 - # as a separate metric rather than reusing num_external_computed_tokens

286 - # for metric directly.

287 - self.local_cache_hit += max(

288 - 0, (num_cached_tokens - num_external_computed_tokens)

289 - )

290 - self.cached_tokens += num_cached_tokens

291 - self.total += prompt_len

306 + self.computed += prefill_stats.num_computed_tokens

307 + self.cached_tokens += prefill_stats.num_cached_tokens

308 + self.total += prefill_stats.num_prompt_tokens

309 +

310 + self.local_cache_hit += prefill_stats.num_local_cached_tokens

311 + self.external_kv_transfer += prefill_stats.num_external_cached_tokens

312

313 def get_by_source(self, source: str) -> int:

314 """Get token count by source label."""

@@ -335,7 +355,6 @@ class IterationStats:

355 output: "EngineCoreOutput",

356 engine_core_timestamp: float,

357 is_prefilling: bool,

338 - prompt_len: int,

358 req_stats: RequestStateStats,

359 lora_states: "LoRARequestStates",

360 lora_name: str | None,

@@ -344,11 +363,8 @@ class IterationStats:

363

364 self.num_generation_tokens += num_new_generation_tokens

365 if is_prefilling:

347 - self.prompt_token_stats.update_from_output(

348 - num_cached_tokens=output.num_cached_tokens,

349 - num_external_computed_tokens=output.num_external_computed_tokens,

350 - prompt_len=prompt_len,

351 - )

366 + if output.prefill_stats is not None:

367 + self.prompt_token_stats.update_from_output(output.prefill_stats)

368

369 first_token_latency = self._time_since(req_stats.arrival_time)

370 self.time_to_first_tokens_iter.append(first_token_latency)

vllm/v1/request.py

+9 -5

@@ -20,6 +20,7 @@ from vllm.v1.engine import (

20 EngineCoreRequest,

21 FinishReason,

22 )

23 +from vllm.v1.metrics.stats import PrefillStats

24 from vllm.v1.structured_output.request import StructuredOutputRequest

25 from vllm.v1.utils import ConstantList

@@ -145,9 +146,6 @@ class Request:

146 self.all_token_ids = ConstantList(self._all_token_ids)

147 # trace_headers

148 self.trace_headers = trace_headers

148 - # State

149 - # The number of tokens with prefix cache hits.

150 - self.num_cached_tokens = -1

149

150 # True if this request is scheduled as a non-final prefill chunk.

151 self.is_prefill_chunk = False

@@ -159,8 +157,7 @@ class Request:

157 # The number of times this request has been preempted by the scheduler.

158 self.num_preemptions = 0

159

162 - # The number of tokens that have been computed remotely.

163 - self.num_external_computed_tokens = 0

160 + self.prefill_stats: PrefillStats | None = PrefillStats()

161

162 self.block_hashes: list[BlockHash] = []

163 # Store the block hasher without binding self to avoid creating a

@@ -278,6 +275,13 @@ class Request:

275 events, self.events = self.events, []

276 return events

277

278 + def take_prefill_stats(self) -> PrefillStats | None:

279 + if self.prefill_stats is None:

280 + return None

281 + prefill_stats = self.prefill_stats

282 + self.prefill_stats = None

283 + return prefill_stats

284 +

285 def __lt__(self, other: "Request") -> bool:

286 """

287 Compare two requests based on priority, arrival time, and request ID.