kvblock

package
v0.9.0 Latest Latest
Warning

This package is not in the latest version of its module.

 Go to latest Published: Jun 16, 2026 License: Apache-2.0 Imports: 27 Imported by: 3

Documentation

Index

Constants

This section is empty.

Variables

This section is empty.

Functions

This section is empty.

Types

type BlockExtraFeatures added in v0.7.0 

type BlockExtraFeatures struct {
	MMHashes []MMHash
}

BlockExtraFeatures holds per-block extra data that taints the block hash. A nil *BlockExtraFeatures means pure text (no taint).

func ComputeBlockExtraFeatures added in v0.7.0 

func ComputeBlockExtraFeatures(
	mmHashes map[string][]string,
	mmPlaceholders map[string][]PlaceholderRange,
	blockSize, numTokens int,
) []*BlockExtraFeatures

ComputeBlockExtraFeatures converts tokenizer-provided multimodal metadata into per-block extra features, matching vLLM's _gen_mm_extra_hash_keys() algorithm.

For each block, it finds overlapping placeholder ranges and emits the identifier (hash) of each overlapping multimodal item. This matches vLLM v0.18.0+ which appends mm_feature.identifier per overlapping item.

func ParseRawExtraKeys added in v0.7.0 

func ParseRawExtraKeys(raw [][]any) ([]*BlockExtraFeatures, error)

ParseRawExtraKeys converts the raw [][]any from BlockStoredEvent.ExtraKeys into typed []*BlockExtraFeatures. Each inner []any element is either:

  • a bare string identifier (vLLM v0.18.0+: mm_feature.identifier), or
  • a 2-element [string, int] tuple (legacy format, offset is ignored).

nil inner slices produce nil entries. Returns nil if raw is nil.

type BlockHash 

type BlockHash uint64

BlockHash struct represents a unique identifier for a KV-cache block. 

const EmptyBlockHash BlockHash = 0

EmptyBlockHash represents an invalid or uninitialized block hash. This serves as the "error value".

func (BlockHash) String 

func (c BlockHash) String() string

String returns a string representation of the Key.

type CostAwareMemoryIndex 

type CostAwareMemoryIndex struct {
	// contains filtered or unexported fields
}

CostAwareMemoryIndex implements the Index interface using Ristretto cache for cost-aware memory management. The two caches below are kept in sync:

  • data: requestKey -> pod cache (cost-bound by Ristretto MaxCost)
  • requestKeys: engineKey -> requestKey (LRU to cap mapping size)

Add always writes both maps; Evict removes pods and, when empty, removes both the requestKey entry and its engineKey mapping to avoid dangling keys.

func NewCostAwareMemoryIndex 

func NewCostAwareMemoryIndex(cfg *CostAwareMemoryIndexConfig) (*CostAwareMemoryIndex, error)

NewCostAwareMemoryIndex creates a new CostAwareMemoryIndex instance.

func (*CostAwareMemoryIndex) Add 

func (m *CostAwareMemoryIndex) Add(ctx context.Context, engineKeys, requestKeys []BlockHash, entries []PodEntry) error

Add adds a set of keys and their associated pod entries to the index backend. If engineKeys is nil, only requestKey -> PodEntry mappings are created (no engineKey -> requestKey mapping). This is used for speculative entries where engine keys are not yet known. When engineKeys is non-nil, the mapping type is inferred from the ratio of array lengths.

func (*CostAwareMemoryIndex) Clear added in v0.9.0 

func (m *CostAwareMemoryIndex) Clear(ctx context.Context, podIdentifier string) error

Clear removes every entry for the pod from the index, across all device tiers. It is O(N) over the index but runs off the Lookup/Add hot path at a coarse cadence. The scan is chunked: it snapshots the request keys, then processes them in fixed-size chunks, each taking mu only briefly, so a Clear never blocks Lookup (which takes mu.RLock) for the whole pass. The trade is atomicity — a concurrent Lookup may see the pod cleared from some keys but not yet from others; that is acceptable since the pod's cache is cold post-reset and a stale hit only costs a cache miss.

The engineKey->requestKey mapping (requestKeys) is intentionally left untouched: it is LRU-bounded, self-heals when the pod re-Adds the same prefixes, and any stale mapping resolves to an emptied request key that correctly breaks the prefix chain in Lookup. Reverse-pruning it would need an O(M) scan for no correctness gain.

func (*CostAwareMemoryIndex) Evict 

func (m *CostAwareMemoryIndex) Evict(ctx context.Context, key BlockHash, keyType KeyType, entries []PodEntry) error

Evict removes a key and its associated pod entries from the index backend. keyType indicates whether the key is an EngineKey (requires engine→request lookup) or a RequestKey (used directly for speculative entries without engineKey mapping).

func (*CostAwareMemoryIndex) GetRequestKey 

func (m *CostAwareMemoryIndex) GetRequestKey(ctx context.Context, engineKey BlockHash) (BlockHash, error)

GetRequestKey returns the last request key (highest index in the chain) associated with the given engineKey. Returns an error if the engineKey is not mapped (e.g., evicted earlier).

func (*CostAwareMemoryIndex) Lookup 

func (m *CostAwareMemoryIndex) Lookup(ctx context.Context, requestKeys []BlockHash,
	podIdentifierSet sets.Set[string],
) (map[BlockHash][]PodEntry, error)

func (*CostAwareMemoryIndex) MaxCost 

func (m *CostAwareMemoryIndex) MaxCost() int64

type CostAwareMemoryIndexConfig 

type CostAwareMemoryIndexConfig struct {
	// Size is the maximum memory size that can be used by the index.
	// Supports human-readable formats like "2GiB", "500MiB", "1GB", etc.
	Size string `json:"size,omitempty"`
}

CostAwareMemoryIndexConfig holds the configuration for the CostAwareMemoryIndex.

func DefaultCostAwareMemoryIndexConfig 

func DefaultCostAwareMemoryIndexConfig() *CostAwareMemoryIndexConfig

type CostPodCache 

type CostPodCache struct {
	// contains filtered or unexported fields
}

CostPodCache wraps a sync.Map of PodEntry and provides cost calculation for memory usage estimation.

func (*CostPodCache) Add 

func (c *CostPodCache) Add(entry PodEntry)

Add adds a PodEntry to the cache.

func (*CostPodCache) CalculateByteSize 

func (c *CostPodCache) CalculateByteSize(keyStr string) int64

CalculateByteSize estimates memory usage for ristretto cost calculation. This is an approximation used for cache eviction decisions.

func (*CostPodCache) Delete added in v0.7.0 

func (c *CostPodCache) Delete(entry PodEntry)

Delete removes a PodEntry from the cache.

func (*CostPodCache) Len 

func (c *CostPodCache) Len() int

Len returns the number of entries in the cache.

type GroupCatalog added in v0.9.0 

type GroupCatalog struct {
	// contains filtered or unexported fields
}

GroupCatalog is a thread-safe catalog of per-pod KV cache group metadata.

func NewGroupCatalog added in v0.9.0 

func NewGroupCatalog() *GroupCatalog

NewGroupCatalog creates a new, empty GroupCatalog.

func (*GroupCatalog) Get added in v0.9.0 

func (c *GroupCatalog) Get(podID string, g GroupID) (GroupMetadata, bool)

Get returns the metadata for a pod group.

func (*GroupCatalog) Learn added in v0.9.0 

func (c *GroupCatalog) Learn(podID string, g GroupID, meta GroupMetadata)

Learn records group metadata for a pod.

type GroupID added in v0.9.0 

type GroupID int

GroupID identifies a vLLM KV cache group.

type GroupMetadata added in v0.9.0 

type GroupMetadata struct {
	Kind              string
	BlockSize         int
	SlidingWindowSize *int
}

GroupMetadata holds per-group KV cache spec info learned from BlockStored events.

type InMemoryIndex 

type InMemoryIndex struct {
	// contains filtered or unexported fields
}

InMemoryIndex is an in-memory implementation of the Index interface.

func NewInMemoryIndex 

func NewInMemoryIndex(cfg *InMemoryIndexConfig) (*InMemoryIndex, error)

NewInMemoryIndex creates a new InMemoryIndex instance.

func (*InMemoryIndex) Add 

func (m *InMemoryIndex) Add(ctx context.Context, engineKeys, requestKeys []BlockHash, entries []PodEntry) error

Add adds a set of engineKeys/requestKeys and their associated pod entries to the index backend. If engineKeys is nil, only requestKey -> PodEntry mappings are created (no engineKey -> requestKey mapping). This is used for speculative entries where engine keys are not yet known. When engineKeys is non-nil, the mapping type is inferred from the ratio of array lengths.

func (*InMemoryIndex) Clear added in v0.9.0 

func (m *InMemoryIndex) Clear(ctx context.Context, podIdentifier string) error

Clear removes every entry for the pod from the index, across all device tiers. O(N) over the index, but Clear is rare and off the Lookup/Add hot path. Reuses evictPodsFromRequestKey for race-safe removal, and holds no global lock — only each PodCache's mu, briefly — so it does not stall Lookup.

The engineKey->requestKey mapping (engineToRequestKeys) is intentionally left untouched: it is LRU-bounded, self-heals when the pod re-Adds the same prefixes, and any stale mapping resolves to an emptied request key that correctly breaks the prefix chain in Lookup.

func (*InMemoryIndex) Evict 

func (m *InMemoryIndex) Evict(ctx context.Context, key BlockHash, keyType KeyType, entries []PodEntry) error

Evict removes a key and its associated pod entries from the index backend. keyType indicates whether the key is an EngineKey (requires engine→request lookup) or a RequestKey (used directly for speculative entries without engineKey mapping).

func (*InMemoryIndex) GetRequestKey 

func (m *InMemoryIndex) GetRequestKey(ctx context.Context, engineKey BlockHash) (BlockHash, error)

GetRequestKey returns the last request key (highest index in the chain) associated with the given engineKey. This is what Pool uses for parent hash resolution. Returns an error if the engineKey mapping is missing (e.g., already evicted).

func (*InMemoryIndex) Lookup 

func (m *InMemoryIndex) Lookup(ctx context.Context, requestKeys []BlockHash,
	podIdentifierSet sets.Set[string],
) (map[BlockHash][]PodEntry, error)

Lookup receives a list of requestKeys and a set of pod identifiers, and retrieves the filtered pods associated with those keys. The filtering is done based on the pod identifiers provided. If the podIdentifierSet is empty, all pods are returned.

It returns: 1. A map where the keys are those in (1) and the values are pod-identifiers. 2. An error if any occurred during the operation.

type InMemoryIndexConfig 

type InMemoryIndexConfig struct {
	// Size is the maximum number of keys that can be stored in the index.
	Size int `json:"size"`
	// PodCacheSize is the maximum number of pod entries per key.
	PodCacheSize int `json:"podCacheSize"`
}

InMemoryIndexConfig holds the configuration for the InMemoryIndex.

func DefaultInMemoryIndexConfig 

func DefaultInMemoryIndexConfig() *InMemoryIndexConfig

DefaultInMemoryIndexConfig returns a default configuration for the InMemoryIndex.

type Index 

type Index interface {
	// Lookup receives a list of keys and a set of pod identifiers,
	// and retrieves the filtered pods associated with those keys.
	// The filtering is done based on the pod identifiers provided.
	// If the podIdentifierSet is empty, all pods are returned.
	//
	// It returns:
	// 1. A map where the keys are those in requestKeys and the values are pod-identifiers.
	// 2. An error if any occurred during the operation.
	Lookup(ctx context.Context, requestKeys []BlockHash, podIdentifierSet sets.Set[string]) (map[BlockHash][]PodEntry, error)
	// Add stores requestKey -> pod entries and (optionally) engineKey -> requestKey
	// mappings. If engineKeys is nil, only requestKey -> pod mappings are created
	// (used for speculative entries where engine keys are not yet known).
	//
	// When engineKeys is non-nil, the backend infers the mapping from the ratio
	// of len(engineKeys) to len(requestKeys). Both lengths derive from the same
	// token count divided by their respective block sizes, so they always divide
	// evenly. Examples with 256 tokens:
	//
	//   1:1   (engine=64, canonical=64)  -> 4 eng, 4 req  -> E0->R0, E1->R1, ...
	//   many:1 (engine=16, canonical=64) -> 16 eng, 4 req -> E0..E3->R0, E4..E7->R1, ...
	//   1:many (engine=128, canonical=64) -> 2 eng, 4 req -> E0->[R0,R1], E1->[R2,R3]
	Add(ctx context.Context, engineKeys, requestKeys []BlockHash, entries []PodEntry) error
	// Evict removes a key and its associated pod entries from the index backend.
	// keyType indicates whether the key is an EngineKey (requires engine→request lookup)
	// or a RequestKey (used directly).
	Evict(ctx context.Context, key BlockHash, keyType KeyType, entries []PodEntry) error
	// GetRequestKey returns the requestKey associated with the given engineKey.
	GetRequestKey(ctx context.Context, engineKey BlockHash) (BlockHash, error)
	// Clear removes all index entries for the given pod, across every device tier.
	// It backs the AllBlocksCleared KV-event (a vLLM prefix-cache reset, e.g. after
	// an RLHF weight update), which is pod-wide — vLLM emits it with no tier. Clear is
	// O(N) over the index but runs off the Lookup/Add hot path, at a coarse cadence
	// (typically once per weight sync).
	Clear(ctx context.Context, podIdentifier string) error
}

Index defines the interface for a backend that manages KV-block indexing.

An index backend is a data store that will aggregate possibly the entire global KV cache block index, and will be used to retrieve pod-localities for a given set of consecutive keys that constitute a prefix-cache hit. The hit may not necessarily be on all keys, but of the longest prefix match.

The index backend allows efficient tracking of which vLLM engines hold which KV-blocks, on what device tier, and when they were last updated.

Index operations are thread-safe and can be performed concurrently.

func NewIndex 

func NewIndex(ctx context.Context, cfg *IndexConfig) (Index, error)

NewIndex creates a new Index instance.

func NewInstrumentedIndex 

func NewInstrumentedIndex(next Index) Index

NewInstrumentedIndex wraps an Index and emits metrics for Add, Evict, and Lookup.

func NewRedisIndex 

func NewRedisIndex(config *RedisIndexConfig) (Index, error)

NewRedisIndex creates a new RedisIndex instance. This constructor supports both Redis and Valkey backends.

func NewTracedIndex added in v0.5.1 

func NewTracedIndex(next Index) Index

NewTracedIndex wraps an Index and emits OpenTelemetry traces for index operations. This encapsulates all tracing logic for the kvblock.Index interface.

func NewValkeyIndex 

func NewValkeyIndex(config *RedisIndexConfig) (Index, error)

NewValkeyIndex creates a new RedisIndex instance configured for Valkey. This is a convenience constructor that sets up Valkey-specific defaults.

type IndexConfig 

type IndexConfig struct {
	// InMemoryConfig holds the configuration for the in-memory index.
	InMemoryConfig *InMemoryIndexConfig `json:"inMemoryConfig"`
	// RedisConfig holds the configuration for the Redis index.
	RedisConfig *RedisIndexConfig `json:"redisConfig"`
	// ValkeyConfig holds the configuration for the Valkey index.
	ValkeyConfig *RedisIndexConfig `json:"valkeyConfig"`
	// CostAwareMemoryConfig holds the configuration for the cost-aware memory index.
	CostAwareMemoryConfig *CostAwareMemoryIndexConfig `json:"costAwareMemoryConfig"`

	// EnableMetrics toggles whether admissions/evictions/hits/misses are
	// recorded.
	EnableMetrics bool `json:"enableMetrics"`
	// MetricsLoggingInterval defines the interval at which metrics are logged.
	// If zero, metrics logging is disabled.
	// Requires `EnableMetrics` to be true.
	MetricsLoggingInterval time.Duration `json:"metricsLoggingInterval"`
}

IndexConfig holds the configuration for the KV-block index. It may configure several backends such as listed within the struct. If multiple backends are configured, only the first one will be used.

func DefaultIndexConfig 

func DefaultIndexConfig() *IndexConfig

DefaultIndexConfig returns a default configuration for the KV-block index.

type KeyType added in v0.7.0 

type KeyType int

KeyType indicates whether a key passed to Evict is an engine key or a request key. 

const (
	// EngineKey means the key is an engine-assigned key that must be resolved
	// to a request key via the engineToRequestKeys mapping.
	EngineKey KeyType = iota
	// RequestKey means the key is a request key and can be used directly.
	// This is used for speculative entries that were added without engineKey mapping.
	RequestKey
)

type MMHash added in v0.7.0 

type MMHash struct {
	Hash string
}

MMHash represents a single multimodal content hash entry. This matches vLLM's per-block extra_keys format where each entry is the mm_feature.identifier string for an overlapping multimodal item.

type PlaceholderRange added in v0.7.0 

type PlaceholderRange struct {
	Offset int // absolute start token index
	Length int // number of placeholder tokens
}

PlaceholderRange describes a contiguous range of placeholder tokens for one multimodal item within the full token sequence.

type PodCache 

type PodCache struct {
	// contains filtered or unexported fields
}

PodCache represents a cache for pod entries.

type PodEntry 

type PodEntry struct {
	// PodIdentifier is the unique identifier for the pod.
	PodIdentifier string
	// DeviceTier is the tier of the device where the KV-block is stored.
	DeviceTier string
	// Speculative indicates the entry was added predictively before a KV event confirmed it.
	Speculative bool
	// HasGroup indicates GroupIdx identifies a vLLM KV cache group.
	HasGroup bool
	// GroupIdx identifies the vLLM KV cache group for HMA events.
	GroupIdx GroupID
}

PodEntry struct represents a pod entry in the KV-block index.

func (*PodEntry) String 

func (e *PodEntry) String() string

String returns a string representation of the PodEntry.

type RedisIndex 

type RedisIndex struct {
	RedisClient *redis.Client
	// BackendType indicates whether this is connecting to "redis" or "valkey"
	BackendType string
	// EnableRDMA indicates if RDMA transport is enabled (for Valkey)
	EnableRDMA bool
}

RedisIndex implements the Index interface using Redis or Valkey as the backend for KV block indexing.

func (*RedisIndex) Add 

func (r *RedisIndex) Add(ctx context.Context, engineKeys, requestKeys []BlockHash, entries []PodEntry) error

Add adds a set of keys and their associated pod entries to the index backend. If engineKeys is nil, only requestKey -> PodEntry mappings are created (no engineKey -> requestKey mapping). This is used for speculative entries where engine keys are not yet known. When engineKeys is non-nil, the mapping type is inferred from the ratio of array lengths.

func (*RedisIndex) Clear added in v0.9.0 

func (r *RedisIndex) Clear(ctx context.Context, podIdentifier string) error

Clear removes every hash field for the pod across all request-key hashes and device tiers. Each field is a JSON-encoded PodEntry, so matching decodes the field and compares PodIdentifier — catching every tier, group, and speculative variant. It pages the keyspace with SCAN (skipping engine: keys), HDELs the pod's fields, and prunes now-empty hashes. O(keyspace), but Clear is rare and off the Lookup/Add hot path. Because it deletes from the shared store, it is correct for multi-replica deployments with no cross-process coordination.

func (*RedisIndex) Evict 

func (r *RedisIndex) Evict(ctx context.Context, key BlockHash, keyType KeyType, entries []PodEntry) error

Evict removes a key and its associated pod entries from the index backend. keyType indicates whether the key is an EngineKey (requires engine→request lookup) or a RequestKey (used directly for speculative entries without engineKey mapping).

func (*RedisIndex) GetRequestKey 

func (r *RedisIndex) GetRequestKey(ctx context.Context, engineKey BlockHash) (BlockHash, error)

GetRequestKey returns the last request key (highest score) associated with the given engineKey.

func (*RedisIndex) Lookup 

func (r *RedisIndex) Lookup(ctx context.Context, requestKeys []BlockHash,
	podIdentifierSet sets.Set[string],
) (map[BlockHash][]PodEntry, error)

Lookup receives a list of keys and a set of pod identifiers, and retrieves the filtered pods associated with those keys. The filtering is done based on the pod identifiers provided. If the podIdentifierSet is empty, all pods are returned.

It returns: 1. A map where the keys are those in (1) and the values are pod-identifiers. 2. An error if any occurred during the operation.

type RedisIndexConfig 

type RedisIndexConfig struct {
	Address string `json:"address,omitempty"` // Redis/Valkey server address
	// BackendType specifies whether to connect to "redis" or "valkey" (optional, defaults to "redis")
	// This is mainly for documentation and future extensibility (e.g., RDMA support)
	BackendType string `json:"backendType,omitempty"`
	// EnableRDMA enables RDMA transport for Valkey when supported (experimental)
	EnableRDMA bool `json:"enableRDMA,omitempty"`
}

RedisIndexConfig holds the configuration for the RedisIndex. This configuration supports both Redis and Valkey backends since they are API-compatible.

func DefaultRedisIndexConfig 

func DefaultRedisIndexConfig() *RedisIndexConfig

func DefaultValkeyIndexConfig 

func DefaultValkeyIndexConfig() *RedisIndexConfig

DefaultValkeyIndexConfig returns a default configuration for Valkey.

type TokenProcessor 

type TokenProcessor interface {
	// TokensToKVBlockKeys converts tokens into kv_block.Keys.
	// It accepts an optional parentKey to continue a hash chain.
	// extraFeatures provides per-block multimodal data that taints the hash;
	// nil means text-only (no taint). When non-nil, its length must match the
	// number of token chunks.
	// It returns a slice of generated Keys.
	TokensToKVBlockKeys(
		parentKey BlockHash, tokens []uint32, modelName string,
		extraFeatures []*BlockExtraFeatures,
	) ([]BlockHash, error)

	// BlockSize returns the number of tokens per block used by this processor.
	BlockSize() int
}

TokenProcessor defines the interface for converting tokens to KVBlockKeys.

func NewChunkedTokenDatabase 

func NewChunkedTokenDatabase(config *TokenProcessorConfig) (TokenProcessor, error)

NewChunkedTokenDatabase creates a new instance with the given config and metadata.

type TokenProcessorConfig 

type TokenProcessorConfig struct {
	// BlockSize is deprecated. Use BlockSizeTokens instead.
	//
	// Deprecated: Use BlockSizeTokens instead.
	BlockSize int `json:"blockSize,omitempty"`
	// BlockSizeTokens is the number of tokens per block.
	// A value of zero is treated as "not set" and resolved to the default (16) by NewChunkedTokenDatabase.
	BlockSizeTokens int `json:"blockSizeTokens"`
	// HashSeed is used to prefix initial hash chunks, similarly to vLLM's NONE_HASH.
	// This should be aligned with vLLM's `PYTHONHASHSEED` environment variable.
	// The system's deployer is responsible for aligning the vLLM deployments
	// with the same seed value.
	HashSeed string `json:"hashSeed"`
	// contains filtered or unexported fields
}

TokenProcessorConfig holds the configuration for the token processor.

func DefaultTokenProcessorConfig 

func DefaultTokenProcessorConfig() *TokenProcessorConfig

DefaultTokenProcessorConfig returns the default configuration for the token processor.

Source Files

View all Source files

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.