[Bugfix] StatLoggers: cache spec decode metrics when they get collected.

[tdoublep]

Fixes #6643

From issue:

After merging #6578, the AsyncMetricsCollector now respects the time interval between collecting speculative decoding metrics. It now collects the metrics every 5 seconds, instead of continuously collecting them.

While this is the desired behaviour, I now see that in practice, we are never logging the speculative decoding stats. This is because, in order to get logged, the speculative decoding stats need to happen to be collected in exactly the same step that we decide to log the output (which is controlled by a different time interval). In practice, this does not seem to happen using the default settings.

This PR caches the spec decode metrics inside the loggers, so that they will generally always get logged, and we don't need to rely to the stats getting collected at precisely the right moment.

I guess an alternative solution could be to enforce that the signal to collect the spec decoding metrics coincides with the signal to log the output somehow (it would actually need to be triggered a bit before we log output due to async implementation) but that sounds unnecessarily complex to me.

cc @cadedaniel @njhill

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.
Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which consists a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of default ones by unblocking the steps in your fast-check build on Buildkite UI.

Once the PR is approved and ready to go, please make sure to run full CI as it is required to merge (or just use auto-merge).

To run full CI, you can do one of these:

• Comment /ready on the PR
• Add ready label to the PR
• Enable auto-merge.

🚀

[comaniac]

The root cause of this issue is we have two logging intervals and they are unlikely to align. I'm wondering whether we could keep just one logging interval to make the process simple. Specifically, record raw spec decode metrics (e.g. number of accepted tokens, etc) every step, and only process them (e.g. compute system efficiency) every logging interval, similar to number of generated tokens and decoding throughput.

[tdoublep]

Specifically, record raw spec decode metrics (e.g. number of accepted tokens, etc) every step, and only process them (e.g. compute system efficiency) every logging interval,

@comaniac I think that the AsyncMetricsCollector was designed specifically so we don't collect the number of accepted tokens at every step, presumably because it incurs some overhead related to device-host copy and/or synchronization. This is just my understanding from reading the code, pls correct me if I'm wrong here.

I agree in principle that a simpler approach might make sense but unsure of the performance implications. I was also wondering if there is at least a subset of spec decode metrics that we could collect at each step directly from CPU (e.g., based on the info we are already copying from device to host).

[cadedaniel]

Thanks for catching this, seems somehow this regressed without me noticing.

Some thoughts:

• For performance, spec decode doesn't want to perform metrics operations every forward pass. So in terms of "single metrics interval for all vLLM", spec decode has specialized concerns due to the latency sensitivity. But if we want to unify the intervals in the Engine level and move spec metrics state into the spec worker, fine by me !
• Can we add a test that metrics are processed by the engine? FYI we have a test that metrics are emitted when interval == 0 (thanks to @comaniac ), but no test when interval > 0.
• FYI (to share context), I wrote this async metrics code in a rush when deploying spec decode to production. I didn't want to add any latency overheads which is why I took the extremely conservative design of async metrics collection. So if it adds a lot of complexity we can go with a simpler design where we collect each time. However I feel we should make a reasonable attempt to keep it async as the spec decode framework pytorch op overheads add up..

screenshot below showing my original latency measurement, it took 130µs to do every time

[comaniac]

Per offline discussion we could let this PR in first and revisit in the future if there are more features that also introduce a separate intervals. Meanwhile, a unit test for this would be great.

[tdoublep]

Will write a test later today

[tdoublep]

/ready

[tdoublep]

Added a test to catch this, and also verified that it fails on main currently.

[tdoublep]

Failing test will be resolved by #6690