OffloadingConnector: Add cpu_bytes_to_use configuration

[orozery]

This PR replaces the OffloadingConnector size configuration from num_cpu_blocks to cpu_bytes_to_use.
This allows for a more intuitive space allocation (per vLLM instance, across workers).

---

Note

Modernizes KV offloading configuration and wiring.

• Replace num_cpu_blocks/kv_bytes_per_rank with instance-wide cpu_bytes_to_use for OffloadingConnector (docs and configs)
• CPUOffloadingSpec now derives num_blocks from KVCacheConfig (page size, tensors, world size) and block_size; requires passing kv_cache_config into OffloadingSpecFactory.create_spec(...) and OffloadingConnector
• Update VllmConfig._post_init_kv_transfer_config to set cpu_bytes_to_use = kv_offloading_size * (1 << 30) (no per-rank split); LMCache path unchanged
• Tests adjusted for new config keys and spec signatures (MockOffloadingSpec, unit/integration offloading tests)

^{Written by Cursor Bugbot for commit 8208b41. This will update automatically on new commits. Configure here.}

[gemini-code-assist]

Warning

Gemini encountered an error creating the review. You can try again by commenting /gemini review.

[ApostaC]

Otherwise LGTM.

[ApostaC]

There is also a knob called offloaded_block_size in #22595. IIUC, it also impacts the calculation of num_cpu_blocks, right? (i.e., if we have larger CPU blocks, we should have less number of CPU blocks)

[orozery]

In v0, the offloading was part of the core.
My suggestion for v1 is to have the offloading as a connector.
I wanted to follow the convention for connectors, where all of their arguments are actually defined in their kv_connector_extra_config.

However, deriving num_cpu_blocks from some kind of a swap_space parameter requires knowledge of kv_bytes_per_block.
So basically, I need my connector (both scheduler-side and worker-side) to be aware of kv_bytes_per_block.
This requires changing things in core, so I tried to make minimal changes and came up with the approach here:

For the scheduler-side connector, report kv_bytes_per_block by setting the existing V0 field num_cpu_blocks.
For the worker-side connector, pass-on kv_cache_configs via the register_kv_caches function (in a follow-up PR).

When the offloading connector gets this num_cpu_blocks (given in GPU block size), it can derive the actual num_cpu_blocks by dividing by block_size_factor.

To sum-up, I'm trying to make minimal changes to the core.
This results in the actual offloading configuration parameters split between vllm_config.cache_config and kv_connector_extra_config.

I'm good with taking a different approach.
Your thoughts?
Perhaps we should ask other relevant folks on their opinion here?

[ApostaC]

Yeah. This is a good point. I think at a high level, there should be two parameters that can be configured by users: (1) total_cpu_buffer_size and (2) cpu_buffer_block_size (how many tokens in each CPU block).

For (1), it's also worth thinking whether it's per rank or per vLLM instance (i.e., summed across all ranks). I feel like if it's per rank, probably it will be better to pass it in the KV connector configs, while it makes more sense to have a "global" cache size when it's configured by global configurations like --swap-space.

For (2), I think it should definitely be put into the KV connector config as it's the current CPU-offloading-connector-specific configuration.

To sum up, I feel like putting all the configs into the KV connector config will probably be better and less confusing. WDYT?

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @orozery.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[orozery]

I'm planning to change this PR next week to share the raw kv_bytes_per_block to the connectors (both scheduler and worker side), and to omit the use of the legacy swap_block_bytes.

[orozery]

Thanks for doing this @orozery !

Generally speaking, I am not a huge fan of the pattern we're using to update the config as a way to share global values that are runtime dependent.

Following your comment, I've spent some time thinking on it.
I still think we need the kv_bytes_per_block field added to KVCacheConfig (and @heheda12345 can sure confirm or prove me wrong, as it is due to hybrid models).
On the other hand, I was able to get rid of the introducing of new lazy-initialized fields to the cache_config.
(By noticing that the connector API on both sides has access to KVCacheConfig via its init.

So PR is modified now and is much leaner.
Basically, the only "intrusive" part is in kv_cache_utils.py (and 1 line in kc_cache_interface.py).

[heheda12345]

For passing in KVCacheConfig, we now have an optional kv_cache_config field in KVConnectorBase_V1. And for page size per block, you can get it from kv_cache_config.kv_cache_groups[0].kv_cache_spec.page_size_bytes. Now the page_size_bytes of all kv_cache_groups are the same and you can also add an assert for this assumption.

[orozery]

For passing in KVCacheConfig, we now have an optional kv_cache_config field in KVConnectorBase_V1. And for page size per block, you can get it from kv_cache_config.kv_cache_groups[0].kv_cache_spec.page_size_bytes. Now the page_size_bytes of all kv_cache_groups are the same and you can also add an assert for this assumption.

Thanks!

This is what I do now:

        page_sizes = {
            kv_cache_group.kv_cache_spec.page_size_bytes
            for kv_cache_group in kv_cache_config.kv_cache_groups
        }
        assert len(page_sizes) == 1
        page_size_bytes = page_sizes.pop()
        kv_bytes_per_block = (
            page_size_bytes
            * len(kv_cache_config.kv_cache_tensors)
            * vllm_config.parallel_config.world_size
        )

@heheda12345 does this look ok to you?

@NickLucche I removed all intrusive changes, so hopefully we're good to go?

[heheda12345]

For passing in KVCacheConfig, we now have an optional kv_cache_config field in KVConnectorBase_V1. And for page size per block, you can get it from kv_cache_config.kv_cache_groups[0].kv_cache_spec.page_size_bytes. Now the page_size_bytes of all kv_cache_groups are the same and you can also add an assert for this assumption.

Thanks!
This is what I do now:

        page_sizes = {
            kv_cache_group.kv_cache_spec.page_size_bytes
            for kv_cache_group in kv_cache_config.kv_cache_groups
        }
        assert len(page_sizes) == 1
        page_size_bytes = page_sizes.pop()
        kv_bytes_per_block = (
            page_size_bytes
            * len(kv_cache_config.kv_cache_tensors)
            * vllm_config.parallel_config.world_size
        )

@heheda12345 does this look ok to you?
@NickLucche I removed all intrusive changes, so hopefully we're good to go?

Yes, sounds good!

[ApostaC]

Thanks for the update. Removing my previous "request changes"

Dismiss my previous outdated reviews

[esmeetu]

@ApostaC @heheda12345 Is this PR ready to be merged? Hopefully can goes into v0.14.0.

[NickLucche]

Leaner, thanks for the work @orozery !

[NickLucche]

nit: we should probably follow up with a more human readable way of expressing the value

[orozery]

I'm open for suggestions :)

[NickLucche]

thinking about unifying with max-model-len format

[orozery]

@NickLucche rebased