[Attention][TokenSpeed MLA] Also warm up prefill kernel from decode impl

The prefill backend may be paired with flash_attn / trtllm in production —
in that case the prefill backend's __init__ never runs and the prefill
kernel's first call pays a 1.5–2 minute JIT cost. Add the same idempotent
`warmup_compile_prefill` invocation to TokenspeedMLAImpl.__init__ (the
decode-side backend, always present when tokenspeed is selected).

The function dedupes by config key, so the double call is a no-op when both
backends are tokenspeed.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>

docs/design/attention_backends.md

@@ -203,6 +203,7 @@ hardware and configuration.
 | `FLASH_ATTN`‡ | FlashAttention varlen (FA2/FA3/FA4) | fp16, bf16 | Any | FA4 on SM100+, FA3 on SM90, FA2 otherwise |
 | `TRTLLM_RAGGED` | TensorRT-LLM ragged attention | fp16, bf16 | 10.x | DeepSeek R1 dims only |
 | `FLASHINFER` | FlashInfer CUTLASS backend | fp16, bf16 | 10.x | DeepSeek R1 dims only |
+| `TOKENSPEED_MLA` | | fp16, bf16 | 10.x | DeepSeek R1 dims only |
 
 > **‡** TRT-LLM Ragged is the default on Blackwell (SM100).
 > On other GPUs, FlashAttention is used as the default.
@@ -223,5 +224,6 @@ 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 |
 | `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

@@ -23,3 +23,6 @@ fastsafetensors >= 0.2.2
 # QuACK and Cutlass DSL for FA4 (cute-DSL implementation)
 nvidia-cutlass-dsl>=4.4.2
 quack-kernels>=0.3.3
+
+# Tokenspeed_MLA for faster mla with spec decode
+tokenspeed-mla==0.1.1

tests/kernels/attention/test_tokenspeed_mla_decode.py

@@ -0,0 +1,128 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Parity: tokenspeed_mla_decode vs flashinfer trtllm_batch_decode_with_kv_cache_mla."""
+
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+
+if not current_platform.has_device_capability(100):
+    pytest.skip(
+        reason="tokenspeed_mla / TRT-LLM MLA decode require Blackwell (SM100+).",
+        allow_module_level=True,
+    )
+
+try:
+    from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla
+except ImportError:
+    pytest.skip(reason="flashinfer not installed", allow_module_level=True)
+
+try:
+    from tokenspeed_mla import get_num_sm, tokenspeed_mla_decode
+except ImportError:
+    pytest.skip(reason="tokenspeed_mla not installed", allow_module_level=True)
+
+
+FLASHINFER_WORKSPACE_BUFFER_SIZE = 128 * 1024 * 1024
+_TS_MAX_Q_LEN = 8
+
+
+def _ts_workspace(device, num_heads, kv_lora_rank):
+    needed = get_num_sm(device) * num_heads * _TS_MAX_Q_LEN * (kv_lora_rank + 1) * 4
+    return torch.empty(needed, dtype=torch.int8, device=device)
+
+
+@pytest.mark.parametrize("bs", [1, 2, 4, 16])
+@pytest.mark.parametrize("block_size", [32, 64])
+@pytest.mark.parametrize("q_len_per_request", [1, 2, 4])
+def test_tokenspeed_vs_trtllm_decode(bs: int, block_size: int, q_len_per_request: int):
+    """Match tokenspeed_mla_decode against TRT-LLM batch decode MLA.
+
+    Both kernels consume the same FP8 KV cache, paged block table, and
+    seq_lens. The only structural difference is rank: TRT-LLM expects 4D
+    (`unsqueeze(1)` for the kv-head dim) while tokenspeed expects 3D. We
+    pass each kernel its preferred shape from the same underlying tensor.
+    """
+    torch.set_default_device("cuda")
+    torch.manual_seed(42)
+
+    # Deepseek R1 dims — both kernels are R1-shape-specialized.
+    num_heads = 128
+    kv_lora_rank = 512
+    qk_nope_head_dim = 128
+    qk_rope_head_dim = 64
+    qk_head_dim = kv_lora_rank + qk_rope_head_dim
+    scale = (qk_nope_head_dim + qk_rope_head_dim) ** -0.5
+
+    MAX_SEQ_LEN = 1024
+
+    seq_lens = [torch.randint(2, MAX_SEQ_LEN, (1,)).item() for _ in range(bs)]
+    seq_lens[-1] = MAX_SEQ_LEN
+    max_seq_len = max(seq_lens)
+    seq_lens_tensor = torch.tensor(seq_lens, dtype=torch.int32)
+
+    blocks_per_seq = (seq_lens_tensor + block_size - 1) // block_size
+    max_num_blocks_per_seq = max(blocks_per_seq.max().item(), 4)
+    total_blocks_needed = sum(blocks_per_seq).item()
+    all_block_ids = torch.randperm(total_blocks_needed, dtype=torch.int32)
+
+    block_tables = torch.zeros((bs, max_num_blocks_per_seq), dtype=torch.int32)
+    block_id = 0
+    for i in range(bs):
+        n = blocks_per_seq[i].item()
+        block_tables[i, :n] = all_block_ids[block_id : block_id + n]
+        block_id += n
+
+    # KV cache: build in BF16 then cast once to FP8 so both kernels see the
+    # exact same quantized values. Shape (num_blocks, block_size, qk_head_dim).
+    kv_cache_bf16 = torch.randn(
+        block_tables.numel(), block_size, qk_head_dim, dtype=torch.bfloat16
+    )
+    kv_cache = kv_cache_bf16.to(torch.float8_e4m3fn)
+
+    # Query: (bs, q_len_per_request, num_heads, qk_head_dim) — same layout as
+    # FlashInferMLAImpl.forward_mqa. Cast to FP8 to match KV.
+    q = torch.randn(
+        bs, q_len_per_request, num_heads, qk_head_dim, dtype=torch.bfloat16
+    ).to(torch.float8_e4m3fn)
+
+    # --- TRT-LLM reference ---
+    fi_workspace = torch.zeros(FLASHINFER_WORKSPACE_BUFFER_SIZE, dtype=torch.uint8)
+    out_ref = trtllm_batch_decode_with_kv_cache_mla(
+        query=q,
+        kv_cache=kv_cache.unsqueeze(1),
+        workspace_buffer=fi_workspace,
+        qk_nope_head_dim=qk_nope_head_dim,
+        kv_lora_rank=kv_lora_rank,
+        qk_rope_head_dim=qk_rope_head_dim,
+        block_tables=block_tables,
+        seq_lens=seq_lens_tensor,
+        max_seq_len=max_seq_len,
+        bmm1_scale=scale,
+    )
+
+    # --- TokenSpeed candidate ---
+    ts_workspace = _ts_workspace(q.device, num_heads, kv_lora_rank)
+    out_ts = tokenspeed_mla_decode(
+        query=q,
+        kv_cache=kv_cache,
+        workspace_buffer=ts_workspace,
+        kv_lora_rank=kv_lora_rank,
+        qk_rope_head_dim=qk_rope_head_dim,
+        block_tables=block_tables,
+        seq_lens=seq_lens_tensor,
+        max_seq_len=max_seq_len,
+        softmax_scale=scale,
+    )
+
+    # Both kernels output v_head_dim=kv_lora_rank=512 per head.
+    # Output dtypes can differ; compare in float32.
+    out_ref_f = out_ref.to(torch.float32)
+    out_ts_f = out_ts.to(torch.float32)
+    assert out_ref_f.shape == out_ts_f.shape, (
+        f"shape mismatch: trtllm={tuple(out_ref_f.shape)} "
+        f"tokenspeed={tuple(out_ts_f.shape)}"
+    )
+
+    torch.testing.assert_close(out_ts_f, out_ref_f, atol=2e-2, rtol=2e-2)

tests/kernels/attention/test_tokenspeed_mla_prefill.py

@@ -0,0 +1,249 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Numeric accuracy parity: tokenspeed_mla_prefill vs trtllm_ragged_attention_deepseek.
+
+Two cases mirror what the vLLM MLA prefill backend does in production:
+- `test_prefill_no_context`: causal Q==KV ragged batch (run_prefill_new_tokens).
+- `test_prefill_with_context`: non-causal Q ragged + KV ragged with
+  per-request kv_len > q_len (run_prefill_context_chunk).
+"""
+
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+
+if not current_platform.has_device_capability(100):
+    pytest.skip(
+        reason="tokenspeed_mla / TRT-LLM ragged require Blackwell (SM100+).",
+        allow_module_level=True,
+    )
+
+try:
+    from flashinfer.prefill import trtllm_ragged_attention_deepseek
+except ImportError:
+    pytest.skip(reason="flashinfer not installed", allow_module_level=True)
+
+try:
+    from tokenspeed_mla import tokenspeed_mla_prefill, warmup_compile_prefill
+except ImportError:
+    pytest.skip(reason="tokenspeed_mla not installed", allow_module_level=True)
+
+
+FLASHINFER_WORKSPACE_BUFFER_SIZE = 384 * 1024 * 1024
+
+
+# Deepseek R1 dimensions — both kernels are shape-specialized for these.
+NUM_HEADS = 128
+KV_LORA_RANK = 512
+QK_NOPE_HEAD_DIM = 128
+QK_ROPE_HEAD_DIM = 64
+V_HEAD_DIM = 128
+QK_HEAD_DIM = QK_NOPE_HEAD_DIM + QK_ROPE_HEAD_DIM  # 192
+SCALE = QK_HEAD_DIM**-0.5
+
+
+def _make_q_kv(
+    seq_lens: list[int],
+    kv_lens: list[int],
+    dtype: torch.dtype,
+):
+    """Build ragged Q (qk_head_dim) and K (qk_head_dim) / V (v_head_dim)."""
+    total_q = sum(seq_lens)
+    total_kv = sum(kv_lens)
+
+    q = torch.randn(total_q, NUM_HEADS, QK_HEAD_DIM, dtype=torch.bfloat16).to(dtype)
+    k = torch.randn(total_kv, NUM_HEADS, QK_HEAD_DIM, dtype=torch.bfloat16).to(dtype)
+    v = torch.randn(total_kv, NUM_HEADS, V_HEAD_DIM, dtype=torch.bfloat16).to(dtype)
+    return q, k, v
+
+
+def _cumsum_int32(lens: list[int]) -> torch.Tensor:
+    out = torch.zeros(len(lens) + 1, dtype=torch.int32)
+    out[1:] = torch.tensor(lens, dtype=torch.int32).cumsum(0)
+    return out
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn])
+@pytest.mark.parametrize("bs", [1, 4, 16])
+@pytest.mark.parametrize("max_q_len", [64, 256, 1024])
+def test_prefill_no_context(dtype: torch.dtype, bs: int, max_q_len: int):
+    """Causal Q==KV ragged: matches the run_prefill_new_tokens code path."""
+    torch.set_default_device("cuda")
+    torch.manual_seed(0)
+
+    if dtype == torch.float8_e4m3fn:
+        warmup_compile_prefill(
+            q_dtype=torch.float8_e4m3fn,
+            d_qk=QK_HEAD_DIM,
+            d_v=V_HEAD_DIM,
+            enable_pdl=False,
+        )
+
+    seq_lens = [int(torch.randint(2, max_q_len + 1, (1,)).item()) for _ in range(bs)]
+    seq_lens[-1] = max_q_len  # pin the last so max_q_len is hit
+
+    seq_lens_tensor = torch.tensor(seq_lens, dtype=torch.int32)
+    cum_seq_lens = _cumsum_int32(seq_lens)
+
+    q, k, v = _make_q_kv(seq_lens, seq_lens, dtype)
+
+    # --- TRT-LLM reference ---
+    workspace = torch.zeros(FLASHINFER_WORKSPACE_BUFFER_SIZE, dtype=torch.uint8)
+    out_ref = torch.empty(q.shape[0], q.shape[1], v.shape[2], dtype=torch.bfloat16)
+    ref_ret = trtllm_ragged_attention_deepseek(
+        query=q,
+        key=k,
+        value=v,
+        workspace_buffer=workspace,
+        seq_lens=seq_lens_tensor,
+        max_q_len=max_q_len,
+        max_kv_len=max_q_len,
+        bmm1_scale=SCALE,
+        bmm2_scale=1.0,
+        o_sf_scale=1.0,
+        batch_size=bs,
+        window_left=-1,
+        cum_seq_lens_q=cum_seq_lens,
+        cum_seq_lens_kv=cum_seq_lens,
+        enable_pdl=False,
+        is_causal=True,
+        return_lse=False,
+        out=out_ref,
+    )
+    out_ref = ref_ret if not isinstance(ref_ret, tuple) else ref_ret[0]
+
+    # --- TokenSpeed candidate ---
+    out_ts = tokenspeed_mla_prefill(
+        query=q,
+        key=k,
+        value=v,
+        seq_lens=seq_lens_tensor,
+        cum_seq_lens=cum_seq_lens,
+        max_seq_len=max_q_len,
+        batch_size=bs,
+        softmax_scale=SCALE,
+        is_causal=True,
+        return_lse=False,
+        enable_pdl=False,
+    )
+    if isinstance(out_ts, tuple):
+        out_ts = out_ts[0]
+
+    out_ref_f = out_ref.to(torch.float32)
+    out_ts_f = out_ts.to(torch.float32)
+    assert out_ref_f.shape == out_ts_f.shape, (
+        f"shape mismatch: trtllm={tuple(out_ref_f.shape)} "
+        f"tokenspeed={tuple(out_ts_f.shape)}"
+    )
+
+    if dtype == torch.float8_e4m3fn:
+        atol, rtol = 5e-2, 5e-2
+    else:
+        atol, rtol = 1e-2, 1e-2
+    torch.testing.assert_close(out_ts_f, out_ref_f, atol=atol, rtol=rtol)
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn])
+@pytest.mark.parametrize("bs", [1, 4, 16])
+def test_prefill_with_context(dtype: torch.dtype, bs: int):
+    """Non-causal Q ragged + KV ragged: run_prefill_context_chunk path.
+
+    Per-request KV length is independent of (and >=) Q length, mimicking the
+    chunked-context call site where KV is the cache chunk and Q is the new tokens.
+    """
+    torch.set_default_device("cuda")
+    torch.manual_seed(1)
+
+    if dtype == torch.float8_e4m3fn:
+        warmup_compile_prefill(
+            q_dtype=torch.float8_e4m3fn,
+            d_qk=QK_HEAD_DIM,
+            d_v=V_HEAD_DIM,
+            enable_pdl=False,
+        )
+
+    q_lens = [int(torch.randint(16, 257, (1,)).item()) for _ in range(bs)]
+    kv_lens = [q_lens[i] + int(torch.randint(0, 1025, (1,)).item()) for i in range(bs)]
+
+    kv_lens_t = torch.tensor(kv_lens, dtype=torch.int32)
+    cum_q = _cumsum_int32(q_lens)
+    cum_kv = _cumsum_int32(kv_lens)
+    max_q_len = max(q_lens)
+    max_kv_len = max(kv_lens)
+
+    q, k, v = _make_q_kv(q_lens, kv_lens, dtype)
+
+    # --- TRT-LLM reference ---
+    workspace = torch.zeros(FLASHINFER_WORKSPACE_BUFFER_SIZE, dtype=torch.uint8)
+    out_ref = torch.empty(q.shape[0], q.shape[1], v.shape[2], dtype=torch.bfloat16)
+    ref_ret = trtllm_ragged_attention_deepseek(
+        query=q,
+        key=k,
+        value=v,
+        workspace_buffer=workspace,
+        seq_lens=kv_lens_t,
+        max_q_len=max_q_len,
+        max_kv_len=max_kv_len,
+        bmm1_scale=SCALE,
+        bmm2_scale=1.0,
+        o_sf_scale=1.0,
+        batch_size=bs,
+        window_left=-1,
+        cum_seq_lens_q=cum_q,
+        cum_seq_lens_kv=cum_kv,
+        enable_pdl=False,
+        is_causal=False,
+        return_lse=True,
+        out=out_ref,
+    )
+    out_ref, lse_ref = ref_ret[0], ref_ret[1]
+
+    # --- TokenSpeed candidate ---
+    ts_ret = tokenspeed_mla_prefill(
+        query=q,
+        key=k,
+        value=v,
+        seq_lens=kv_lens_t,
+        cum_seq_lens=cum_kv,
+        max_seq_len=max_kv_len,
+        batch_size=bs,
+        softmax_scale=SCALE,
+        is_causal=False,
+        return_lse=True,
+        cum_seq_lens_q=cum_q,
+        max_seq_len_q=max_q_len,
+        enable_pdl=False,
+    )
+    out_ts, lse_ts = ts_ret[0], ts_ret[1]
+
+    if dtype == torch.float8_e4m3fn:
+        atol, rtol = 5e-2, 5e-2
+    else:
+        atol, rtol = 1e-2, 1e-2
+    torch.testing.assert_close(
+        out_ts.to(torch.float32),
+        out_ref.to(torch.float32),
+        atol=atol,
+        rtol=rtol,
+    )
+
+    # LSE: trtllm returns (q_len, num_heads). Tokenspeed convention should
+    # match shape-by-shape — if it doesn't, the LSE transpose contract that
+    # merge_attn_states relies on is broken and this assert surfaces it.
+    assert lse_ref.shape == lse_ts.shape, (
+        f"LSE shape mismatch: trtllm={tuple(lse_ref.shape)} "
+        f"tokenspeed={tuple(lse_ts.shape)}"
+    )
+    # Log-base normalization: trtllm returns LSE in log2, tokenspeed and
+    # vLLM's merge_attn_states (triton_merge_attn_states.py:138) both use
+    # natural-log. Convert trtllm's log2 LSE to natural log before
+    # comparison, otherwise we'd be comparing different bases (factor ln 2).
+    import math
+
+    torch.testing.assert_close(
+        lse_ts.to(torch.float32),
+        lse_ref.to(torch.float32) * math.log(2),
+        atol=5e-3,
+        rtol=5e-3,
+    )

vllm/model_executor/layers/attention/mla_attention.py

@@ -1367,6 +1367,7 @@ def backend_supports_prefill_query_quantization() -> bool:
     return backend_cls.get_name() in (
         "FLASHINFER",
         "TRTLLM_RAGGED",
+        "TOKENSPEED_MLA",
     )
 
 

vllm/v1/attention/backends/mla/prefill/registry.py

@@ -43,6 +43,10 @@ class MLAPrefillBackendEnum(Enum, metaclass=_MLAPrefillBackendEnumMeta):
         "vllm.v1.attention.backends.mla.prefill.trtllm_ragged."
         "TrtllmRaggedPrefillBackend"
     )
+    TOKENSPEED_MLA = (
+        "vllm.v1.attention.backends.mla.prefill.tokenspeed_mla."
+        "TokenspeedMLAPrefillBackend"
+    )
 
     def get_path(self) -> str:
         """Get the fully qualified class path for this backend."""

vllm/v1/attention/backends/mla/prefill/selector.py

@@ -67,6 +67,7 @@ def _get_mla_prefill_backend_priorities(
             MLAPrefillBackendEnum.FLASH_ATTN,
             MLAPrefillBackendEnum.TRTLLM_RAGGED,
             MLAPrefillBackendEnum.FLASHINFER,
+            MLAPrefillBackendEnum.TOKENSPEED_MLA,
         ]
     else:  # Hopper (SM90) and older
         return [

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

@@ -0,0 +1,180 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""TokenSpeed CuTe DSL backend for MLA prefill."""
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+    from vllm.model_executor.layers.attention.mla_attention import (
+        MLACommonPrefillMetadata,
+    )
+    from vllm.platforms.interface import DeviceCapability
+
+
+class TokenspeedMLAPrefillBackend(MLAPrefillBackend):
+    """TokenSpeed CuTe DSL backend for MLA prefill."""
+
+    requires_r1_mla_dimensions = True
+
+    @staticmethod
+    def get_name() -> str:
+        return "TOKENSPEED_MLA"
+
+    @classmethod
+    def supports_compute_capability(cls, device_capability: "DeviceCapability") -> bool:
+        return device_capability.major == 10
+
+    _INSTALL_HINT = (
+        "tokenspeed_mla package is not installed. "
+        "Install it with: `uv pip install tokenspeed-mla`"
+    )
+
+    @classmethod
+    def is_available(cls) -> bool:
+        try:
+            from tokenspeed_mla import (
+                tokenspeed_mla_prefill,  # noqa: F401
+            )
+
+            return True
+        except ImportError:
+            return False
+
+    @classmethod
+    def validate_configuration(
+        cls,
+        device_capability,
+        selector_config,
+    ) -> list[str]:
+        # Replace the generic "required dependencies not available" message
+        # from the base class with a specific install hint so users know
+        # exactly which package to install when they explicitly select this
+        # backend without having tokenspeed_mla installed.
+        reasons = super().validate_configuration(device_capability, selector_config)
+        return [
+            cls._INSTALL_HINT if r == "required dependencies not available" else r
+            for r in reasons
+        ]
+
+    def __init__(
+        self,
+        num_heads: int,
+        scale: float,
+        kv_lora_rank: int,
+        qk_nope_head_dim: int,
+        qk_rope_head_dim: int,
+        v_head_dim: int,
+        vllm_config: "VllmConfig",
+    ) -> None:
+        super().__init__(
+            num_heads=num_heads,
+            scale=scale,
+            kv_lora_rank=kv_lora_rank,
+            qk_nope_head_dim=qk_nope_head_dim,
+            qk_rope_head_dim=qk_rope_head_dim,
+            v_head_dim=v_head_dim,
+            vllm_config=vllm_config,
+        )
+
+        # Pre-JIT BF16 and FP8 prefill kernels. Idempotent — also called from
+        # TokenspeedMLAImpl.__init__; second call is a no-op.
+        from tokenspeed_mla import warmup_compile_prefill
+
+        for q_dtype in (torch.bfloat16, torch.float8_e4m3fn):
+            warmup_compile_prefill(
+                q_dtype=q_dtype,
+                d_qk=qk_nope_head_dim + qk_rope_head_dim,
+                d_v=v_head_dim,
+                enable_pdl=False,
+            )
+
+    def prepare_metadata(
+        self,
+        prefill_metadata: "MLACommonPrefillMetadata",
+    ) -> None:
+        super().prepare_metadata(prefill_metadata)
+        # Kernel signature requires `seq_lens` but the implementation never reads
+        # it (per-batch lengths are derived from `cum_seq_lens` diffs); compute
+        # for parity with trtllm_ragged. cuda-graph padding in
+        # `query_start_loc` is saturated to `total_num_tokens`
+        # (gpu_model_runner.py:1905), so trailing diffs are 0 and padded batches
+        # are kernel no-ops — same reason trtllm passes the padded length as
+        # batch_size directly.
+        self._query_seq_lens = (
+            prefill_metadata.query_start_loc[1:] - prefill_metadata.query_start_loc[:-1]
+        )
+
+    def run_prefill_new_tokens(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        return_softmax_lse: bool,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        from tokenspeed_mla import tokenspeed_mla_prefill
+
+        # `v` arrives as the second half of `kv_nope.split(...)` in
+        # mla_attention.forward_mha — a non-contiguous view of `kv_nope` along
+        # dim=-1. The kernel does `v.reshape(1, total_kv, h_k, 1, d_v)` which
+        # would silently copy on a non-contiguous tensor; force contiguity here
+        # so the copy (if any) happens once outside the kernel call.
+        v = v.contiguous()
+
+        ret = tokenspeed_mla_prefill(
+            query=q,
+            key=k,
+            value=v,
+            seq_lens=self._query_seq_lens,
+            cum_seq_lens=self._prefill_metadata.query_start_loc,
+            max_seq_len=self._prefill_metadata.max_query_len,
+            batch_size=self._query_seq_lens.shape[0],
+            softmax_scale=self.scale,
+            is_causal=True,
+            return_lse=return_softmax_lse,
+            enable_pdl=False,
+        )
+
+        if isinstance(ret, tuple):
+            # Convert from (q_len, num_heads) to (num_heads, q_len)
+            return ret[0], ret[1].transpose(0, 1).contiguous()
+        return ret
+
+    def run_prefill_context_chunk(
+        self,
+        chunk_idx: int,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        from tokenspeed_mla import tokenspeed_mla_prefill
+
+        assert self._prefill_metadata.chunked_context is not None
+        chunked = self._prefill_metadata.chunked_context
+
+        # See note in run_prefill_new_tokens — `v` is a split-view of `kv_nope`
+        # in `_compute_prefill_context` and arrives non-contiguous.
+        v = v.contiguous()
+
+        attn_out, lse = tokenspeed_mla_prefill(
+            query=q,
+            key=k,
+            value=v,
+            seq_lens=chunked.seq_lens[chunk_idx],
+            cum_seq_lens=chunked.cu_seq_lens[chunk_idx],
+            max_seq_len=chunked.max_seq_lens[chunk_idx],
+            batch_size=chunked.seq_lens[chunk_idx].shape[0],
+            softmax_scale=self.scale,
+            is_causal=False,
+            return_lse=True,
+            cum_seq_lens_q=self._prefill_metadata.query_start_loc,
+            max_seq_len_q=self._prefill_metadata.max_query_len,
+            enable_pdl=False,
+        )
+
+        # Convert from (q_len, num_heads) to (num_heads, q_len)
+        return attn_out, lse.transpose(0, 1).contiguous()

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

@@ -0,0 +1,277 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""TokenSpeed CuTe DSL MLA decode backend (Blackwell, FP8 KV cache only)."""
+
+from typing import ClassVar
+
+import torch
+
+from vllm.config.cache import CacheDType
+from vllm.logger import init_logger
+from vllm.model_executor.layers.attention.mla_attention import (
+    MLACommonBackend,
+    MLACommonImpl,
+    MLACommonMetadata,
+    MLACommonMetadataBuilder,
+    QueryLenSupport,
+)
+from vllm.platforms.interface import DeviceCapability
+from vllm.utils.torch_utils import is_quantized_kv_cache
+from vllm.v1.attention.backend import (
+    AttentionCGSupport,
+    AttentionLayer,
+    AttentionType,
+    MultipleOf,
+)
+from vllm.v1.attention.backends.utils import KVCacheLayoutType
+
+logger = init_logger(__name__)
+
+# Workspace upper bound for tokenspeed_mla_decode (per-device, lazy):
+#   num_sms * num_heads * MAX_Q_LEN * (kv_lora_rank + 1) * sizeof(float32)
+# Matches the kernel's `get_workspace_size` formula. MAX_Q_LEN=8 covers up to
+# EAGLE3 / MTP-2 spec decoding query lengths; larger q_len fails the kernel's
+# own buffer check.
+_TOKENSPEED_MAX_Q_LEN = 8
+
+_g_workspace: dict[torch.device, torch.Tensor] = {}
+
+
+def _get_workspace(
+    device: torch.device, num_heads: int, kv_lora_rank: int
+) -> torch.Tensor:
+    from tokenspeed_mla import get_num_sm
+
+    needed = (
+        get_num_sm(device) * num_heads * _TOKENSPEED_MAX_Q_LEN * (kv_lora_rank + 1) * 4
+    )
+    existing = _g_workspace.get(device)
+    if existing is None or existing.numel() < needed:
+        _g_workspace[device] = torch.empty(needed, dtype=torch.int8, device=device)
+    return _g_workspace[device]
+
+
+class TokenspeedMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
+    _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
+    query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.UNIFORM
+
+
+class TokenspeedMLABackend(MLACommonBackend):
+    supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
+    supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
+        "fp8",
+        "fp8_e4m3",
+    ]
+
+    @staticmethod
+    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+        return [32, 64]
+
+    @staticmethod
+    def get_name() -> str:
+        return "TOKENSPEED_MLA"
+
+    @staticmethod
+    def get_impl_cls() -> type["TokenspeedMLAImpl"]:
+        return TokenspeedMLAImpl
+
+    @staticmethod
+    def get_builder_cls() -> type["TokenspeedMLAMetadataBuilder"]:
+        return TokenspeedMLAMetadataBuilder
+
+    @classmethod
+    def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
+        return capability.major == 10
+
+    @classmethod
+    def supports_combination(
+        cls,
+        head_size: int,
+        dtype: torch.dtype,
+        kv_cache_dtype: CacheDType | None,
+        block_size: int | None,
+        use_mla: bool,
+        has_sink: bool,
+        use_sparse: bool,
+        device_capability: DeviceCapability,
+    ) -> str | None:
+        # Surface a clear install hint up front rather than letting a raw
+        # ModuleNotFoundError fire deep inside `forward_mqa` at first request.
+        try:
+            import tokenspeed_mla  # noqa: F401
+        except ImportError:
+            return (
+                "tokenspeed_mla package is not installed. "
+                "Install it with: `uv pip install tokenspeed-mla`"
+            )
+
+        # tokenspeed_mla CuTe DSL kernel is shape-specialized for DeepSeek R1
+        # MLA dimensions (qk_nope=128, qk_rope=64, v=128). Reject anything else.
+        from vllm.config import get_current_vllm_config
+
+        vllm_config = get_current_vllm_config()
+        if vllm_config.model_config is not None:
+            hf_text_config = vllm_config.model_config.hf_text_config
+            qk_nope_head_dim = getattr(hf_text_config, "qk_nope_head_dim", 0)
+            qk_rope_head_dim = getattr(hf_text_config, "qk_rope_head_dim", 0)
+            v_head_dim = getattr(hf_text_config, "v_head_dim", 0)
+            if qk_nope_head_dim != 128 or qk_rope_head_dim != 64 or v_head_dim != 128:
+                return (
+                    "tokenspeed_mla requires DeepSeek R1 MLA dimensions "
+                    "(qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128), "
+                    f"got ({qk_nope_head_dim}, {qk_rope_head_dim}, {v_head_dim})"
+                )
+        return None
+
+    @classmethod
+    def get_required_kv_cache_layout(cls) -> "KVCacheLayoutType | None":
+        return "HND"
+
+
+class TokenspeedMLAImpl(MLACommonImpl[MLACommonMetadata]):
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        scale: float,
+        num_kv_heads: int,
+        alibi_slopes: list[float] | None,
+        sliding_window: int | None,
+        kv_cache_dtype: str,
+        logits_soft_cap: float | None,
+        attn_type: str,
+        kv_sharing_target_layer_name: str | None,
+        # MLA Specific Arguments
+        **mla_args,
+    ) -> None:
+        super().__init__(
+            num_heads,
+            head_size,
+            scale,
+            num_kv_heads,
+            alibi_slopes,
+            sliding_window,
+            kv_cache_dtype,
+            logits_soft_cap,
+            attn_type,
+            kv_sharing_target_layer_name,
+            **mla_args,
+        )
+
+        unsupported_features = [alibi_slopes, sliding_window, logits_soft_cap]
+        if any(unsupported_features):
+            raise NotImplementedError(
+                "TokenspeedMLAImpl does not support one of the following: "
+                "alibi_slopes, sliding_window, logits_soft_cap"
+            )
+
+        if attn_type != AttentionType.DECODER:
+            raise NotImplementedError(
+                "Encoder self-attention and "
+                "encoder/decoder cross-attention "
+                "are not implemented for "
+                "TokenspeedMLAImpl"
+            )
+
+        if not is_quantized_kv_cache(self.kv_cache_dtype):
+            raise NotImplementedError(
+                "TokenspeedMLAImpl requires an FP8 KV cache "
+                "(--kv-cache-dtype fp8 or fp8_e4m3); "
+                f"got kv_cache_dtype={self.kv_cache_dtype!r}."
+            )
+
+        # Allocate (or fetch the cached) workspace lazily on first forward —
+        # __init__ runs before the device is necessarily set on the worker;
+        # we know it for sure at forward time when we see the input tensor.
+        self._workspace_buffer: torch.Tensor | None = None
+        self.softmax_scale: float | None = None
+        self.output_scale: float | None = None
+
+        # Pre-JIT BF16 and FP8 prefill kernels here too — decode impl always
+        # runs when tokenspeed is selected, prefill backend may not (user can
+        # pair with flash_attn / trtllm). Idempotent.
+        from tokenspeed_mla import warmup_compile_prefill
+
+        for q_dtype in (torch.bfloat16, torch.float8_e4m3fn):
+            warmup_compile_prefill(
+                q_dtype=q_dtype,
+                d_qk=self.qk_nope_head_dim + self.qk_rope_head_dim,
+                d_v=self.v_head_dim,
+                enable_pdl=False,
+            )
+
+    def forward_mqa(
+        self,
+        q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
+        kv_c_and_k_pe_cache: torch.Tensor,
+        attn_metadata: MLACommonMetadata,
+        layer: AttentionLayer,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        from tokenspeed_mla import tokenspeed_mla_decode
+
+        assert kv_c_and_k_pe_cache.numel() > 0
+        assert attn_metadata.decode is not None
+
+        if isinstance(q, tuple):
+            q_nope, q_pe = q
+            q = torch.cat([q_nope, q_pe], dim=-1)
+
+        # supports_quant_query_input=True (set in MLACommonImpl) tells the
+        # pipeline to concat+FP8-quantize Q upstream via _decode_concat_quant_fp8_op.
+        # The kernel is shape-specialized for FP8 Q + FP8 KV, so anything else
+        # here means the upstream quant didn't run and the kernel will produce
+        # garbage.
+        assert q.dtype == torch.float8_e4m3fn, (
+            f"TokenspeedMLAImpl expected FP8 query (supports_quant_query_input=True), "
+            f"got {q.dtype}. Pipeline isinstance(q, tuple)={isinstance(q, tuple)}, "
+            f"q_scale={layer._q_scale_float}, k_scale={layer._k_scale_float}."
+        )
+
+        # tokenspeed_mla_decode expects query shape
+        # (num_decodes, q_len_per_request, num_heads, head_dim).
+        if attn_metadata.num_decode_tokens % attn_metadata.num_decodes != 0:
+            logger.warning_once(
+                """TokenspeedMLAImpl got a query of uneven length.
+                This usually indicates an issue in batch reordering
+                or incorrect setup in dummy_run."""
+            )
+            q = q.unsqueeze(1)
+        else:
+            q = q.view(attn_metadata.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
+            # always apply. softmax_scale is bmm1; output_scale is bmm2 — both
+            # required to recover the correct attention output from the FP8
+            # KV cache (V is stored as V_real/k_scale).
+            self.softmax_scale = (
+                self.scale * layer._q_scale_float * layer._k_scale_float
+            )
+            self.output_scale = layer._k_scale_float
+
+        if self._workspace_buffer is None:
+            self._workspace_buffer = _get_workspace(
+                q.device, self.num_heads, self.kv_lora_rank
+            )
+
+        # vLLM kv_c_and_k_pe_cache is already (num_blocks, block_size, head_size).
+        # tokenspeed_mla_decode wants 3D — pass as-is (no unsqueeze, unlike trtllm).
+        o = tokenspeed_mla_decode(
+            query=q,
+            kv_cache=kv_c_and_k_pe_cache,
+            workspace_buffer=self._workspace_buffer,
+            kv_lora_rank=self.kv_lora_rank,
+            qk_rope_head_dim=self.qk_rope_head_dim,
+            block_tables=attn_metadata.decode.block_table,
+            seq_lens=attn_metadata.decode.seq_lens,
+            max_seq_len=attn_metadata.max_seq_len,
+            softmax_scale=self.softmax_scale,
+            output_scale=self.output_scale,
+            enable_pdl=False,
+        )
+
+        # Flatten the output for consistent shape
+        o = o.view(-1, o.shape[-2], o.shape[-1])
+
+        # tokenspeed_mla_decode does not return LSE.
+        return o, None

vllm/v1/attention/backends/registry.py

@@ -63,6 +63,9 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
     FLASHINFER_MLA = (
         "vllm.v1.attention.backends.mla.flashinfer_mla.FlashInferMLABackend"
     )
+    TOKENSPEED_MLA = (
+        "vllm.v1.attention.backends.mla.tokenspeed_mla.TokenspeedMLABackend"
+    )
     FLASHINFER_MLA_SPARSE = (
         "vllm.v1.attention.backends.mla.flashinfer_mla_sparse."
         "FlashInferMLASparseBackend"