iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Bugfix][KVConnector][NIXL] Size per-region replicate flags by each region's block count (#60107)

Browse Source 
Signed-off-by: Lucas Wilkinson <lwilkinson@neuralmagic.com>
Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>
This commit is contained in:
Lucas Wilkinson
2026-10-06 01:27:50 +00:00
committed by GitHub
co-authored by Claude Opus 5.5 Matthew Bonanni
parent 0e8f8ecba8
commit 2537629eb2
1 changed files with 11 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+11 -8
View File
@@ -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}"
)