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/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
@@ -784,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/utils/flashinfer.py

@@ -439,14 +439,6 @@ def use_trtllm_attention(
     if force_use_trtllm is not None and not force_use_trtllm:
         return False
 
-    # Decode context parallel is not supported
-    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"
-        )
-        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

@@ -734,12 +734,25 @@ def __init__(
             if can_use_xqa_or_trtllm_gen_decode
             else None
         )
+        if (
+            self.use_dcp
+            and self.flashinfer_trtllm_api_decode_kernel == FlashInferDecodeKernel.XQA
+        ):
+            logger.warning_once(
+                "FlashInfer XQA decode does not support returning LSE and "
+                "therefore does not support DCP, reverting to native FlashInfer "
+                "decode."
+            )
+            self.use_trtllm_decode_attention = False
+            self.flashinfer_trtllm_api_decode_kernel = None
         supports_spec_as_decode = (
             self.flashinfer_trtllm_api_decode_kernel
             == FlashInferDecodeKernel.TRTLLM_GEN
         )
         self._init_reorder_batch_threshold(
-            1, supports_spec_as_decode=supports_spec_as_decode
+            1,
+            supports_spec_as_decode=supports_spec_as_decode,
+            supports_dcp_with_varlen=supports_spec_as_decode,
         )
 
         self._cascade_wrapper = None  # Wrapper for cascade attention
@@ -1109,9 +1122,24 @@ def build(
             has_sinks=self.has_sinks,
             has_spec=uses_spec_reorder,
         )
-        decode_with_flashinfer_trtllm_api = (
-            causal and self.use_trtllm_decode_attention and self.dcp_world_size <= 1
-        )
+        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:
             raise NotImplementedError(
@@ -1212,7 +1240,13 @@ def build(
         # Compute paged_kv_indices if necessary
         # paged_kv_indices is only needed for FlashInfer native paths;
         # XQA/trtllm-gen paths use block_tables directly on GPU.
-        needs_paged_kv_indices = use_cascade or not all_uses_trtllm
+        needs_native_paged_prefill = num_prefills > 0 and not prefill_use_trtllm
+        needs_native_paged_decode = (
+            num_decodes > 0 and not decode_with_flashinfer_trtllm_api
+        )
+        needs_paged_kv_indices = (
+            use_cascade or needs_native_paged_prefill or needs_native_paged_decode
+        )
         if needs_paged_kv_indices:
             assert num_blocks_np is not None
             assert seq_lens_np is not None
@@ -1296,6 +1330,13 @@ 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
@@ -1314,7 +1355,7 @@ def build(
                 max_q_len_prefill = int(query_lens_prefill_cpu.max().item())
                 attn_metadata.prefill = TRTLLMPrefill(
                     block_tables=block_table_tensor[prefill_start:],
-                    seq_lens=seq_lens[prefill_start:],
+                    seq_lens=prefill_seq_lens,
                     cum_seq_lens_q=qo_indptr_prefill_gpu,
                     cum_seq_lens_kv=paged_kv_indptr_prefill_gpu,
                     max_q_len=max_q_len_prefill,
@@ -1391,10 +1432,16 @@ def build(
                     f"Got {num_decode_tokens=} and {num_decodes=}."
                 )
                 assert self.flashinfer_trtllm_api_decode_kernel is not None
+                seq_lens_decode = seq_lens[:num_decodes]
+                if self.use_dcp:
+                    assert common_attn_metadata.dcp_local_seq_lens is not None
+                    seq_lens_decode = common_attn_metadata.dcp_local_seq_lens[
+                        :num_decodes
+                    ]
                 attn_metadata.decode = FlashInferTrtllmAPIDecode(
                     kernel=self.flashinfer_trtllm_api_decode_kernel,
                     block_tables=block_table_tensor[:num_decodes],
-                    seq_lens=seq_lens[:num_decodes],
+                    seq_lens=seq_lens_decode,
                     max_seq_len=max_seq_len,
                 )
             else:
@@ -2038,6 +2085,18 @@ def forward(
                     f"contiguous, got strides {kv_strides}"
                 )
 
+                if use_dcp:
+                    assert decode_with_trtllm_gen
+                    if output.dtype == FP4_DTYPE:
+                        raise NotImplementedError(
+                            "DCP decode with FlashInfer trtllm-gen does not support "
+                            "FP4 attention output yet."
+                        )
+                    decode_query = get_dcp_group().all_gather(
+                        decode_query.contiguous(), dim=-2
+                    )
+                    decode_query = canonicalize_singleton_dim_strides(decode_query)
+
                 if output.dtype == FP4_DTYPE:
                     assert self.o_sf_scale is not None
                     out = FP4Tensor(
@@ -2076,6 +2135,19 @@ def forward(
                     else self.bmm1_scale
                 )
 
+                lse = None
+                if use_dcp:
+                    out = torch.empty(
+                        decode_query.shape,
+                        dtype=output.dtype,
+                        device=decode_query.device,
+                    )
+                    lse = torch.empty(
+                        (decode_query.size(0), decode_query.size(1)),
+                        dtype=torch.float32,
+                        device=decode_query.device,
+                    )
+
                 trtllm_batch_decode_with_kv_cache(
                     query=decode_query,
                     kv_cache=(
@@ -2097,9 +2169,19 @@ def forward(
                     kv_cache_sf=(
                         nvfp4_kv_block_scales if self.is_kvcache_nvfp4 else None
                     ),
+                    lse=lse,
+                    return_lse=self.need_to_return_lse_for_decode,
                 )
 
-                if needs_fp8_out:
+                if use_dcp:
+                    assert isinstance(out, torch.Tensor)
+                    assert lse is not None
+                    output[:num_decode_tokens] = self.dcp_combine(
+                        out,
+                        lse,
+                        get_dcp_group(),
+                    )
+                elif needs_fp8_out:
                     output[:num_decode_tokens].copy_(out.to(output.dtype))
         return output_padded