A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 13, 02:00 UTC
[Bugfix][NIXL] Fix `block_size_ratio` when logical !=physical blocks (#28925)
Signed-off-by: NickLucche <nlucches@redhat.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Nicolò Lucchesi committed
Nov 18, 2025 at 15:07 UTC
184b12fdc6dce87485e3bd793e13e90421f93924
1 file changed
+12
-6
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+12
-6
@@ -677,12 +677,13 @@ class NixlConnectorWorker:
677
mapping between local and remote TP workers.
678
"""
679
680
- tp_size: int
680
tp_rank: int
681
remote_tp_size: dict[EngineId, int]
682
is_mla: bool
683
total_num_kv_heads: int
684
attn_backend: type[AttentionBackend]
685
+ engine_id: EngineId
686
+ remote_block_size: dict[EngineId, int]
687
688
def __post_init__(self):
689
# Figure out whether the first dimension of the cache is K/V
@@ -710,8 +711,13 @@ class NixlConnectorWorker:
711
self.is_mla or self._use_pallas or self.is_kv_layout_blocks_first
712
)
713
713
- block_size: int
714
- remote_block_size: dict[EngineId, int]
714
+ @property
715
+ def tp_size(self) -> int:
716
+ return self.remote_tp_size[self.engine_id]
717
+
718
+ @property
719
+ def block_size(self) -> int:
720
+ return self.remote_block_size[self.engine_id]
721
722
def tp_ratio(
723
self,
@@ -957,13 +963,12 @@ class NixlConnectorWorker:
963
self.xfer_stats = NixlKVConnectorStats()
964
965
self.kv_topo = self.TpKVTopology(
960
- tp_size=self.world_size,
966
tp_rank=self.tp_rank,
967
+ engine_id=self.engine_id,
968
remote_tp_size=self._tp_size, # shared state
969
+ remote_block_size=self._block_size, # shared state
970
is_mla=self.use_mla,
971
total_num_kv_heads=self.model_config.get_total_num_kv_heads(),
965
- block_size=self.block_size,
966
- remote_block_size=self._block_size,
972
attn_backend=backend,
973
)
974
self._use_pallas = self.kv_topo._use_pallas
@@ -1185,6 +1190,7 @@ class NixlConnectorWorker:
1190
self.block_size // kernel_block_size
1191
)
1192
self.block_size = kernel_block_size
1193
+ self._block_size[self.engine_id] = kernel_block_size
1194
1195
seen_base_addresses.append(base_addr)
1196
curr_tensor_size_bytes = cache.numel() * cache.element_size()