[Bugfix][KV Connector] Mooncake: transfer GLM-5.3-Flash's kpool indexer pages and tail like NIXL

tests/v1/kv_connector/unit/test_mooncake_connector.py

@@ -50,6 +50,7 @@
 )
 from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
+    KpoolTailSpec,
     KVCacheConfig,
     KVCacheGroupSpec,
     KVCacheLayout,
@@ -1910,6 +1911,81 @@ def test_register_kv_caches_collapses_shared_mla_storage(
         assert worker.region_shared_groups == [(0, 1)]
 
 
+@pytest.mark.parametrize("kernel_blocks_per_block", [1, 4])
+def test_register_kv_caches_pages_kpool_indexer_and_folds_its_tail(
+    kernel_blocks_per_block,
+):
+    """A kpool indexer row holding several kernel blocks registers per page, and
+    the kpool tail inside its allocation is read through that region.
+
+    With ``kernel_blocks_per_block > 1`` each tensor row is one logical block
+    split into kernel pages, as on GLM-5.3-Flash (2304-token blocks over
+    64-token kernel blocks)."""
+    num_blocks, rows_per_tensor_row = 8, 4
+    indexer_spec = MLAAttentionSpec(
+        block_size=16 * kernel_blocks_per_block,
+        num_kv_heads=1,
+        head_size=132,
+        dtype=torch.uint8,
+        tokens_per_state=4,
+    )
+    tail_spec = KpoolTailSpec(
+        block_size=4,
+        num_kv_heads=2,
+        head_size=8,
+        head_size_v=0,
+        dtype=torch.bfloat16,
+        sliding_window=4,
+    )
+    page = indexer_spec.page_size_bytes // kernel_blocks_per_block
+    state = indexer_spec.state_content_size_bytes
+    if kernel_blocks_per_block > 1:
+        # One tensor row per logical block.
+        num_blocks, rows_per_tensor_row = 2, 1
+    num_kernel_blocks = num_blocks * kernel_blocks_per_block
+    pages_per_row = rows_per_tensor_row * kernel_blocks_per_block
+    config = KVCacheConfig(
+        num_blocks=num_blocks,
+        kv_cache_tensors=[],
+        kv_cache_groups=[
+            KVCacheGroupSpec(["model.layers.0.indexer"], indexer_spec),
+            KVCacheGroupSpec(["model.layers.0.tail"], tail_spec),
+        ],
+    )
+    with mooncake_register_worker(config) as (worker, reg):
+        worker._physical_blocks_per_logical_kv_block = kernel_blocks_per_block
+        backing = torch.zeros(num_kernel_blocks * page, dtype=torch.uint8)
+        indexer = backing.view(
+            num_kernel_blocks // pages_per_row, 1, pages_per_row * page // state, state
+        )
+        tail_bytes = num_blocks * tail_spec.page_size_bytes
+        tail = backing[:tail_bytes].view(torch.bfloat16).view(num_blocks, 2, 4, 8)
+        worker.register_kv_caches(
+            {"model.layers.0.indexer": indexer, "model.layers.0.tail": tail}
+        )
+
+    reg.assert_called_once()
+    assert reg.call_args[0][0] == [backing.data_ptr()]
+    # Kernel block j is the j-th page, not the j-th tensor row.
+    assert worker.kv_caches_base_addr == [indexer.data_ptr()]
+    assert worker.block_len_per_layer == [page]
+    assert worker.kv_block_len_per_layer == [page]
+    assert worker.registered_layer_names == ["model.layers.0.indexer"]
+    assert worker.registered_group_indices == [_SHARED_REGION_GROUP_ID]
+    assert worker.region_shared_groups == [(0, 1)]
+    # Each kernel block of logical block 1 addresses its own page of the row.
+    kernel_ids = worker._logical_to_kernel_block_ids([[1], [1]])[0]
+    assert kernel_ids == list(
+        range(kernel_blocks_per_block, 2 * kernel_blocks_per_block)
+    )
+    for k in kernel_ids:
+        row, page_in_row = divmod(k, pages_per_row)
+        first_state = indexer[row, 0, page_in_row * page // state]
+        assert first_state.data_ptr() == (
+            worker.kv_caches_base_addr[0] + k * worker.block_len_per_layer[0]
+        )
+
+
 def test_block_ids_for_region_flattens_shared_groups():
     ids = [[10, 11], [12]]
     assert _block_ids_for_region(ids, 0) == [10, 11]

vllm/distributed/kv_transfer/kv_connector/utils.py

@@ -20,6 +20,7 @@
 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import async_tensor_h2d
 from vllm.v1.attention.backend import AttentionBackend
+from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec
 from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
 
 if TYPE_CHECKING:
@@ -43,6 +44,51 @@ def clip_ssm_state_blocks(blocks: list[int], num_spec_blocks: int) -> list[int]:
     return blocks[-1:]
 
 
+def tensor_byte_span_end(cache: torch.Tensor) -> int:
+    """Return the exclusive end address touched by a nonnegative-stride view."""
+    if cache.numel() == 0:
+        return cache.data_ptr()
+    if any(stride < 0 for stride in cache.stride()):
+        raise ValueError("KV cache views must have nonnegative strides")
+    max_element_offset = sum(
+        (size - 1) * stride for size, stride in zip(cache.shape, cache.stride())
+    )
+    return cache.data_ptr() + (max_element_offset + 1) * cache.element_size()
+
+
+def uses_dense_virtual_transfer_pages(
+    layer_spec: KVCacheSpec,
+    cache: torch.Tensor,
+    physical_page_size: int,
+    num_blocks: int,
+) -> bool:
+    """Return whether a compressed kernel view can be split into transfer pages.
+
+    A compressed cache (``tokens_per_state > 1``, such as GLM-5.3-Flash's kpool
+    indexer) can store several kernel blocks in one tensor row. Its kernel block
+    ``j`` is then the ``j``-th ``physical_page_size`` page of the tensor.
+    """
+    if not (
+        isinstance(layer_spec, MLAAttentionSpec)
+        and layer_spec.tokens_per_state > 1
+        and cache.ndim == 4
+        and cache.shape[1] == 1
+        and cache.is_contiguous()
+        and physical_page_size > 0
+        and layer_spec.state_content_size_bytes > 0
+    ):
+        return False
+
+    block_stride = cache.stride(0) * cache.element_size()
+    return (
+        block_stride > physical_page_size
+        and block_stride % physical_page_size == 0
+        and physical_page_size % layer_spec.state_content_size_bytes == 0
+        and cache.shape[0] * (block_stride // physical_page_size) == num_blocks
+        and cache.nbytes == num_blocks * physical_page_size
+    )
+
+
 def get_kv_connector_cache_layout(vllm_config: VllmConfig | None = None):
     # NOTE (NickLucche) When running disaggregated PD with NIXL, LBHNC layout is
     # used for faster transfer.

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

@@ -26,6 +26,8 @@
     EngineId,
     TransferTopology,
     get_current_attn_backends,
+    tensor_byte_span_end,
+    uses_dense_virtual_transfer_pages,
 )
 from vllm.distributed.kv_transfer.kv_connector.v1.base import (
     KVConnectorBase_V1,
@@ -1954,6 +1956,10 @@ def _send_blocks(
             )
         return ret_value
 
+    def _physical_page_size(self, layer_spec: KVCacheSpec) -> int:
+        """Bytes per kernel block of an attention layer's page."""
+        return layer_spec.page_size_bytes // self._physical_blocks_per_logical_kv_block
+
     def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
         """Register the KV Cache data in mooncake."""
         logger.info("Registering KV_Caches. use_mla: %s", self.use_mla)
@@ -1978,6 +1984,22 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
         num_blocks = (
             self.kv_cache_config.num_blocks * self._physical_blocks_per_logical_kv_block
         )
+        # As in NIXL, a compressed cache that packs several kernel blocks per
+        # tensor row (GLM-5.3-Flash's kpool indexer) is registered page by page,
+        # and a kpool tail inside its allocation moves with it.
+        dense_page_layers = {
+            layer_name
+            for layer_name, cache in kv_caches.items()
+            if (spec := self._layer_specs.get(layer_name)) is not None
+            and uses_dense_virtual_transfer_pages(
+                spec, cache, self._physical_page_size(spec), num_blocks
+            )
+        }
+        dense_page_owners = {
+            kv_caches[name].data_ptr(): kv_caches[name] for name in dense_page_layers
+        }
+        dense_page_regions: dict[int, int] = {}
+        covered_tail_groups: dict[int, set[int]] = {}
 
         for layer_name, cache in kv_caches.items():
             layer_index = extract_layer_index(layer_name)
@@ -1999,6 +2021,37 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
                 kv_data_ptrs.append(storage_addr)
                 kv_data_lens.append(storage.nbytes())
 
+            if isinstance(layer_spec, KpoolTailSpec):
+                owner = dense_page_owners.get(cache.data_ptr())
+                if owner is not None:
+                    if not (
+                        owner.untyped_storage().data_ptr() == storage_addr
+                        and tensor_byte_span_end(cache)
+                        <= owner.data_ptr() + owner.nbytes
+                    ):
+                        raise AssertionError(
+                            "Kpool tail cache is not fully covered by its "
+                            f"compressed indexer region: layer={layer_name}, "
+                            f"tail_shape={tuple(cache.shape)}, "
+                            f"owner_shape={tuple(owner.shape)}"
+                        )
+                    covered_tail_groups.setdefault(cache.data_ptr(), set()).add(
+                        group_index
+                    )
+                    continue
+            if layer_name in dense_page_layers:
+                page = self._physical_page_size(layer_spec)
+                dense_page_regions[cache.data_ptr()] = len(region_base_addresses)
+                region_base_addresses.append(cache.data_ptr())
+                self.block_len_per_layer.append(page)
+                self.kv_block_len_per_layer.append(page)
+                self.registered_layer_names.append(layer_name)
+                self.registered_layer_indices.append(layer_index)
+                self.registered_group_indices.append(group_index)
+                self.region_shared_groups.append((group_index,))
+                self.region_row_offsets.append(-1)
+                continue
+
             is_mla_region = isinstance(
                 layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec)
             )
@@ -2134,6 +2187,16 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
                 # Not a view into a shared packed row, so do not promote it.
                 self.region_row_offsets.append(-1)
 
+        # The tail group's blocks are read through the indexer region it shares.
+        for owner_ptr, tail_groups in covered_tail_groups.items():
+            region_idx = dense_page_regions[owner_ptr]
+            shared = tuple(
+                dict.fromkeys((*self.region_shared_groups[region_idx], *tail_groups))
+            )
+            self.region_shared_groups[region_idx] = shared
+            if len(shared) > 1:
+                self.registered_group_indices[region_idx] = _SHARED_REGION_GROUP_ID
+
         self.kv_caches_base_addr = region_base_addresses
         self.seen_base_addresses = kv_data_ptrs
 

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

@@ -32,6 +32,8 @@
     kv_postprocess_blksize_and_layout_on_receive,
     kv_postprocess_blksize_on_receive,
     kv_postprocess_layout_on_receive,
+    tensor_byte_span_end,
+    uses_dense_virtual_transfer_pages,
 )
 from vllm.distributed.kv_transfer.kv_connector.v1.base import (
     CopyBlocksOp,
@@ -117,46 +119,6 @@ def _share_storage_and_block_stride(caches: list[torch.Tensor]) -> bool:
     return len(block_strides) == len(storage_ptrs) == 1
 
 
-def _tensor_byte_span_end(cache: torch.Tensor) -> int:
-    """Return the exclusive end address touched by a nonnegative-stride view."""
-    if cache.numel() == 0:
-        return cache.data_ptr()
-    if any(stride < 0 for stride in cache.stride()):
-        raise ValueError("NIXL cache views must have nonnegative strides")
-    max_element_offset = sum(
-        (size - 1) * stride for size, stride in zip(cache.shape, cache.stride())
-    )
-    return cache.data_ptr() + (max_element_offset + 1) * cache.element_size()
-
-
-def _uses_dense_virtual_transfer_pages(
-    layer_spec: KVCacheSpec,
-    cache: torch.Tensor,
-    physical_page_size: int,
-    num_blocks: int,
-) -> bool:
-    """Return whether a compressed kernel view can be split into NIXL pages."""
-    if not (
-        isinstance(layer_spec, MLAAttentionSpec)
-        and layer_spec.tokens_per_state > 1
-        and cache.ndim == 4
-        and cache.shape[1] == 1
-        and cache.is_contiguous()
-        and physical_page_size > 0
-        and layer_spec.state_content_size_bytes > 0
-    ):
-        return False
-
-    block_stride = cache.stride(0) * cache.element_size()
-    return (
-        block_stride > physical_page_size
-        and block_stride % physical_page_size == 0
-        and physical_page_size % layer_spec.state_content_size_bytes == 0
-        and cache.shape[0] * (block_stride // physical_page_size) == num_blocks
-        and cache.nbytes == num_blocks * physical_page_size
-    )
-
-
 class NixlBaseConnectorWorker:
     """Base implementation of Worker side methods shared by pull and push."""
 
@@ -1479,7 +1441,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
                 if isinstance(layer_spec, MambaSpec)
                 else self.num_blocks
             )
-            if _uses_dense_virtual_transfer_pages(
+            if uses_dense_virtual_transfer_pages(
                 layer_spec, cache, physical_page_size, num_blocks
             ):
                 compressed_region_owners.setdefault(cache.data_ptr(), cache)
@@ -1576,7 +1538,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
                     tail_is_covered = (
                         compressed_owner.is_contiguous()
                         and owner_storage.data_ptr() == storage_addr
-                        and _tensor_byte_span_end(cache) <= owner_end
+                        and tensor_byte_span_end(cache) <= owner_end
                     )
                     if not tail_is_covered:
                         raise AssertionError(
@@ -1628,7 +1590,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
                     and cache[0].is_contiguous()
                     and cache[0].nbytes == physical_page_size
                 )
-                virtual_transfer_pages = _uses_dense_virtual_transfer_pages(
+                virtual_transfer_pages = uses_dense_virtual_transfer_pages(
                     layer_spec, cache, physical_page_size, num_blocks
                 )
                 # Push workers address the pages of packed MLA rows by layer.