From f81865e14949755e6108fd0af727c8ddd60cc9a4 Mon Sep 17 00:00:00 2001 From: Pz1116 Date: Mon, 3 Aug 2026 15:20:39 +0800 Subject: [PATCH] fix(kv_pool): transfer registered MTP layers Build layerwise execution and group offsets from the KV caches registered on each worker. This includes MTP caches while keeping PP-local task indices and dense per-group storage indices for hybrid models. Signed-off-by: Pz1116 --- .../ascend_store/test_pool_worker.py | 51 ++++++++ .../kv_pool/ascend_store/pool_worker.py | 113 ++++++++---------- 2 files changed, 103 insertions(+), 61 deletions(-) diff --git a/tests/ut/distributed/ascend_store/test_pool_worker.py b/tests/ut/distributed/ascend_store/test_pool_worker.py index 09d98465108a..27de0fae24c4 100644 --- a/tests/ut/distributed/ascend_store/test_pool_worker.py +++ b/tests/ut/distributed/ascend_store/test_pool_worker.py @@ -194,6 +194,57 @@ def test_lookup_reuses_grouped_hashes_for_hit_resolution(self): self.assertEqual(hit, 32) self.assertEqual(rehash.call_count, 2) + def test_registered_layer_layout_includes_mtp_multi_group_and_pp(self): + cls = self._make_worker_class() + worker = object.__new__(cls) + worker.num_layers = 2 + worker.hf_config = MagicMock(num_hidden_layers=4) + caches = [MagicMock(), MagicMock(), MagicMock()] + for address, cache in zip((1000, 2000, 3000), caches): + cache.data_ptr.return_value = address + worker.kv_caches = { + "model.layers.2.self_attn": (caches[0],), + "model.layers.3.self_attn": (caches[1],), + "model.mtp.0.self_attn": (caches[2],), + } + worker._get_cache_block_metadata = MagicMock(return_value=(160, 160, 160, 1)) + worker.group_kv_caches_base_addr = {} + worker.group_block_len = {} + worker.group_block_stride = {} + worker.group_num_layers = {} + + worker._infer_cache_group_metadata(0, list(worker.kv_caches)) + self.assertEqual(worker.group_kv_caches_base_addr[0], [1000, 2000, 3000]) + self.assertEqual(worker.group_num_layers[0], 3) + worker.layer_load_tasks = [[], []] + worker.layer_save_tasks = [[], []] + + worker._configure_registered_layerwise_layers( + [ + ( + 0, + [ + "model.layers.2.self_attn", + "model.layers.3.self_attn", + "model.mtp.0.self_attn", + ], + ), + (1, ["model.layers.2.indexer", "model.layers.3.indexer"]), + ] + ) + + self.assertEqual(worker.num_layers, 3) + self.assertEqual( + worker.local_layer_to_group_layers, + { + 0: [(0, 0), (1, 0)], + 1: [(0, 1), (1, 1)], + 2: [(0, 2)], + }, + ) + self.assertEqual(worker._extract_physical_layer_index("model.layers.4.self_attn"), 4) + self.assertEqual(len(worker.layer_load_tasks), 3) + class TestKVPoolWorkerInit(unittest.TestCase): """Test KVPoolWorker initialization with mocked dependencies.""" diff --git a/vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/pool_worker.py b/vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/pool_worker.py index 89a4d83efed5..7fde9441b67e 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/pool_worker.py +++ b/vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/pool_worker.py @@ -275,38 +275,10 @@ def _init_state_vars(self) -> None: self._allocated_gvas: dict[str, int] = {} def _init_layerwise_config(self) -> None: - # Build mapping: physical_layer -> [(group_id, layer_idx_in_group), ...] - # layer_idx_in_group is the index of the physical layer within the - # group (not the index in layer_names). Multiple layer_names at the - # same physical layer (e.g. indexer.k_cache + attn) are treated as - # multiple cache tensors of ONE layer (caches_per_layer > 1). - self.physical_layer_to_group_layers: dict[int, list[tuple[int, int]]] = {} - - if self.kv_cache_config is not None and self.num_kv_cache_groups > 1: - for group_id, group_spec in enumerate(self.kv_cache_config.kv_cache_groups): - # Map each unique physical layer to a sequential layer_idx_in_group - phys_to_layer_idx: dict[int, int] = {} - for layer_name in group_spec.layer_names: - physical_layer = self._extract_physical_layer_index(layer_name) - if physical_layer >= getattr(self.hf_config, "num_hidden_layers", self.num_layers): - continue - if physical_layer not in phys_to_layer_idx: - phys_to_layer_idx[physical_layer] = len(phys_to_layer_idx) - - # Add one entry per unique physical layer (no duplicates) - for physical_layer, layer_idx_in_group in phys_to_layer_idx.items(): - existing = self.physical_layer_to_group_layers.setdefault(physical_layer, []) - entry = (group_id, layer_idx_in_group) - if entry not in existing: - existing.append(entry) - - logger.info( - "layerwise group %d: %d layer_names, %d unique physical layers, caches_per_layer=%d", - group_id, - len(group_spec.layer_names), - len(phys_to_layer_idx), - len(group_spec.layer_names) // max(1, len(phys_to_layer_idx)), - ) + # Registration provides the authoritative layer layout. In + # particular, model_config.get_num_layers() only reports the local PP + # target layers and does not include registered MTP draft layers. + self.local_layer_to_group_layers: dict[int, list[tuple[int, int]]] = {} self.layer_load_tasks: list[list[LayerTransferTask]] = [[] for _ in range(self.num_layers)] self.layer_save_tasks: list[list[LayerTransferTask]] = [[] for _ in range(self.num_layers)] @@ -318,10 +290,9 @@ def _init_layerwise_config(self) -> None: self.sync_save_events: list[torch.npu.Event] | None = None logger.info( - "layerwise config: num_layers=%d num_groups=%d physical_layer_to_group_layers_sample=%s", + "layerwise config: num_layers=%d num_groups=%d", self.num_layers, self.num_kv_cache_groups, - {k: v for k, v in list(self.physical_layer_to_group_layers.items())[:3]}, ) def _build_group_layer_builders(self) -> list[LayerBatchBuilder]: @@ -613,8 +584,6 @@ def _infer_cache_group_metadata(self, group_id: int, layer_names: list[str]): physical_layers = set() for layer_name in layer_names: phys = self._extract_physical_layer_index(layer_name) - if phys >= getattr(self.hf_config, "num_hidden_layers", self.num_layers): - continue physical_layers.add(phys) cache_or_caches = self.kv_caches[layer_name] for cache in self._as_cache_tuple(cache_or_caches): @@ -628,6 +597,38 @@ def _infer_cache_group_metadata(self, group_id: int, layer_names: list[str]): self.group_block_stride[group_id] = group_block_strides self.group_num_layers[group_id] = len(physical_layers) + def _configure_registered_layerwise_layers(self, cache_groups: list[tuple[int, list[str]]]) -> None: + """Map registered physical layers to dense local and group indices.""" + groups_by_physical: dict[int, list[tuple[int, int]]] = {} + + for group_id, layer_names in cache_groups: + group_physical_layers: list[int] = [] + seen_in_group: set[int] = set() + for layer_name in layer_names: + physical_layer = self._extract_physical_layer_index(layer_name) + if physical_layer not in seen_in_group: + seen_in_group.add(physical_layer) + group_physical_layers.append(physical_layer) + + for layer_idx_in_group, physical_layer in enumerate(group_physical_layers): + groups_by_physical.setdefault(physical_layer, []).append((group_id, layer_idx_in_group)) + + physical_layers = sorted(groups_by_physical) + self.local_layer_to_group_layers = { + local_layer: groups_by_physical[physical_layer] + for local_layer, physical_layer in enumerate(physical_layers) + } + + original_num_layers = self.num_layers + self.num_layers = len(physical_layers) + self.layer_load_tasks = [[] for _ in range(self.num_layers)] + self.layer_save_tasks = [[] for _ in range(self.num_layers)] + logger.info( + "KVPoolWorker: configured %d registered layers (was %d).", + self.num_layers, + original_num_layers, + ) + def _align_kv_ptrs(self, registered_regions: dict[int, tuple[int, int]]): """ In hybrid scenario, where a KVCacheTensor is shared by multiple layers, @@ -704,27 +705,17 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): ptrs = [start for start, _ in registered_regions.values()] lengths = [end - start for start, end in registered_regions.values()] + cache_groups: list[tuple[int, list[str]]] if self.kv_cache_config is not None and self.use_hybrid: - for group_id, group_spec in enumerate(self.kv_cache_config.kv_cache_groups): - self._infer_cache_group_metadata(group_id, group_spec.layer_names) + cache_groups = [ + (group_id, [layer_name for layer_name in group_spec.layer_names if layer_name in kv_caches]) + for group_id, group_spec in enumerate(self.kv_cache_config.kv_cache_groups) + ] else: - self._infer_cache_group_metadata(0, list(kv_caches.keys())) - - # group_num_layers is computed from the actual kv_caches dict which - # includes ALL attention layers (main + MTP). For single-group models, - # sum(group_num_layers.values()) equals the physical layer count - # (including MTP). For multi-group models, it counts (group, layer) - # pairs which is NOT the physical layer count — keep the original - # num_layers (physical layers) in that case. - original_num_layers = self.num_layers - new_num_layers = sum(self.group_num_layers.values()) - if self.num_kv_cache_groups == 1 and new_num_layers != original_num_layers: - self.num_layers = new_num_layers - logger.info( - "KVPoolWorker: updated num_layers %d -> %d (includes MTP/spec-decode draft layers).", - original_num_layers, - self.num_layers, - ) + cache_groups = [(0, list(kv_caches.keys()))] + for group_id, layer_names in cache_groups: + self._infer_cache_group_metadata(group_id, layer_names) + self._configure_registered_layerwise_layers(cache_groups) self.page_size_bytes = sum(self.block_len) self.token_database.set_group_buffers( @@ -1344,17 +1335,17 @@ def _build_shared_load_data(self) -> None: def process_layer_data(self, requests: list[ReqMeta]) -> None: if not requests: return - for physical_layer in range(self.num_layers): - group_layers = self.physical_layer_to_group_layers.get(physical_layer, [(0, physical_layer)]) + for local_layer in range(self.num_layers): + group_layers = self.local_layer_to_group_layers.get(local_layer, [(0, local_layer)]) for group_id, layer_idx_in_group in group_layers: - self._process_save_for_layer_batch(requests, physical_layer, group_id, layer_idx_in_group) + self._process_save_for_layer_batch(requests, local_layer, group_id, layer_idx_in_group) self._alloc_gvas_for_save(requests) self._build_shared_save_data() self._prepare_load_gvas(requests) - for physical_layer in range(self.num_layers): - group_layers = self.physical_layer_to_group_layers.get(physical_layer, [(0, physical_layer)]) + for local_layer in range(self.num_layers): + group_layers = self.local_layer_to_group_layers.get(local_layer, [(0, local_layer)]) for group_id, layer_idx_in_group in group_layers: - self._process_load_for_layer_batch(requests, physical_layer, group_id, layer_idx_in_group) + self._process_load_for_layer_batch(requests, local_layer, group_id, layer_idx_in_group) self._build_shared_load_data() def _submit_ready_layer_loads(self) -> None: