llmgw/internal/kvevents/decode.go

package kvevents

import (
	"bytes"
	"errors"
	"fmt"
	"strings"

	"github.com/vmihailenco/msgpack/v5"
)

// Kind is what an event says happened to the cache.
type Kind int

// The event kinds vLLM publishes.
const (
	Stored Kind = iota + 1
	Removed
	Cleared
)

// Event is one cache event, reduced to what the index needs.
type Event struct {
	Kind Kind
	// Hashes are the blocks stored or removed, in external form.
	Hashes []uint64
	// Parent is the hash of the block before the first stored one; HasParent
	// is false for the first block of a prompt.
	Parent    uint64
	HasParent bool
	// Tokens are the stored blocks' token ids, and BlockSize how many make a
	// block.
	Tokens    []int
	BlockSize int
	// Extra says the blocks were hashed with more than their tokens: a LoRA
	// adapter, an image, a cache salt. Their hashes are indexed like any
	// other, but say nothing about how plain tokens hash.
	Extra bool
	// Elsewhere says the event is about a copy outside this engine's GPU
	// memory — offloaded to CPU or storage, or another engine's — or about a
	// second KV cache group. Such events are ignored; see Mirror.Apply.
	Elsewhere bool
}

// Batch is one published message: the events of one engine step.
type Batch struct {
	TS     float64
	Events []Event
}

// DecodeBatch reads a msgpack-encoded KVEventBatch.
//
// Two encodings are accepted, because vLLM has used both. A batch is always
// an array, [ts, events, data_parallel_rank]. An event was an array too —
// ["BlockStored", block_hashes, parent_block_hash, token_ids, block_size,
// lora_id, medium, ...] — in the releases that introduced the events, and is
// now a map tagged by its "type" key, which lets fields be added without
// breaking anyone. A hash is an integer, or 32 bytes with
// VLLM_KV_EVENTS_USE_INT_BLOCK_HASHES off.
//
// An event of a kind not listed here is skipped rather than refused: a newer
// engine adding one must not stop the index.
func DecodeBatch(payload []byte) (Batch, error) {
	d := msgpack.NewDecoder(bytes.NewReader(payload))
	d.UseLooseInterfaceDecoding(true)
	v, err := d.DecodeInterface()
	if err != nil {
		return Batch{}, fmt.Errorf("kv events: not msgpack: %w", err)
	}
	arr, ok := v.([]any)
	if !ok || len(arr) < 2 {
		return Batch{}, errors.New("kv events: a batch is an array of at least ts and events")
	}
	var b Batch
	b.TS, _ = toFloat(arr[0])
	evs, ok := arr[1].([]any)
	if !ok && arr[1] != nil {
		return Batch{}, errors.New("kv events: the events of a batch are not an array")
	}
	for i, raw := range evs {
		e, ok, err := decodeEvent(raw)
		if err != nil {
			return Batch{}, fmt.Errorf("kv events: event %d: %w", i, err)
		}
		if ok {
			b.Events = append(b.Events, e)
		}
	}
	return b, nil
}

// decodeEvent turns either encoding into a map by field name, then reads it.
func decodeEvent(raw any) (Event, bool, error) {
	fields := map[string]any{}
	var tag string
	switch v := raw.(type) {
	case map[string]any:
		fields = v
		tag, _ = v["type"].(string)
	case []any:
		if len(v) == 0 {
			return Event{}, false, errors.New("empty event")
		}
		tag, _ = v[0].(string)
		var names []string
		switch tag {
		case "BlockStored":
			names = []string{"block_hashes", "parent_block_hash", "token_ids", "block_size",
				"lora_id", "medium", "lora_name", "extra_keys", "group_idx"}
		case "BlockRemoved":
			names = []string{"block_hashes", "medium", "group_idx"}
		}
		for i, name := range names {
			if i+1 < len(v) {
				fields[name] = v[i+1]
			}
		}
	default:
		return Event{}, false, fmt.Errorf("an event is a map or an array, not %T", raw)
	}

	var e Event
	switch tag {
	case "BlockStored":
		e.Kind = Stored
	case "BlockRemoved":
		e.Kind = Removed
	case "AllBlocksCleared":
		e.Kind = Cleared
		return e, true, nil
	default:
		return Event{}, false, nil
	}

	hs, _ := fields["block_hashes"].([]any)
	for _, h := range hs {
		x, err := hashOf(h)
		if err != nil {
			return Event{}, false, err
		}
		e.Hashes = append(e.Hashes, x)
	}
	if medium, ok := fields["medium"].(string); ok && !strings.EqualFold(medium, "GPU") {
		e.Elsewhere = true
	}
	if loc, ok := fields["locality"].(string); ok && !strings.EqualFold(loc, "LOCAL") {
		e.Elsewhere = true
	}
	if g, ok := toInt(fields["group_idx"]); ok && g != 0 {
		e.Elsewhere = true
	}
	if e.Kind == Removed {
		return e, true, nil
	}

	if p := fields["parent_block_hash"]; p != nil {
		x, err := hashOf(p)
		if err != nil {
			return Event{}, false, err
		}
		e.Parent, e.HasParent = x, true
	}
	ts, _ := fields["token_ids"].([]any)
	e.Tokens = make([]int, 0, len(ts))
	for _, t := range ts {
		n, ok := toInt(t)
		if !ok {
			return Event{}, false, fmt.Errorf("token id %v is not an integer", t)
		}
		e.Tokens = append(e.Tokens, n)
	}
	e.BlockSize, _ = toInt(fields["block_size"])
	if fields["lora_id"] != nil || fields["lora_name"] != nil {
		e.Extra = true
	}
	if ks, ok := fields["extra_keys"].([]any); ok {
		for _, k := range ks {
			if k != nil {
				e.Extra = true
			}
		}
	}
	return e, true, nil
}

func hashOf(v any) (uint64, error) {
	switch h := v.(type) {
	case []byte:
		return External(h), nil
	case string:
		// msgpack has no separate bytes type in its oldest form, and a
		// publisher using that writes bytes as a string.
		return External([]byte(h)), nil
	case int64:
		return uint64(h), nil
	case uint64:
		return h, nil
	}
	return 0, fmt.Errorf("a block hash is an integer or bytes, not %T", v)
}

func toInt(v any) (int, bool) {
	switch n := v.(type) {
	case int64:
		return int(n), true
	case uint64:
		return int(n), true
	}
	return 0, false
}

func toFloat(v any) (float64, bool) {
	switch n := v.(type) {
	case float64:
		return n, true
	case float32:
		return float64(n), true
	case int64:
		return float64(n), true
	case uint64:
		return float64(n), true
	}
	return 0, false
}