[KV Connector] Canonical KV Cache Allocation for HMA Models by Etelis · Pull Request #37885 · vllm-project/vllm

Etelis · 2026-03-23T11:56:40Z

This is the first phase of a multi-phase effort to enable contiguous KV cache allocation for all model architectures. Currently, only single-group (uniform) models benefit from contiguous cross-layer blocks. This PR extends that to HMA models with uniform page sizes. Future phases will broaden support to models with varying page sizes and additional architectures.

The existing allocate_uniform_kv_caches path only supports single-group models (all layers identical). HMA models like Gemma 3 have multiple KV cache groups (full attention + sliding window) with different eviction policies but the same page size. Previously, these models fell back to per-layer allocation, which scatters block data across non-contiguous memory regions, making RDMA transfers inefficient.

This PR extends contiguous KV cache allocation to HMA models where all KV cache groups share the same page size.

Test plan

Unit tests: pytest tests/v1/kv_connector/unit/test_canonical_kv_caches.py -v -s
- Happy path for use_canonical_kv_caches
- Parametrized rejection cases (single group, no connector, no HMA, mamba layers, no stride order)
- Allocation correctness: shapes, memory sharing, physical contiguity, group refs, page sizes

Related PRs

#34373 — Original KVCacheTopology PR (closed, too complex).
#37339 — WorkerConnectorInitializationData pattern. We adopt their interface design -- Hopefully to be merged after that PR.

mergify · 2026-03-23T11:57:34Z

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

Signed-off-by: Itay Etelis <itay.etelis@ibm.com>

gemini-code-assist

Code Review

This pull request introduces canonical KV cache allocation for Hybrid Multi-Attention (HMA) models, specifically targeting those with uniform page sizes. This is a significant improvement as it enables contiguous cross-layer block allocation, which was previously limited to single-group models. The changes involve new data structures to represent canonical KV caches and their references, along with modifications to the KV cache allocation logic within the gpu_model_runner and kv_connector_model_runner_mixin. A comprehensive unit test suite has been added to validate the new allocation strategy under various conditions, including happy paths and rejection cases. The implementation appears well-considered and robust, addressing the stated goal of improving RDMA transfer efficiency by ensuring memory contiguity for HMA models.

mergify · 2026-03-23T12:09:50Z

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

uv pip install pre-commit>=4.5.1
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 failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

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

Signed-off-by: Itay Etelis <itay.etelis@ibm.com>

orozery · 2026-03-23T12:15:06Z

+        # The connector must support HMA
+        if not supports_hma(get_kv_transfer_group()):
+            return False
+        if len(kv_cache_config.kv_cache_groups) <= 1:


Suggested change

if len(kv_cache_config.kv_cache_groups) <= 1:

if len(kv_cache_config.kv_cache_groups) < 1:

Done,
makes sense.

orozery · 2026-03-23T12:21:11Z

+        if len(kv_cache_config.kv_cache_groups) <= 1:
+            return False
+
+        # All groups must use AttentionSpec with uniform page size


Suggested change

# All groups must use AttentionSpec with uniform page size

# Currently, all groups must use AttentionSpec with uniform page size

# We plan to gradually relax this requirement to support other cases

Thanks, sorry.

orozery · 2026-03-23T12:27:31Z

+        spec = kv_cache_config.kv_cache_groups[0].kv_cache_spec
+        assert isinstance(spec, AttentionSpec)


Can we remove this and use the spec inside the loop per each group?

orozery · 2026-03-23T12:30:37Z

            )
            self.cross_layers_kv_cache = cross_layers_kv_cache
            self.cross_layers_attn_backend = attn_backend
+        elif self.use_canonical_kv_caches(


Let's move this check before checking use_uniform_kv_cache.

orozery · 2026-03-23T13:28:33Z

+        kernel_num_blocks = num_blocks * num_blocks_per_kv_block
+
+        # prepend a group_size dimension into the shape
+        kv_cache_shape = attn_backend.get_kv_cache_shape(


Can we move this logic AFTER we allocate the single tensor?

Then, inside the layer loop, we can reshape?
I think we can also remove assert len(unique_kernel_bs) == 1.

I think it's better to also build the group_data_refs inside the same loop.

orozery · 2026-03-23T13:29:57Z

    @property
    def needs_kv_cache_zeroing(self) -> bool:
        return self.has_mamba_layers
+


These classes are currently specific to connector usage.
I think we should move them to base.py.

Signed-off-by: Itay Etelis <itay.etelis@ibm.com>

…nector Signed-off-by: Itay Etelis <itay.etelis@ibm.com>

orozery · 2026-03-23T16:49:10Z

+                    WorkerConnectorInitializationData,
+                )
+
+                kv_transfer_group.initialize_worker_connector(


Actually, initialize_worker_connector is needed for the CacheBlend use-case.
Let's try to call it exactly as in #37339.
But keep this if here and simply pass, commenting that the canonical kv caches will be registered below.

I thought they'd add it themselves afterwards,
nvm I will fix it.

Combine the kv_caches population, block tensor splitting, and layer-to-position mapping into a single pass over positions. Remove the unique kernel block size assertion. Signed-off-by: Itay Etelis <itay.etelis@ibm.com>

Call initialize_worker_connector unconditionally so connectors like CacheBlend can use it regardless of the allocation path taken. Signed-off-by: Itay Etelis <itay.etelis@ibm.com>

orozery · 2026-03-24T08:22:20Z

+                canonical_kv_caches is the CanonicalKVCaches wrapping
+                    for the connector.
+        """
+        # all tensors have the same size (validated by use_canonical_kv_caches)


Where did we validate this?

@orozery sharp eye

fixed.

Move the uniform tensor size check into use_canonical_kv_caches so the precondition is validated before entering the allocation path, keeping the assert in allocate_canonical_kv_caches as a safety net. Signed-off-by: Itay Etelis <itay.etelis@ibm.com>

orozery · 2026-05-05T17:09:26Z

Tested performance with gpt-oss-20b on H100:
Time to complete 128K token prefill loading from CPU (offloading connector):
main: 286ms
this PR: 56ms

NickLucche

I guess one main question, why aren't we folding this into the cross-layer functionality as a natural extension to hma?
It feels like things are tied in the current formulation (canonical implies cross-layer) and yet there ares separate branches in shared code path like the runner

NickLucche · 2026-05-06T07:43:04Z

+class CanonicalKVCaches:
+    """
+    Canonicalized block-level representation of the KV caches.
+
+    Composed of:
+        - Unique list of KV cache data tensors,
+          each with shape (num_blocks, page_size_in_bytes) and int8 dtype.
+        - Per-group data references of the tensors.
+          i.e. how each KV cache group maps to the tensors.
+    """
+
+    # Ordered list of unique block tensors, each with shape
+    # (num_blocks, ...).
+    tensors: list[KVCacheBlockTensor]
+    # Per-KV-cache-group list of data references that map each layer
+    # in the group to the appropriate entry in the tensors list.
+    group_data_refs: list[list[KVCacheBlockDataRef]]


I am not sure these dataclasses about tensors belong here with kv_connector interface. They look a lot more related to whats in kv_cache_manager.py.

I'd rather keep this file lean for the actual interface.

Connectors need a way to know how to access the KV cache tensors.
Currently, connectors have 2 tasks:

Determine the topology for each KV cache tensor

Determine how each group maps to each KV cache tensor (using KVCacheConfig)

Using the canonical KV caches saves connectors these 2 tasks:

All tensors are (num_blocks, ) first

group_data_refs describes how each group maps to tensors.

With cross-layers layout you cannot use KVCacheConfig as the tensors (single one) do not match kv_cache_config.kv_cache_tensors.

NickLucche · 2026-05-06T07:49:41Z

+        if len(page_sizes) != 1:
+            return False


isn't this case unexpected if we take UniformTypeKVCacheSpecs out of the equation?
We should probably assert or at least log

This is the use_canonical_kv_caches function.
The purpose of this function is to determine if should allocate cross-layers or use the regular allocation.
If the spec is UniformTypeKVCacheSpecs we should return False to disable cross layers allocation (for now).

NickLucche · 2026-05-06T07:52:32Z

+                # num_blocks must be the leading physical dimension.
+                # +1 accounts for the prepended group_size dimension.
+                if stride_order[0] != kv_cache_shape.index(1234) + 1:
+                    return False


can we unify to use get_kv_cache_block_dim to get the block dim?

We already get the kv_cache_shape to check if cross-layers is supported, so it seems easier and more efficient to leave this check here.

NickLucche · 2026-05-06T07:56:53Z

+                if len(stride_order) != len(kv_cache_shape) + 1:
+                    return False


I feel like some of this invariants could be asserted by use_uniform_kv_cache per-group?

use_uniform_kv_cache should actually be deprecated, along with register_cross_layers_kv_cache.

I suggest that we remove it.
Connectors using prefers_cross_layers_block will simply fall back to the regular register_kv_caches.

What do you think?

@orozery I am ok with removing use_uniform_kv_cache, just in a separate PR

orozery · 2026-05-13T04:36:22Z

@Etelis I think it's better to move the implementation from MRV1 to MRV2, for 2 reasons:

MRV2 will soon be the default model runner.
It will look cleaner in MRV2 since we'll only have the new cross layers API.

… (V4) OffloadingConnectorWorker.register_kv_caches assumed every layer inside a single KVCacheTensor.shared_by group had the same physical layout (same tensor count, same data_ptr, same stride). That is true on uniform-attention models, but DeepSeek V4's _get_kv_cache_config_deepseek_v4 (cherry-pick of PR vllm-project#37885) intentionally builds one KVCacheTensor whose shared_by spans the three V4 attention groups (MLA, Compress-4, Compress-128). Those layers view the same physical memory through different attn-backend shapes -- a 1-tuple vs a 2-tuple, with different strides -- so the three uniformity asserts fire during engine init: AssertionError at worker.py:197 len({len(tensors_per_block[n]) for n in tensor_layer_names}) == 1 Fix: bucket layers by (tuple_arity, stride). If all layers share the same signature, behave exactly as before. Otherwise (V4 hybrid), walk each layer separately -- register one canonical tensor per layer's own layout and map its CanonicalKVCacheRef to that canonical entry. The homogeneous branch is bit-for-bit equivalent to the previous code, so non-V4 models are unaffected. This unblocks V4-Flash + OffloadingConnector + --no-disable-hybrid-kv-cache-manager (the 6.3x GPU KV pool gain). Companion fix to PR vllm-project#42992 (math.gcd hash_block_size derivation); both are required end-to-end. Signed-off-by: Kfir Toledo <kfir.toledo@ibm.com>

mergify · 2026-05-23T07:32:53Z

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

…ical KV path PR vllm-project#37885's gpu_model_runner branch skips register_kv_caches() when the canonical KV cache layout is active, routing setup through initialize_worker_connector() instead. The PR added the abstract method but did not implement it on OffloadingConnector, so transfer handlers were never registered and the first transfer_async asserted with an empty transfer_type_to_handler map: worker.py:133: assert handler is not None AssertionError Wire initialize_worker_connector to call connector_worker._register_handlers with the canonical_kv_caches from the initialization data so the SharedStorage handler registers as expected. Also add an info log when the canonical KV path activates, for easier runtime confirmation during benchmarking.

LucasWilkinson · 2026-06-02T04:33:30Z

@Etelis @orozery @NickLucche Thanks for all the hard work! In the interest in getting this perf improvements into a usable release I think it may make sense to land this for 0.23 and then land #42374 quickly after the 0.23 cut giving use some time to retest and harden this with the refactored layout

orozery · 2026-06-02T05:01:05Z

@Etelis @orozery @NickLucche Thanks for all the hard work! In the interest in getting this perf improvements into a usable release I think it may make sense to land this for 0.23 and then land #42374 quickly after the 0.23 cut giving use some time to retest and harden this with the refactored layout

This PR introduces a new API for registering kv caches, while #42374 takes a different approach.
The problem I see is we don't want to play back-and-forth on the connector API.

# Conflicts: # vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

LucasWilkinson · 2026-06-03T03:37:20Z

@Etelis @orozery @NickLucche Thanks for all the hard work! In the interest in getting this perf improvements into a usable release I think it may make sense to land this for 0.23 and then land #42374 quickly after the 0.23 cut giving use some time to retest and harden this with the refactored layout

This PR introduces a new API for registering kv caches, while #42374 takes a different approach. The problem I see is we don't want to play back-and-forth on the connector API.

The idea would be just to have this in interim as it may take a while #42374 to land, we can label the API as deprecated if that helps?

orozery · 2026-06-03T04:12:33Z

@Etelis @orozery @NickLucche Thanks for all the hard work! In the interest in getting this perf improvements into a usable release I think it may make sense to land this for 0.23 and then land #42374 quickly after the 0.23 cut giving use some time to retest and harden this with the refactored layout

This PR introduces a new API for registering kv caches, while #42374 takes a different approach. The problem I see is we don't want to play back-and-forth on the connector API.

The idea would be just to have this in interim as it may take a while #42374 to land, we can label the API as deprecated if that helps?

@NickLucche @markmc what do you guys think?

NickLucche

My view: I like this as an extension/replacement to existing cross layer methods.

What I like less is handling things like

                # prepend a group_size dimension into the shape
                full_shape = (group_size,) + kv_cache_shape

or similarly

                if len(stride_order) != len(kv_cache_shape) + 1:

which feel brittle and makes the feature feel like a separate thing/not properly integrated in the backend.
Which I believe #42082 could help in addressing, moving things to attn backend under the "BHLNC" semantics.

Second thing, I am not sure how to handle heteroTP here for day0 model releases. #42082 is making an effort to move the logic into pieces that one could extend and provide along model definition.
I feel using canonical view here wouldn't provide much extensibility in this scenario.

Also not a fan of interim solution, so I think we should try to either integrate it with #42082 or keep it as a physical allocation optimization for offloading.
API changes are lgtm, I think we need initialize_worker_connector anyways, so second option still remains valid in my view.

NickLucche · 2026-06-03T08:11:02Z

+        if len(kv_cache_config.kv_cache_groups) < 1:
+            return False


how can it have no groups?

NickLucche · 2026-06-03T08:13:06Z

+                if len(stride_order) != len(kv_cache_shape) + 1:
+                    return False


@orozery I am ok with removing use_uniform_kv_cache, just in a separate PR

NickLucche · 2026-06-03T08:14:57Z

+                except (AttributeError, NotImplementedError):
+                    return False


we should log the fallback on unexpected cases

mergify · 2026-06-21T15:00:01Z

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

Etelis requested review from ApostaC, NickLucche, heheda12345, njhill and orozery as code owners March 23, 2026 11:56

mergify Bot added the v1 label Mar 23, 2026

mergify Bot added needs-rebase kv-connector labels Mar 23, 2026

EtelisIBM added 5 commits March 23, 2026 13:57

Add CanonicalKVCaches data classes for HMA KV cache representation

52688de

Signed-off-by: Itay Etelis <itay.etelis@ibm.com>

Add WorkerConnectorInitializationData and initialize_worker_connector

e25e020

Signed-off-by: Itay Etelis <itay.etelis@ibm.com>

Add canonical KV cache allocation for HMA models

01b7897

Signed-off-by: Itay Etelis <itay.etelis@ibm.com>

Wire up canonical KV cache allocation in gpu_model_runner

03903f1

Signed-off-by: Itay Etelis <itay.etelis@ibm.com>

Add unit tests for canonical KV cache allocation

9697d1c

Signed-off-by: Itay Etelis <itay.etelis@ibm.com>

Etelis force-pushed the canonical-kv-caches branch from aa5d414 to 9697d1c Compare March 23, 2026 11:58

mergify Bot removed the needs-rebase label Mar 23, 2026

gemini-code-assist Bot reviewed Mar 23, 2026

View reviewed changes

Fix mypy error: rename shadowed variable in use_canonical_kv_caches

d9f7203

Signed-off-by: Itay Etelis <itay.etelis@ibm.com>

orozery requested changes Mar 23, 2026

View reviewed changes

EtelisIBM added 3 commits March 23, 2026 15:53

Move canonical KV cache dataclasses to connector base

f39ae32

Signed-off-by: Itay Etelis <itay.etelis@ibm.com>

Address CR: relax group count check and use per-group spec

68ce39a

Signed-off-by: Itay Etelis <itay.etelis@ibm.com>

Address CR: prioritize canonical path and scope initialize_worker_con…

5b2b3bc

…nector Signed-off-by: Itay Etelis <itay.etelis@ibm.com>

Etelis force-pushed the canonical-kv-caches branch from db277a8 to 5b2b3bc Compare March 23, 2026 13:54

orozery reviewed Mar 23, 2026

View reviewed changes

EtelisIBM added 2 commits March 23, 2026 19:23

Address CR: merge loops in allocate_canonical_kv_caches

3f90424

Combine the kv_caches population, block tensor splitting, and layer-to-position mapping into a single pass over positions. Remove the unique kernel block size assertion. Signed-off-by: Itay Etelis <itay.etelis@ibm.com>

Address CR: always call initialize_worker_connector

432d002

Call initialize_worker_connector unconditionally so connectors like CacheBlend can use it regardless of the allocation path taken. Signed-off-by: Itay Etelis <itay.etelis@ibm.com>

Etelis force-pushed the canonical-kv-caches branch from 30cca9f to 432d002 Compare March 23, 2026 17:27

orozery reviewed Mar 24, 2026

View reviewed changes

NickLucche reviewed May 6, 2026

View reviewed changes

markmc changed the title ~~Canonical KV Cache Allocation for HMA Models~~ [KV Connector] Canonical KV Cache Allocation for HMA Models May 6, 2026

This was referenced May 10, 2026

[P/D][Mooncake] Add cross-layer KV cache support to MooncakeConnector #41093

Open

[RFC]: Standardize KV-cache Layouts #42082

Open

Merge branch 'main' into canonical-kv-caches

afd03b6

orozery mentioned this pull request May 14, 2026

[Core][WIP][1/N] Standardize kv layout #42374

Open

mergify Bot added the needs-rebase label May 23, 2026

LucasWilkinson added 2 commits June 2, 2026 17:22

Merge remote-tracking branch 'origin/main' into canonical-kv-caches

73ab712

# Conflicts: # vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

fixes

76cefbd

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

mergify Bot removed the needs-rebase label Jun 2, 2026

LucasWilkinson added the ready ONLY add when PR is ready to merge/full CI is needed label Jun 3, 2026

Etelis added 5 commits June 3, 2026 07:24

Merge branch 'main' into canonical-kv-caches

f567c43

Merge branch 'main' into canonical-kv-caches

6220b0b

Merge branch 'main' into canonical-kv-caches

99c1c97

Merge branch 'main' into canonical-kv-caches

73edcc8

Merge branch 'main' into canonical-kv-caches

fba712c

NickLucche reviewed Jun 3, 2026

View reviewed changes

Merge branch 'main' into canonical-kv-caches

757e0fc

LucasWilkinson mentioned this pull request Jun 20, 2026

[KV Offload] Support packed HMA KV cache layout #46205

Merged

mergify Bot added the needs-rebase label Jun 21, 2026

	if len(kv_cache_config.kv_cache_groups) <= 1:
	if len(kv_cache_config.kv_cache_groups) < 1:

	# All groups must use AttentionSpec with uniform page size
	# Currently, all groups must use AttentionSpec with uniform page size
	# We plan to gradually relax this requirement to support other cases

		spec = kv_cache_config.kv_cache_groups[0].kv_cache_spec
		assert isinstance(spec, AttentionSpec)

		if len(stride_order) != len(kv_cache_shape) + 1:
		return False

Uh oh!

Uh oh!

Conversation

Etelis commented Mar 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Test plan

Related PRs

Uh oh!

mergify Bot commented Mar 23, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

mergify Bot commented Mar 23, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

orozery commented May 5, 2026

Uh oh!

NickLucche left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

orozery commented May 13, 2026

Uh oh!

mergify Bot commented May 23, 2026

Uh oh!

LucasWilkinson commented Jun 2, 2026

Uh oh!

orozery commented Jun 2, 2026

Uh oh!

LucasWilkinson commented Jun 3, 2026

Uh oh!

Etelis commented Mar 23, 2026 •

edited

Loading