Profile the different `kvblock.Index` implementations

[vMaroon]

As of writing this issue, we have three kvblock.Index implementations. An index contains KV-block locality information and is updated when KV-Events are digested from the vLLM servers, each containing a set of block-hashes, admission or removal type and an identifier of the sender.

The kvblock.Index is the heart of the kv-cache-manager and its performance is essential for low-latency and accurate scheduling.

The implementations are:

1. InMemoryIndex (default) - an in-memory index backed by a hashicorp LRU cache implementation, with relaxed locking
2. CostAwareMemoryIndex - contributed by @yankay, a hypermodeinc/ristretto memory bound high-performance in-memory cache
3. Redis - an index backed by a Redis server

Profiling the three options in realistic/high-scale workloads inputs would be helpful for drawing configuration recommendations and general transparency. The promotion of the 2nd option to default would also be desirable given good profiling data.

[kartikx]

@vMaroon I would like to work on this issue.

[vMaroon]

@kartikx thank you for your interest! Please note the milestone deadline and feel free to reach out on the llm-d slack anytime.

[kartikx]

@vMaroon Do you have any suggestions on what a "realistic" workload looks like? Currently, I have the following setup

1. Add 100,000 keys to the index at once, and profile the time for the Add.
2. Lookup all keys and profile the time for the Lookup.

All of these are added to the same pod. I repeat experiments 10 times for each index type, the keys are different on each run.

Current set of results for this:

[Cost] Add: 22.271475ms Lookup 21.758358ms
[Redis] Add: 22.533808ms Lookup 22.307492ms
[In memory] Add: 22.59525ms Lookup 22.8048ms

The numbers are pretty close, and the ordering doesn't stay consistent. Perhaps a better profiling workload can help make differences clear.

I can share a cleaned up version of my script for you to review if that helps.

[kartikx]

@vMaroon I found a bug in my implementation which wasn't initializing the index correctly, which is why the numbers were all so close together.

Here are the numbers (again) for adding and looking up 100,000 keys. The reported numbers are averaged over 5 trials.

[Cost] Add: 82.093341ms Lookup 8.742659583s
[Redis] Add: 160.122641ms Lookup 82.22675ms
[In memory] Add: 69.18185ms Lookup 9.4533644s

This aligns more with our expectations.

Here is my profiling script that you can use to test out.

I would be happy to create a better version of this script as a PR. Please let me know if you have any high-level comments I could address.

[vMaroon]

Thanks @kartikx - the script looks OK - the TODOs you outlined are a good improvement start, but we can start with a straightforward profiling test and iterate on it in PRs.

[kartikx]

@vMaroon I am sorry for the delay as I got caught up with other work. I will raise a PR this week.

[samber]

I would suggest a W-TinyLFU (2023) eviction algo. It efficiently balances recency and frequency to maximize hit rate while remaining lightweight and fast. Otter is cool.