diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 0e791d57b728..92d8b35e2c0e 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -1683,6 +1683,164 @@ def test_resolve_kv_cache_block_sizes_mixed_dcp_replicated_groups(): assert hash_block_size == 64 +def test_replicated_mla_uses_lockstep_pool_capacity_and_contiguous_tensors(): + vllm_config = SimpleNamespace( + model_config=SimpleNamespace(max_model_len=262144), + parallel_config=SimpleNamespace( + decode_context_parallel_size=4, + prefill_context_parallel_size=1, + ), + cache_config=SimpleNamespace(num_gpu_blocks_override=None), + kv_transfer_config=None, + ) + specs: dict[str, KVCacheSpec] = { + f"target.{i}": MLAAttentionSpec( + block_size=64, + num_kv_heads=1, + head_size=576, + dtype=torch.uint8, + cache_dtype_str="nvfp4_ds_mla", + ) + for i in range(3) + } + specs.update( + { + f"indexer.{i}": MLAAttentionSpec( + block_size=256, + num_kv_heads=1, + head_size=132, + dtype=torch.uint8, + dcp_replicated=True, + ) + for i in range(2) + } + ) + + grouped_specs = kv_cache_utils.group_and_unify_kv_cache_specs(specs, 4, 1) + assert grouped_specs is not None + groups = kv_cache_utils._get_kv_cache_groups_uniform_groups(grouped_specs) + assert all( + isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs) for group in groups + ) + assert kv_cache_utils._use_lockstep_mla_allocation(groups, 4, 1) + + bytes_per_pool_block = 3 * (64 * 432) + 2 * (256 * 132) + request_blocks = 262144 // 256 + required_memory = bytes_per_pool_block * request_blocks + kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups( + vllm_config, + groups, + available_memory=required_memory * 2, + ) + + assert kv_cache_config.num_blocks == request_blocks * 2 + assert all(tensor.block_stride == 0 for tensor in kv_cache_config.kv_cache_tensors) + assert sum(tensor.size for tensor in kv_cache_config.kv_cache_tensors) == ( + required_memory * 2 + ) + assert ( + kv_cache_utils._max_memory_usage_bytes_from_groups(vllm_config, groups) + == required_memory + ) + assert get_max_concurrency_for_kv_cache_config( + vllm_config, kv_cache_config + ) == pytest.approx(2.0) + + +def test_lockstep_mla_predicate_rejects_nonmatching_layouts(): + sharded = MLAAttentionSpec( + block_size=64, + num_kv_heads=1, + head_size=128, + dtype=torch.float32, + ) + replicated = MLAAttentionSpec( + block_size=256, + num_kv_heads=1, + head_size=128, + dtype=torch.float32, + dcp_replicated=True, + ) + groups = [ + KVCacheGroupSpec(["target"], sharded), + KVCacheGroupSpec(["indexer"], replicated), + ] + assert not kv_cache_utils._use_lockstep_mla_allocation(groups, 1, 1) + + mismatched = MLAAttentionSpec( + block_size=64, + num_kv_heads=1, + head_size=128, + dtype=torch.float32, + dcp_replicated=True, + ) + mismatched_groups = [groups[0], KVCacheGroupSpec(["indexer"], mismatched)] + assert not kv_cache_utils._use_lockstep_mla_allocation(mismatched_groups, 4, 1) + assert ( + kv_cache_utils.group_and_unify_kv_cache_specs( + {"target": sharded, "indexer": mismatched}, 4, 1 + ) + is None + ) + + non_mla = FullAttentionSpec( + block_size=256, + num_kv_heads=1, + head_size=128, + dtype=torch.float32, + dcp_replicated=True, + ) + assert not kv_cache_utils._use_lockstep_mla_allocation( + [groups[0], KVCacheGroupSpec(["draft"], non_mla)], 4, 1 + ) + + +def test_lockstep_mla_equal_page_sizes_use_distinct_tensors(): + sharded = MLAAttentionSpec( + block_size=64, + num_kv_heads=1, + head_size=128, + dtype=torch.bfloat16, + ) + replicated = MLAAttentionSpec( + block_size=256, + num_kv_heads=1, + head_size=32, + dtype=torch.bfloat16, + dcp_replicated=True, + ) + assert sharded.page_size_bytes == replicated.page_size_bytes + groups = [ + KVCacheGroupSpec(["target"], sharded), + KVCacheGroupSpec(["indexer"], replicated), + ] + vllm_config = SimpleNamespace( + cache_config=SimpleNamespace(num_gpu_blocks_override=None), + parallel_config=SimpleNamespace( + decode_context_parallel_size=4, + prefill_context_parallel_size=1, + ), + kv_transfer_config=None, + ) + bytes_per_pool_block = sharded.page_size_bytes + replicated.page_size_bytes + + kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups( + vllm_config, + groups, + available_memory=3 * bytes_per_pool_block, + ) + + assert kv_cache_config.num_blocks == 3 + assert [tensor.shared_by for tensor in kv_cache_config.kv_cache_tensors] == [ + ["target"], + ["indexer"], + ] + assert [tensor.size for tensor in kv_cache_config.kv_cache_tensors] == [ + 3 * sharded.page_size_bytes, + 3 * replicated.page_size_bytes, + ] + + def test_dsv4_engine_capacity_uses_worker_kv_cache_config(): from vllm.v1.engine.core import EngineCore diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index dd89769315b9..74fdcfccb56f 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -177,6 +177,183 @@ def make_kv_cache_config_hybrid_model( ) +def make_lockstep_mla_manager(num_blocks: int = 5) -> KVCacheManager: + global_block_size = 256 + kv_cache_config = KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["target"], + MLAAttentionSpec( + block_size=global_block_size // 4, + num_kv_heads=1, + head_size=432, + dtype=torch.uint8, + ), + ), + KVCacheGroupSpec( + ["indexer"], + MLAAttentionSpec( + block_size=global_block_size, + num_kv_heads=1, + head_size=132, + dtype=torch.uint8, + dcp_replicated=True, + ), + ), + ], + ) + return KVCacheManager( + kv_cache_config=kv_cache_config, + max_model_len=4 * global_block_size, + scheduler_block_size=global_block_size, + hash_block_size=global_block_size, + enable_caching=True, + dcp_world_size=4, + ) + + +def test_mixed_mla_groups_share_block_ids_hashes_and_eviction_order(): + block_size = 256 + manager = make_lockstep_mla_manager() + assert manager.coordinator.lockstep_mla_allocations + request = make_request( + "lockstep", + list(range(4 * block_size)), + block_size, + sha256, + ) + + new_blocks = manager.allocate_slots( + request, + num_new_tokens=4 * block_size, + full_sequence_must_fit=True, + ) + assert new_blocks is not None + target_ids, indexer_ids = manager.get_block_ids(request.request_id) + assert target_ids == indexer_ids + assert new_blocks.get_block_ids() == (target_ids, target_ids) + assert manager.block_pool.get_num_free_blocks() == 0 + assert all( + manager.block_pool.blocks[block_id].ref_cnt == 2 for block_id in target_ids + ) + + for block_hash, block_id in zip(request.block_hashes, target_ids): + group_keys = [ + make_block_hash_with_group_id(block_hash, group_id) for group_id in range(2) + ] + assert group_keys[0] != group_keys[1] + cached = manager.block_pool.get_cached_block(block_hash, [0, 1]) + assert cached is not None + assert [block.block_id for block in cached] == [block_id, block_id] + + manager.free(request) + assert manager.block_pool.get_num_free_blocks() == 4 + assert all( + manager.block_pool.blocks[block_id].ref_cnt == 0 for block_id in target_ids + ) + free_ids = [ + block.block_id + for block in manager.block_pool.free_block_queue.get_all_free_blocks() + ] + assert free_ids == target_ids[::-1] + + replay = make_request( + "lockstep-replay", + list(range(4 * block_size)), + block_size, + sha256, + ) + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(replay) + assert num_computed_tokens == 3 * block_size + assert computed_blocks.get_block_ids() == ( + target_ids[:3], + target_ids[:3], + ) + + replay_new_blocks = manager.allocate_slots( + replay, + num_new_tokens=block_size, + num_new_computed_tokens=num_computed_tokens, + new_computed_blocks=computed_blocks, + ) + assert replay_new_blocks is not None + assert manager.get_block_ids(replay.request_id) == (target_ids, target_ids) + assert all( + manager.block_pool.blocks[block_id].ref_cnt == 2 for block_id in target_ids + ) + manager.free(replay) + assert [ + block.block_id + for block in manager.block_pool.free_block_queue.get_all_free_blocks() + ] == target_ids[::-1] + + evicted = manager.block_pool.get_new_blocks(1)[0] + assert evicted.block_id == target_ids[-1] + assert manager.block_pool.get_cached_block(request.block_hashes[-1], [0]) is None + assert manager.block_pool.get_cached_block(request.block_hashes[-1], [1]) is None + + +def test_lockstep_mla_rejects_external_computed_blocks(): + manager = make_lockstep_mla_manager() + + with pytest.raises(NotImplementedError, match="External KV loads"): + manager.coordinator.allocate_new_computed_blocks( + "external", + ([], []), + num_local_computed_tokens=0, + num_external_computed_tokens=256, + ) + + assert manager.get_block_ids("external") == ([], []) + + +def test_lockstep_group_hashes_promote_partial_block_together(): + hash_block_size = 64 + block_size = 4 * hash_block_size + pool = BlockPool( + num_gpu_blocks=2, + enable_caching=True, + hash_block_size=hash_block_size, + ) + block = pool.get_new_blocks(1)[0] + request = make_request( + "promotion", + list(range(block_size)), + hash_block_size, + sha256, + ) + + for group_id in range(2): + pool.cache_partial_block( + request=request, + block=block, + num_tokens=2 * hash_block_size, + kv_cache_group_id=group_id, + block_size=block_size, + ) + partial_hash = request.block_hashes[1] + partial_cached = pool.get_cached_block(partial_hash, [0, 1]) + assert partial_cached == [block, block] + + for group_id in range(2): + pool.cache_full_blocks( + request=request, + blocks=[block], + num_cached_blocks=0, + num_full_blocks=1, + block_size=block_size, + kv_cache_group_id=group_id, + ) + + assert pool.get_cached_block(partial_hash, [0]) is None + assert pool.get_cached_block(partial_hash, [1]) is None + full_hash = request.block_hashes[-1] + assert pool.get_cached_block(full_hash, [0, 1]) == [block, block] + assert block.block_hash_num_tokens == block_size + + def make_kv_cache_config_three_types( block_size: int, num_blocks: int, third_spec_type: str = "mamba" ) -> KVCacheConfig: diff --git a/tests/v1/worker/test_attn_utils.py b/tests/v1/worker/test_attn_utils.py index 7e65d650f7ed..2591a426913d 100644 --- a/tests/v1/worker/test_attn_utils.py +++ b/tests/v1/worker/test_attn_utils.py @@ -1,10 +1,20 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + import torch -from vllm.v1.kv_cache_interface import FullAttentionSpec, KVQuantMode -from vllm.v1.worker.gpu.attn_utils import _reshape_kv_cache +from vllm.v1.core.kv_cache_utils import _get_kv_cache_config_packed +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheConfig, + KVCacheGroupSpec, + KVQuantMode, + MLAAttentionSpec, + UniformTypeKVCacheSpecs, +) +from vllm.v1.worker.gpu.attn_utils import _allocate_kv_cache, _reshape_kv_cache from vllm.v1.worker.utils import AttentionGroup @@ -149,6 +159,104 @@ def get_kv_cache_stride_order( return (0, 1, 2, 3) +class FakeSingleKVBackend: + @staticmethod + def get_kv_cache_shape( + num_blocks: int, + block_size: int, + num_kv_heads: int, + head_size: int, + cache_dtype_str: str = "auto", + ) -> tuple[int, ...]: + assert num_kv_heads == 1 + return (num_blocks, block_size, head_size) + + @staticmethod + def get_kv_cache_stride_order( + include_num_layers_dimension: bool = False, + ) -> tuple[int, ...]: + assert not include_num_layers_dimension + return (0, 1, 2) + + @staticmethod + def get_kv_cache_block_dim( + block_size: int, + num_kv_heads: int, + head_size: int, + cache_dtype_str: str = "auto", + ) -> int: + return 0 + + +def test_lockstep_manager_pages_reshape_to_contiguous_kernel_pages(): + manager_blocks = 3 + kernel_block_size = 64 + main_layer = MLAAttentionSpec( + block_size=kernel_block_size, + num_kv_heads=1, + head_size=328, + dtype=torch.bfloat16, + ) + indexer_layer = MLAAttentionSpec( + block_size=kernel_block_size * 4, + num_kv_heads=1, + head_size=132, + dtype=torch.bfloat16, + dcp_replicated=True, + ) + main_spec = UniformTypeKVCacheSpecs( + block_size=main_layer.block_size, + kv_cache_specs={"main": main_layer}, + ) + indexer_spec = UniformTypeKVCacheSpecs( + block_size=indexer_layer.block_size, + kv_cache_specs={"indexer": indexer_layer}, + ) + groups = [ + KVCacheGroupSpec(["main"], main_spec), + KVCacheGroupSpec(["indexer"], indexer_spec), + ] + config = SimpleNamespace( + cache_config=SimpleNamespace(num_gpu_blocks_override=None), + parallel_config=SimpleNamespace( + decode_context_parallel_size=4, + prefill_context_parallel_size=1, + ), + kv_transfer_config=None, + ) + bytes_per_manager_block = main_layer.page_size_bytes + indexer_layer.page_size_bytes + num_blocks, tensors = _get_kv_cache_config_packed( + config, + groups, + available_memory=bytes_per_manager_block * manager_blocks, + ) + kv_cache_config = KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=tensors, + kv_cache_groups=groups, + ) + raw_tensors = _allocate_kv_cache(kv_cache_config, {}, torch.device("cpu")) + attn_groups = [ + AttentionGroup(FakeSingleKVBackend, ["main"], main_layer, 0), + AttentionGroup(FakeSingleKVBackend, ["indexer"], indexer_layer, 1), + ] + + kv_caches = _reshape_kv_cache( + attn_groups, + raw_tensors, + "auto", + [kernel_block_size, kernel_block_size], + {}, + kv_cache_config, + ) + + assert kv_caches["main"].shape == (manager_blocks, 64, 328) + assert kv_caches["indexer"].shape == (manager_blocks * 4, 64, 132) + assert kv_caches["main"].is_contiguous() + assert kv_caches["indexer"].is_contiguous() + assert all(tensor.block_stride == 0 for tensor in tensors) + + def test_reshape_padded_diff_kv_cache_does_not_infer_kv_dim(): num_blocks = 3 spec = FullAttentionSpec( diff --git a/tests/v1/worker/test_gpu_block_table.py b/tests/v1/worker/test_gpu_block_table.py index 4c15915912ab..954f788781a5 100644 --- a/tests/v1/worker/test_gpu_block_table.py +++ b/tests/v1/worker/test_gpu_block_table.py @@ -177,3 +177,49 @@ def test_compute_slot_mappings_applies_padding_mask(): assert slot_mappings.cpu().tolist() == [ [32, PAD_SLOT_ID, 34, 48, PAD_SLOT_ID, PAD_SLOT_ID, PAD_SLOT_ID, PAD_SLOT_ID] ] + + +def test_compute_slot_mappings_mixed_sharded_and_replicated_groups(): + device = torch.device("cuda") + block_tables = BlockTables( + block_sizes=[64, 256], + max_num_reqs=1, + max_num_batched_tokens=8, + max_num_blocks_per_group=[1, 1], + device=device, + kernel_block_sizes=[64, 64], + cp_size=4, + cp_rank=2, + group_cp_sizes=[4, 1], + ) + block_tables.append_block_ids( + req_index=0, + new_block_ids=([5], [5]), + overwrite=True, + ) + block_tables.apply_staged_writes() + + idx_mapping = torch.tensor([0], dtype=torch.int32, device=device) + query_start_loc = torch.tensor([0, 8], dtype=torch.int32, device=device) + positions = torch.arange(8, dtype=torch.int64, device=device) + slot_mappings = block_tables.compute_slot_mappings( + idx_mapping, + query_start_loc, + positions, + num_tokens_padded=8, + ) + torch.accelerator.synchronize() + + assert slot_mappings.cpu().tolist() == [ + [ + PAD_SLOT_ID, + PAD_SLOT_ID, + 320, + PAD_SLOT_ID, + PAD_SLOT_ID, + PAD_SLOT_ID, + 321, + PAD_SLOT_ID, + ], + list(range(1280, 1288)), + ] diff --git a/vllm/v1/core/block_pool.py b/vllm/v1/core/block_pool.py index b42c7662d2ba..7e241ec14035 100644 --- a/vllm/v1/core/block_pool.py +++ b/vllm/v1/core/block_pool.py @@ -282,14 +282,16 @@ def cache_full_blocks( block_hash, kv_cache_group_id ) if blk.block_hash is not None: - # The only valid case where a "new full block" already has a - # hash is partial->full promotion of the same cache block. - assert ( - blk.block_hash_num_tokens is not None - and blk.block_hash_num_tokens < num_hash_tokens - ) - removed_hashes = self._remove_cached_block_hashes(blk) - self._emit_block_removed_events(removed_hashes) + assert blk.block_hash_num_tokens is not None + if blk.block_hash_num_tokens < num_hash_tokens: + removed_hashes = self._remove_cached_block_hashes(blk) + self._emit_block_removed_events(removed_hashes) + else: + # Lockstep groups attach group-specific keys to one block + # at the same token boundary. + assert blk.block_hash_num_tokens == num_hash_tokens + assert get_block_hash(blk.block_hash) == block_hash + assert get_group_id(blk.block_hash) != kv_cache_group_id self._insert_block_hash( block_hash_with_group_id, blk, diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index aebd694bb854..efc22ad7c9e1 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -11,6 +11,7 @@ from vllm.v1.core.kv_cache_utils import ( BlockHash, KVCacheBlock, + _use_lockstep_mla_allocation, is_deepseek_v4_hybrid_kv_cache_config, ) from vllm.v1.core.single_type_kv_cache_manager import ( @@ -119,6 +120,11 @@ def __init__( ) for i, kv_cache_group in enumerate(self.kv_cache_config.kv_cache_groups) ) + self.lockstep_mla_allocations = _use_lockstep_mla_allocation( + kv_cache_config.kv_cache_groups, + dcp_world_size, + pcp_world_size, + ) # A positive retention interval must be a multiple of the base hit granularity # (``scheduler_block_size``) to land on real cache-hit boundaries. @@ -164,31 +170,37 @@ def get_num_blocks_to_allocate( Returns: The number of blocks to allocate. """ - num_blocks_to_allocate = 0 + blocks_by_group: list[int] = [] for i, manager in enumerate(self.single_type_managers): if isinstance(manager, CrossAttentionManager): # For cross-attention, we issue a single static allocation # of blocks based on the number of encoder input tokens. - num_blocks_to_allocate += manager.get_num_blocks_to_allocate( - request_id, - num_encoder_tokens, - [], - 0, - 0, - num_encoder_tokens, - apply_admission_cap=apply_admission_cap, + blocks_by_group.append( + manager.get_num_blocks_to_allocate( + request_id, + num_encoder_tokens, + [], + 0, + 0, + num_encoder_tokens, + apply_admission_cap=apply_admission_cap, + ) ) else: - num_blocks_to_allocate += manager.get_num_blocks_to_allocate( - request_id, - num_tokens, - new_computed_blocks[i], - total_computed_tokens, - num_local_computed_tokens, - num_tokens_main_model, - apply_admission_cap=apply_admission_cap, + blocks_by_group.append( + manager.get_num_blocks_to_allocate( + request_id, + num_tokens, + new_computed_blocks[i], + total_computed_tokens, + num_local_computed_tokens, + num_tokens_main_model, + apply_admission_cap=apply_admission_cap, + ) ) - return num_blocks_to_allocate + if self.lockstep_mla_allocations: + return max(blocks_by_group, default=0) + return sum(blocks_by_group) def allocate_new_computed_blocks( self, @@ -208,6 +220,11 @@ def allocate_new_computed_blocks( num_local_computed_tokens: The number of local computed tokens. num_external_computed_tokens: The number of external computed tokens. """ + if self.lockstep_mla_allocations and num_external_computed_tokens: + raise NotImplementedError( + "External KV loads are not supported with lockstep MLA allocations." + ) + # A running request is already tracked in num_cached_block and won't # have new prefix-cache hits, so this is a no-op for it. if any( @@ -236,6 +253,17 @@ def allocate_new_computed_blocks( num_external_computed_tokens, ) + if self.lockstep_mla_allocations: + group_blocks = [ + manager.req_to_blocks[request_id] + for manager in self.single_type_managers + ] + block_ids = [block.block_id for block in group_blocks[0]] + assert all( + [block.block_id for block in blocks] == block_ids + for blocks in group_blocks[1:] + ) + def allocate_new_blocks( self, request_id: str, @@ -260,16 +288,38 @@ def allocate_new_blocks( Returns: The new allocated blocks. """ - return tuple( - manager.allocate_new_blocks( - request_id, - num_encoder_tokens - if isinstance(manager, CrossAttentionManager) - else num_tokens, - num_tokens_main_model, + if not self.lockstep_mla_allocations: + return tuple( + manager.allocate_new_blocks( + request_id, + num_encoder_tokens + if isinstance(manager, CrossAttentionManager) + else num_tokens, + num_tokens_main_model, + ) + for manager in self.single_type_managers ) - for manager in self.single_type_managers + + managers = self.single_type_managers + primary_ids = [ + block.block_id for block in managers[0].req_to_blocks[request_id] + ] + assert all( + [block.block_id for block in manager.req_to_blocks[request_id]] + == primary_ids + for manager in managers[1:] + ) + + new_blocks = managers[0].allocate_new_blocks( + request_id, + num_tokens, + num_tokens_main_model, ) + for manager in managers[1:]: + if new_blocks: + self.block_pool.touch(new_blocks) + manager.req_to_blocks[request_id].extend(new_blocks) + return tuple(list(new_blocks) for _ in managers) def cache_blocks(self, request: Request, num_computed_tokens: int) -> None: """ diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 01aa1196d714..fbd71d5bb232 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -1002,6 +1002,16 @@ def _pool_bytes_per_block( `available_memory` into `num_blocks`. Used to compute the effective KV cache capacity once `num_gpu_blocks_override` is applied. """ + parallel_config = vllm_config.parallel_config + if _use_lockstep_mla_allocation( + kv_cache_groups, + parallel_config.decode_context_parallel_size, + parallel_config.prefill_context_parallel_size, + ): + return sum( + page_size + for page_size, _ in _get_lockstep_mla_tensor_slots(kv_cache_groups) + ) if len(kv_cache_groups) == 1 and isinstance( kv_cache_groups[0].kv_cache_spec, UniformTypeKVCacheSpecs ): @@ -1313,10 +1323,91 @@ def _bucket_layers_by_page_size( return buckets +def _is_lockstep_mla_spec_layout( + kv_cache_specs: Iterable[KVCacheSpec], + dcp_world_size: int, + pcp_world_size: int, +) -> bool: + specs = list(kv_cache_specs) + if ( + len(specs) < 2 + or not isinstance(dcp_world_size, int) + or not isinstance(pcp_world_size, int) + ): + return False + cp_world_size = dcp_world_size * pcp_world_size + if cp_world_size <= 1: + return False + if not all(isinstance(spec, MLAAttentionSpec) for spec in specs): + return False + + replication_modes = {bool(getattr(spec, "dcp_replicated", False)) for spec in specs} + global_block_sizes = { + spec.block_size + if getattr(spec, "dcp_replicated", False) + else spec.block_size * cp_world_size + for spec in specs + } + return replication_modes == {False, True} and len(global_block_sizes) == 1 + + +def _use_lockstep_mla_allocation( + kv_cache_groups: list[KVCacheGroupSpec], + dcp_world_size: int, + pcp_world_size: int, +) -> bool: + """Return whether mixed MLA groups can share one block-ID namespace.""" + if len(kv_cache_groups) < 2: + return False + + layer_specs: list[KVCacheSpec] = [] + for group in kv_cache_groups: + group_spec = group.kv_cache_spec + group_layer_specs = ( + list(group_spec.kv_cache_specs.values()) + if isinstance(group_spec, UniformTypeKVCacheSpecs) + else [group_spec] + ) + group_modes = { + bool(getattr(spec, "dcp_replicated", False)) for spec in group_layer_specs + } + if len(group_modes) != 1: + return False + layer_specs.extend(group_layer_specs) + + return _is_lockstep_mla_spec_layout(layer_specs, dcp_world_size, pcp_world_size) + + +def _get_lockstep_mla_tensor_slots( + kv_cache_groups: list[KVCacheGroupSpec], +) -> list[tuple[int, list[str]]]: + slots: list[tuple[int, list[str]]] = [] + for group in kv_cache_groups: + group_spec = group.kv_cache_spec + if isinstance(group_spec, UniformTypeKVCacheSpecs): + slots.extend( + (group_spec.kv_cache_specs[layer_name].page_size_bytes, [layer_name]) + for layer_name in group.layer_names + ) + else: + slots.extend( + (group_spec.page_size_bytes, [layer_name]) + for layer_name in group.layer_names + ) + return slots + + def _use_packed_kv_cache_config( vllm_config: VllmConfig, kv_cache_groups: list[KVCacheGroupSpec], ) -> bool: + parallel_config = vllm_config.parallel_config + if _use_lockstep_mla_allocation( + kv_cache_groups, + parallel_config.decode_context_parallel_size, + parallel_config.prefill_context_parallel_size, + ): + return True is_dsv4 = all( isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs) for group in kv_cache_groups @@ -1346,11 +1437,27 @@ def _get_kv_cache_config_packed( groups at the same slot share a tensor (they have independent block tables so block-id namespaces never collide). Each emitted tensor aliases one physical backing allocation, with per-block data laid out contiguously. - """ + Lockstep MLA groups instead get distinct contiguous tensors because their + block IDs intentionally coincide. + """ + parallel_config = vllm_config.parallel_config + if _use_lockstep_mla_allocation( + kv_cache_groups, + parallel_config.decode_context_parallel_size, + parallel_config.prefill_context_parallel_size, + ): + tensor_slots = _get_lockstep_mla_tensor_slots(kv_cache_groups) + total_num_bytes_per_block = sum(page_size for page_size, _ in tensor_slots) + num_blocks = available_memory // total_num_bytes_per_block + num_blocks = may_override_num_blocks(vllm_config, num_blocks) + return num_blocks, [ + KVCacheTensor(size=page_size * num_blocks, shared_by=slot) + for page_size, slot in tensor_slots + ] + # buckets = {page_size: [[layer_names], [layer_names], ...]} buckets = _bucket_layers_by_page_size(kv_cache_groups) total_num_bytes_per_block = sum(ps * len(slots) for ps, slots in buckets.items()) - num_blocks = available_memory // total_num_bytes_per_block num_blocks = may_override_num_blocks(vllm_config, num_blocks) @@ -1557,6 +1664,8 @@ def unify_hybrid_kv_cache_specs(kv_cache_spec: dict[str, KVCacheSpec]): def group_and_unify_kv_cache_specs( kv_cache_spec: dict[str, KVCacheSpec], + dcp_world_size: int = 1, + pcp_world_size: int = 1, ) -> list[UniformTypeKVCacheSpecs] | None: """ Group the KV cache specs and unify each group into one UniformTypeKVCacheSpecs. @@ -1565,29 +1674,33 @@ def group_and_unify_kv_cache_specs( has_swa = any( isinstance(spec, SlidingWindowMLASpec) for spec in kv_cache_spec.values() ) - # DFlash-under-DCP draft: full-attention layers replicated on every DCP - # rank. They have a different page size than the MLA target and need their - # own group, but the DeepseekV4 multi-group allocator (with page-size - # padding) handles exactly that, so route them through here too. + lockstep_mla_layout = _is_lockstep_mla_spec_layout( + kv_cache_spec.values(), dcp_world_size, pcp_world_size + ) + # Replicated layers need a group separate from sharded layers because they + # have different token ownership and block-table semantics. has_repl = any( getattr(spec, "dcp_replicated", False) - and not isinstance(spec, MLAAttentionSpec) + and (not isinstance(spec, MLAAttentionSpec) or lockstep_mla_layout) for spec in kv_cache_spec.values() ) if not (has_swa or has_repl): return None - # SlidingWindowMLASpec models with uniform page sizes don't need tuple packing. + # Other uniform page layouts do not need tuple packing. page_sizes = {spec.page_size_bytes for spec in kv_cache_spec.values()} - if len(page_sizes) <= 1: + if len(page_sizes) <= 1 and not lockstep_mla_layout: return None mla_specs: dict[str, KVCacheSpec] = {} grouped_swa_mla_specs: dict[tuple[int, int, bool, bool], dict[str, KVCacheSpec]] = ( defaultdict(dict) ) - # dcp_replicated non-MLA groups (e.g. the DFlash draft), keyed by block_size. - grouped_repl_specs: dict[tuple[int], dict[str, KVCacheSpec]] = defaultdict(dict) + # Keep incompatible replicated types and page layouts separate. This + # includes DFlash drafts and replicated MLA indexer caches. + grouped_repl_specs: dict[ + tuple[int] | tuple[str, int, int], dict[str, KVCacheSpec] + ] = defaultdict(dict) # NOTE: Here we group SWA layers by (block_size, sliding_window, # dcp_replicated, dcp_sharded), which separates SWA layers, C4I+C4A # layers, C128A layers, and replicated compressor-state groups. @@ -1601,10 +1714,17 @@ def group_and_unify_kv_cache_specs( spec.dcp_sharded, ) ][name] = spec + elif getattr(spec, "dcp_replicated", False) and ( + not isinstance(spec, MLAAttentionSpec) or lockstep_mla_layout + ): + key = ( + (type(spec).__name__, spec.block_size, spec.page_size_bytes) + if isinstance(spec, MLAAttentionSpec) + else (spec.block_size,) + ) + grouped_repl_specs[key][name] = spec elif isinstance(spec, MLAAttentionSpec): mla_specs[name] = spec - elif getattr(spec, "dcp_replicated", False): - grouped_repl_specs[(spec.block_size,)][name] = spec if len(mla_specs) == 0: # No full-MLA group to anchor the DeepseekV4 layout; let the generic @@ -1698,11 +1818,17 @@ def _get_kv_cache_groups_uniform_groups( ] swa_mla_specs = grouped_specs[1:] - # Non-first groups are SWA-MLA, full-attention dcp_replicated drafts, or - # sliding-window dcp_replicated drafts. All are padded to MLA buckets - # identically. + # Non-first groups are SWA-MLA or DCP-replicated caches. assert all( - isinstance(spec, (SlidingWindowMLASpec, FullAttentionSpec, SlidingWindowSpec)) + isinstance( + spec, + ( + MLAAttentionSpec, + SlidingWindowMLASpec, + FullAttentionSpec, + SlidingWindowSpec, + ), + ) for group in swa_mla_specs for spec in group.kv_cache_specs.values() ) @@ -1716,19 +1842,30 @@ def _get_kv_cache_groups_uniform_groups( for sm_spec in swa_mla_specs: sm_page_sizes = sm_spec.get_page_sizes() layers_per_size: dict[int, list[str]] = defaultdict(list) - assert max(sm_page_sizes) <= max(all_page_sizes) + is_replicated_mla = all( + isinstance(spec, MLAAttentionSpec) + and getattr(spec, "dcp_replicated", False) + for spec in sm_spec.kv_cache_specs.values() + ) + if not is_replicated_mla: + assert max(sm_page_sizes) <= max(all_page_sizes) # Unify page size by padding layers' page_size to the nearest larger page_size. # Compute candidate (nearest larger page_size) for each unique page size. size_to_candidate: dict[int, int] = {} for ps in sm_page_sizes: - size_to_candidate[ps] = min(x for x in all_page_sizes if x >= ps) + if ps <= max(all_page_sizes): + size_to_candidate[ps] = min(x for x in all_page_sizes if x >= ps) # Pad and collect layer names per page size. for layer_name, layer_spec in sm_spec.kv_cache_specs.items(): current_size = layer_spec.page_size_bytes - candidate = size_to_candidate[current_size] - if current_size < candidate: - object.__setattr__(layer_spec, "page_size_padded", candidate) + if is_replicated_mla: + candidate = current_size + else: + assert current_size <= max(all_page_sizes) + candidate = size_to_candidate[current_size] + if current_size < candidate: + object.__setattr__(layer_spec, "page_size_padded", candidate) layers_per_size[candidate].append(layer_name) # NOTE(yifan): for now, inside a UniformKV group, each page_size should # have the same number of layers. This also means we don't need to pad layers @@ -1816,7 +1953,11 @@ def get_kv_cache_groups( # full attention, or all layers are sliding window attention with the # same window size). Put all layers into one group. return _get_kv_cache_groups_uniform_type(uniform_spec) - elif grouped_specs := group_and_unify_kv_cache_specs(kv_cache_spec): + elif grouped_specs := group_and_unify_kv_cache_specs( + kv_cache_spec, + vllm_config.parallel_config.decode_context_parallel_size, + vllm_config.parallel_config.prefill_context_parallel_size, + ): # DeepseekV4 case: All layers need the same number of token slots, # yet some layers are full attention while others are sliding window # attention in different sizes. Need to group layers into multiple @@ -1903,6 +2044,21 @@ def _max_memory_usage_bytes_from_groups( if not kv_cache_groups: return 0 + parallel_config = vllm_config.parallel_config + if _use_lockstep_mla_allocation( + kv_cache_groups, + parallel_config.decode_context_parallel_size, + parallel_config.prefill_context_parallel_size, + ): + request_blocks = max( + cdiv( + group.kv_cache_spec.max_memory_usage_bytes(vllm_config), + group.kv_cache_spec.page_size_bytes, + ) + for group in kv_cache_groups + ) + return request_blocks * _pool_bytes_per_block(vllm_config, kv_cache_groups) + if len(kv_cache_groups) == 1 and isinstance( kv_cache_groups[0].kv_cache_spec, UniformTypeKVCacheSpecs ):