@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 57a96e26c9
Fork 0

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 7, 11:30 UTC

Revert "[Bugfix] Disable TRTLLM attention with KV transfer enabled (#33192)" (#34832)

Signed-off-by: Zhanqiu Hu <zh338@cornell.edu>

zhanqiuhu committed Mar 1, 2026 at 17:32 UTC 57a96e26c913cb9fae96c9e600fa4ff10dc40a1a
1 file changed -17
vllm/v1/attention/backends/flashinfer.py
-17
@@ -575,20 +575,6 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
575 # try to use fp8 q if kv cache is fp8, and will fall back to model dtype
576 # if TRTLLM attention kernel is not used when building attn metadata
577 can_use_trtllm = can_use_trtllm_attention(self.num_qo_heads, self.num_kv_heads)
578 -
579 - # TRTLLM attention requires strictly contiguous KV cache tensors.
580 - # When KV transfer (P/D disaggregation) is enabled, the KV cache may be
581 - # permuted into non-contiguous views, which causes assertion failures.
582 - self._kv_transfer_enabled = vllm_config.kv_transfer_config is not None
583 - if can_use_trtllm and self._kv_transfer_enabled:
584 - logger.info_once(
585 - "TRTLLM attention is disabled because KV transfer "
586 - "(P/D disaggregation) is enabled. TRTLLM attention requires "
587 - "strictly contiguous KV cache tensors which may not be "
588 - "guaranteed with KV transfer."
589 - )
590 - can_use_trtllm = False
591 -
578 if (
579 can_use_trtllm
580 and not vllm_config.attention_config.disable_flashinfer_q_quantization
@@ -865,9 +851,6 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
851 has_sinks=self.has_sinks,
852 has_spec=uses_spec_reorder,
853 )
868 - # KV transfer requires non-contiguous KV cache views, incompatible with TRTLLM
869 - if self._kv_transfer_enabled:
870 - prefill_use_trtllm = False
854 decode_use_trtllm = (
855 self.use_trtllm_decode_attention and self.dcp_world_size <= 1
856 )