indexer.go

/*
Copyright 2025 The llm-d Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package kvcache

import (
	"context"
	"fmt"

	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/codes"
	"go.opentelemetry.io/otel/trace"
	"k8s.io/apimachinery/pkg/util/sets"
	"sigs.k8s.io/controller-runtime/pkg/log"

	"github.com/llm-d/llm-d-kv-cache/pkg/kvcache/kvblock"
	"github.com/llm-d/llm-d-kv-cache/pkg/telemetry"
	"github.com/llm-d/llm-d-kv-cache/pkg/tokenization"
	"github.com/llm-d/llm-d-kv-cache/pkg/tokenization/types"
	"github.com/llm-d/llm-d-kv-cache/pkg/utils/logging"
)

// AttentionGroupConfig defines the attention window size and type for a group.
type AttentionGroupConfig struct {
	WindowSize    int    `json:"windowSize"`    // Attention window size (0 or omit for full attention = no constraint)
	AttentionType string `json:"attentionType"` // Attention type (e.g., "full", "sliding", "block_sparse")
	BlockSize     int    `json:"blockSize"`     // KV block size for this group (if 0, falls back to ModelConfig.BlockSize)
}

// ModelConfig holds model-specific configuration including block size and attention groups.
type ModelConfig struct {
	Name            string                 `json:"name"`            // Model name
	BlockSize       int                    `json:"blockSize"`       // Default KV block size (used when AttentionGroupConfig.BlockSize is 0)
	AttentionGroups []AttentionGroupConfig `json:"attentionGroups"` // Multiple attention groups with different window sizes and block sizes
}

// Config holds the configuration for the Indexer module.
// The configuration cover the different components found in the Indexer
// module.
type Config struct {
	KVBlockIndexConfig   *kvblock.IndexConfig    `json:"kvBlockIndexConfig"`
	KVBlockScorerConfig  *KVBlockScorerConfig    // not exported
	TokenizersPoolConfig *tokenization.Config    `json:"tokenizersPoolConfig"`
	BackendConfigs       []*KVCacheBackendConfig `json:"kvCacheBackendConfigs"`
	ModelConfigs         []*ModelConfig
}

// NewDefaultConfig returns a default configuration for the Indexer module.
func NewDefaultConfig() (*Config, error) {
	tokenizerPoolConfig, err := tokenization.DefaultConfig()
	if err != nil {
		return &Config{}, fmt.Errorf("failed to get default tokenizer pool config: %w", err)
	}

	return &Config{
		KVBlockIndexConfig:   kvblock.DefaultIndexConfig(),
		KVBlockScorerConfig:  DefaultKVBlockScorerConfig(),
		TokenizersPoolConfig: tokenizerPoolConfig,
		BackendConfigs:       DefaultKVCacheBackendConfig(),
		ModelConfigs:         nil, // No default model configs - must be explicitly configured
	}, nil
}

// Indexer is a concrete implementation of the KVCacheIndex interface.
type Indexer struct {
	config *Config

	tokenProcessor kvblock.TokenProcessor // turns tokens to kv block keys
	kvBlockIndex   kvblock.Index          // looks up pods for block keys
	kvBlockScorer  KVBlockScorer          // scores pods based on block hits

	tokenizersPool TokenizersPool
}

// NewKVCacheIndexer creates a KVCacheIndex given a Config.
func NewKVCacheIndexer(ctx context.Context, config *Config, tokenProcessor kvblock.TokenProcessor) (*Indexer, error) {
	if config == nil {
		return nil, fmt.Errorf("config cannot be nil")
	}
	if tokenProcessor == nil {
		return nil, fmt.Errorf("tokenProcessor cannot be nil")
	}

	kvBlockIndex, err := kvblock.NewIndex(ctx, config.KVBlockIndexConfig)
	if err != nil {
		return nil, fmt.Errorf("failed to create RedisKVBlockIndexer: %w", err)
	}

	// Wrap index with tracing instrumentation.
	// When tracing is not configured, otel.Tracer() returns a no-op implementation.
	kvBlockIndex = kvblock.NewTracedIndex(kvBlockIndex)

	// override backend configs with the ones from the config, if the defaults are not used.
	config.KVBlockScorerConfig.BackendConfigs = config.BackendConfigs
	scorer, err := NewKVBlockScorer(config.KVBlockScorerConfig)
	if err != nil {
		return nil, fmt.Errorf("failed to create KVBlockScorer: %w", err)
	}

	// Wrap scorer with tracing instrumentation.
	// When tracing is not configured, otel.Tracer() returns a no-op implementation.
	scorer = NewTracedScorer(scorer)

	tokenizersPool, err := tokenization.NewTokenizationPool(ctx, config.TokenizersPoolConfig)
	if err != nil {
		return nil, fmt.Errorf("failed to create tokenizers pool: %w", err)
	}

	return &Indexer{
		config:         config,
		tokenProcessor: tokenProcessor,
		kvBlockIndex:   kvBlockIndex,
		kvBlockScorer:  scorer,
		tokenizersPool: tokenizersPool,
	}, nil
}

// Run starts the indexer.
func (k *Indexer) Run(ctx context.Context) {
	k.tokenizersPool.Run(ctx)
}

// KVBlockIndex returns the kvblock.Index used by the Indexer.
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, features := 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. Compute per-block extra features from multimodal metadata (if present).
	var extraFeatures []*kvblock.BlockExtraFeatures
	if features != nil {
		extraFeatures = kvblock.ComputeBlockExtraFeatures(
			features.MMHashes, features.MMPlaceholders,
			k.blockSize(), len(tokens))
	}

	// 4. get block keys
	blockKeys, err := k.tokenProcessor.TokensToKVBlockKeys(kvblock.EmptyBlockHash, tokens, modelName, extraFeatures)
	if err != nil {
		traceLogger.Error(err, "blockKey conversion failed")
		return nil, fmt.Errorf("blockKey conversion failed: %w", err)
	}
	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.
// If the set of pod identifiers is empty, the function assumes all pods are
// relevant.
//
// The function returns a map of pod identifiers to scores.
func (k *Indexer) GetPodScores(ctx context.Context, renderReq *types.RenderChatRequest, prompt, modelName string,
	podIdentifiers []string,
) (map[string]float64, error) {
	// 1. tokenize prompt
	tokens, features := 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. Compute per-block extra features from multimodal metadata (if present).
	var extraFeatures []*kvblock.BlockExtraFeatures
	if features != nil {
		extraFeatures = kvblock.ComputeBlockExtraFeatures(
			features.MMHashes, features.MMPlaceholders,
			k.blockSize(), len(tokens))
	}

	return k.ScoreTokens(ctx, tokens, modelName, podIdentifiers, extraFeatures)
}

// ScoreTokens computes pod scores for the given tokens and model.
// It converts tokens into KV block keys, looks up which pods hold
// matching blocks in the index, and scores each pod based on cache hits.
//
// extraFeatures provides per-block multimodal data that taints the hash;
// nil means text-only. podIdentifiers limits scoring to the given pod addresses.
// If empty, all pods are considered.
func (k *Indexer) ScoreTokens(
	ctx context.Context,
	tokens []uint32,
	modelName string,
	podIdentifiers []string,
	extraFeatures []*kvblock.BlockExtraFeatures,
) (map[string]float64, error) {
	tracer := otel.Tracer(telemetry.InstrumentationName)
	ctx, span := tracer.Start(ctx, "llm_d.kv_cache.score_tokens",
		trace.WithSpanKind(trace.SpanKindInternal),
	)
	defer span.End()

	traceLogger := log.FromContext(ctx).V(logging.TRACE).WithName("kvcache.ScoreTokens")

	blockKeys, err := k.tokenProcessor.TokensToKVBlockKeys(kvblock.EmptyBlockHash, tokens, modelName, extraFeatures)
	if err != nil {
		return nil, fmt.Errorf("blockKey conversion failed: %w", err)
	}

	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)),
	)

	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)

	keyToPods, err := k.kvBlockIndex.Lookup(ctx, blockKeys, sets.New(podIdentifiers...))
	if err != nil {
		span.SetStatus(codes.Error, err.Error())
		return nil, fmt.Errorf("failed to query kvblock indexer: %w", err)
	}
	traceLogger.Info("found block keys", "block-keys", blockKeys,
		"pods", podsPerKeyPrintHelper(keyToPods))

	// Calculate block-level hit ratio (blocks found / blocks requested).
	blocksFound := 0
	for _, pods := range keyToPods {
		if len(pods) > 0 {
			blocksFound++
		}
	}
	blockHitRatio := 0.0
	if len(blockKeys) > 0 {
		blockHitRatio = float64(blocksFound) / float64(len(blockKeys))
	}
	span.SetAttributes(
		attribute.Float64("llm_d.kv_cache.block_hit_ratio", blockHitRatio),
		attribute.Int("llm_d.kv_cache.blocks_found", blocksFound),
	)

	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)
	}

	return podScores, nil
}

// podsPerKeyPrintHelper formats a map of keys to pod entries for printing.
func podsPerKeyPrintHelper(ks map[kvblock.BlockHash][]kvblock.PodEntry) string {
	flattened := ""
	for k, v := range ks {
		entries := make([]string, len(v))
		for i, entry := range v {
			entries[i] = entry.String()
		}
		flattened += fmt.Sprintf("%s: %v\n", k.String(), entries)
	}

	return flattened
}

func (k *Indexer) SetTokenizer(tokenizer tokenization.Tokenizer, modelName string) {
	k.tokenizersPool.SetTokenizer(tokenizer, modelName)
}

// blockSize returns the block size from the injected token processor.
func (k *Indexer) blockSize() int {
	return k.tokenProcessor.BlockSize()
}

// GetModelConfig returns the model configuration for the given model name.
// Returns nil if no configuration is found for the model.
func (c *Config) GetModelConfig(modelName string) *ModelConfig {
	if c.ModelConfigs == nil {
		return nil
	}

	for _, mc := range c.ModelConfigs {
		if mc.Name == modelName {
			return mc
		}
	}

	return nil
}