class TokenspeedMLAPrefillBackend(MLAPrefillBackend):
"""TokenSpeed CuTe DSL backend for MLA prefill."""
supported_mla_dimensions: ClassVar[list[MLADimensions]] = [
MLADimensions(
qk_nope_head_dim=128,
qk_rope_head_dim=64,
v_head_dim=128,
),
]
@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 supports_out(self) -> bool:
# Output head dim is v_head_dim (unpadded); `out` supported since 0.1.8.
return True
def run_prefill_new_tokens(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
return_softmax_lse: bool,
out: torch.Tensor | None = None,
output_scale: torch.Tensor | None = None,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
from tokenspeed_mla import tokenspeed_mla_prefill
# `v` arrives as the second half of `kv_nope.split(...)` — a
# non-contiguous view of `kv_nope` along dim=-1. The kernel wraps inputs
# via `from_dlpack`, which preserves strides, but the FMHA kernel reads
# `v` assuming contiguous storage and silently produces wrong output on
# a strided view (verified on tokenspeed-mla 0.1.8). Force contiguity.
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,
out=out,
)
if isinstance(ret, tuple):
# Convert from (q_len, num_heads) to (num_heads, q_len)
return ret[0], ret[1].transpose(0, 1)
return ret
def run_prefill_context_chunk(
self,
chunk: "MLACommonPrefillMetadata.ContextChunk",
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
out: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
from tokenspeed_mla import tokenspeed_mla_prefill
# See note in run_prefill_new_tokens — `v` is a split-view of `kv_nope`
# in the generic `_compute_prefill_context` and arrives non-contiguous.
# (K3's fused fp8 K/V pack already hands over a contiguous `v`.)
v = v.contiguous()
attn_out, lse = tokenspeed_mla_prefill(
query=q,
key=k,
value=v,
seq_lens=chunk.seq_lens,
cum_seq_lens=chunk.cu_seq_lens,
max_seq_len=chunk.max_seq_len,
batch_size=chunk.num_requests,
softmax_scale=self.scale,
is_causal=False,
return_lse=True,
cum_seq_lens_q=chunk.query_start_loc,
max_seq_len_q=chunk.max_query_len,
enable_pdl=False,
out=out,
)
# Convert from (q_len, num_heads) to (num_heads, q_len)
return attn_out, lse.transpose(0, 1)