From 2537629eb22a5e3134cb8b09d367e07efbbfdf00 Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Mon, 5 Oct 2026 21:27:50 -0400 Subject: [PATCH] [Bugfix][KVConnector][NIXL] Size per-region replicate flags by each region's block count (#60107) Signed-off-by: Lucas Wilkinson Co-authored-by: Claude Opus 5.5 (1M context) Co-authored-by: Matthew Bonanni --- .../kv_connector/v1/nixl/base_worker.py | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) 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 a9a0d1980a..dade5056ee 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 @@ -340,7 +340,7 @@ class NixlBaseConnectorWorker: ) # Per-FA-descriptor replicate flag, in _build_fa_local emission order. - fa_desc_replicated = self._fa_desc_replicated(num_fa_descs) + fa_desc_replicated = self._fa_desc_replicated(num_fa_descs, block_size_ratio) sharded_desc_end = len(src_blocks_data) - ( self._logical_num_blocks if self._ple_region_index is not None else 0 ) @@ -383,19 +383,22 @@ class NixlBaseConnectorWorker: """ return tp_ratio < 0 and (not self.use_mla or len(plan.all_source_ranks) > 1) - def _fa_desc_replicated(self, num_fa_descs: int) -> list[bool]: - """Per-FA-descriptor replicate flag, in _build_fa_local emission order - (region-major; one desc per block, with K/V packed). Length ``num_fa_descs``. + def _fa_desc_replicated( + self, num_fa_descs: int, block_size_ratio: int = 1 + ) -> list[bool]: + """Per-FA-descriptor replicate flag, in _build_fa_local emission order: + region-major, ``region_num_blocks[i] * block_size_ratio`` descs per + region. Length ``num_fa_descs``. """ assert self.transfer_topo is not None n_regions = len(self.block_len_per_layer) if n_regions == 0 or self.num_regions == 0: return [False] * num_fa_descs - nblk = num_fa_descs // self.num_regions + region_indices = self._transfer_layer_region_indices or range(n_regions) flags: list[bool] = [] - for i in range(n_regions): - replicated = self._is_region_replicated(i) - flags.extend([replicated] * nblk) + for i in region_indices: + num_descs = self.region_num_blocks[i] * block_size_ratio + flags.extend([self._is_region_replicated(i)] * num_descs) assert len(flags) == num_fa_descs, ( f"FA desc flags {len(flags)} != num_fa_descs {num_fa_descs}" )