llmgw/internal/kvevents/hash.go

// Package kvevents reads the KV cache events a vLLM engine publishes, and
// computes the block hashes those events are keyed by (plan §7.5 tier 2).
//
// Both halves live together because they must agree byte for byte. The agent
// reads the events and uses the hashing to check that the engine hashes the
// way llmgw expects; the router uses the same hashing to ask which replicas
// hold the leading blocks of a prompt. A difference between the two is not an
// error anything would report — it is a cache index that never matches.
package kvevents

import (
	"crypto/sha256"
	"encoding/binary"
)

// The hashing scheme llmgw starts vLLM with, and assumes.
//
// sha256_cbor because it is the one of vLLM's algorithms defined outside
// Python: canonical CBOR of the tuple (parent hash, token ids, extra keys),
// then SHA-256. The default, "sha256", hashes a pickle, which only Python can
// reproduce; the builtin hash vLLM once offered changes from process to
// process.
//
// HashSeed is what PYTHONHASHSEED is set to. vLLM derives the hash every
// chain starts from — its NONE_HASH — from that variable, and without it
// either from a random value (older releases) or a fixed string (newer ones).
// Setting it pins the start of the chain in every release. It has to be an
// integer for Python to start at all; which one does not matter, only that
// the engine and the router agree, and it is a constant so they cannot
// disagree.
const (
	HashAlgo = "sha256_cbor"
	HashSeed = "1729"
)

// BlockHashes returns the external hashes of the full blocks of tokens, as
// vLLM would publish them for a request carrying exactly these token ids and
// nothing that adds extra keys (no LoRA adapter, no images, no cache salt).
//
// A partial last block has no hash: vLLM caches only full blocks. At most max
// blocks are hashed, and max <= 0 means all of them.
func BlockHashes(tokens []int, blockSize, max int) []uint64 {
	if blockSize <= 0 {
		return nil
	}
	n := len(tokens) / blockSize
	if max > 0 && n > max {
		n = max
	}
	out := make([]uint64, 0, n)
	parent := noneHash()
	buf := make([]byte, 0, 64+blockSize*5)
	for i := range n {
		buf = blockInput(buf[:0], parent[:], tokens[i*blockSize:(i+1)*blockSize])
		parent = sha256.Sum256(buf)
		out = append(out, External(parent[:]))
	}
	return out
}

// FirstBlockHash is the hash of a chain's first block, which needs nothing
// but its tokens. It is how the agent checks an engine hashes the way llmgw
// expects: the first block of every prompt is published with its tokens and
// no parent.
func FirstBlockHash(tokens []int) uint64 {
	parent := noneHash()
	sum := sha256.Sum256(blockInput(nil, parent[:], tokens))
	return External(sum[:])
}

// External is the form vLLM publishes a hash in: the low 64 bits of the
// digest read big-endian, which is what VLLM_KV_EVENTS_USE_INT_BLOCK_HASHES
// (on by default, and set by llmgw) makes of the 32 bytes. A publisher with it
// off sends the bytes, and they are reduced the same way, so both forms index
// alike.
func External(digest []byte) uint64 {
	if len(digest) < 8 {
		var pad [8]byte
		copy(pad[8-len(digest):], digest)
		return binary.BigEndian.Uint64(pad[:])
	}
	return binary.BigEndian.Uint64(digest[len(digest)-8:])
}

// noneHash is vLLM's NONE_HASH under PYTHONHASHSEED=HashSeed:
// sha256_cbor of the seed, which Python reads as a string.
func noneHash() [32]byte {
	return sha256.Sum256(cborText(nil, HashSeed))
}

// blockInput is canonical CBOR of the tuple vLLM hashes for a block:
// (parent_block_hash, tuple(token_ids), extra_keys) with no extra keys, as
// cbor2.dumps(..., canonical=True) writes it. A tuple is an array, bytes a
// byte string, a token id an unsigned integer, None the simple value null.
func blockInput(buf, parent []byte, tokens []int) []byte {
	buf = cborHead(buf, 4, 3) // array of three
	buf = cborHead(buf, 2, uint64(len(parent)))
	buf = append(buf, parent...)
	buf = cborHead(buf, 4, uint64(len(tokens)))
	for _, t := range tokens {
		if t < 0 {
			// Token ids are never negative; a negative one would be a major
			// type 1 integer, and is encoded as such rather than wrapped.
			buf = cborHead(buf, 1, uint64(-1-t))
			continue
		}
		buf = cborHead(buf, 0, uint64(t))
	}
	return append(buf, 0xf6) // None
}

func cborText(buf []byte, s string) []byte {
	buf = cborHead(buf, 3, uint64(len(s)))
	return append(buf, s...)
}

// cborHead writes a CBOR item head in its shortest form, which is what
// canonical encoding requires.
func cborHead(buf []byte, major byte, n uint64) []byte {
	m := major << 5
	switch {
	case n < 24:
		return append(buf, m|byte(n))
	case n <= 0xff:
		return append(buf, m|24, byte(n))
	case n <= 0xffff:
		return binary.BigEndian.AppendUint16(append(buf, m|25), uint16(n))
	case n <= 0xffffffff:
		return binary.BigEndian.AppendUint32(append(buf, m|26), uint32(n))
	default:
		return binary.BigEndian.AppendUint64(append(buf, m|27), n)
	}
}