[Feature] Enable uniform KV cache allocation for multi-group HMA models

[Etelis]

use_uniform_kv_cache() currently rejects any model with more than one KV cache group, which means hybrid-attention models (alternating full + sliding-window layers) cannot use the contiguous cross-layer layout for efficient KV transfers.

This PR relaxes the single-group gate: instead of requiring exactly one group, we loop over all groups and check that they share the same backend shape and stride order.

Test Plan

Unit tests (tests/v1/kv_connector/unit/test_uniform_kv_cache.py) — 4 tests:

test_uniform_kv_cache.py::test_multi_group_compatible PASSED
test_uniform_kv_cache.py::test_multi_group_incompatible PASSED
test_uniform_kv_cache.py::test_allocate_multi_group_shared_tensors PASSED
test_uniform_kv_cache.py::test_allocate_rejects_mismatched_kernel_block_sizes PASSED

Test Result

E2E on an H100 with google/gemma-2-2b:

Single-group regression (HMA disabled, OffloadingConnector):

INFO  Allocating a cross layer KV cache of shape (7842, 26, 2, 16, 4, 256)
                                                         ^^
                                                    all 26 layers in 1 group
Prompt: 'The capital of France is'
Output: ' a city of contrasts. It is a city of art, culture, and history. It is a'

Multi-group (HMA enabled, SupportsHMA test connector):

INFO  Allocating a cross layer KV cache of shape (15685, 13, 2, 16, 4, 256)
                                                          ^^
                                                     13 layers per group (2 groups)
[HMATestConnector] register_cross_layers_kv_cache:
  tensor shape=torch.Size([15685, 13, 2, 16, 4, 256]), dtype=torch.bfloat16
  backend=FlashAttentionBackend

Prompt: 'The capital of France is'
Output: ' a city of contrasts. It is a city of art, culture, and history. It is a'

[gemini-code-assist]

Code Review

This pull request enables uniform KV cache allocation for models with multiple attention groups, such as Hybrid-Attention Models (HMA). This is achieved by relaxing the single-group constraint and instead checking for compatibility (e.g., same shape and stride order) across all groups. The changes in use_uniform_kv_cache are logical and well-supported by a comprehensive new test suite. I have one suggestion to make the compatibility check more robust by also ensuring all attention groups use the same backend, which is implied by the docstring and the subsequent allocation logic.

[gemini-code-assist]

The docstring for use_uniform_kv_cache states that for a uniform layout, all KV cache groups must have the same backend. However, the current implementation only checks for compatible kv_cache_shape and kv_cache_stride_order, but not that the attn_backend is the same across all groups.

The subsequent allocate_uniform_kv_caches function uses the backend from the first attention group, which could lead to incorrect behavior or runtime errors if other groups use a different backend.

To prevent this potential issue and align with the documentation, I suggest also checking that all attention groups share the same backend instance.

        if not attn_groups:
            return False

        reference_shape = None
        reference_stride_order = None
        reference_backend = None

        for subgroups in attn_groups:
            if len(subgroups) != 1:
                return False

            attn_group = subgroups[0]
            kv_cache_spec = attn_group.kv_cache_spec
            if not isinstance(kv_cache_spec, AttentionSpec):
                return False

            attn_backend = attn_group.backend
            kv_cache_shape = attn_backend.get_kv_cache_shape(
                1234,
                kv_cache_spec.block_size,
                kv_cache_spec.num_kv_heads,
                kv_cache_spec.head_size,
                cache_dtype_str=cache_dtype,
            )

            try:
                kv_cache_stride_order = attn_backend.get_kv_cache_stride_order(
                    include_num_layers_dimension=True
                )
            except (AttributeError, NotImplementedError):
                return False

            # check that attention backend includes a layers dimension
            if len(kv_cache_stride_order) != len(kv_cache_shape) + 1:
                return False

            if reference_backend is None:
                reference_shape = kv_cache_shape
                reference_stride_order = kv_cache_stride_order
                reference_backend = attn_backend
            elif (
                kv_cache_shape != reference_shape
                or kv_cache_stride_order != reference_stride_order
                or attn_backend is not reference_backend
            ):
                return False

        return True

[orozery]

Looks great for a start! Thanks @Etelis !
I also tested this on my side and it seems to work well.
As expected, KV transfer performance is greatly improved!

IIUC right now you handle the case of multiple groups, but requiring:

1. A single attention group per KV group.
2. Consistent KV cache shape and stride order across all layers
   I think we can try and relax the constraints for using cross-layer blocks even further.

Take a look at the current options for defining KV cache tensors:
https://github.com/vllm-project/vllm/blob/8a798be929d62a6467fd079c03c83632f8231b11/vllm/v1/core/kv_cache_utils.py#L1095-1140

There are 2 cases:

1. A single group containing all layers, but layers can have different page (hidden) size.
2. Multiple groups but all layers have the same page size.

For case 1 (which is less common I think), we can group layers by their page size and create a single tensor per each group of layers with the same page size.
A group of layers having the same page size (and the same number of blocks, which is always true) can always share a tensor, even if they have different KV cache shapes / stride order.
The only requirement in this case is the cache shape to start with (num_blocks, ...).
In that case, you can set a physical layout of (num_blocks, num_layers, page_size).
Then, each layer will create a different view of it: (num_blocks, num_layers, rest_of_unique_logical_shape...)
And finally, create a new view that fixates the layer dimension to the layer_idx.
Then, each layer can permute according to its stride_order (setting with_layers_dim=False).

Does that make sense?

[Etelis]

Looks great for a start! Thanks @Etelis ! I also tested this on my side and it seems to work well. As expected, KV transfer performance is greatly improved!

IIUC right now you handle the case of multiple groups, but requiring:

1. A single attention group per KV group.

2. Consistent KV cache shape and stride order across all layers
   I think we can try and relax the constraints for using cross-layer blocks even further.

Take a look at the current options for defining KV cache tensors: https://github.com/vllm-project/vllm/blob/8a798be929d62a6467fd079c03c83632f8231b11/vllm/v1/core/kv_cache_utils.py#L1095-1140

There are 2 cases:

1. A single group containing all layers, but layers can have different page (hidden) size.

2. Multiple groups but all layers have the same page size.

For case 1 (which is less common I think), we can group layers by their page size and create a single tensor per each group of layers with the same page size. A group of layers having the same page size (and the same number of blocks, which is always true) can always share a tensor, even if they have different KV cache shapes / stride order. The only requirement in this case is the cache shape to start with (num_blocks, ...). In that case, you can set a physical layout of (num_blocks, num_layers, page_size). Then, each layer will create a different view of it: (num_blocks, num_layers, rest_of_unique_logical_shape...) And finally, create a new view that fixates the layer dimension to the layer_idx. Then, each layer can permute according to its stride_order (setting with_layers_dim=False).

Does that make sense?

Thanks @orozery! Went ahead and implemented this.

The main idea is what you described — layers are grouped by page_size_bytes, and each group gets a contiguous int8 tensor of shape (num_blocks, num_layers_in_group, page_size_bytes). Both your cases are covered: different hidden sizes within one group produce separate sub-groups by page size, and multiple groups with the same page size share one tensor.

For per-layer views I'm doing the view → slice → permute pipeline you suggested:

raw.view(dtype).view(kernel_num_blocks, num_layers, kernel_page_elements)
[:, layer_idx]
permute(...)  # each layer uses its own stride order

I also extended this to handle MambaSpec layers — those use as_strided with cross-layer-aware strides so the per-block contiguity is preserved for transfers. (hopefully I understood it correctly)

Tested on H100 with a hybrid setup (4 attention + 2 Mamba layers, 2 page_size groups), all passing. No connector changes yet — will handle nixl_connector/utils.py for multi-group registration as a follow-up.

[Etelis]

Checked with GPT-OSS-20B, alternating 12 SW + 12 FA(layer_types: [sliding_attention, full_attention, ...]).

All 24 layers share: (num_kv_heads=8, head_size=64, block_size=16),

so they should merge into a single cross-layer group.

H100 80GB — TP=1 (default layout)

Group 0:
  tp_layout:       False
  tensor shape:    (136756, 12, 32768)
  page_size_bytes: 32768
  num layers:      24  (12 sliding + 12 full, paired into 12 shared tensors)
  spec type:       SlidingWindowSpec
  backend:         FlashAttentionBackend
  num_kv_heads:    8
  head_size:       64
  strides:         (393216, 32768, 1)
  per-layer view:  (2, 136756, 16, 8, 64) bfloat16

Model loaded and generated text.

8x RTX 3090 — TP=2 (TP layout)

With tensor_parallel_size=2, (4 heads per GPU). The grouping key becomes ("tp", 4, 4096)

Group 0:
  tp_layout:       True
  tensor shape:    (66511, 4, 12, 4096)
  page_size_bytes: 16384
  num layers:      24  (12 sliding + 12 full, paired into 12 shared tensors)
  spec type:       SlidingWindowSpec
  backend:         TritonAttentionBackend (had an issue with FLASHINFER on this setup)
  num_kv_heads:    4   (8 total / TP=2)
  head_size:       64
  strides:         (196608, 49152, 4096, 1)
  per-layer view:  (66511, 2, 16, 4, 64) bfloat16

[mergify]

Hi @Etelis, 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 @Etelis, 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 @Etelis, 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

[Etelis]

Resolved all CRs @orozery.
Hope it looks fine now.

[mergify]

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

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