fix(kvcache): decode vLLM map-format KV events; tier-aware prefix scoring - #2587
fix(kvcache): decode vLLM map-format KV events; tier-aware prefix scoring#2587wuyan-zs wants to merge 1 commit into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request adds compatibility with newer vLLM wire formats (including msgspec map encoding) and introduces tier-aware scoring (GPU, CPU, STORAGE) for prefix cache routing. Feedback on the changes highlights a potential decoding failure in decodeBatchArray due to a strict float64 type assertion on the timestamp, as well as issues where a pod's storage medium is not updated when the pod already exists in the prefix store, leading to stale tier information.
| tsFloat, ok := arr[0].(float64) | ||
| if !ok { | ||
| return nil, fmt.Errorf("invalid batch timestamp type: %T", rawBatch[0]) | ||
| return nil, fmt.Errorf("invalid batch timestamp type: %T", arr[0]) | ||
| } |
There was a problem hiding this comment.
In decodeEventArray, the switch statement accepts various numeric types for the timestamp (e.g., int, int64, float32, etc.). However, inside decodeBatchArray, there is a strict type assertion to float64 (tsFloat, ok := arr[0].(float64)). If the timestamp is decoded as any other numeric type, this assertion will fail, resulting in a decoding error. We should use a type switch to robustly convert any of the accepted numeric types to float64.
var tsFloat float64
switch x := arr[0].(type) {
case float64:
tsFloat = x
case float32:
tsFloat = float64(x)
case int:
tsFloat = float64(x)
case int8:
tsFloat = float64(x)
case int16:
tsFloat = float64(x)
case int32:
tsFloat = float64(x)
case int64:
tsFloat = float64(x)
case uint:
tsFloat = float64(x)
case uint8:
tsFloat = float64(x)
case uint16:
tsFloat = float64(x)
case uint32:
tsFloat = float64(x)
case uint64:
tsFloat = float64(x)
default:
return nil, fmt.Errorf(\"invalid batch timestamp type: %T\", arr[0])
}| } else { | ||
| pods[podName] = &PodInfo{ | ||
| SourcePod: podName, | ||
| Medium: medium, |
There was a problem hiding this comment.
When the pod already exists in the pods map (the if block of this conditional, which is outside this diff), we only update its LastAccessTime but do not update its Medium. If a prefix is migrated to a different storage tier, the tier information will remain stale. Please ensure that podInfo.Medium = medium is also updated when the pod already exists.
| } else { | ||
| pods[podName] = &PodInfo{ | ||
| SourcePod: podName, | ||
| Medium: medium, |
There was a problem hiding this comment.
Similarly, when the pod already exists in the pods map (the if block of this conditional, which is outside this diff), we should also update podInfo.Medium = medium to prevent stale tier information when a prefix is migrated to a different storage tier.
39be9ed to
059fe5e
Compare
There was a problem hiding this comment.
Pull request overview
This PR restores compatibility with current vLLM KV-event msgpack wire formats (including the post-#42892 map encoding) and improves prefix-routing correctness by carrying tier/adaptor/group metadata end-to-end and weighting prefix match scores by storage tier.
Changes:
- Extend KV-event decoding to accept map-encoded single events, legacy tag-first arrays, and
[ts, events]batch wrappers; decode additional vLLM fields (medium,lora_name,group_idx,extra_keys, etc.). - Add tier-aware prefix scoring by recording
mediumper (prefix, pod) and weightingMatchPrefixscores (GPU/CPU/STORAGE). - Plumb
medium/lora_name/group_idxthrough the kvevent handler and store-provider adapter into the sync prefix indexer; add focused compatibility and scoring tests.
Reviewed changes
Copilot reviewed 9 out of 9 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| pkg/utils/syncprefixcacheindexer/sync_hash.go | Record per-pod storage tier for prefixes and weight MatchPrefix scores by tier; add APIs to add prefixes with medium. |
| pkg/utils/syncprefixcacheindexer/medium_scoring_test.go | New tests validating tier weighting via both direct prefix insertion and event ingestion. |
| pkg/utils/syncprefixcacheindexer/events.go | Extend sync event structs to carry medium, lora_name, and group_idx. |
| pkg/kvevent/interfaces.go | Extend kvevent interfaces to include tier/adaptor/group metadata. |
| pkg/kvevent/handler.go | Populate and forward medium/lora_name/group_idx from decoded KV events into sync events. |
| pkg/cache/store_providers.go | Plumb new metadata through the sync indexer adapter. |
| pkg/cache/kvcache/msgpack_decoder.go | Decode vLLM map-format events and legacy single-event arrays; broaden decoding to all supported shapes and fields. |
| pkg/cache/kvcache/event_types.go | Extend KV event types with additional optional fields retained during decode. |
| pkg/cache/kvcache/decoder_vllm_compat_test.go | New byte-exact fixtures ensuring decoder compatibility with real vLLM payload shapes. |
Suppressed comments (1)
pkg/utils/syncprefixcacheindexer/sync_hash.go:583
- addPrefixToPodLocked has the same issue as AddPrefixWithMedium: when the pod entry already exists, it doesn't update PodInfo.Medium. That can leave tier info stale across repeated BlockStored events and skew MatchPrefix weighting.
if podInfo, exists := pods[podName]; exists {
podInfo.LastAccessTime.Store(now)
} else {
pods[podName] = &PodInfo{
SourcePod: podName,
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| func toExtraKeys(v any) ([][]interface{}, error) { | ||
| if v == nil { | ||
| return nil, nil | ||
| } | ||
| raw, ok := v.([]interface{}) | ||
| if !ok { | ||
| return nil, fmt.Errorf("expected []interface{}, got %T", v) | ||
| } | ||
| out := make([][]interface{}, len(raw)) | ||
| for i, x := range raw { | ||
| if x == nil { | ||
| out[i] = nil | ||
| continue | ||
| } | ||
| entry, ok := x.([]interface{}) | ||
| if !ok { | ||
| return nil, fmt.Errorf("extra_keys[%d]: expected []interface{}, got %T", i, x) | ||
| } | ||
| out[i] = entry | ||
| } | ||
| return out, nil | ||
| } |
| // 0: batch timestamp | ||
| tsFloat, ok := rawBatch[0].(float64) | ||
| tsFloat, ok := arr[0].(float64) | ||
| if !ok { | ||
| return nil, fmt.Errorf("invalid batch timestamp type: %T", rawBatch[0]) | ||
| return nil, fmt.Errorf("invalid batch timestamp type: %T", arr[0]) | ||
| } |
| if podInfo, exists := pods[podName]; exists { | ||
| podInfo.LastAccessTime.Store(now) | ||
| } else { | ||
| pods[podName] = &PodInfo{ | ||
| SourcePod: podName, |
|
Thanks for the PR. The tier-aware scoring and format compatibility improvements look solid, but there are a few edge cases involving state updates and logic duplication that need addressing:
|
|
1. Stale Currently, This breaks the feature this PR is adding. ZMQ is at-least-once, and vLLM CPU offload can emit another Please update the medium on the existing path, ensuring an empty value doesn't overwrite a known tier: if podInfo, exists := pods[podName]; exists {
podInfo.LastAccessTime.Store(now)
if medium != "" {
podInfo.Medium = medium
}
}Suggestion: Add a test that stores the same 2. "Current vLLM" fixtures do not match the real ZMQ shape ( The live The new tests cover:
They never cover Suggestion: Add a fixture encoded as a real |
…ring - DecodeEventBatch now accepts the msgspec map encoding vLLM switched to in vllm-project/vllm#42892 (single event per ZMQ message, flat map with a `type` key), the legacy tag-first array encoding, and the [ts, events] batch wrapper. Previously every map payload failed to unmarshal, so kv-events prefix routing was broken against current vLLM. - BlockStored/BlockRemoved now parse all fields: lora_id, medium, lora_name, extra_keys, group_idx, kv_cache_spec_kind, kv_cache_spec_sliding_window, locality (map and array forms). - syncprefixcacheindexer records the storage medium per (prefix, pod) and weights MatchPrefix scores (GPU=1.0, CPU=0.5, STORAGE=0.25; unspecified keeps full score) so pods holding a prefix only on a slower tier are not over-credited. - Plumb medium/lora_name/group_idx through the kvevent handler and the store provider adapter (previously dropped at the handler). Fixes vllm-project#2285 Signed-off-by: wuyan-zs <91841343+wuyan-zs@users.noreply.github.com>
|
Bumping this. Running into the same issue |
Summary
Fixes #2285. While investigating, I found a larger compatibility break:
vLLM switched KV event encoding from positional arrays to msgpack maps in
vllm-project/vllm#42892
(merged 2026-06-09).
DecodeEventBatchstill expected[ts, events]arrays,so every KV event payload from current vLLM fails to unmarshal
(
msgpack: invalid code=0x88 decoding array length) — kv-events prefixrouting is effectively dead against current vLLM. The legacy publisher also
sent one event per ZMQ message (tag-first array), which the batch-shaped
decoder never handled either.
Changes
pkg/cache/kvcache/msgpack_decoder.go: accept all three wire shapes —map single event (current vLLM), tag-first array single event (legacy
vLLM), and
[ts, events]batch (older code paths / test encoder).BlockStored/BlockRemovedfield:lora_id,medium,lora_name,extra_keys,group_idx,kv_cache_spec_kind,kv_cache_spec_sliding_window,locality.pkg/utils/syncprefixcacheindexer: record the storagemediumper(prefix, pod) and weight
MatchPrefixscores (GPU=1.0, CPU=0.5,STORAGE=0.25; unspecified keeps full score for backward compatibility)
so pods holding a prefix only on a slower tier are not over-credited.
medium/lora_name/group_idxthrough thekveventhandler andstore_providersadapter (previously dropped at the handler).Tests
New
decoder_vllm_compat_test.goverifies decoding against byte-exactpayloads generated with the real msgspec encoder for both current map and
legacy array formats (single-event and batch shapes). New
medium_scoring_test.gocovers tier weighting viaAddPrefixWithMediumandthe event ingestion path. All existing
kvcacheandsyncprefixcacheindexertests keep passing.Follow-up
The
(model, lora_name, group_idx)index keying part of #2285 needs arequest-side adapter id (
RoutingContexthas no LoRA field today); see theissue thread.