Skip to content

[None][perf] optimize native V2 KV event production - #16876

Draft
alec-flowers wants to merge 2 commits into
NVIDIA:mainfrom
alec-flowers:agent/native-kv-events-v2-fastpath
Draft

[None][perf] optimize native V2 KV event production#16876
alec-flowers wants to merge 2 commits into
NVIDIA:mainfrom
alec-flowers:agent/native-kv-events-v2-fastpath

Conversation

@alec-flowers

Copy link
Copy Markdown
Collaborator

Summary

Depends on #16869.

This stacked draft optimizes the native TensorRT-LLM V2 KV-event production path introduced by #16869. Until the base PR merges, GitHub will show both commits; this branch will be rebased onto main afterward.

  • replace the generic V2 event-manager path with a scheduler-local manager that accumulates only publishable full-block store/remove events from the maximum attention-window lifecycle
  • reuse each radix-tree block's existing chained SHA-256 Block.key, converting its low 64 bits to vLLM's signed integer wire representation
  • preserve the legacy pull-event path and its configured hash behavior
  • continue to perform no KV-event allgather in native publishing mode

Root cause

The generic V2 manager constructs events for cache creation, storage-tier updates, sliding-window lifecycles, partial blocks, and other mutations that the vLLM wire adapter later discards. In the distributed AgentX run, 3.249 million native events produced only 54,491 wire events, so 98.3% of the generic events were filtered after construction.

For the blocks that survived filtering, the native path also translated the existing V2 radix key back into the legacy V1 token hash on the scheduler thread. vLLM instead reuses the SHA-256 hash already owned by its prefix cache.

The new manager filters at the cache mutation hook and directly accumulates the events the router consumes. It reuses Block.key for both stores and removals, avoiding token re-hashing.

Validation

The final implementation was validated against the exact TensorRT-LLM 1.3.0rc21 benchmark image.

AgentX configuration Job Req/s Total tok/s/GPU vs. baseline Cached-read ratio
Reproduced baseline 2484714 16.95 79,289.09 baseline 95.6464%
Legacy KV-event publishing 2488135 14.88 69,452.68 -12.4058% 95.6573%
Native V2 fast path + ZMQ 2496292 16.9709 78,853.00 -0.5500% 95.6218%

Job 2496292 completed the full 600-second warmup and 3,600-second AgentX profile on 8 nodes / 32 GPUs with 61,774 successful requests and zero errors. It improved throughput 13.5349% over legacy event publishing, with no event-gap or event-drop markers and zero KV-event allgathers.

Local exact-image long points:

Mode Total tok/s Delta vs. adjacent events-off
Events off 163,767.29 control
Native fast path, null publisher 163,884.72 +0.07%
Native fast path, ZMQ publisher 161,968.02 -1.10%
  • exact rc21 real-socket tests passed, covering signed radix-hash conversion, maximum-window filtering, partial-block suppression, vLLM msgpack framing, removals, and idempotent shutdown
  • end-to-end smoke passed from the TRT-LLM V2 manager through three-frame ZMQ into Dynamo's Rust event publisher, local indexer, and KV router
  • the current-main forward-port passes Python compilation and git diff --check; focused CI remains pending on this draft

Regression protected by the focused test: publishing could revert to legacy token re-hashing or emit incorrect lifecycle/removal identities, causing scheduler throughput loss or stale KV-router state; the test checks the observable real-ZMQ wire boundary.

Signed-off-by: Alec Flowers <aflowers@nvidia.com>
Signed-off-by: Alec Flowers <aflowers@nvidia.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant