[Hybrid Allocator] Support full attention with different hidden size

[heheda12345]

Purpose

Hybrid allocator requires all layers have the same kv_hidden_size. But some model breaks this assumption, e.g.,

1. spec decode, the MTP layer and the main model, like [Feature]: Support Eagle Draft Model with different number of KV heads #22432
2. recent native sparse attention like minicpm 4.1. It needs to run attention.0 on a smaller KV cache to select tokens, and then do attention.1 only on the selected tokens. Both attention.0 and attention.1 can be regarded as full attention, but with different hidden size.

If the model only have one type of kv cache (e.g., all layers are full attention / all layers are sliding window with the same window size), KVCacheManager can regard this model as only having one layer of that type (i.e., only one KVCacheGroup with all layers in it).

This PR supports this kind of KVCacheGroup by introducing a UniformTypeKVCacheSpecs

The model runner is updated to support different kv_cache_specs within one group. The layers will be split to different attn_groups and have different attention metadata builder based on their kv_cache_spec.

Test Plan

Initialize the model with

llm = LLM(
    model="mistralai/Mixtral-8x7B-Instruct-v0.1",
    enforce_eager=True,
    speculative_config={"model": "yuhuili/EAGLE-mixtral-instruct-8x7B", "num_speculative_tokens": 3, "method":"eagle"},
    tensor_parallel_size=4
)

and run basic.py

Test Result

Can generate meaningful result

Fix #22432

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request introduces support for models with varying hidden sizes across layers by adding UniformTypeKVCacheSpecs. This is a valuable feature, particularly for speculative decoding and certain sparse attention models. The changes are well-structured, with necessary refactoring in the model runner and the addition of relevant tests. I have identified one issue where a configuration option is not being applied in the new code path, which could lead to incorrect behavior.

[gemini-code-assist]

The num_gpu_blocks_override configuration is not being respected in this new code path for UniformTypeKVCacheSpecs. The calculation for num_blocks should use the get_num_blocks helper function, similar to the else branch, to ensure that user-provided overrides are applied correctly.

        page_size = kv_cache_groups[0].kv_cache_spec.page_size_bytes
        num_blocks = get_num_blocks(vllm_config,
                                    1,
                                    available_memory=available_memory,
                                    page_size=page_size)

[heheda12345]

@LucasWilkinson Why didn't we add kv_cache_spec as a key to split kv cache groups to attention groups?

[heheda12345]

this assert is in generate_scheduler_kv_cache_config now.