[Feat][Core] Support multiple KV cache groups in Hybrid KV Coordinator

[ivanium]

Purpose

The current hybrid KV cache coordinator supports at most two attention types (full attention + another sliding-window/mamba attention). However, emerging models may need more flexible support. For example, full attention + sliding window attn with various window sizes, etc., as required in #31592 and #30263.

Since prefix caching for sliding window and mamba does not have the monotonic prefix cache hit property, viz., a cache hit at position i does not imply a cache hit at position j where j < i, find_longest_cache_hit needs to check all attention groups until a prefix gets cache hits from all of them. This is what implemented by this PR.

Test Plan

pytest tests/v1/core/test_prefix_caching.py -q

Test Result

Passed

---

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.

---

Note

Expands hybrid KV caching beyond the prior 2-type (full+other) limit and introduces a unified cache-hit algorithm across arbitrary attention groups.

• Refactors HybridKVCacheCoordinator to group KV cache groups by identical KVCacheSpec, prioritize FullAttentionSpec, and compute LCM across all block sizes for alignment

• Implements an iterative fixed-point find_longest_cache_hit that iteratively constrains hit length across attention types and reuses cached full-attention hits when possible

• Updates imports to use SingleTypeKVCacheManager and removes the hardcoded 2-group/full-attn assumptions

• Keeps divisibility validation and DCP/PCP constraints; returns per-group hit blocks as a tuple aligned to group indices

• Tests: adds helpers to build mixed-spec configs and a parameterized test_prefill_hybrid_model_combinations covering 2–4 groups, interleaving, sliding-window variants, and Mamba; updates existing hybrid tests and fixtures accordingly

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

---

Note

Expands hybrid KV caching beyond the prior 2-type limit and introduces a unified cache-hit algorithm across arbitrary attention groups.

• Refactors HybridKVCacheCoordinator to group KV cache groups by identical KVCacheSpec, prioritize FullAttentionSpec, and compute LCM across all group block sizes for alignment
• Replaces the 2-group/full-attn assumption with a generic iterative fixed-point find_longest_cache_hit that reconciles hits across all groups and reuses full-attn hits when shrinking
• Updates imports to use SingleTypeKVCacheManager; keeps divisibility validation and DCP/PCP constraints; returns per-group hit blocks aligned to original group indices
• Tests: adds hybrid helpers and a parameterized test_prefill_hybrid_model_combinations covering 2–4 groups, interleaving, sliding-window variants, and Mamba; updates existing hybrid tests accordingly (tests/v1/core/test_prefix_caching.py)

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

[gemini-code-assist]

Code Review

This pull request introduces a significant improvement to the HybridKVCacheCoordinator by enabling support for multiple, interleaved KV cache groups, moving beyond the previous two-type limitation. The core logic in find_longest_cache_hit has been thoughtfully re-implemented with an iterative approach to correctly identify common cache prefixes across various attention mechanisms, including those with non-monotonic hit properties like sliding window and Mamba. The new tests are thorough, covering configurations with three attention types and interleaved group IDs, which validates the increased flexibility. The changes are well-structured and appear robust. Overall, this is a solid enhancement to support more complex model architectures.

[heheda12345]

Per offline discussion, we'll also consider models without full attention in this PR.

[mergify]

Hi @ivanium, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[mergify]

Hi @ivanium, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[mergify]

Hi @ivanium, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[heheda12345]

use kv_cache_spec as key?

[ivanium]

But we need hit_blocks_by_group as a list for return values anyway

[heheda12345]

I think trim full attention in every iteration is more clear and should have similar efficiency

[heheda12345]

What about this flow?

We only need to compute the cache hit length for full attention once. Starting from the second iteration, we can simply keep the first hit_length // block_size blocks of the last iteration where the hit_length is reducing in each step.

[heheda12345]

Should we add a break here? IMO it makes sense to iterate over all groups to get the minimum hit length in every while-loop iteration.

[heheda12345]

LGTM!

[mergify]

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

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