diff --git a/tests/v1/kv_connector/unit/test_mooncake_connector.py b/tests/v1/kv_connector/unit/test_mooncake_connector.py index d326615aad8f..39933374d027 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,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] 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..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 @@ -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 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.