feat: add speculative indexing support for PodEntry and Indexer (#369)

* feat: speculative indexing support

* refactor: flatten nested cleanup logic in InMemoryIndex.Evict

* refactor: replace Annotation string with Annotations struct in PodEntry

* fix: lint

* refactor: use strconv for BlockHash string representation

* feat: add speculative indexing support to RedisIndex

* refactor: simplify nested if in Evict() with early returns

* refactor: simplify PodEntry Annotations to Speculative bool

* feat: add explicit KeyType parameter to Evict() interface

* fix: suppress nilerr lint in Redis Evict for missing engine key

* fix: handle nil engineKeys in CostAwareMemoryIndex.Add and add common test

* refactor: delegate GetPodScores to ScoreTokens

GetPodScores duplicated the tokenize→lookup→score logic that ScoreTokens
already provides. Reduce it to a thin wrapper that tokenizes, truncates,
and delegates to ScoreTokens. Fixes rebase regression from PR 415.

* fix: restore traceLogger lines in ScoreTokens

* fix: restore block hit ratio span attributes in ScoreTokens

Author: bongwoobak

examples/valkey_example/main.go

@@ -188,7 +188,7 @@ func demonstrateValkeyOperations(ctx context.Context, indexer *kvcache.Indexer)
 
 	// Demonstrate eviction
 	logger.Info("Demonstrating cache eviction")
-	err = indexer.KVBlockIndex().Evict(ctx, promptKeys[0], podEntries[:1])
+	err = indexer.KVBlockIndex().Evict(ctx, promptKeys[0], kvblock.EngineKey, podEntries[:1])
 	if err != nil {
 		return fmt.Errorf("failed to evict cache entry: %w", err)
 	}

pkg/kvcache/indexer.go

@@ -123,6 +123,36 @@ func (k *Indexer) KVBlockIndex() kvblock.Index {
 	return k.kvBlockIndex
 }
 
+// ComputeBlockKeys computes the KV-block keys for a given prompt and model name.
+// This method extracts the tokenization and block key computation logic so that
+// callers (e.g., IGW::EPP::PrepareDataPlugin) can compute block keys once and reuse them
+// across multiple extension points without re-tokenizing.
+func (k *Indexer) ComputeBlockKeys(ctx context.Context, renderReq *types.RenderChatRequest, prompt, modelName string,
+) ([]kvblock.BlockHash, error) {
+	traceLogger := log.FromContext(ctx).V(logging.TRACE).WithName("kvcache.ComputeBlockKeys")
+
+	// 1. tokenize prompt
+	tokens := k.tokenizersPool.Tokenize(renderReq, prompt)
+
+	// 2. Truncate prompt (if set in the request)
+	if renderReq != nil && renderReq.TruncatePromptTokens != nil {
+		limit := *renderReq.TruncatePromptTokens
+		if limit > 0 && len(tokens) > limit {
+			tokens = tokens[len(tokens)-limit:]
+		}
+	}
+
+	// 3. get block keys
+	blockKeys := k.tokenProcessor.TokensToKVBlockKeys(kvblock.EmptyBlockHash, tokens, modelName)
+	if len(blockKeys) == 0 {
+		traceLogger.Info("no block keys found")
+		return nil, nil
+	}
+	traceLogger.Info("computed block keys", "tokens", tokens, "block-keys", blockKeys)
+
+	return blockKeys, nil
+}
+
 // GetPodScores retrieves the pod scores for a given prompt and model name.
 // The function receives the mentioned information and a list of relevant pod
 // identifiers. A Pod identifier should be its address.
@@ -159,9 +189,8 @@ func (k *Indexer) ScoreTokens(
 	modelName string,
 	podIdentifiers []string,
 ) (map[string]float64, error) {
-	// Start tracing span for main operation
 	tracer := otel.Tracer(telemetry.InstrumentationName)
-	ctx, span := tracer.Start(ctx, "llm_d.kv_cache.get_scores",
+	ctx, span := tracer.Start(ctx, "llm_d.kv_cache.score_tokens",
 		trace.WithSpanKind(trace.SpanKindInternal),
 	)
 	defer span.End()
@@ -170,21 +199,20 @@ func (k *Indexer) ScoreTokens(
 
 	blockKeys := k.tokenProcessor.TokensToKVBlockKeys(kvblock.EmptyBlockHash, tokens, modelName)
 
-	// Set initial attributes
 	span.SetAttributes(
 		attribute.String("gen_ai.request.model", modelName),
 		attribute.Int("llm_d.kv_cache.pod_count", len(podIdentifiers)),
 		attribute.Int("llm_d.kv_cache.token_count", len(tokens)),
+		attribute.Int("llm_d.kv_cache.block_keys.count", len(blockKeys)),
 	)
-	span.SetAttributes(attribute.Int("llm_d.kv_cache.block_keys.count", len(blockKeys)))
+
 	if len(blockKeys) == 0 {
 		traceLogger.Info("no block keys found, returning empty scores")
 		//nolint:nilnil // no need to return an error
 		return nil, nil
 	}
 	traceLogger.Info("found tokens", "tokens", tokens, "block-keys", blockKeys)
 
-	// query kvblock indexer for pods
 	keyToPods, err := k.kvBlockIndex.Lookup(ctx, blockKeys, sets.New(podIdentifiers...))
 	if err != nil {
 		span.SetStatus(codes.Error, err.Error())
@@ -209,13 +237,11 @@ func (k *Indexer) ScoreTokens(
 		attribute.Int("llm_d.kv_cache.blocks_found", blocksFound),
 	)
 
-	// 5. score pods
 	podScores, err := k.kvBlockScorer.Score(ctx, blockKeys, keyToPods)
 	if err != nil {
 		span.SetStatus(codes.Error, err.Error())
 		return nil, fmt.Errorf("failed to query kvblock scorer: %w", err)
 	}
-	traceLogger.Info("found pod scores", "pod-scores", podScores)
 
 	return podScores, nil
 }

pkg/kvcache/kvblock/cost_aware_memory.go

@@ -166,24 +166,26 @@ func (c *CostPodCache) CalculateByteSize(keyStr string) int64 {
 var _ Index = &CostAwareMemoryIndex{}
 
 // 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.
 func (m *CostAwareMemoryIndex) Add(ctx context.Context, engineKeys, requestKeys []BlockHash, entries []PodEntry) error {
 	m.mu.Lock()
 	defer m.mu.Unlock()
 
-	if len(engineKeys) == 0 || len(requestKeys) == 0 || len(entries) == 0 {
+	if len(requestKeys) == 0 || len(entries) == 0 {
 		return fmt.Errorf("no keys or entries provided for adding to index")
 	}
-	if len(engineKeys) != len(requestKeys) {
+	if engineKeys != nil && len(engineKeys) != len(requestKeys) {
 		return fmt.Errorf("mismatch between engine keys and request keys length")
 	}
 
 	traceLogger := log.FromContext(ctx).V(logging.TRACE).WithName("kvblock.CostAwareMemoryIndex.Add")
 
 	for i, requestKey := range requestKeys {
-		engineKey := engineKeys[i]
-
-		// Store engineKey -> requestKey mapping
-		m.requestKeys.Add(engineKey, requestKey)
+		// Store engineKey -> requestKey mapping (only if engineKeys provided)
+		if engineKeys != nil {
+			m.requestKeys.Add(engineKeys[i], requestKey)
+		}
 
 		keyStr := requestKey.String()
 		podCache, found := m.data.Get(keyStr)
@@ -198,7 +200,7 @@ func (m *CostAwareMemoryIndex) Add(ctx context.Context, engineKeys, requestKeys
 		// Calculate the actual cost for this cache entry
 		cost := podCache.CalculateByteSize(keyStr)
 		m.data.Set(keyStr, podCache, cost)
-		traceLogger.Info("added pods to key", "requestKey", requestKey, "engineKey", engineKey, "pods", entries, "cost-bytes", cost)
+		traceLogger.Info("added pods to key", "requestKey", requestKey, "pods", entries, "cost-bytes", cost)
 	}
 	m.data.Wait()
 	return nil
@@ -260,7 +262,9 @@ func (m *CostAwareMemoryIndex) Lookup(ctx context.Context, requestKeys []BlockHa
 }
 
 // Evict removes a key and its associated pod entries from the index backend.
-func (m *CostAwareMemoryIndex) Evict(ctx context.Context, engineKey BlockHash, entries []PodEntry) error {
+// keyType indicates whether the key is an EngineKey (requires engine→request lookup)
+// or a RequestKey (used directly for speculative entries without engineKey mapping).
+func (m *CostAwareMemoryIndex) Evict(ctx context.Context, key BlockHash, keyType KeyType, entries []PodEntry) error {
 	m.mu.Lock()
 	defer m.mu.Unlock()
 
@@ -270,17 +274,33 @@ func (m *CostAwareMemoryIndex) Evict(ctx context.Context, engineKey BlockHash, e
 
 	traceLogger := log.FromContext(ctx).V(logging.TRACE).WithName("kvblock.CostAwareMemoryIndex.Evict")
 
-	requestKey, found := m.requestKeys.Get(engineKey)
-	if !found {
-		traceLogger.Info("engineKey not found in index, nothing to evict", "engineKey", engineKey)
-		return nil
+	var requestKey BlockHash
+	hasEngineKeyMapping := false
+
+	switch keyType {
+	case EngineKey:
+		rk, found := m.requestKeys.Get(key)
+		if !found {
+			traceLogger.Info("engineKey not found in mapping, nothing to evict", "engineKey", key)
+			return nil
+		}
+		requestKey = rk
+		hasEngineKeyMapping = true
+	case RequestKey:
+		requestKey = key
+	default:
+		return fmt.Errorf("unknown key type: %d", keyType)
 	}
 
 	keyStr := requestKey.String()
 	podCache, found := m.data.Get(keyStr)
 	if !found || podCache == nil {
-		traceLogger.Info("requestKey not found in index, cleaning up engineKey", "requestKey", requestKey, "engineKey", engineKey)
-		m.requestKeys.Remove(engineKey)
+		if hasEngineKeyMapping {
+			traceLogger.Info("requestKey not found in index, cleaning up engineKey", "requestKey", requestKey, "engineKey", key)
+			m.requestKeys.Remove(key)
+		} else {
+			traceLogger.Info("key not found in index, nothing to evict", "key", key)
+		}
 		return nil
 	}
 
@@ -292,12 +312,13 @@ func (m *CostAwareMemoryIndex) Evict(ctx context.Context, engineKey BlockHash, e
 
 	if podCache.Len() == 0 {
 		m.data.Del(keyStr)
-		m.requestKeys.Remove(engineKey)
-
-		traceLogger.Info("removed requestKey from index as no pods remain", "requestKey", requestKey, "engineKey", engineKey)
+		if hasEngineKeyMapping {
+			m.requestKeys.Remove(key)
+		}
+		traceLogger.Info("removed requestKey from index as no pods remain", "requestKey", requestKey, "key", key)
 	} else if podCacheLenBefore != podCache.Len() {
 		m.data.Set(keyStr, podCache, podCache.CalculateByteSize(keyStr))
-		traceLogger.Info("evicted pods from engineKey", "requestKey", requestKey, "engineKey", engineKey, "pods", entries)
+		traceLogger.Info("evicted pods from key", "requestKey", requestKey, "key", key, "keyType", keyType, "pods", entries)
 	}
 	m.data.Wait()
 	return nil

pkg/kvcache/kvblock/in_memory.go

@@ -147,21 +147,23 @@ func (m *InMemoryIndex) Lookup(ctx context.Context, requestKeys []BlockHash,
 }
 
 // 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.
 func (m *InMemoryIndex) Add(ctx context.Context, engineKeys, requestKeys []BlockHash, entries []PodEntry) error {
-	if len(engineKeys) == 0 || len(requestKeys) == 0 || len(entries) == 0 {
+	if len(requestKeys) == 0 || len(entries) == 0 {
 		return fmt.Errorf("no keys or entries provided for adding to index")
 	}
-	if len(engineKeys) != len(requestKeys) {
+	if engineKeys != nil && len(engineKeys) != len(requestKeys) {
 		return fmt.Errorf("mismatch between engine keys and request keys length")
 	}
 
 	traceLogger := log.FromContext(ctx).V(logging.TRACE).WithName("kvblock.InMemoryIndex.Add")
 
 	for i, requestKey := range requestKeys {
-		engineKey := engineKeys[i]
-
-		// 1. Store engineKey -> requestKey mapping
-		m.engineToRequestKeys.Add(engineKey, requestKey)
+		// 1. Store engineKey -> requestKey mapping (only if engineKeys provided)
+		if engineKeys != nil {
+			m.engineToRequestKeys.Add(engineKeys[i], requestKey)
+		}
 
 		// 2. Store requestKey -> PodCache mapping
 		var podCache *PodCache
@@ -203,30 +205,48 @@ func (m *InMemoryIndex) Add(ctx context.Context, engineKeys, requestKeys []Block
 		}
 		podCache.mu.Unlock()
 
-		traceLogger.Info("added pods to key", "requestKey", requestKey, "engineKey", engineKey, "pods", entries)
+		traceLogger.Info("added pods to key", "requestKey", requestKey, "pods", entries)
 	}
 
 	return nil
 }
 
-// Evict removes a engineKey and its associated pod entries from the index backend.
-func (m *InMemoryIndex) Evict(ctx context.Context, engineKey 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 for speculative entries without engineKey mapping).
+func (m *InMemoryIndex) Evict(ctx context.Context, key BlockHash, keyType KeyType, entries []PodEntry) error {
 	if len(entries) == 0 {
 		return fmt.Errorf("no entries provided for eviction from index")
 	}
 
 	traceLogger := log.FromContext(ctx).V(logging.TRACE).WithName("kvblock.InMemoryIndex.Evict")
 
-	requestKey, found := m.engineToRequestKeys.Get(engineKey)
-	if !found {
-		traceLogger.Info("engineKey not found in index, nothing to evict", "engineKey", engineKey)
-		return nil
+	var requestKey BlockHash
+	hasEngineKeyMapping := false
+
+	switch keyType {
+	case EngineKey:
+		rk, found := m.engineToRequestKeys.Get(key)
+		if !found {
+			traceLogger.Info("engineKey not found in mapping, nothing to evict", "engineKey", key)
+			return nil
+		}
+		requestKey = rk
+		hasEngineKeyMapping = true
+	case RequestKey:
+		requestKey = key
+	default:
+		return fmt.Errorf("unknown key type: %d", keyType)
 	}
 
 	podCache, found := m.data.Get(requestKey)
 	if !found || podCache == nil {
-		traceLogger.Info("requestKey not found in index, cleaning up engineKey", "requestKey", requestKey, "engineKey", engineKey)
-		m.engineToRequestKeys.Remove(engineKey)
+		if hasEngineKeyMapping {
+			traceLogger.Info("requestKey not found in index, cleaning up engineKey", "requestKey", requestKey, "engineKey", key)
+			m.engineToRequestKeys.Remove(key)
+		} else {
+			traceLogger.Info("key not found in index, nothing to evict", "key", key)
+		}
 		return nil
 	}
 
@@ -238,21 +258,28 @@ func (m *InMemoryIndex) Evict(ctx context.Context, engineKey BlockHash, entries
 	isEmpty := podCache.cache.Len() == 0
 	podCache.mu.Unlock()
 
-	traceLogger.Info("evicted pods from key", "requestKey", requestKey, "engineKey", engineKey, "pods", entries)
-
-	// Remove key from main cache if empty
-	if isEmpty {
-		// Re-fetch and hold the lock through removal to prevent racing with Add
-		if currentCache, stillExists := m.data.Get(requestKey); stillExists && currentCache != nil {
-			currentCache.mu.Lock()
-			if currentCache.cache.Len() == 0 {
-				m.data.Remove(requestKey)
-				m.engineToRequestKeys.Remove(engineKey)
-				traceLogger.Info("removed requestKey from index as no pods remain", "requestKey", requestKey, "engineKey", engineKey)
-			}
-			currentCache.mu.Unlock()
+	traceLogger.Info("evicted pods from key", "requestKey", requestKey, "key", key, "keyType", keyType, "pods", entries)
+
+	// Remove key from main cache if empty.
+	// Re-fetch and hold the lock through removal to prevent racing with Add.
+	if !isEmpty {
+		return nil
+	}
+
+	currentCache, stillExists := m.data.Get(requestKey)
+	if !stillExists || currentCache == nil {
+		return nil
+	}
+
+	currentCache.mu.Lock()
+	if currentCache.cache.Len() == 0 {
+		m.data.Remove(requestKey)
+		if hasEngineKeyMapping {
+			m.engineToRequestKeys.Remove(key)
 		}
+		traceLogger.Info("removed requestKey from index as no pods remain", "requestKey", requestKey, "key", key)
 	}
+	currentCache.mu.Unlock()
 
 	return nil
 }