[V1][Metrics] add support for kv event publishing

[alec-flowers]

RFC: KVBlocks and Metrics Publishing In Inference Frameworks

• Added KVCacheEvent, BlockStored, BlockRemoved, and AllBlocksCleared msgspec classes
• Created a queue in the BlockPool and write these events in the appropriate functions
• Bubble the events up to the scheduler where they are appended to EngineCoreOutputs
• Add kv_cach_events to EngineCoreOutputs
• Wrote unit tests at the BlockManager level to test basic functionality and at the EngineCore level testing correct propagation and serializing over zmq.

API

• add enable_kv_cache_events to engineArgs
  - add external_stat_loggers field to AsyncLLM API Covered by [V1][Metrics] Allow V1 AsyncLLM to use custom logger #14661

With #14661 and this PR a 3rd party can write a custom stat logger to consume both engine Stats and Events and publish them elsewhere.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[mergify]

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

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

[njhill]

@alec-flowers re the custom stats loggers, there's another PR for this which is hopefully close to being merged: #14661. As you alluded to, we need to take factories/constructors for the multi-engine case (or otherwise change the stats logger API to include an an engine id).

[alec-flowers]

@alec-flowers re the custom stats loggers, there's another PR for this which is hopefully close to being merged: #14661. As you alluded to, we need to take factories/constructors for the multi-engine case (or otherwise change the stats logger API to include an an engine id).

@njhill Awesome! Looks like there has been some good discussion around the API. I will remove my WIP stuff from this PR and reference that PR.

[robertgshaw2-redhat]

cc @markmc - FYI, this feature it to enable creating a global KV cache view in NVIDIA Dynamo

[robertgshaw2-redhat]

This generally looks very reasonable to me.

It would be useful to add an example of the Logger which uses the the events, for testing purposes. Otherwise this will break.

Is there a reason why you don't want to implement the KVCacheEventsLogger inside VLLM?

[alec-flowers]

This generally looks very reasonable to me.

It would be useful to add an example of the Logger which uses the the events, for testing purposes. Otherwise this will break.

Is there a reason why you don't want to implement the KVCacheEventsLogger inside VLLM?

Main reason is to avoid adding dynamo (or pieces of dynamo) as a dependency to vLLM. For our logger, we would need to add in a rust based publisher (with python bindings) from Dynamo that sends the events over NATS to another component.
https://github.com/ai-dynamo/dynamo/blob/main/lib/llm/src/kv_router/publisher.rs
https://github.com/ai-dynamo/dynamo/blob/main/lib/bindings/python/src/dynamo/_core.pyi#L497

I'm definitely flexible and open to ideas here.

And good reminder, I will add in an integration test with a mock StatLogger that will break if the API changes too much.

[ApostaC]

Quick question: how is KVCacheEvent published? I saw it's passed into the metrics logger, but seems like it's not used?

[ApostaC]

Should those events be defined in a source file that is more related to KV cache management? For example, kv_cache_utils.py

[markmc]

Main reason is to avoid adding dynamo (or pieces of dynamo) as a dependency to vLLM. For our logger, we would need to add in a rust based publisher (with python bindings) from Dynamo that sends the events over NATS to another component. https://github.com/ai-dynamo/dynamo/blob/main/lib/llm/src/kv_router/publisher.rs https://github.com/ai-dynamo/dynamo/blob/main/lib/bindings/python/src/dynamo/_core.pyi#L497

I'm definitely flexible and open to ideas here.

re my other comment #16669 (comment)

If this is a hook to capture stored/removed KV cache events so they can be sent elsewhere, wouldn't the KV connector API be a better fit - no need to route these events to a logger in the frontend for this?

[alec-flowers]

Quick question: how is KVCacheEvent published? I saw it's passed into the metrics logger, but seems like it's not used?

@ApostaC It would be published by a 3rd party defined Logger. You would be free to write your own that would publish it how you want. That is why in this case we would be reliant on this issue - #14661, which adds support for passing in a logger at runtime to the AysncLLM class.

[njhill]

@alec-flowers it looks like there may be an issue with the test_kv_cache_events test timing out: https://buildkite.com/vllm/ci/builds/18784#01967c2d-91f5-42bb-b3ae-ce520c3ddb11/210-1613. It might help to set daemon=True when starting the publishing thread.

[alec-flowers]

@alec-flowers it looks like there may be an issue with the test_kv_cache_events test timing out: https://buildkite.com/vllm/ci/builds/18784#01967c2d-91f5-42bb-b3ae-ce520c3ddb11/210-1613. It might help to set daemon=True when starting the publishing thread.

I'm not able to replicate locally

When I start the publishing thread the daemon=True is set

 self._thread = threading.Thread(target=self._publisher_thread,
                                        daemon=True,
                                        name="zmq-publisher")

will try a few things out.

[mergify]

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

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

[mergify]

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

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

[alec-flowers]

Took me some time, but isolated the problem to be the previous test_engine_core_client_asyncio which was leaving a thread open with a zmq socket that was causing problems with my test. I believe I've fixed it.

@njhill I'm having some problems with CI. It seems the last run failed in the build stage. I'm also constantly having to re-base due to a large import block. Hopefully this run goes smooth.

[njhill]

Thanks @alec-flowers, sorry you had to take extra time to debug that. I think it's a lingering issue where the client destructor doesn't always get run. In theory the explicit shutdown shouldn't be needed but makes sense to include it here for now.