From deedcffa59a222fc97be7b5711fe2985600f29ef Mon Sep 17 00:00:00 2001 From: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Date: Sat, 3 Oct 2026 07:22:07 -0700 Subject: [PATCH 1/3] [Bugfix][KV Connector] Mooncake: transfer GLM-5.3-Flash's kpool indexer like NIXL GLM-5.3-Flash's kpool indexer stores several kernel blocks in one tensor row, and its kpool tail lives inside that allocation. #53906 taught NIXL both; Mooncake registered the indexer per row and the tail as its own region, then expanded both groups' block ids to kernel blocks. Tail and indexer descriptors landed past their pages, so P/D transfers failed with "Memory region not registered" or wrote the wrong bytes. Register a compressed cache whose rows hold several kernel blocks page by page, and fold a kpool tail inside it into that region as a shared group, as NIXL does. The two NIXL helpers move to kv_connector/utils.py so both connectors use them. Co-authored-by: Claude Opus 5.5 (1M context) Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> --- .../unit/test_mooncake_connector.py | 58 ++++++++++++++++ .../kv_transfer/kv_connector/utils.py | 46 +++++++++++++ .../v1/mooncake/mooncake_connector.py | 68 +++++++++++++++++++ .../kv_connector/v1/nixl/base_worker.py | 48 ++----------- 4 files changed, 177 insertions(+), 43 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_mooncake_connector.py b/tests/v1/kv_connector/unit/test_mooncake_connector.py index d326615aad8f..c36813b0e45b 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_connector.py +++ b/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,63 @@ def test_register_kv_caches_collapses_shared_mla_storage( assert worker.region_shared_groups == [(0, 1)] +def test_register_kv_caches_pages_kpool_indexer_and_folds_its_tail(): + """A kpool indexer row holding several kernel blocks registers per page, and + the kpool tail inside its allocation is read through that region.""" + num_blocks, pages_per_row = 8, 4 + indexer_spec = MLAAttentionSpec( + block_size=16, + 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 + state = indexer_spec.state_content_size_bytes + 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): + assert worker._physical_blocks_per_logical_kv_block == 1 + backing = torch.zeros(num_blocks * page, dtype=torch.uint8) + indexer = backing.view( + num_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)] + assert _block_ids_for_region([[1, 2], [5]], _SHARED_REGION_GROUP_ID, (0, 1)) == [ + 1, + 2, + 5, + ] + + def test_block_ids_for_region_flattens_shared_groups(): ids = [[10, 11], [12]] assert _block_ids_for_region(ids, 0) == [10, 11] diff --git a/vllm/distributed/kv_transfer/kv_connector/utils.py b/vllm/distributed/kv_transfer/kv_connector/utils.py index a0ead27c2ff1..e55b8ddfd73e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/utils.py +++ b/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. diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py index 172d41efdb45..85c71c2dcf2e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py +++ b/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,17 @@ 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 _uses_dense_transfer_pages( + self, layer_spec: KVCacheSpec | None, cache: torch.Tensor, num_blocks: int + ) -> bool: + return layer_spec is not None and uses_dense_virtual_transfer_pages( + layer_spec, cache, self._physical_page_size(layer_spec), num_blocks + ) + 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 +1991,18 @@ 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_owners = { + cache.data_ptr(): cache + for layer_name, cache in kv_caches.items() + if self._uses_dense_transfer_pages( + self._layer_specs.get(layer_name), cache, num_blocks + ) + } + 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 +2024,39 @@ 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 dense_page_owners.get(cache.data_ptr(), None) is not None and ( + dense_page_owners[cache.data_ptr()] is cache + ): + 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 +2192,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 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 0eaa1aadf34f..61738fe78a81 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/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. From ee660f5683e9e3823285bcfc12a5de339a72206f Mon Sep 17 00:00:00 2001 From: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Date: Sat, 3 Oct 2026 15:58:36 -0700 Subject: [PATCH 2/3] [Bugfix][KV Connector] Mooncake: track dense-page owners by layer name Inline the one-use predicate wrapper and drop a test assertion that only restated _block_ids_for_region. Co-authored-by: Claude Opus 5.5 (1M context) Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> --- .../unit/test_mooncake_connector.py | 5 ---- .../v1/mooncake/mooncake_connector.py | 23 ++++++++----------- 2 files changed, 9 insertions(+), 19 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_mooncake_connector.py b/tests/v1/kv_connector/unit/test_mooncake_connector.py index c36813b0e45b..b1900a341d28 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_connector.py +++ b/tests/v1/kv_connector/unit/test_mooncake_connector.py @@ -1961,11 +1961,6 @@ def test_register_kv_caches_pages_kpool_indexer_and_folds_its_tail(): 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)] - assert _block_ids_for_region([[1, 2], [5]], _SHARED_REGION_GROUP_ID, (0, 1)) == [ - 1, - 2, - 5, - ] def test_block_ids_for_region_flattens_shared_groups(): diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py index 85c71c2dcf2e..27cd830648cf 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py @@ -1960,13 +1960,6 @@ 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 _uses_dense_transfer_pages( - self, layer_spec: KVCacheSpec | None, cache: torch.Tensor, num_blocks: int - ) -> bool: - return layer_spec is not None and uses_dense_virtual_transfer_pages( - layer_spec, cache, self._physical_page_size(layer_spec), num_blocks - ) - 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) @@ -1994,13 +1987,17 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): # 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_owners = { - cache.data_ptr(): cache + dense_page_layers = { + layer_name for layer_name, cache in kv_caches.items() - if self._uses_dense_transfer_pages( - self._layer_specs.get(layer_name), cache, num_blocks + 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]] = {} @@ -2042,9 +2039,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): group_index ) continue - if dense_page_owners.get(cache.data_ptr(), None) is not None and ( - dense_page_owners[cache.data_ptr()] is cache - ): + 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()) From 970e91faa85df833a78ff0a3bc9f4a5ee1eca615 Mon Sep 17 00:00:00 2001 From: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Date: Tue, 6 Oct 2026 05:38:51 -0700 Subject: [PATCH 3/3] [Bugfix][KV Connector] Mooncake: test kpool paging with several kernel blocks per block The kpool registration test pinned one kernel block per logical block, but GLM-5.3-Flash runs 2304-token blocks over 64-token kernel blocks, one tensor row per logical block. Parametrize the test over the ratio and check that every kernel block of a logical block addresses its own page of the indexer tensor. Co-authored-by: Claude Opus 5.5 (1M context) Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> --- .../unit/test_mooncake_connector.py | 39 +++++++++++++++---- 1 file changed, 31 insertions(+), 8 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_mooncake_connector.py b/tests/v1/kv_connector/unit/test_mooncake_connector.py index b1900a341d28..39933374d027 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_connector.py +++ b/tests/v1/kv_connector/unit/test_mooncake_connector.py @@ -1911,12 +1911,19 @@ def test_register_kv_caches_collapses_shared_mla_storage( assert worker.region_shared_groups == [(0, 1)] -def test_register_kv_caches_pages_kpool_indexer_and_folds_its_tail(): +@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.""" - num_blocks, pages_per_row = 8, 4 + 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, + block_size=16 * kernel_blocks_per_block, num_kv_heads=1, head_size=132, dtype=torch.uint8, @@ -1930,8 +1937,13 @@ def test_register_kv_caches_pages_kpool_indexer_and_folds_its_tail(): dtype=torch.bfloat16, sliding_window=4, ) - page = indexer_spec.page_size_bytes + 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=[], @@ -1941,10 +1953,10 @@ def test_register_kv_caches_pages_kpool_indexer_and_folds_its_tail(): ], ) with mooncake_register_worker(config) as (worker, reg): - assert worker._physical_blocks_per_logical_kv_block == 1 - backing = torch.zeros(num_blocks * page, dtype=torch.uint8) + 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_blocks // pages_per_row, 1, pages_per_row * page // state, state + 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) @@ -1961,6 +1973,17 @@ def test_register_kv_caches_pages_kpool_indexer_and_folds_its_tail(): 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():