Fix MLA DCP spec decode metadata

Signed-off-by: Pavani Majety <pmajety@nvidia.com>

Author: pavanimajety

tests/v1/attention/test_mla_backends.py

@@ -882,6 +882,59 @@ def fake_get_mla_metadata_dense_fp8(
         assert fp8_call is None
 
 
+def test_flashinfer_mla_dcp_spec_decode_keeps_reorder_threshold():
+    builder_cls, _ = try_get_attention_backend(AttentionBackendEnum.FLASHINFER_MLA)
+
+    from vllm.config import SpeculativeConfig
+
+    class _DummyPrefillBackend:
+        def clone(self):
+            return self
+
+    vllm_config = create_vllm_config(
+        model_name="deepseek-ai/DeepSeek-R1",
+        tensor_parallel_size=1,
+        max_model_len=1024,
+        max_num_batched_tokens=512,
+        hf_config_override={
+            "num_attention_heads": 128,
+            "num_key_value_heads": 128,
+            "hidden_size": 7168,
+            "qk_nope_head_dim": 128,
+            "qk_rope_head_dim": 64,
+            "v_head_dim": 128,
+            "kv_lora_rank": 512,
+        },
+    )
+    vllm_config.parallel_config.decode_context_parallel_size = 4
+    vllm_config.speculative_config = SpeculativeConfig(
+        method="ngram", num_speculative_tokens=3
+    )
+    vllm_config.compilation_config.static_forward_context["layer.0"] = SimpleNamespace(
+        prefill_backend=_DummyPrefillBackend()
+    )
+
+    kv_cache_spec = MLAAttentionSpec(
+        block_size=32,
+        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,
+        sliding_window=vllm_config.model_config.get_sliding_window(),
+        cache_dtype_str="fp8",
+    )
+
+    builder = builder_cls(
+        kv_cache_spec,
+        ["layer.0"],
+        vllm_config,
+        torch.device("cpu"),
+    )
+
+    assert builder.reorder_batch_threshold == 4
+
+
 def run_attention_backend(
     backend: AttentionBackendEnum,
     kv_cache_spec: MLAAttentionSpec,

vllm/v1/attention/backends/mla/flashinfer_mla.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-from typing import ClassVar
+from typing import TYPE_CHECKING, ClassVar
 
 import torch
 from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla
@@ -25,6 +25,10 @@
 )
 from vllm.v1.attention.backends.utils import KVCacheLayoutType
 
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+    from vllm.v1.kv_cache_interface import AttentionSpec
+
 logger = init_logger(__name__)
 
 FLASHINFER_MLA_WORKSPACE_BUFFER_SIZE = 128 * 1024 * 1024
@@ -52,6 +56,22 @@ class FlashInferMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
     _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
     query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.UNIFORM
 
+    def __init__(
+        self,
+        kv_cache_spec: "AttentionSpec",
+        layer_names: list[str],
+        vllm_config: "VllmConfig",
+        device: torch.device,
+    ) -> None:
+        super().__init__(
+            kv_cache_spec,
+            layer_names,
+            vllm_config,
+            device,
+            MLACommonMetadata,
+            supports_dcp_with_varlen=True,
+        )
+
 
 class FlashInferMLABackend(MLACommonBackend):
     supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]

vllm/v1/attention/backends/mla/tokenspeed_mla.py

@@ -23,10 +23,7 @@
     AttentionType,
     MultipleOf,
 )
-from vllm.v1.attention.backends.utils import (
-    KVCacheLayoutType,
-    get_dcp_local_seq_lens,
-)
+from vllm.v1.attention.backends.utils import KVCacheLayoutType
 
 if TYPE_CHECKING:
     from vllm.config import VllmConfig
@@ -263,9 +260,7 @@ def forward_mqa(
         causal_seqs = attn_metadata.decode.dcp_tot_seq_lens
 
         # tokenspeed_mla_decode expects query shape
-        # (num_decodes, q_len_per_request, num_heads, head_dim). Under DCP,
-        # multi-token speculative decode must be split into one-token rows:
-        # local(G-k) is not local(G)-k in interleaved DCP layout.
+        # (num_decodes, q_len_per_request, num_heads, head_dim).
         if num_decode_tokens % num_decodes != 0:
             logger.warning_once(
                 """TokenspeedMLAImpl got a query of uneven length.
@@ -274,33 +269,7 @@ def forward_mqa(
             )
             q = q.unsqueeze(1)
         else:
-            tokens_per_req = num_decode_tokens // num_decodes
-            if self.dcp_world_size > 1 and tokens_per_req > 1:
-                assert causal_seqs is not None
-                offsets = torch.arange(
-                    1 - tokens_per_req,
-                    1,
-                    device=causal_seqs.device,
-                    dtype=causal_seqs.dtype,
-                )
-                causal_seqs = (
-                    causal_seqs[:num_decodes].unsqueeze(1) + offsets
-                ).flatten()
-                seq_lens = get_dcp_local_seq_lens(
-                    causal_seqs,
-                    self.dcp_world_size,
-                    self.dcp_rank,
-                    self.cp_kv_cache_interleave_size,
-                )
-                block_tables = (
-                    block_tables[:num_decodes]
-                    .unsqueeze(1)
-                    .expand(-1, tokens_per_req, -1)
-                    .reshape(num_decode_tokens, -1)
-                )
-                q = q.view(num_decode_tokens, 1, q.shape[-2], q.shape[-1])
-            else:
-                q = q.view(num_decodes, -1, q.shape[-2], q.shape[-1])
+            q = q.view(num_decodes, -1, q.shape[-2], q.shape[-1])
 
         if self.softmax_scale is None:
             # FP8 KV cache is mandatory for this backend, so q_scale/k_scale