Add DCP + Eagle support for Tokenspeed MLA backends

docs/design/attention_backends.md

@@ -232,7 +232,7 @@ MLA decode backends are selected using the standard
 | `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %1 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
 | `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 1, 64 | Any | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
 | `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
-| `TOKENSPEED_MLA` | fp16, bf16 | `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
+| `TOKENSPEED_MLA` | fp16, bf16 | `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
 | `TRITON_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
 | `XPU_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | Any | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | Any |
 

requirements/cuda.txt

@@ -26,7 +26,7 @@ nvidia-cutlass-dsl[cu13]==4.5.2
 quack-kernels>=0.3.3
 
 # Tokenspeed_MLA for faster mla with spec decode
-tokenspeed-mla==0.1.2; platform_system == "Linux"
+tokenspeed-mla==0.1.8; platform_system == "Linux"
 
 # Humming kernels for quantization gemm
 humming-kernels[cu13]==0.1.10

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,73 @@
+# 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, set_current_vllm_config
+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,
+    )
+    with set_current_vllm_config(vllm_config):
+        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

tests/v1/attention/test_mla_backends.py

@@ -8,6 +8,9 @@
   test_backend_correctness[small_prefill], but passes when run alone.
 """
 
+import sys
+from types import SimpleNamespace
+
 import pytest
 import torch
 
@@ -32,6 +35,7 @@
 from vllm.v1.attention.backend import CommonAttentionMetadata
 from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla
 from vllm.v1.attention.backends.mla import flashmla as flashmla_module
+from vllm.v1.attention.backends.mla import tokenspeed_mla as tokenspeed_mla_module
 from vllm.v1.attention.backends.mla.prefill import (
     MLAPrefillBackendEnum,
     get_mla_prefill_backend,
@@ -693,6 +697,100 @@ class _AttnMeta:
     )
 
 
+def test_tokenspeed_mla_dcp_single_token_decode_contract(monkeypatch):
+    decode_call = None
+    num_decodes = 2
+    tokens_per_decode = 1
+    num_decode_tokens = num_decodes * tokens_per_decode
+    dcp_world_size = 2
+    dcp_rank = 1
+    num_heads = 128
+    kv_lora_rank = 512
+    qk_rope_head_dim = 64
+    head_size = kv_lora_rank + qk_rope_head_dim
+    block_size = 64
+    num_blocks = 4
+    max_seq_len = 24
+
+    def fake_decode(**kwargs):
+        nonlocal decode_call
+        decode_call = kwargs
+        q = kwargs["query"]
+        out = torch.empty(
+            q.shape[0],
+            q.shape[1],
+            q.shape[2],
+            kv_lora_rank,
+            dtype=torch.bfloat16,
+        )
+        lse = torch.empty(q.shape[0], q.shape[1], q.shape[2], dtype=torch.float32)
+        return out, lse
+
+    monkeypatch.setitem(
+        sys.modules,
+        "tokenspeed_mla",
+        SimpleNamespace(tokenspeed_mla_decode=fake_decode),
+    )
+
+    impl = object.__new__(tokenspeed_mla_module.TokenspeedMLAImpl)
+    impl.dcp_world_size = dcp_world_size
+    impl.dcp_rank = dcp_rank
+    impl.cp_kv_cache_interleave_size = 1
+    impl.need_to_return_lse_for_decode = True
+    impl.kv_lora_rank = kv_lora_rank
+    impl.qk_rope_head_dim = qk_rope_head_dim
+    impl.num_heads = num_heads
+    impl.scale = 1.0
+    impl.softmax_scale = 1.0
+    impl.output_scale = 1.0
+    impl._workspace_buffer = torch.empty(1, dtype=torch.int8)
+
+    metadata = SimpleNamespace(
+        num_decodes=num_decodes,
+        num_decode_tokens=num_decode_tokens,
+        max_seq_len=max_seq_len,
+        decode=SimpleNamespace(
+            block_table=torch.empty((num_decodes, 1), dtype=torch.int32),
+            seq_lens=torch.tensor([16, max_seq_len], dtype=torch.int32),
+            dcp_tot_seq_lens=torch.tensor([16, max_seq_len], dtype=torch.int32),
+        ),
+    )
+    q = torch.empty(num_decode_tokens, num_heads, head_size, dtype=torch.float8_e4m3fn)
+    kv_cache = torch.empty(
+        num_blocks,
+        block_size,
+        head_size,
+        dtype=torch.float8_e4m3fn,
+    )
+
+    out, lse = impl.forward_mqa(
+        q,
+        kv_cache,
+        metadata,
+        SimpleNamespace(_q_scale_float=2.0, _k_scale_float=3.0),
+    )
+
+    assert out.shape == (num_decode_tokens, num_heads, kv_lora_rank)
+    assert lse is not None
+    assert lse.shape == (num_decode_tokens, num_heads)
+
+    assert decode_call is not None
+    assert decode_call["query"].shape == (
+        num_decodes,
+        tokens_per_decode,
+        num_heads,
+        head_size,
+    )
+    torch.testing.assert_close(decode_call["seq_lens"], metadata.decode.seq_lens)
+    torch.testing.assert_close(decode_call["block_tables"], metadata.decode.block_table)
+    torch.testing.assert_close(
+        decode_call["causal_seqs"], metadata.decode.dcp_tot_seq_lens
+    )
+    assert decode_call["return_lse"] is True
+    assert decode_call["cp_world"] == dcp_world_size
+    assert decode_call["cp_rank"] == dcp_rank
+
+
 @pytest.mark.parametrize("is_fp8_kvcache", [False, True], ids=["bf16", "fp8"])
 def test_flashmla_dcp_decode_metadata_uses_gathered_query_heads(
     monkeypatch, is_fp8_kvcache

vllm/utils/flashinfer.py

@@ -439,11 +439,13 @@ def use_trtllm_attention(
     if force_use_trtllm is not None and not force_use_trtllm:
         return False
 
-    # Decode context parallel is not supported
+    # 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 does not support returning LSE and as a result "
-            "does not support DCP, reverting to FlashInfer"
+            "TRTLLM prefill does not support DCP, reverting to FlashInfer"
         )
         return False