feat: profiling kv cache index implementations (#108) #175

kartikx · 2025-11-22T17:59:58Z

This PR adds a profiling test to compare the various KVIndex implementations.

It adds and looks up keys for each implementation, and reports the average over a specified number of trials.

go run -ldflags="-extldflags '-L$(pwd)/lib'" ./tests/profiling/kv_cache_index/main.go -trials 10 -keys 10000

[Cost Aware] Add: 7.516833ms Lookup 141.652299ms
[Redis] Add: 10.614779ms Lookup 6.190304ms
[InMemory] Add: 3.515912ms Lookup 127.265687ms

Requested Feedback:

Would it be better if I set this up as an actual test? I'm not sure if we want this to run in CI and on every build, so I set it up as a script for now.
The current workload simply generates a given number of keys. I'm open to suggestions on whether we should improve this.

Issue: #108

Copilot

Pull request overview

This PR introduces a profiling tool to benchmark and compare the performance of three KVIndex implementations: Cost Aware, Redis, and InMemory. The tool measures Add and Lookup operation times over multiple trials and reports averaged results, enabling developers to make informed decisions about which index implementation to use based on performance characteristics.

Key Changes:

Adds a standalone profiling tool at tests/profiling/kv_cache_index/main.go that benchmarks Add and Lookup operations
Implements configurable trials and key counts via command-line flags
Generates random workload keys and measures averaged performance across multiple runs

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2025-11-22T18:03:57Z

tests/profiling/kv_cache_index/main.go

+	if err != nil {
+		fmt.Printf("Failed to profile cost index: %v\n", err)
+	} else {
+		fmt.Printf("[Cost Aware] Add: %v Lookup %v \n", result.AddTime, result.LookupTime)


Missing colon after "Lookup" in the output format string. The output format should be consistent with lines 169 and 176 which have a colon after "Lookup".

Suggested change

fmt.Printf("[Cost Aware] Add: %v Lookup %v \n", result.AddTime, result.LookupTime)

fmt.Printf("[Cost Aware] Add: %v Lookup: %v \n", result.AddTime, result.LookupTime)

Copilot · 2025-11-22T18:03:58Z

tests/profiling/kv_cache_index/main.go

+	if err != nil {
+		fmt.Printf("Failed to profile redis index: %v\n", err)
+	} else {
+		fmt.Printf("[Redis] Add: %v Lookup %v \n", result.AddTime, result.LookupTime)


Missing colon after "Lookup" in the output format string. The output format should be consistent with line 176 which has a colon after "Lookup".

Suggested change

fmt.Printf("[Redis] Add: %v Lookup %v \n", result.AddTime, result.LookupTime)

fmt.Printf("[Redis] Add: %v Lookup: %v \n", result.AddTime, result.LookupTime)

Copilot · 2025-11-22T18:03:58Z

tests/profiling/kv_cache_index/main.go

+		fmt.Printf("[Cost Aware] Add: %v Lookup %v \n", result.AddTime, result.LookupTime)
+	}
+
+	result, err = profileRedisIndex(*numTrials, *numKeys)
+	if err != nil {
+		fmt.Printf("Failed to profile redis index: %v\n", err)
+	} else {
+		fmt.Printf("[Redis] Add: %v Lookup %v \n", result.AddTime, result.LookupTime)
+	}
+
+	result, err = profileInMemoryIndex(*numTrials, *numKeys)
+	if err != nil {
+		fmt.Printf("Failed to profile in memory index: %v\n", err)
+	} else {
+		fmt.Printf("[InMemory] Add: %v Lookup %v \n", result.AddTime, result.LookupTime)


Extra space before the newline. This should be consistent with the format: \n not \n.

Suggested change

fmt.Printf("[Cost Aware] Add: %v Lookup %v \n", result.AddTime, result.LookupTime)

}

result, err = profileRedisIndex(*numTrials, *numKeys)

if err != nil {

fmt.Printf("Failed to profile redis index: %v\n", err)

} else {

fmt.Printf("[Redis] Add: %v Lookup %v \n", result.AddTime, result.LookupTime)

}

result, err = profileInMemoryIndex(*numTrials, *numKeys)

if err != nil {

fmt.Printf("Failed to profile in memory index: %v\n", err)

} else {

fmt.Printf("[InMemory] Add: %v Lookup %v \n", result.AddTime, result.LookupTime)

fmt.Printf("[Cost Aware] Add: %v Lookup %v\n", result.AddTime, result.LookupTime)

}

result, err = profileRedisIndex(*numTrials, *numKeys)

if err != nil {

fmt.Printf("Failed to profile redis index: %v\n", err)

} else {

fmt.Printf("[Redis] Add: %v Lookup %v\n", result.AddTime, result.LookupTime)

}

result, err = profileInMemoryIndex(*numTrials, *numKeys)

if err != nil {

fmt.Printf("Failed to profile in memory index: %v\n", err)

} else {

fmt.Printf("[InMemory] Add: %v Lookup %v\n", result.AddTime, result.LookupTime)

Copilot · 2025-11-22T18:03:58Z

tests/profiling/kv_cache_index/main.go

+// TODO @kartikx: Use a more realistic workload if possible.
+func generateWorkloadKeys(numKeys int) []kvblock.Key {
+	// Uses time as seed to ensure that different profiling runs get different keys.
+	randGen := rand.New(rand.NewPCG(42, uint64(time.Now().UnixNano())))


The random number generator is seeded with a fixed seed (42) combined with the current time, which contradicts the comment on line 39. The fixed seed of 42 means that if the same time value is obtained (or if profiling is run very quickly in succession), the same random keys could be generated. Consider removing the fixed seed and using only the time-based seed: rand.New(rand.NewPCG(uint64(time.Now().UnixNano()), uint64(time.Now().UnixNano()))) or simply use two different time-based values.

Suggested change

randGen := rand.New(rand.NewPCG(42, uint64(time.Now().UnixNano())))

randGen := rand.New(rand.NewPCG(uint64(time.Now().UnixNano()), uint64(time.Now().UnixNano())))

Copilot · 2025-11-22T18:03:58Z

tests/profiling/kv_cache_index/main.go

+	for i := range numTrials {
+		ctx := context.Background()
+
+		indexConfig := createConfig()
+		index, err := kvblock.NewIndex(ctx, indexConfig)
+		if err != nil {
+			return IndexProfileResult{}, fmt.Errorf("failed to create index: %w", err)
+		}
+
+		result, err := measureIndexRun(ctx, index, "pod1", numKeys)
+		if err != nil {
+			return IndexProfileResult{}, fmt.Errorf("failed to profile index: %w", err)
+		}
+


[nitpick] Creating a new index instance for each trial may skew profiling results due to initialization overhead. Consider clarifying whether this is intentional (to measure cold-start performance) or if the index should be created once outside the loop for warmed-up performance measurement. The current approach profiles both initialization and operation costs together.

Suggested change

for i := range numTrials {

ctx := context.Background()

indexConfig := createConfig()

index, err := kvblock.NewIndex(ctx, indexConfig)

if err != nil {

return IndexProfileResult{}, fmt.Errorf("failed to create index: %w", err)

}

result, err := measureIndexRun(ctx, index, "pod1", numKeys)

if err != nil {

return IndexProfileResult{}, fmt.Errorf("failed to profile index: %w", err)

}

ctx := context.Background()

indexConfig := createConfig()

index, err := kvblock.NewIndex(ctx, indexConfig)

if err != nil {

return IndexProfileResult{}, fmt.Errorf("failed to create index: %w", err)

}

for i := range numTrials {

result, err := measureIndexRun(ctx, index, "pod1", numKeys)

if err != nil {

return IndexProfileResult{}, fmt.Errorf("failed to profile index: %w", err)

}

Copilot · 2025-11-22T18:03:59Z

tests/profiling/kv_cache_index/main.go

+func measureIndexRun(ctx context.Context, index kvblock.Index, podName string, numKeys int) (IndexProfileResult, error) {
+	keys := generateWorkloadKeys(numKeys)
+
+	podEntries := []kvblock.PodEntry{{PodIdentifier: podName, DeviceTier: "gpu"}}


[nitpick] The empty podIdentifierSet means all pods will be returned during lookup (as documented in the Index interface). Consider adding a comment to clarify this is intentional, as it may be confusing why an empty set is used instead of including "pod1" to measure filtered lookup performance.

Suggested change

podEntries := []kvblock.PodEntry{{PodIdentifier: podName, DeviceTier: "gpu"}}

podEntries := []kvblock.PodEntry{{PodIdentifier: podName, DeviceTier: "gpu"}}

// Intentionally use an empty podIdentifierSet to return all pods during lookup,

// as documented in the Index interface. This measures unfiltered lookup performance.

sagiahrac

Thanks for the contribution @kartikx!
Maybe we can refactor this into a standard Go benchmark using the testing package. This approach is better for project stability as it automatically adjusts b.N (the number of keys) to achieve stable timing, which allows us to remove the manual numTrials and easily integrate with CI. It also handles warmups automatically and dynamically.

You can split Add and Lookup into 2 separate benchmarks for simplicity, or reset the timer after the Add phase.

sagiahrac · 2025-12-01T15:46:53Z

tests/profiling/kv_cache_index/main.go

+// TODO @kartikx: Use a more realistic workload if possible.
+func generateWorkloadKeys(numKeys int) []kvblock.Key {
+	// Uses time as seed to ensure that different profiling runs get different keys.
+	randGen := rand.New(rand.NewPCG(42, uint64(time.Now().UnixNano())))


Can you generate the same workload keys for all profiling sessions? That way, the only difference will be the indexer implementation.

feat: profiling kv cache index implementations (llm-d#108)

2547bfa

Copilot AI review requested due to automatic review settings November 22, 2025 17:59

kartikx requested review from dannyharnik, elevran, kfirtoledo and vMaroon as code owners November 22, 2025 17:59

Copilot started reviewing on behalf of kartikx November 22, 2025 18:00 View session

Copilot finished reviewing on behalf of kartikx November 22, 2025 18:02

Copilot AI reviewed Nov 22, 2025

View reviewed changes

sagiahrac suggested changes Dec 1, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

feat: profiling kv cache index implementations (#108) #175

feat: profiling kv cache index implementations (#108) #175

Uh oh!

kartikx commented Nov 22, 2025 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI Nov 22, 2025

Uh oh!

Copilot AI Nov 22, 2025

Uh oh!

Copilot AI Nov 22, 2025

Uh oh!

Copilot AI Nov 22, 2025

Uh oh!

Copilot AI Nov 22, 2025

Uh oh!

Copilot AI Nov 22, 2025

Uh oh!

sagiahrac left a comment

Uh oh!

sagiahrac Dec 1, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

	fmt.Printf("[Cost Aware] Add: %v Lookup %v \n", result.AddTime, result.LookupTime)
	fmt.Printf("[Cost Aware] Add: %v Lookup: %v \n", result.AddTime, result.LookupTime)

	fmt.Printf("[Redis] Add: %v Lookup %v \n", result.AddTime, result.LookupTime)
	fmt.Printf("[Redis] Add: %v Lookup: %v \n", result.AddTime, result.LookupTime)

	randGen := rand.New(rand.NewPCG(42, uint64(time.Now().UnixNano())))
	randGen := rand.New(rand.NewPCG(uint64(time.Now().UnixNano()), uint64(time.Now().UnixNano())))

feat: profiling kv cache index implementations (#108) #175

Are you sure you want to change the base?

feat: profiling kv cache index implementations (#108) #175

Uh oh!

Conversation

kartikx commented Nov 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI Nov 22, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Nov 22, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Nov 22, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Nov 22, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Nov 22, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Nov 22, 2025

Choose a reason for hiding this comment

Uh oh!

sagiahrac left a comment

Choose a reason for hiding this comment

Uh oh!

sagiahrac Dec 1, 2025

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

kartikx commented Nov 22, 2025 •

edited

Loading