[V1] Move more control of kv cache initialization from model_executor to EngineCore

[heheda12345]

This pr changes the workflow of EngineCore._initialize_kv_caches to enable more flexible control of kv cache format in the future.
It is splitted from #11938 and is a preparation for #11382
Original workflow:

num_gpu_blocks, _ = self.model_executor.determine_num_available_blocks()
self.model_executor.initialize(num_gpu_blocks)

New workflow:

# Get all kv cache tensor needed by the model
kv_cache_spec = self.model_executor.get_kv_cache_spec()

# Profiles the peak memory usage of the model to determine how much
# memory can be allocated for kv cache.
availble_gpu_memory = self.model_executor.get_available_memory()

# Get the kv cache tensor size
kv_cache_config, num_gpu_blocks = get_kv_cache_config(
    vllm_config, kv_cache_spec, availble_gpu_memory)


# Initialize kv cache and warmup the execution
self.model_executor.initialize(kv_cache_config)

This pr introduces 2 new concepts:

1. KVCacheSpec, a data structure to represent the kv cache needed by each attention layer, which is constructed by asking the model runner to analyze all Attention modules. Will add more types of Spec in the future, e.g., SlidingWindowSpec, MLASpec
2. KVCacheConfig, a class to represent how to allocate the kv cache Tensor. It is quite simple now, i.e., tensors with the same size. But it may be extended to the following cases:
   1. tensors with different sizes, to support MLA & spec decode
   2. allocate a global buffer, and make the kv_cache tensors to point to different offsets, to support multiple types of layer sharing the same memory pool.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.
Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can do one of these:

• Add ready label to the PR
• Enable auto-merge.

🚀

[comaniac]

Overall LGTM. Most comments are style changes. One important thing: Please docstring all functions and files in the desired format:

def func_name(arg1, arg2):
    """What does this function do?
    
    Args:
        arg1: ...
        arg2: ...
    
    Returns:
        ... (skip if return None) ...
    """

[comaniac]

Can we try to make this name more informative?

[heheda12345]

KVCacheSpecBase.key -> KVCacheSpecBase.type_id
is_same_key -> is_same_type

[comaniac]

Should this be in the kv_cache_utils.py as it is kv cache related?

[heheda12345]

bind_kv_cache is called by GPUModelRunner. I think it is strange to let GPUModelRunner call a function inside core. These two parts should be independent.

[comaniac]

I don't understand [group_id][layer_name in the group]. What's group ID? It might be better to just show an example.

# A list of kv-cache groups. Each group includes a set of layers with
# the same kv-cache spec. For example: ...

[heheda12345]

Comments updated. Do you feel better now?

[heheda12345]

@comaniac Thank you for the review. I've updated the code based on your suggestions. Can you take another look?

[comaniac]

Should be the last batch of comments. Approve to unblock the PR first.

[comaniac]

I feel this function name is a bit unclear. It actually checks whether the "kv cache specs" of "all" layers are the same, so it should be informative about "spec", "all" and "same". Maybe "is_uniformed_kv_cache_type" or something like that would be better.

[heheda12345]

changed to is_kv_cache_type_uniform

[comaniac]

Again this function name (also the callee _get_kv_cache_config_same_type) doesn't illustrate the number of GPU blocks. It's weird to see config, num_blocks = get_kv_cache_config(...).

One possibility:

def get_kv_cache_config_and_available_blocks(...):
    check_enough_kv_cache_memory(...)

   # Later maybe you can introduce a registry when you have more policies.
    if is_uniformed_kv_cache_type(...):
        return _get_kv_cache_config_and_blocks_for_unifiemd_type(...)
    return _get_kv_cache_config_and_blocks_for_xxx(...)

[heheda12345]

Changed num_blocks to an attribute of KVCacheConfig

[comaniac]

How would this be extended later if you have different specs?

[heheda12345]

It won't be extended. kv_cache_spec[i] is for all layers of one GPU. Different TP GPUs always have the same spec though the spec of each layer inside one GPU can be different.
PP executors of different stages can have different specs.

[comaniac]

Please elaborate more, because the concept of "binding" kv-cache is not common for most people. For example, bind what kv-cache to what, and what's the purpose.

[comaniac]

Add a comment if you're going to extend this logic for xxx; otherwise it looks a bit weird.

[heheda12345]

Changed to raise an error if multiple attention layers have the same layer_index.
Layer_name and layer_index are defined by model implementation instead of this pr. There will be no conflict of layer_index in decoder-only models.

[comaniac]

LGTM. Please also sync the main branch. The SPMD PR may have incompatible changes to this one

[comaniac]

Indent

[heheda12345]

@comaniac Thanks for the review. I've update this PR.

[heheda12345]

Trying to fix the CI failure with #12138

[youkaichao]

this is not really v1-specific, I think we should just make it for v0 and v1 directly. I really hate code duplication, it makes later bugfix and v0-v1 agnostic features difficult to develop.

[heheda12345]

Will move it to v0 & v1 common code if we need these interfaces for v0. But as I don't plan to support Hybrid allocator on v0, I prefer not to introduce these concepts in v0 at this moment.

[youkaichao]

this is heavy code duplication.

[heheda12345]

Running the same test as v0, but the detail workflow is very different due to the dict->list and list->dict difference.

[youkaichao]

this executor is not used. you should add it in vllm/executor/executor_base.py , and do not use _run_workers . use collective_rpc instead.

[heheda12345]

This have been done in your PR #12171 . Thanks for it.