[Core][Metrics][BugFix] Replace num_cached_tokens/num_external_computed_tokens with PrefillStats

[markmc]

Related to the discussion in #36859 and the Counters can only be incremented by non-negative amounts error with the vllm:prompt_tokens_by_source_total metric.

In OutputProcessor, we take the first EngineCoreOutput as a signal that prefill has completed, and record certain statistics about it.

On the scheduler side, because of preemption, we might have prefills that are scheduled but never completed, or we might need to recompute an already completed prefill. To add clarity, we use PrefillStats to track the first scheduled prefill so that the stats can be returned to the frontend via EngineCoreOutput.

num_cached_tokens was previously used for KV transfer failure recovery, but this is no longer true as of #38096. We also no longer attempt to correct these prefill metrics if KV transfers failed, since this introduced unjustified brittleness to an already brittle code path.

[markmc]

I've just realized an obvious gap - the prefill stats aren't updated on the failed-KV-transfer path

[gemini-code-assist]

Code Review

This pull request introduces PrefillStats to more accurately track and report prefill-related metrics, which is a valuable improvement for clarity and correctness, especially in scenarios involving preemption. The refactoring to use a dedicated PrefillStats object instead of separate numerical values is a good design choice. The changes are consistently applied across the scheduler, engine outputs, request objects, and tests. I have found one critical issue in the metrics calculation logic that leads to double-counting of tokens, for which I've provided a detailed comment and a suggested fix.

[orozery]

num_cached_tokens and num_external_computed_tokens remain in place and untouched for now - next we will look at what the remaining uses of these are.

Right now we basically duplicate these fields into prefill stats.
But we already have problems with these original fields, so I think it's better we set prefill stats directly, removing any connection to the old fields.

And I guess once we do that, we can think whether we want to keep the old fields (probably not?).

[markmc]

num_cached_tokens and num_external_computed_tokens remain in place and untouched for now - next we will look at what the remaining uses of these are.

Right now we basically duplicate these fields into prefill stats. But we already have problems with these original fields, so I think it's better we set prefill stats directly, removing any connection to the old fields.

IMO, the PR basically does that (set the stats directly) ?

And I guess once we do that, we can think whether we want to keep the old fields (probably not?).

I'm going to take a look at this, and I expect I can remove/rework these fields without affecting PrefillStats

[orozery]

MO, the PR basically does that (set the stats directly) ?

Right, my bad.

[markmc]

@orozery from #36859 (comment)

The failure handling code use of this field is hacky IMO and I can think of alternative more robust ways to go without it. (I actually implemented a fix in one of the revisions of #35223: f02a5c8).

I've incorporated this approach from that earlier revision now, PTAL

[orozery]

I've incorporated this approach from that earlier revision now, PTAL

I think it's better to create a separate PR for these changes in the failure recovery flow.
They stand for themselves, and don't relay on the other changes in this PR.

[markmc]

Moved to draft since #38096 needs to merge first now

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @markmc.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[markmc]

FTR, this is a significant change in the most recent update:

We now track stats from the first time the prefill is scheduled (previously: first time prefill is completed) since the proven-error-prone effort to update the stats in the KV transfer error-handling path seems unjustified

[orozery]

LGTM

[ZhanqiuHu]

Note: This PR changed the prefill stats reported in the case of multi kv connector (NIXL connector + CPUOffloading connector). It impacts a couple of test cases in tests/v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh, which uses the metrics reported to check if the transferred kv cache and offloading kv cache are expected.

Will post a PR that update the test cases accordingly.

[NickLucche]

yep it broke CI for v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh