@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 6700813f86

[3/N][KV-Cache Layout Refactor] Standardize Mamba cache; drop `get_transfer_cache_regions` (#44456)

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: OpenAI Codex <codex@openai.com>

Lucas Wilkinson committed Jul 21, 2026 at 11:16 UTC 6700813f86552a2e0787989dd9f468fddcb88bac

10 files changed +115 -225

tests/v1/kv_connector/unit/offloading_connector/test_worker.py

+1 -2

@@ -125,8 +125,7 @@ def test_register_kv_caches(backend):

125 own dedicated tensors.

126

127 Uses the real GPUModelRunner.initialize_kv_cache_tensors to produce

128 - kv_caches, which automatically applies

129 - _update_hybrid_attention_mamba_layout for hybrid models.

128 + the raw per-layer kv_caches registered by the connector.

129

130 Verifies that the canonicalized CanonicalKVCaches has the correct

131 block tensors, tensor_idx references, and page sizes across all groups.

vllm/distributed/kv_transfer/kv_connector/utils.py

-28

@@ -21,12 +21,10 @@ from vllm.logger import init_logger

21 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase

22 from vllm.platforms import current_platform

23 from vllm.v1.attention.backend import AttentionBackend

24 -from vllm.v1.kv_cache_interface import MambaSpec

24 from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput

26 if TYPE_CHECKING:

27 from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBase

29 - from vllm.v1.kv_cache_interface import KVCacheSpec

29 logger = init_logger(__name__)

@@ -594,32 +592,6 @@ class TransferTopology:

592 abs_ratio = -tp_ratio

593 return [self.tp_rank * abs_ratio + i for i in range(abs_ratio)]

594

597 - def get_transfer_cache_regions(

598 - self, cache: torch.Tensor, layer_spec: "KVCacheSpec"

599 - ) -> list[torch.Tensor] | torch.Tensor:

600 - """Return the cache tensor(s) to register as NIXL memory regions,

601 - also accounting for hybrid SSM models specificities.

602 - """

603 - if isinstance(layer_spec, MambaSpec):

604 - # Register the whole kv cache shared tensor, including

605 - # SSM/Conv.

606 - conv, ssm = cache

607 - return [conv]

608 -

609 - # Check may be hacky but it's matching

610 - # `_update_hybrid_attention_mamba_layout`.

611 - if self.is_mamba and cache.shape[0] == 2:

612 - # When MAMBA is present, all backends are blocks first, so

613 - # that blocks can be shared between attention layers and mamba

614 - # layers. Runner already adjusted strides for FlashAttn-like

615 - # backends so its num_blocks first.

616 - # Swap [2<>num_blocks] dims for hybrid SSM layout.

617 - cache = cache.transpose(0, 1)

618 -

619 - # K and V are packed into one tensor (content dim), so each layer

620 - # registers as a single region.

621 - return [cache]

622 -

595 def describe(self, remote_engine_id: EngineId, remote_pp_rank: int = 0) -> str:

596 """One-line summary of transfer config for logging."""

597 info = self._engines[(remote_engine_id, remote_pp_rank)]

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py

+3 -3

@@ -1678,9 +1678,9 @@ class MooncakeConnectorWorker:

1678 conv, _ = cache_or_caches

1679 cache_list = [conv]

1680 else:

1681 - cache_list = self.transfer_topo.get_transfer_cache_regions(

1682 - cache_or_caches, layer_spec

1683 - )

1681 + # K and V are packed into one blocks-first tensor per layer,

1682 + # so each layer registers as a single region.

1683 + cache_list = [cache_or_caches]

1684

1685 logger.debug(

1686 "registering layer %s with %d cache tensor(s)",

vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py

+51 -61

@@ -1092,7 +1092,7 @@ class NixlBaseConnectorWorker:

1092 # to better exploit the memory layout (ie num_blocks is the first dim).

1093 tensor_size_bytes = None

1094

1095 - for layer_name, cache_or_caches in xfer_buffers.items():

1095 + for layer_name, cache in xfer_buffers.items():

1096 # NOTE (NickLucche) Hybrid SSM models assume a layout that is similar to

1097 # that of FI, with block laid out as in `get_backend_aware_kv_block_len`.

1098 # However, physical page_size may differ when kernel requires a specific

@@ -1109,9 +1109,6 @@ class NixlBaseConnectorWorker:

1109 if isinstance(layer_spec, UniformTypeKVCacheSpecs):

1110 # MLA DSv32 Indexer case: UniformTypeKVCacheSpecs merges kv_cache_specs

1111 layer_spec = layer_spec.kv_cache_specs[layer_name]

1112 - cache_list = self.transfer_topo.get_transfer_cache_regions(

1113 - cache_or_caches, layer_spec

1114 - )

1112 # `layer_spec.page_size_bytes` only accounts for logical page_size, that is

1113 # the page_size assuming constant `self._logical_num_blocks`.

1114 physical_page_size = (

@@ -1120,8 +1117,6 @@ class NixlBaseConnectorWorker:

1117 else layer_spec.page_size_bytes

1118 // self._physical_blocks_per_logical_kv_block

1119 )

1123 - # For when registering multiple tensors eg K/V in separate regions.

1124 - physical_page_size = physical_page_size // len(cache_list)

1120 if self.transfer_topo._cross_layers_blocks:

1121 # When cross-layers blocks are used, multiply by number of layers

1122 physical_page_size = physical_page_size * len(

@@ -1136,66 +1131,61 @@ class NixlBaseConnectorWorker:

1131 # [`num_blocks` * `page_size`]

1132 curr_tensor_size_bytes = num_blocks * physical_page_size

1133

1139 - # TODO (NickLucche) we could eventually unify how we handle FA/FI regions,

1140 - # registering a single tensor for both K/V and splitting logically like FI.

1141 - for cache in cache_list:

1142 - base_addr = cache.data_ptr()

1143 - if base_addr in seen_base_addresses:

1144 - # NOTE (NickLucche) HMA employs memory pooling to share tensors

1145 - # across groups. This results in skipping all tensors but the ones

1146 - # pointed to by group0. Also, generally we will have more blocks

1147 - # per tensor but fewer regions.

1148 - logger.debug("Skipping %s because it's already seen", layer_name)

1149 - continue

1150 - logger.debug(

1151 - "Registering layer %s with cache shape: %s", layer_name, cache.shape

1152 - )

1153 - seen_base_addresses.append(base_addr)

1154 - # Only record non-Mamba page sizes.

1155 - if isinstance(layer_spec, MambaSpec):

1156 - self.block_len_per_layer.append(

1157 - physical_page_size // self._physical_blocks_per_logical_kv_block

1158 - )

1159 - else:

1160 - self.block_len_per_layer.append(physical_page_size)

1161 - is_mla_region = isinstance(

1162 - layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec)

1134 + base_addr = cache.data_ptr()

1135 + if base_addr in seen_base_addresses:

1136 + # NOTE (NickLucche) HMA employs memory pooling to share tensors

1137 + # across groups. This results in skipping all tensors but the ones

1138 + # pointed to by group0. Also, generally we will have more blocks

1139 + # per tensor but fewer regions.

1140 + logger.debug("Skipping %s because it's already seen", layer_name)

1141 + continue

1142 + logger.debug(

1143 + "Registering layer %s with cache shape: %s", layer_name, cache.shape

1144 + )

1145 + seen_base_addresses.append(base_addr)

1146 + # Only record non-Mamba page sizes.

1147 + if isinstance(layer_spec, MambaSpec):

1148 + self.block_len_per_layer.append(

1149 + physical_page_size // self._physical_blocks_per_logical_kv_block

1150 )

1164 - self._region_is_mla.append(is_mla_region)

1165 -

1166 - if not is_mla_region:

1167 - if tensor_size_bytes is None:

1168 - tensor_size_bytes = curr_tensor_size_bytes

1169 - assert tensor_size_bytes == curr_tensor_size_bytes, (

1170 - "All non-MLA kv cache tensors must have the same size"

1171 - )

1151 + else:

1152 + self.block_len_per_layer.append(physical_page_size)

1153 + is_mla_region = isinstance(

1154 + layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec)

1155 + )

1156 + self._region_is_mla.append(is_mla_region)

1157

1173 - # When there's a mismatch between kbs<>bs, we rely on HMA to ensure

1174 - # caches are either [NB, PS] or [NB*r, PS/r] where r is bs/kbs.

1175 - if (

1176 - self._physical_blocks_per_logical_kv_block == 1

1177 - and cache.shape[0] != num_blocks

1178 - ):

1179 - raise AssertionError(

1180 - "All kv cache tensors must have the same number of "

1181 - f"blocks; layer={layer_name}, "

1182 - f"expected_num_blocks={num_blocks}, "

1183 - f"cache_shape={tuple(cache.shape)}, "

1184 - f"cache_stride={tuple(cache.stride())}, "

1185 - f"layer_spec={type(layer_spec).__name__}, "

1186 - f"backend={self.backend_name}, "

1187 - "all_backends="

1188 - f"{[backend.get_name() for backend in self.attn_backends]}, "

1189 - f"kv_cache_layout={self.kv_cache_layout}"

1190 - )

1158 + if not is_mla_region:

1159 + if tensor_size_bytes is None:

1160 + tensor_size_bytes = curr_tensor_size_bytes

1161 + assert tensor_size_bytes == curr_tensor_size_bytes, (

1162 + "All non-MLA kv cache tensors must have the same size"

1163 + )

1164

1192 - # Need to make sure the device ID is non-negative for NIXL,

1193 - # Torch uses -1 to indicate CPU tensors.

1194 - self.device_id = max(cache.get_device(), 0)

1195 - caches_data.append(

1196 - (base_addr, curr_tensor_size_bytes, self.device_id, "")

1165 + # When there's a mismatch between kbs<>bs, we rely on HMA to ensure

1166 + # caches are either [NB, PS] or [NB*r, PS/r] where r is bs/kbs.

1167 + if (

1168 + self._physical_blocks_per_logical_kv_block == 1

1169 + and cache.shape[0] != num_blocks

1170 + ):

1171 + raise AssertionError(

1172 + "All kv cache tensors must have the same number of "

1173 + f"blocks; layer={layer_name}, "

1174 + f"expected_num_blocks={num_blocks}, "

1175 + f"cache_shape={tuple(cache.shape)}, "

1176 + f"cache_stride={tuple(cache.stride())}, "

1177 + f"layer_spec={type(layer_spec).__name__}, "

1178 + f"backend={self.backend_name}, "

1179 + "all_backends="

1180 + f"{[backend.get_name() for backend in self.attn_backends]}, "

1181 + f"kv_cache_layout={self.kv_cache_layout}"

1182 )

1183

1184 + # Need to make sure the device ID is non-negative for NIXL,

1185 + # Torch uses -1 to indicate CPU tensors.

1186 + self.device_id = max(cache.get_device(), 0)

1187 + caches_data.append((base_addr, curr_tensor_size_bytes, self.device_id, ""))

1188 +

1189 logger.debug(

1190 "Different block lengths collected: %s", set(self.block_len_per_layer)

1191 )

vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py

+7 -20

@@ -56,9 +56,7 @@ class OffloadingConnectorWorker:

56 def _init_worker(self, kv_caches: CanonicalKVCaches) -> None:

57 self.worker = self.spec.get_worker(kv_caches)

59 - def register_kv_caches(

60 - self, kv_caches: dict[str, torch.Tensor | list[torch.Tensor]]

61 - ):

59 + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):

60 kv_cache_config = self.kv_cache_config

61 num_blocks = kv_cache_config.num_blocks

@@ -120,24 +118,13 @@ class OffloadingConnectorWorker:

118 )

119

120 elif isinstance(layer_kv_cache_spec, MambaSpec):

123 - state_tensors = kv_caches[layer_name]

124 - assert isinstance(state_tensors, list)

125 -

126 - # re-construct the raw (num_blocks, page_size) tensor

127 - # from the first state tensor

128 - assert len(state_tensors) > 0

129 - first_state_tensor = state_tensors[0]

130 - assert first_state_tensor.storage_offset() == 0

131 - tensor = (

132 - torch.tensor(

133 - [],

134 - dtype=torch.int8,

135 - device=first_state_tensor.device,

136 - )

137 - .set_(first_state_tensor.untyped_storage())

138 - .view((num_blocks, layer_kv_cache_spec.page_size_bytes))

121 + layer_kv_cache = kv_caches[layer_name]

122 + assert layer_kv_cache.dtype == torch.int8

123 + tensors_per_block[layer_name] = (

124 + layer_kv_cache.view(

125 + num_blocks, layer_kv_cache_spec.page_size_bytes

126 + ),

127 )

140 - tensors_per_block[layer_name] = (tensor,)

128

129 page_size_bytes[layer_name] = layer_kv_cache_spec.page_size_bytes

130 unpadded_page_size_bytes[layer_name] = replace(

vllm/model_executor/layers/attention_layer_base.py

+10

@@ -4,6 +4,8 @@

5 from abc import ABC, abstractmethod

7 +import torch

8 +

9 from vllm.config import VllmConfig

10 from vllm.v1.attention.backend import AttentionBackend, AttentionImpl

11 from vllm.v1.kv_cache_interface import KVCacheSpec

@@ -21,6 +23,14 @@ class AttentionLayerBase(ABC):

23 impl: "AttentionImpl"

24 supports_dcp: bool = True

26 + def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:

27 + """Bind the allocated KV cache tensor to this layer.

28 +

29 + The default stores the cache view as-is; subclasses (e.g. Mamba)

30 + override this to unpack the raw buffer into per-state views.

31 + """

32 + self.kv_cache = kv_cache

33 +

34 @abstractmethod

35 def get_attn_backend(self) -> type[AttentionBackend]:

36 """Get the attention backend class for this layer."""

vllm/model_executor/layers/mamba/abstract.py

+18

@@ -2,11 +2,13 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from abc import abstractmethod

4 from collections.abc import Iterable

5 +from math import prod

7 import torch

9 from vllm.config import VllmConfig

10 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase

11 +from vllm.utils.torch_utils import get_dtype_size

12 from vllm.v1.attention.backend import AttentionBackend

13 from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum

14 from vllm.v1.attention.selector import get_mamba_attn_backend

@@ -24,6 +26,22 @@ class MambaBase(AttentionLayerBase):

26 kv_cache: tuple[torch.Tensor, ...]

27 supports_dcp: bool = False

29 + def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:

30 + """Unpack a raw ``[B, 1, 1, C]`` int8 page view into per-state views.

31 +

32 + Each block's ``C`` bytes hold the layer's states (e.g. conv, ssm)

33 + packed contiguously; slice them out and reinterpret per dtype/shape.

34 + """

35 + pages = kv_cache.squeeze(dim=(1, 2))

36 + states: list[torch.Tensor] = []

37 + offset = 0

38 + for shape, dtype in zip(self.get_state_shape(), self.get_state_dtype()):

39 + nbytes = prod(shape) * get_dtype_size(dtype)

40 + state = pages[:, offset : offset + nbytes].view(dtype)

41 + states.append(state.view(-1, *shape))

42 + offset += nbytes

43 + self.kv_cache = tuple(states)

44 +

45 @abstractmethod

46 def get_state_shape(self) -> Iterable[tuple[int, ...]]:

47 """

vllm/v1/worker/gpu/attn_utils.py

+11 -87

@@ -262,7 +262,7 @@ def _reshape_kv_cache(

262 kv_cache_config: "KVCacheConfig | None" = None,

263 ) -> dict[str, Any]:

264 kv_caches: dict[str, Any] = {}

265 - has_attn, has_mamba = False, False

265 + has_attn = False

266

267 layer_packing: dict[str, tuple[int, int]] = {}

268 if kv_cache_config is not None:

@@ -340,38 +340,21 @@ def _reshape_kv_cache(

340 )

341

342 elif isinstance(kv_cache_spec, MambaSpec):

343 - has_mamba = True

344 - state_tensors = []

345 - storage_offset_bytes = 0

346 - for shape, dtype in zip(kv_cache_spec.shapes, kv_cache_spec.dtypes):

347 - dtype_size = get_dtype_size(dtype)

348 - num_element_per_page = kv_cache_spec.page_size_bytes // dtype_size

349 - target_shape = (num_blocks, *shape)

350 - stride = torch.empty(target_shape).stride()

351 - target_stride = (num_element_per_page, *stride[1:])

352 - assert storage_offset_bytes % dtype_size == 0

353 - tensor = torch.as_strided(

354 - kv_raw_tensor.view(dtype),

355 - size=target_shape,

356 - stride=target_stride,

357 - storage_offset=storage_offset_bytes // dtype_size,

358 - )

359 - state_tensors.append(tensor)

360 - storage_offset_bytes += stride[0] * dtype_size

361 - kv_caches[layer_name] = state_tensors

343 + page_size_bytes = kv_cache_spec.page_size_bytes

344 + # Hold a single contiguous [num_blocks, 1, 1, page_size_bytes]

345 + # int8 page view per layer; the layer's bind_kv_cache unpacks

346 + # each block's bytes into its conv/ssm state views. Keeping

347 + # one tensor per layer lets the KV connector register it

348 + # without special-casing Mamba.

349 + kv_caches[layer_name] = kv_raw_tensor[

350 + : num_blocks * page_size_bytes

351 + ].view(num_blocks, 1, 1, page_size_bytes)

352 else:

353 raise NotImplementedError(

354 f"Unsupported KV cache spec type: {type(kv_cache_spec)}"

355 )

356

367 - if has_attn and has_mamba:

368 - _update_hybrid_attention_layout(

369 - attn_groups=attn_groups,

370 - kv_caches=kv_caches,

371 - kernel_block_sizes=kernel_block_sizes,

372 - cache_dtype=cache_dtype,

373 - )

374 - elif has_attn and kv_cache_config is not None:

357 + if has_attn and kv_cache_config is not None:

358 _align_mixed_attention_kv_cache_views(

359 attn_groups=attn_groups,

360 kv_caches=kv_caches,

@@ -458,65 +441,6 @@ def _restride_blocks_first_kv_cache_to_kv_first_storage(

441 )

442

443

461 -def _update_hybrid_attention_layout(

462 - attn_groups: Iterable[AttentionGroup],

463 - kv_caches: dict[str, Any],

464 - kernel_block_sizes: list[int],

465 - cache_dtype: str,

466 -) -> None:

467 - for group in attn_groups:

468 - if group.kv_cache_group_id >= len(kernel_block_sizes):

469 - continue

470 -

471 - kv_cache_spec = group.kv_cache_spec

472 - if not isinstance(kv_cache_spec, AttentionSpec):

473 - continue

474 - # Mirror the per-layer dtype selection used when building the shape

475 - # above. The block-dim index is dtype-independent for current backends

476 - # (quantization only changes the last dim), so this is a no-op today,

477 - # but it keeps both call sites consistent for skip layers.

478 - layer_cache_dtype = (

479 - "auto"

480 - if kv_cache_spec.kv_quant_mode == KVQuantMode.NONE

481 - and not isinstance(kv_cache_spec, TQFullAttentionSpec)

482 - else cache_dtype

483 - )

484 - block_dim = group.backend.get_kv_cache_block_dim(

485 - kernel_block_sizes[group.kv_cache_group_id],

486 - kv_cache_spec.num_kv_heads,

487 - kv_cache_spec.head_size,

488 - cache_dtype_str=layer_cache_dtype,

489 - )

490 - # if the first dim of the kvcache's layout is already num_blocks, continue

491 - if block_dim == 0:

492 - continue

493 -

494 - assert block_dim == 1, (

495 - "Expected the dim `num_blocks` at the second dim when updating"

496 - " the kvcache's layout of full attention layer"

497 - )

498 -

499 - for layer_name in group.layer_names:

500 - if layer_name not in kv_caches:

501 - # Shared layer — will be aliased to its target after this pass.

502 - continue

503 -

504 - kv_cache = kv_caches[layer_name]

505 - if kv_cache.shape[0] == 2:

506 - assert kv_cache.shape[1] != 2, (

507 - f"Cannot determine layout for tensor of shape {kv_cache.shape}"

508 - )

509 - hidden_size = kv_cache.shape[2:].numel()

510 - kv_cache.as_strided_(

511 - size=kv_cache.shape,

512 - stride=(

513 - hidden_size,

514 - 2 * hidden_size,

515 - *kv_cache.stride()[2:],

516 - ),

517 - )

518 -

519 -

444 def init_kv_cache(

445 runner_kv_caches: list[torch.Tensor | list[torch.Tensor]],

446 forward_context: dict[str, Any],

vllm/v1/worker/gpu_model_runner.py

+9 -22

@@ -126,7 +126,6 @@ from vllm.utils.torch_utils import (

126 PIN_MEMORY,

127 async_tensor_h2d,

128 current_stream,

129 - get_dtype_size,

129 is_quantized_kv_cache,

130 kv_cache_dtype_str_to_dtype,

131 )

@@ -7381,27 +7380,15 @@ class GPUModelRunner(

7380 elif isinstance(kv_cache_spec, MambaSpec):

7381 has_mamba = True

7382 raw_tensor = kv_cache_raw_tensors[layer_name]

7384 - state_tensors = []

7385 - storage_offset_bytes = 0

7386 - for shape, dtype in zip(kv_cache_spec.shapes, kv_cache_spec.dtypes):

7387 - dtype_size = get_dtype_size(dtype)

7388 - num_element_per_page = (

7389 - kv_cache_spec.page_size_bytes // dtype_size

7390 - )

7391 - target_shape = (num_blocks, *shape)

7392 - stride = torch.empty(target_shape).stride()

7393 - target_stride = (num_element_per_page, *stride[1:])

7394 - assert storage_offset_bytes % dtype_size == 0

7395 - tensor = torch.as_strided(

7396 - raw_tensor.view(dtype),

7397 - size=target_shape,

7398 - stride=target_stride,

7399 - storage_offset=storage_offset_bytes // dtype_size,

7400 - )

7401 - state_tensors.append(tensor)

7402 - storage_offset_bytes += stride[0] * dtype_size

7403 -

7404 - kv_caches[layer_name] = state_tensors

7383 + page_size_bytes = kv_cache_spec.page_size_bytes

7384 + # Hold a single contiguous [num_blocks, 1, 1, page_size_bytes]

7385 + # int8 page view per layer; the layer's bind_kv_cache unpacks

7386 + # each block's bytes into its conv/ssm state views. Keeping

7387 + # one tensor per layer lets the KV connector register it

7388 + # without special-casing Mamba.

7389 + kv_caches[layer_name] = raw_tensor[

7390 + : num_blocks * page_size_bytes

7391 + ].view(num_blocks, 1, 1, page_size_bytes)

7392 else:

7393 raise NotImplementedError

7394

vllm/v1/worker/utils.py

+5 -2

@@ -533,9 +533,12 @@ def bind_kv_cache(

533 for layer_name in layer_names:

534 runner_kv_caches.append(kv_caches[layer_name])

535

536 - # Bind kv_caches to forward context

536 + # Bind kv_caches to forward context. Each layer's bind_kv_cache unpacks

537 + # its raw allocation into the per-layer view(s) it needs (e.g. Mamba

538 + # splits conv/ssm), so the kv_caches dict can hold a single tensor per

539 + # layer for the KV connector to register.

540 for layer_name, kv_cache in kv_caches.items():

538 - forward_context[layer_name].kv_cache = kv_cache

541 + forward_context[layer_name].bind_kv_cache(kv_cache)

542

543

544 def copy_kv_cache_blocks_inplace(