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 15, 04:45 UTC
[P/D][NixlConnector] Support `tp_size > num_kv_heads` deployments (#19691)
Signed-off-by: NickLucche <nlucches@redhat.com> Co-authored-by: Nick Hill <nhill@redhat.com>
Nicolò Lucchesi committed
Jun 23, 2025 at 07:41 UTC
2ebff5b77c049b3e620d5f79f02acbcbbc09bade
1 file changed
+10
-5
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+10
-5
@@ -22,6 +22,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
22
from vllm.distributed.parallel_state import (
23
get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size,
24
get_tp_group)
25
+from vllm.distributed.utils import divide
26
from vllm.logger import init_logger
27
from vllm.platforms import _Backend
28
from vllm.utils import make_zmq_path, make_zmq_socket, round_down
@@ -679,11 +680,15 @@ class NixlConnectorWorker:
680
681
# Number of D TP workers reading from a single P TP worker. This is
682
# 1 when P and D `--tensor-parallel-size` match.
682
- assert self._tp_size[self.engine_id] % self._tp_size[engine_id] == 0, (
683
- "Local TP size must be divisible by remote TP size.")
684
- tp_ratio = self._tp_size[self.engine_id] // self._tp_size[engine_id]
683
+ tp_ratio = divide(self._tp_size[self.engine_id],
684
+ self._tp_size[engine_id])
685
assert tp_ratio > 0, "Decode TP cannot be smaller than prefill TP"
686
- if self.use_mla:
686
+
687
+ # Handle tp_size>num_kv_heads: replicate KV cache.
688
+ total_num_kv_heads = self.model_config.get_total_num_kv_heads()
689
+ is_kv_replicated = self._tp_size[engine_id] // total_num_kv_heads >= 1
690
+
691
+ if self.use_mla or is_kv_replicated:
692
# With MLA the only difference is in the number of blocks.
693
remote_block_size = nixl_agent_meta.block_len // (
694
self.slot_size_bytes)
@@ -720,7 +725,7 @@ class NixlConnectorWorker:
725
self.kv_caches_base_addr[
726
engine_id] = nixl_agent_meta.kv_caches_base_addr
727
rank_offset = self.tp_rank % tp_ratio * self.block_len \
723
- if not self.use_mla else 0
728
+ if not (self.use_mla or is_kv_replicated) else 0
729
# Register all remote blocks, but only the corresponding kv heads.
730
for base_addr in nixl_agent_meta.kv_caches_base_addr:
731
for block_id in range(nixl_agent_meta.num_blocks):