Merge pull request #5 from GirasoleY/fix/pr48180-dcp-prefill-no-trtllm

[Bugfix] Gate DCP prefill and varlen spec decode on CP-aware paths (for #48180)

Author: pavanimajety

tests/kernels/attention/test_use_trtllm_attention.py

@@ -154,6 +154,17 @@ def test_use_dcp_fallback(_mock):
     assert _call(dcp_world_size=2) is False
 
 
+@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
+def test_use_dcp_fallback_prefill(_mock):
+    # TRTLLM prefill has no cross-rank LSE combine for DCP-sharded KV.
+    assert _call(dcp_world_size=2, is_prefill=True) is False
+
+
+@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
+def test_use_dcp_fallback_prefill_force_on_still_false(_mock):
+    assert _call(dcp_world_size=2, is_prefill=True, force_use_trtllm=True) is False
+
+
 @patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=False)
 def test_use_platform_unsupported(_mock):
     assert _call() is False

tests/v1/attention/test_flashinfer_dcp_spec_reorder.py

@@ -0,0 +1,72 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""FlashInfer GQA builder: reorder threshold under DCP with spec decode."""
+
+import pytest
+import torch
+
+from tests.v1.attention.utils import create_vllm_config
+from vllm.config import SpeculativeConfig
+from vllm.platforms import current_platform
+from vllm.v1.attention.backends import flashinfer as flashinfer_backend
+from vllm.v1.attention.backends.flashinfer import (
+    FlashInferDecodeKernel,
+    FlashInferMetadataBuilder,
+)
+from vllm.v1.attention.backends.utils import PerLayerParameters
+from vllm.v1.kv_cache_interface import FullAttentionSpec
+
+if not current_platform.is_cuda():
+    pytest.skip("FlashInfer backend requires a CUDA platform.", allow_module_level=True)
+
+
+def test_flashinfer_gqa_dcp_spec_decode_clamps_reorder_threshold(monkeypatch):
+    """trtllm-gen decode receives no cp_rank/global-seq-len information, so its
+    end-aligned causal mask is wrong for q_len > 1 over the DCP-interleaved
+    local KV shard. The builder must keep reorder_batch_threshold at 1 under
+    DCP so spec queries take the (DCP-aware) prefill path instead.
+    """
+    vllm_config = create_vllm_config(max_model_len=1024)
+    vllm_config.parallel_config.decode_context_parallel_size = 2
+    vllm_config.speculative_config = SpeculativeConfig(
+        method="ngram", num_speculative_tokens=3
+    )
+
+    monkeypatch.setattr(
+        flashinfer_backend, "can_use_trtllm_attention", lambda *args, **kwargs: True
+    )
+    monkeypatch.setattr(
+        FlashInferMetadataBuilder,
+        "_get_flashinfer_trtllm_api_decode_kernel",
+        staticmethod(lambda: FlashInferDecodeKernel.TRTLLM_GEN),
+    )
+    monkeypatch.setattr(
+        flashinfer_backend,
+        "get_per_layer_parameters",
+        lambda *args, **kwargs: {
+            "layer.0": PerLayerParameters(
+                window_left=-1, logits_soft_cap=None, sm_scale=0.1, has_sinks=False
+            )
+        },
+    )
+
+    kv_cache_spec = FullAttentionSpec(
+        block_size=16,
+        num_kv_heads=vllm_config.model_config.get_num_kv_heads(
+            vllm_config.parallel_config
+        ),
+        head_size=vllm_config.model_config.get_head_size(),
+        dtype=vllm_config.model_config.dtype,
+    )
+    builder = FlashInferMetadataBuilder(
+        kv_cache_spec,
+        ["layer.0"],
+        vllm_config,
+        torch.device("cpu"),
+    )
+
+    # Guard against passing vacuously with the kernel disabled.
+    assert (
+        builder.flashinfer_trtllm_api_decode_kernel == FlashInferDecodeKernel.TRTLLM_GEN
+    )
+    assert builder.reorder_batch_threshold == 1

vllm/utils/flashinfer.py

@@ -439,6 +439,16 @@ def use_trtllm_attention(
     if force_use_trtllm is not None and not force_use_trtllm:
         return False
 
+    # TRTLLM prefill attends only the DCP-local KV shard and has no
+    # cross-rank LSE combine, so it cannot be used with DCP; fall back to
+    # FlashInfer's DCP prefill path. TRTLLM decode under DCP is selected
+    # separately (all-gathered query heads + LSE combine in forward).
+    if dcp_world_size > 1:
+        logger.warning_once(
+            "TRTLLM prefill does not support DCP, reverting to FlashInfer"
+        )
+        return False
+
     # The platform is not supported
     if not supports_trtllm_attention(is_prefill=is_prefill):
         if force_use_trtllm:

vllm/v1/attention/backends/flashinfer.py

@@ -752,7 +752,14 @@ def __init__(
         self._init_reorder_batch_threshold(
             1,
             supports_spec_as_decode=supports_spec_as_decode,
-            supports_dcp_with_varlen=supports_spec_as_decode,
+            # trtllm-gen decode receives no cp_rank/global-seq-len information,
+            # so its end-aligned causal mask is wrong for q_len > 1 over the
+            # DCP-interleaved local KV shard (spec token i misses up to
+            # (dcp_world_size - 1) * (q_len - 1 - i) KV entries, including its
+            # own). Keep the threshold at 1 under DCP so spec queries take the
+            # DCP-aware prefill path, until the kernel is CP-aware (compare
+            # flash_attn_varlen_func's cp_world_size/cp_rank/cp_tot_seqused_k).
+            supports_dcp_with_varlen=False,
         )
 
         self._cascade_wrapper = None  # Wrapper for cascade attention
@@ -1122,23 +1129,6 @@ def build(
             has_sinks=self.has_sinks,
             has_spec=uses_spec_reorder,
         )
-        if (
-            causal
-            and self.use_dcp
-            and not prefill_use_trtllm
-            and self.flashinfer_trtllm_api_decode_kernel
-            == FlashInferDecodeKernel.TRTLLM_GEN
-            and can_use_trtllm_attention(
-                self.num_qo_heads,
-                self.num_kv_heads,
-                is_prefill=True,
-            )
-        ):
-            logger.info_once(
-                "Using TRTLLM prefill attention with DCP because trtllm-gen "
-                "attention is available."
-            )
-            prefill_use_trtllm = True
         decode_with_flashinfer_trtllm_api = causal and self.use_trtllm_decode_attention
 
         if not causal and self.use_dcp:
@@ -1322,6 +1312,9 @@ def build(
             assert qo_indptr_prefill_cpu.shape[0] == num_prefills + 1
 
             if prefill_use_trtllm:
+                # TRTLLM prefill has no cross-rank combine for DCP-sharded KV;
+                # use_trtllm_attention never selects it when DCP is enabled.
+                assert not self.use_dcp
                 # Create GPU versions
                 qo_indptr_prefill_gpu = (
                     qo_indptr[prefill_start:] - qo_indptr[prefill_start]
@@ -1330,13 +1323,6 @@ def build(
                 # This is the cumulative sum of the number of KV cache
                 # blocks per prefill request.
                 prefill_seq_lens = seq_lens[prefill_start:]
-                if self.use_dcp:
-                    prefill_seq_lens = get_dcp_local_seq_lens(
-                        prefill_seq_lens,
-                        self.dcp_world_size,
-                        self.dcp_rank,
-                        self.dcp_kv_cache_interleave_size,
-                    )
                 num_blocks_per_req = (prefill_seq_lens + page_size - 1) // page_size
                 paged_kv_indptr_prefill_gpu = self.paged_kv_indptr.gpu[
                     prefill_start : num_reqs + 1