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 8, 21:30 UTC
[Bugfix] Disable TRTLLM attention when KV transfer is enabled (#33192)
Signed-off-by: Zhanqiu Hu <zh338@cornell.edu>
zhanqiuhu committed
Feb 4, 2026 at 19:49 UTC
bbe0574d8e51c1c5935aeff9e92040c61d1d59c5
1 file changed
+17
vllm/v1/attention/backends/flashinfer.py
+17
@@ -573,6 +573,20 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
573
# try to use fp8 q if kv cache is fp8, and will fall back to model dtype
574
# if TRTLLM attention kernel is not used when building attn metadata
575
can_use_trtllm = can_use_trtllm_attention(self.num_qo_heads, self.num_kv_heads)
576
+
577
+ # TRTLLM attention requires strictly contiguous KV cache tensors.
578
+ # When KV transfer (P/D disaggregation) is enabled, the KV cache may be
579
+ # permuted into non-contiguous views, which causes assertion failures.
580
+ self._kv_transfer_enabled = vllm_config.kv_transfer_config is not None
581
+ if can_use_trtllm and self._kv_transfer_enabled:
582
+ logger.info_once(
583
+ "TRTLLM attention is disabled because KV transfer "
584
+ "(P/D disaggregation) is enabled. TRTLLM attention requires "
585
+ "strictly contiguous KV cache tensors which may not be "
586
+ "guaranteed with KV transfer."
587
+ )
588
+ can_use_trtllm = False
589
+
590
if (
591
can_use_trtllm
592
and not vllm_config.attention_config.disable_flashinfer_q_quantization
@@ -822,6 +836,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
836
has_sinks=self.has_sinks,
837
has_spec=uses_spec_reorder,
838
)
839
+ # KV transfer requires non-contiguous KV cache views, incompatible with TRTLLM
840
+ if self._kv_transfer_enabled:
841
+ prefill_use_trtllm = False
842
decode_use_trtllm = (
843
self.use_trtllm_decode_attention and self.dcp_world_size <= 1
844
)