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 5, 16:00 UTC
[KVTransfer] Fix TpKVTopology.is_kv_replicated equality case (#38179)
Signed-off-by: JianDan0212 <zhangyj0212@gmail.com> Co-authored-by: Nicolò Lucchesi <nlucches@redhat.com>
yjz committed
Apr 1, 2026 at 18:41 UTC
992368522f69aa2df15a3f9539d138498e52c5bf
1 file changed
+3
-1
vllm/distributed/kv_transfer/kv_connector/utils.py
+3
-1
@@ -457,9 +457,11 @@ class TpKVTopology:
457
"""
458
Whether the KV cache is replicated across TP workers due to the
459
number of TP workers being greater than the number of KV heads.
460
+ When they are equal, each TP rank still owns one distinct KV head,
461
+ so this is not considered replication.
462
"""
463
tp_size = self.remote_tp_size[engine_id]
462
- return tp_size // self.total_num_kv_heads >= 1
464
+ return tp_size > self.total_num_kv_heads
465
466
def replicates_kv_cache(self, remote_engine_id: EngineId) -> bool:
467
# MLA is always replicated as the hidden dim can't be split.