fix(kvcache): decode lora_id/medium/lora_name/group_idx in BlockStored (#2285)

Author: pjdurden

pkg/cache/kvcache/event_types.go

@@ -28,8 +28,11 @@ import (
 //   None,               # parent_block_hash
 //   [55, 11, 22],       # token_ids
 //   32,                 # block_size
-//   None,               # lora_id
-//   "cuda:0"            # medium
+//   None,               # lora_id (deprecated)
+//   "GPU",              # medium
+//   "adapter-a",        # lora_name
+//   None,               # extra_keys (may be omitted)
+//   0                   # group_idx (may be omitted)
 // ]
 //
 // [
@@ -74,15 +77,18 @@ type KVEvent interface {
 
 // BlockStoredEvent represents blocks being stored in KV cache
 // ------------------------------------------------------------
-// BlockStored (msgspec encoding order)
-// Python fields:
+// BlockStored (msgspec encoding order, vllm/distributed/kv_events.py)
+// Python fields, positions after the tag:
 //
-//	[tag, block_hashes, parent_block_hash, token_ids, block_size, lora_id, medium]
+//	[tag, block_hashes, parent_block_hash, token_ids, block_size,
+//	 lora_id, medium, lora_name, extra_keys, group_idx, ...]
 //
+// The struct is msgspec array_like with omit_defaults, so trailing fields that
+// equal their default (extra_keys, group_idx and newer spec fields) may be
+// omitted from the wire array. The decoder therefore reads positions 5+ by
+// index with bounds checks and tolerates shorter arrays from older vLLM builds.
 // ------------------------------------------------------------
 //
-// lora_id and medium are unused for now.
-//
 // Note: BlockHashes are converted at decode time:
 // - vLLM legacy format (int64) → stored as-is
 // - vLLM new format (32-byte SHA-256 from PR #23673) → first 8 bytes converted to int64
@@ -94,6 +100,15 @@ type BlockStoredEvent struct {
 	ParentBlockHash *int64    // Decoded from vLLM, supports both old and new formats
 	TokenIDs        [][]byte
 
+	// Optional fields carried by newer vLLM builds (positions 5-9). Decoded
+	// best-effort; nil when the connected vLLM does not emit them. These are
+	// exposed here so prefix-cache routing can key and score on them; see
+	// issue #2285. Consumption is handled separately.
+	LoraID   *int64  // position 5; deprecated upstream in favor of LoraName
+	Medium   *string // position 6; storage tier, e.g. "GPU"/"cpu" (nil = unspecified)
+	LoraName *string // position 7; canonical adapter id
+	GroupIdx *int64  // position 9; KV-cache group for hybrid-attention models
+
 	// NOTE: These are NOT part of msgpack
 	Timestamp time.Time `msgpack:"-"`
 	ModelName string    `msgpack:"-"`

pkg/cache/kvcache/msgpack_decoder.go

@@ -145,12 +145,39 @@ func parseEventArray(arr []interface{}) (KVEvent, error) {
 			return nil, err
 		}
 
-		return &BlockStoredEvent{
+		ev := &BlockStoredEvent{
 			Type:            EventTypeBlockStored,
 			BlockHashes:     blockHashes,
 			ParentBlockHash: parentHash,
 			TokenIDs:        tokens,
-		}, nil
+		}
+
+		// Optional fields added by newer vLLM builds. msgspec omit_defaults may
+		// drop trailing ones, so read by position with bounds checks and leave
+		// the rest nil. Position 8 (extra_keys) is skipped here; group_idx is
+		// still indexed at its fixed position 9.
+		if len(arr) > 5 {
+			if ev.LoraID, err = toInt64Ptr(arr[5]); err != nil {
+				return nil, fmt.Errorf("invalid lora_id: %w", err)
+			}
+		}
+		if len(arr) > 6 {
+			if ev.Medium, err = toStringPtr(arr[6]); err != nil {
+				return nil, fmt.Errorf("invalid medium: %w", err)
+			}
+		}
+		if len(arr) > 7 {
+			if ev.LoraName, err = toStringPtr(arr[7]); err != nil {
+				return nil, fmt.Errorf("invalid lora_name: %w", err)
+			}
+		}
+		if len(arr) > 9 {
+			if ev.GroupIdx, err = toInt64Ptr(arr[9]); err != nil {
+				return nil, fmt.Errorf("invalid group_idx: %w", err)
+			}
+		}
+
+		return ev, nil
 
 	case EventTypeBlockRemoved:
 		if len(arr) < 2 {
@@ -348,6 +375,24 @@ func toInt64Ptr(v any) (*int64, error) {
 	return &val, nil
 }
 
+// toStringPtr converts a nullable msgpack string field to *string.
+// A nil input (Python None) yields a nil pointer. msgpack may decode a string
+// as either string or []byte, so both are accepted.
+func toStringPtr(v any) (*string, error) {
+	if v == nil {
+		return nil, nil
+	}
+	switch s := v.(type) {
+	case string:
+		return &s, nil
+	case []byte:
+		str := string(s)
+		return &str, nil
+	default:
+		return nil, fmt.Errorf("expected string, got %T", v)
+	}
+}
+
 func parseUint32(v any) (uint32, error) {
 	switch x := v.(type) {
 

pkg/cache/kvcache/msgpack_decoder_test.go

@@ -372,3 +372,129 @@ func TestBlockRemovedEventWithBytesHashes(t *testing.T) {
 	assert.Equal(t, "removed-model", removed.ModelName)
 	assert.Equal(t, "removed-pod", removed.PodName)
 }
+
+// --- Issue #2285: decoder must read lora_id/medium/lora_name/group_idx ---
+
+// rawBlockStoredBatch marshals a single BlockStored event using the exact
+// positional layout vLLM emits (msgspec array_like, tag first), so these tests
+// assert against the real wire contract rather than Go-encoder symmetry.
+//
+// vLLM BlockStored positions (vllm/distributed/kv_events.py):
+//
+//	[0]tag [1]block_hashes [2]parent_block_hash [3]token_ids [4]block_size
+//	[5]lora_id [6]medium [7]lora_name [8]extra_keys [9]group_idx ...
+func rawBlockStoredBatch(t *testing.T, eventFields []interface{}) []byte {
+	t.Helper()
+	event := append([]interface{}{string(EventTypeBlockStored)}, eventFields...)
+	batch := []interface{}{
+		float64(1_700_000_000), // batch timestamp (seconds)
+		[]interface{}{event},
+	}
+	data, err := msgpack.Marshal(batch)
+	require.NoError(t, err)
+	return data
+}
+
+// blockStoredCore returns the four always-present fields after the tag:
+// block_hashes, parent_block_hash, token_ids, block_size.
+func blockStoredCore() []interface{} {
+	return []interface{}{
+		[]interface{}{int64(101), int64(102)}, // block_hashes
+		nil,                                   // parent_block_hash
+		[]interface{}{int64(1), int64(2), int64(3), int64(4)}, // token_ids
+		int64(4), // block_size
+	}
+}
+
+func decodeSingleBlockStored(t *testing.T, data []byte) *BlockStoredEvent {
+	t.Helper()
+	batch, err := DecodeEventBatch(data, "m", "p")
+	require.NoError(t, err)
+	require.Len(t, batch.Events, 1)
+	ev, ok := batch.Events[0].(*BlockStoredEvent)
+	require.True(t, ok, "expected *BlockStoredEvent")
+	return ev
+}
+
+func TestBlockStored_DecodesAllOptionalFields(t *testing.T) {
+	fields := append(blockStoredCore(),
+		int64(7),    // lora_id
+		"GPU",       // medium
+		"adapter-a", // lora_name
+		nil,         // extra_keys
+		int64(3),    // group_idx
+	)
+	ev := decodeSingleBlockStored(t, rawBlockStoredBatch(t, fields))
+
+	require.NotNil(t, ev.LoraID)
+	assert.Equal(t, int64(7), *ev.LoraID)
+	require.NotNil(t, ev.Medium)
+	assert.Equal(t, "GPU", *ev.Medium)
+	require.NotNil(t, ev.LoraName)
+	assert.Equal(t, "adapter-a", *ev.LoraName)
+	require.NotNil(t, ev.GroupIdx)
+	assert.Equal(t, int64(3), *ev.GroupIdx)
+}
+
+// group_idx sits at position 9, after extra_keys at 8. Verify it is read by
+// position even when extra_keys carries a non-nil payload.
+func TestBlockStored_GroupIdxReadPastExtraKeys(t *testing.T) {
+	fields := append(blockStoredCore(),
+		nil,                               // lora_id
+		"cpu",                             // medium
+		nil,                               // lora_name
+		[]interface{}{[]interface{}{"k"}}, // extra_keys (non-nil)
+		int64(5),                          // group_idx
+	)
+	ev := decodeSingleBlockStored(t, rawBlockStoredBatch(t, fields))
+
+	require.NotNil(t, ev.Medium)
+	assert.Equal(t, "cpu", *ev.Medium)
+	require.NotNil(t, ev.GroupIdx)
+	assert.Equal(t, int64(5), *ev.GroupIdx)
+}
+
+// None values (nil medium / lora_name) must decode to nil pointers, not "".
+func TestBlockStored_NilOptionalValues(t *testing.T) {
+	fields := append(blockStoredCore(),
+		nil, // lora_id
+		nil, // medium
+		nil, // lora_name
+		nil, // extra_keys
+		nil, // group_idx
+	)
+	ev := decodeSingleBlockStored(t, rawBlockStoredBatch(t, fields))
+
+	assert.Nil(t, ev.LoraID)
+	assert.Nil(t, ev.Medium)
+	assert.Nil(t, ev.LoraName)
+	assert.Nil(t, ev.GroupIdx)
+}
+
+// Older vLLM builds send only the first five positions. Decoding must still
+// succeed and leave the new fields nil (no panic, no error).
+func TestBlockStored_BackwardCompatShortArray(t *testing.T) {
+	ev := decodeSingleBlockStored(t, rawBlockStoredBatch(t, blockStoredCore()))
+
+	assert.Equal(t, []int64{101, 102}, ev.BlockHashes)
+	assert.Nil(t, ev.LoraID)
+	assert.Nil(t, ev.Medium)
+	assert.Nil(t, ev.LoraName)
+	assert.Nil(t, ev.GroupIdx)
+}
+
+// Intermediate length: through lora_name (position 7), no extra_keys/group_idx.
+func TestBlockStored_PartialOptionalFields(t *testing.T) {
+	fields := append(blockStoredCore(),
+		int64(2),    // lora_id
+		"GPU",       // medium
+		"adapter-b", // lora_name
+	)
+	ev := decodeSingleBlockStored(t, rawBlockStoredBatch(t, fields))
+
+	require.NotNil(t, ev.LoraID)
+	assert.Equal(t, int64(2), *ev.LoraID)
+	require.NotNil(t, ev.LoraName)
+	assert.Equal(t, "adapter-b", *ev.LoraName)
+	assert.Nil(t, ev.GroupIdx)
+}