Add prefix caching

[maxdebayser]

Description

Note: this PR is stacked on #585

This implementation uses the KV Cache Manager to find matching prefixes that satisfy our chunking constraints. Since in the initial version we're going for scheduler agnostic caching, the execute_model() function just executes dummy steps if the chunk was loaded from cache. This is required because the scheduler controls the num_computed_tokens.

[github-actions]

👋 Hi! Thank you for contributing to vLLM support on Spyre.
Just a reminder: Make sure that your code passes all the linting checks, otherwise your PR won't be able to be merged. To do so, first install the linting requirements, then run format.sh and commit the changes. This can be done with uv directly:

uv sync --frozen --group lint --active --inexact

Or this can be done with pip:

uv pip compile --group lint > requirements-lint.txt
pip install -r requirements-lint.txt
bash format.sh

Now you are good to go 🚀

[maxdebayser]

This PR is now in a state where it is running correctly in a very simple test case where I send the same long prompt twice, (on CPU).

[maxdebayser]

The next steps are:

• test on the spyre device
• figure out what to do about the enable-prefix-caching setting which is on by default in vllm
• write comprehensive tests

[yannicks1]

looks very good! Put some comments. (I did not run any of the code yet)

[yannicks1]

seems a bit odd to introduce the enable_prefix_caching arg in the ContinuousBatchingSpyreModelRunner and not in the ChunkedPrefillModelRunner only as CB cannot do prefix caching. But I understand you need it currently in _set_blocks().

A possible solution could be to have an additional arg in _set_blocks e.g. _set_blocks(..., enable_caching=False) which defaults to False. You can then introduce enable_prefix_caching only for the ChunkedPrefillModelRunner and let the default value handle CB.

[sducouedic]

I like this idea @yannicks1

[sducouedic]

but actually set_blocks is called only in ContinuousBatchingSpyreModelrunner

[maxdebayser]

Considering that this PR also disabled prefix caching by default and that the only real impact that setting it to enabled while running with the continuous batching model runner is that the block pool will do a bit of [extra work[(https://github.com/vllm-project/vllm/blob/e858bfe05167a3bbb064e283da5a1a7709dee24e/vllm/v1/core/block_pool.py#L326) when chunks are freed, I've removed the weird constructor argument.

[yannicks1]

Is this also updating vllm's internal counter for the cache hit rate stats?

[yannicks1]

if not we should also keep these stats up to date for complete benchmark outputs reporting cache hit rates

[maxdebayser]

The BlockPool keeps track of kv cache events if enabled. I didn't do it in the first pass in this PR, but we should definitely reuse this.

[maxdebayser]

bot:test

[maxdebayser]

Caveat: in this current PR we also cache blocks that become full during decoding. But there could be significant numerical differences between blocks generated during prefill and decode.

[sducouedic]

It looks great @maxdebayser, especially that everything is reusing upstream KV cache manager. I'll just give it a run after dinner before approving

[sducouedic]

I believe we can substract the number of reserved blocks here using computed_blocks.
Just did an execution, and I noticed that allocate_new_blocks called after _maybe_load_prefix_from_cache takes into account the number of blocks assigned during save_new_computed_blocks.

But probably we want to do this in another PR and add steps tests for that

[maxdebayser]

Yes, I think you're right, it makes sense to subtract that. But actually @yannicks1 was also proposing to remove the reserved blocks tracking altogether because the volumetric constraint is a tighter constraint anyway.

[sducouedic]

True, I'm okay with that

[sducouedic]

LGTM!!