@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 0d2732dd91
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 3, 16:30 UTC

[MLA Attention Backend] Add TOKENSPEED_MLA backend for DSR1/Kimi K25 prefill + decode on Blackwell (#41778)

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Signed-off-by: Roger Wang <hey@rogerw.io> Co-authored-by: Roger Wang <hey@rogerw.io>

Yongye Zhu committed May 14, 2026 at 02:48 UTC 0d2732dd919b2cdb979923144180bda32d6bf57f
14 files changed +640 -89
benchmarks/attention_benchmarks/configs/mla_decode.yaml
+1
@@ -53,6 +53,7 @@ backends:
53 - FLASHINFER_MLA
54 - FLASH_ATTN_MLA # Hopper only
55 - FLASHMLA # Hopper only
56 + - TOKENSPEED_MLA # Blackwell + R1 dims + FP8 KV (use --kv-cache-dtype fp8)
57
58 device: "cuda:0"
59 repeats: 100
benchmarks/attention_benchmarks/configs/mla_prefill.yaml
+2
@@ -3,6 +3,7 @@
3 # Compares all available MLA prefill backends:
4 # FA backends: fa2, fa3, fa4 (FlashAttention versions)
5 # Non-FA: flashinfer, cudnn, trtllm (Blackwell-only, require flashinfer)
6 +# CuTe DSL: tokenspeed (Blackwell + R1 dims, requires tokenspeed_mla)
7 #
8 # Uses cutlass_mla as the decode backend for impl construction
9 # (only the prefill path is exercised).
@@ -120,6 +121,7 @@ prefill_backends:
121 - flashinfer
122 - cudnn
123 - trtllm
124 + - tokenspeed
125
126 device: "cuda:0"
127 repeats: 20
benchmarks/attention_benchmarks/mla_runner.py
+67 -63
@@ -179,19 +179,27 @@ def create_minimal_vllm_config(
179
180 if prefill_backend is not None:
181 prefill_cfg = get_prefill_backend_config(prefill_backend)
182 - if prefill_cfg["flash_attn_version"] is not None:
183 - vllm_config.attention_config.flash_attn_version = prefill_cfg[
184 - "flash_attn_version"
182 + if prefill_cfg.get("mla_prefill_backend_enum") is not None:
183 + # Registry-based backends bypass the deprecated boolean flags.
184 + from vllm.v1.attention.backends.mla.prefill import MLAPrefillBackendEnum
185 +
186 + vllm_config.attention_config.mla_prefill_backend = MLAPrefillBackendEnum[
187 + prefill_cfg["mla_prefill_backend_enum"]
188 ]
186 - vllm_config.attention_config.disable_flashinfer_prefill = prefill_cfg[
187 - "disable_flashinfer_prefill"
188 - ]
189 - vllm_config.attention_config.use_cudnn_prefill = prefill_cfg[
190 - "use_cudnn_prefill"
191 - ]
192 - vllm_config.attention_config.use_trtllm_ragged_deepseek_prefill = prefill_cfg[
193 - "use_trtllm_ragged_deepseek_prefill"
194 - ]
189 + else:
190 + if prefill_cfg["flash_attn_version"] is not None:
191 + vllm_config.attention_config.flash_attn_version = prefill_cfg[
192 + "flash_attn_version"
193 + ]
194 + vllm_config.attention_config.disable_flashinfer_prefill = prefill_cfg[
195 + "disable_flashinfer_prefill"
196 + ]
197 + vllm_config.attention_config.use_cudnn_prefill = prefill_cfg[
198 + "use_cudnn_prefill"
199 + ]
200 + vllm_config.attention_config.use_trtllm_ragged_deepseek_prefill = (
201 + prefill_cfg["use_trtllm_ragged_deepseek_prefill"]
202 + )
203
204 return vllm_config
205
@@ -223,22 +231,17 @@ _PREFILL_BACKEND_CONFIG: dict[str, dict] = {
231 "use_trtllm_ragged_deepseek_prefill": False,
232 },
233 "flashinfer": {
226 - "flash_attn_version": None,
227 - "disable_flashinfer_prefill": False,
228 - "use_cudnn_prefill": False,
229 - "use_trtllm_ragged_deepseek_prefill": False,
234 + "mla_prefill_backend_enum": "FLASHINFER",
235 },
236 "cudnn": {
232 - "flash_attn_version": None,
233 - "disable_flashinfer_prefill": True,
234 - "use_cudnn_prefill": True,
235 - "use_trtllm_ragged_deepseek_prefill": False,
237 + # cuDNN prefill backend was removed; AttentionConfig raises on use.
238 + "mla_prefill_backend_enum": "FLASHINFER",
239 },
240 "trtllm": {
238 - "flash_attn_version": None,
239 - "disable_flashinfer_prefill": True,
240 - "use_cudnn_prefill": False,
241 - "use_trtllm_ragged_deepseek_prefill": True,
241 + "mla_prefill_backend_enum": "TRTLLM_RAGGED",
242 + },
243 + "tokenspeed": {
244 + "mla_prefill_backend_enum": "TOKENSPEED_MLA",
245 },
246 }
247
@@ -625,6 +628,21 @@ def _create_backend_impl(
628 # Create mock layer
629 layer = MockLayer(device, impl=impl, kv_cache_spec=kv_cache_spec)
630
631 + # Attach a prefill backend (MLAAttention does this in __init__; the metadata
632 + # builder reads layer.prefill_backend from static_forward_context).
633 + from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend
634 +
635 + prefill_backend_cls = get_mla_prefill_backend(vllm_config)
636 + layer.prefill_backend = prefill_backend_cls(
637 + num_heads=mla_dims["num_q_heads"],
638 + scale=(mla_dims["qk_nope_head_dim"] + mla_dims["qk_rope_head_dim"]) ** -0.5,
639 + kv_lora_rank=mla_dims["kv_lora_rank"],
640 + qk_nope_head_dim=mla_dims["qk_nope_head_dim"],
641 + qk_rope_head_dim=mla_dims["qk_rope_head_dim"],
642 + v_head_dim=mla_dims["v_head_dim"],
643 + vllm_config=vllm_config,
644 + )
645 +
646 # Create builder instance if needed
647 builder_instance = None
648 if builder_class:
@@ -961,19 +979,6 @@ def _run_mla_benchmark_batched(
979 results = []
980
981 with set_current_vllm_config(vllm_config):
964 - # Clear cached prefill backend detection functions so they re-evaluate
965 - # with the current VllmConfig. These are @functools.cache decorated and
966 - # would otherwise return stale results from a previous backend's config.
967 - from vllm.model_executor.layers.attention.mla_attention import (
968 - use_cudnn_prefill,
969 - use_flashinfer_prefill,
970 - use_trtllm_ragged_deepseek_prefill,
971 - )
972 -
973 - use_flashinfer_prefill.cache_clear()
974 - use_cudnn_prefill.cache_clear()
975 - use_trtllm_ragged_deepseek_prefill.cache_clear()
976 -
982 # Create backend impl, layer, builder, and indexer (reused across benchmarks)
983 impl, layer, builder_instance, indexer = _create_backend_impl(
984 backend_cfg,
@@ -985,36 +990,35 @@ def _run_mla_benchmark_batched(
990 kv_cache_dtype=kv_cache_dtype,
991 )
992
988 - # Verify the actual prefill backend matches what was requested
993 + # Verify the actual prefill backend matches what was requested. The
994 + # selector + impl construction already raise on misuse; here we just
995 + # check the resolved class against the requested name as a sanity guard.
996 if prefill_backend is not None:
990 - prefill_cfg = get_prefill_backend_config(prefill_backend)
991 - fa_version = prefill_cfg["flash_attn_version"]
992 -
993 - if fa_version is not None:
994 - # FA backend: verify the impl's FA version
995 - actual_fa_version = getattr(impl, "vllm_flash_attn_version", None)
997 + expected_class = {
998 + "fa2": "FlashAttnPrefillBackend",
999 + "fa3": "FlashAttnPrefillBackend",
1000 + "fa4": "FlashAttnPrefillBackend",
1001 + "flashinfer": "FlashInferPrefillBackend",
1002 + "trtllm": "TrtllmRaggedPrefillBackend",
1003 + "tokenspeed": "TokenspeedMLAPrefillBackend",
1004 + }.get(prefill_backend)
1005 + actual_class = type(getattr(layer, "prefill_backend", None)).__name__
1006 + if expected_class and actual_class != expected_class:
1007 + raise RuntimeError(
1008 + f"Prefill backend '{prefill_backend}' requested "
1009 + f"{expected_class}, got {actual_class}. Check "
1010 + f"attention_config plumbing or installed deps."
1011 + )
1012 + if prefill_backend in {"fa2", "fa3", "fa4"}:
1013 + fa_version = int(prefill_backend[2:])
1014 + actual_fa_version = getattr(
1015 + layer.prefill_backend, "vllm_flash_attn_version", None
1016 + )
1017 if actual_fa_version != fa_version:
1018 raise RuntimeError(
1019 f"Prefill backend '{prefill_backend}' requested FA "
999 - f"version {fa_version}, but the impl is using FA "
1000 - f"version {actual_fa_version}. Check "
1001 - f"vllm/v1/attention/backends/fa_utils.py."
1002 - )
1003 - else:
1004 - # Non-FA backend: verify the builder picked the right path
1005 - expected_flags = {
1006 - "flashinfer": "_use_fi_prefill",
1007 - "cudnn": "_use_cudnn_prefill",
1008 - "trtllm": "_use_trtllm_ragged_prefill",
1009 - }
1010 - flag_name = expected_flags.get(prefill_backend)
1011 - if flag_name and not getattr(builder_instance, flag_name, False):
1012 - raise RuntimeError(
1013 - f"Prefill backend '{prefill_backend}' was requested "
1014 - f"but the metadata builder did not enable it. This "
1015 - f"usually means a dependency is missing (e.g., "
1016 - f"flashinfer not installed) or the platform doesn't "
1017 - f"support it."
1020 + f"version {fa_version}, got "
1021 + f"{actual_fa_version} on {actual_class}."
1022 )
1023
1024 # Run each benchmark with the shared impl
docs/design/attention_backends.md
+9 -6
@@ -125,12 +125,13 @@ Priority is **1 = highest** (tried first).
125 | Priority | Backend |
126 | -------- | ------- |
127 | 1 | `FLASHINFER_MLA` |
128 -| 2 | `CUTLASS_MLA` |
129 -| 3 | `FLASH_ATTN_MLA` |
130 -| 4 | `FLASHMLA` |
131 -| 5 | `TRITON_MLA` |
132 -| 6 | `FLASHINFER_MLA_SPARSE`**\*** |
133 -| 7 | `FLASHMLA_SPARSE` |
128 +| 2 | `TOKENSPEED_MLA` |
129 +| 3 | `CUTLASS_MLA` |
130 +| 4 | `FLASH_ATTN_MLA` |
131 +| 5 | `FLASHMLA` |
132 +| 6 | `TRITON_MLA` |
133 +| 7 | `FLASHINFER_MLA_SPARSE`**\*** |
134 +| 8 | `FLASHMLA_SPARSE` |
135
136 **Ampere/Hopper (SM 8.x-9.x):**
137
@@ -202,6 +203,7 @@ hardware and configuration.
203 | `FLASH_ATTN`‡ | FlashAttention varlen (FA2/FA3/FA4) | fp16, bf16 | Any | FA4 on SM100+, FA3 on SM90, FA2 otherwise |
204 | `TRTLLM_RAGGED` | TensorRT-LLM ragged attention | fp16, bf16 | 10.x | DeepSeek R1 dims only |
205 | `FLASHINFER` | FlashInfer CUTLASS backend | fp16, bf16 | 10.x | DeepSeek R1 dims only |
206 +| `TOKENSPEED_MLA` | | fp16, bf16 | 10.x | DeepSeek R1 dims only |
207
208 > **‡** TRT-LLM Ragged is the default on Blackwell (SM100).
209 > On other GPUs, FlashAttention is used as the default.
@@ -222,5 +224,6 @@ MLA decode backends are selected using the standard
224 | `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %1 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
225 | `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 1, 64 | Any | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
226 | `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
227 +| `TOKENSPEED_MLA` | fp16, bf16 | `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
228 | `TRITON_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
229 | `XPU_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | Any | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | Any |
requirements/cuda.txt
+3
@@ -23,3 +23,6 @@ fastsafetensors >= 0.2.2
23 # QuACK and Cutlass DSL for FA4 (cute-DSL implementation)
24 nvidia-cutlass-dsl[cu13]>=4.4.2
25 quack-kernels>=0.3.3
26 +
27 +# Tokenspeed_MLA for faster mla with spec decode
28 +tokenspeed-mla==0.1.2
\ No newline at end of file
tests/conftest.py
+22 -13
@@ -224,19 +224,28 @@ def init_test_http_connection():
224 def dist_init():
225 from tests.utils import ensure_current_vllm_config
226
227 - temp_file = tempfile.mkstemp()[1]
228 -
229 - with ensure_current_vllm_config():
230 - init_distributed_environment(
231 - world_size=1,
232 - rank=0,
233 - distributed_init_method=f"file://{temp_file}",
234 - local_rank=0,
235 - backend="nccl",
236 - )
237 - initialize_model_parallel(1, 1)
238 - yield
239 - cleanup_dist_env_and_memory()
227 + # Close the fd returned by mkstemp; FileStore opens the path itself.
228 + # Leaving it open leaks one FD per test and eventually exhausts the
229 + # ulimit, causing FileStore's destructor to throw c10::DistStoreError
230 + # ("Too many open files") during gc and abort the process.
231 + fd, temp_file = tempfile.mkstemp()
232 + os.close(fd)
233 +
234 + try:
235 + with ensure_current_vllm_config():
236 + init_distributed_environment(
237 + world_size=1,
238 + rank=0,
239 + distributed_init_method=f"file://{temp_file}",
240 + local_rank=0,
241 + backend="nccl",
242 + )
243 + initialize_model_parallel(1, 1)
244 + yield
245 + cleanup_dist_env_and_memory()
246 + finally:
247 + with contextlib.suppress(OSError):
248 + os.unlink(temp_file)
249
250
251 @pytest.fixture
tests/v1/attention/test_mla_backends.py
+66 -7
@@ -20,17 +20,20 @@ from tests.v1.attention.utils import (
20 from vllm import _custom_ops as ops
21 from vllm.config.vllm import set_current_vllm_config
22 from vllm.model_executor.layers.attention.mla_attention import (
23 + MLAAttention,
24 QueryLenSupport,
25 _DecodeConcatQuantFP8,
26 )
26 -from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
27 from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
28 from vllm.platforms import current_platform
29 from vllm.utils.math_utils import cdiv
30 from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
31 from vllm.v1.attention.backend import CommonAttentionMetadata
32 from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla
33 -from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend
33 +from vllm.v1.attention.backends.mla.prefill import (
34 + MLAPrefillBackendEnum,
35 + get_mla_prefill_backend,
36 +)
37 from vllm.v1.attention.backends.registry import AttentionBackendEnum
38 from vllm.v1.attention.ops.flashmla import is_flashmla_dense_supported
39 from vllm.v1.kv_cache_interface import MLAAttentionSpec
@@ -41,6 +44,7 @@ BACKENDS_TO_TEST = [
44 AttentionBackendEnum.FLASH_ATTN_MLA,
45 AttentionBackendEnum.FLASHINFER_MLA,
46 AttentionBackendEnum.TRITON_MLA,
47 + AttentionBackendEnum.TOKENSPEED_MLA,
48 ]
49
50 DEVICE_TYPE = current_platform.device_type
@@ -49,6 +53,7 @@ DEVICE_TYPE = current_platform.device_type
53 if not torch.cuda.is_available() or torch.cuda.get_device_properties(0).major < 10:
54 BACKENDS_TO_TEST.remove(AttentionBackendEnum.CUTLASS_MLA)
55 BACKENDS_TO_TEST.remove(AttentionBackendEnum.FLASHINFER_MLA)
56 + BACKENDS_TO_TEST.remove(AttentionBackendEnum.TOKENSPEED_MLA)
57
58 # Remove FLASH_ATTN_MLA from the list if not supported
59 if not flash_attn_supports_mla():
@@ -58,6 +63,22 @@ if not flash_attn_supports_mla():
63 if not is_flashmla_dense_supported()[0]:
64 BACKENDS_TO_TEST.remove(AttentionBackendEnum.FLASHMLA)
65
66 +# Remove TOKENSPEED_MLA if the optional package is not installed
67 +if AttentionBackendEnum.TOKENSPEED_MLA in BACKENDS_TO_TEST:
68 + try:
69 + import tokenspeed_mla # noqa: F401
70 + except ImportError:
71 + BACKENDS_TO_TEST.remove(AttentionBackendEnum.TOKENSPEED_MLA)
72 +
73 +
74 +# Filtered per-test via validate_configuration (capability/deps/dims).
75 +PREFILL_BACKENDS_TO_TEST = [
76 + MLAPrefillBackendEnum.FLASH_ATTN,
77 + MLAPrefillBackendEnum.FLASHINFER,
78 + MLAPrefillBackendEnum.TRTLLM_RAGGED,
79 + MLAPrefillBackendEnum.TOKENSPEED_MLA,
80 +]
81 +
82
83 SPEC_DECODE_BACKENDS = []
84 for backend in BACKENDS_TO_TEST:
@@ -389,14 +410,18 @@ class MockSparseMLAAttentionLayer:
410 return output
411
412
392 -class MockMLAAttentionLayer(AttentionLayerBase):
413 +class MockMLAAttentionLayer(MLAAttention):
414 """A mock MLA attention layer for testing.
415
416 This replicates the forward_impl logic from MLAAttention to allow
417 testing MLA backends without the full layer infrastructure.
418
398 - The W_UK_T and W_UV weight matrices are created on the layer (like in
399 - MLAAttention.process_weights_after_loading), not on the impl.
419 + Subclasses MLAAttention so that backends that filter
420 + `static_forward_context` by `isinstance(layer, MLAAttention)` (e.g.
421 + FlashInfer prefill, which reads sm_scale through that filter) see the
422 + mock as a real MLA layer. MLAAttention.__init__ is intentionally
423 + skipped — it would create its own impl/prefill_backend and self-register
424 + in static_forward_context, which fights what the test sets up below.
425 """
426
427 def __init__(
@@ -412,6 +437,7 @@ class MockMLAAttentionLayer(AttentionLayerBase):
437 q_scale: float,
438 k_scale: float,
439 ):
440 + torch.nn.Module.__init__(self)
441 self.impl = impl
442 self.num_heads = num_heads
443 self.qk_nope_head_dim = qk_nope_head_dim
@@ -562,11 +588,15 @@ def run_attention_backend(
588 q_scale: float,
589 k_scale: float,
590 kv_cache_dtype: str = "auto",
591 + prefill_backend: MLAPrefillBackendEnum | None = None,
592 ) -> torch.Tensor:
593 """Run attention computation using the specified backend's AttentionImpl."""
594
595 builder_cls, impl_cls = try_get_attention_backend(backend)
596
597 + # Force the prefill backend selection (None means auto-select).
598 + vllm_config.attention_config.mla_prefill_backend = prefill_backend
599 +
600 # Set the current vllm config so that get_current_vllm_config() works
601 # in the backend implementations
602 with set_current_vllm_config(vllm_config):
@@ -578,7 +608,11 @@ def run_attention_backend(
608 vllm_config.parallel_config
609 )
610 head_size = vllm_config.model_config.get_head_size()
581 - scale = 1.0 / (head_size**0.5)
611 + # Production MLA passes 1/sqrt(qk_head_dim) (the prefill scale) to the
612 + # impl and forwards the same value to the prefill backend. FLASHINFER
613 + # prefill reads sm_scale back from impl.scale via global_hyperparameters
614 + # at plan() time, so impl.scale must agree with prefill_backend.scale.
615 + scale = (qk_nope_head_dim + qk_rope_head_dim) ** -0.5
616 impl = impl_cls(
617 num_heads=num_heads,
618 head_size=head_size,
@@ -683,6 +717,7 @@ def run_attention_backend(
717 @pytest.mark.parametrize("tensor_parallel_size", [1, 4, 8, 16])
718 @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8", "fp8_e4m3"])
719 @pytest.mark.parametrize(("q_scale", "k_scale"), [(1.0, 1.0), (2.0, 3.0)])
720 +@pytest.mark.parametrize("prefill_backend", PREFILL_BACKENDS_TO_TEST)
721 def test_backend_correctness(
722 default_vllm_config,
723 dist_init,
@@ -693,6 +728,7 @@ def test_backend_correctness(
728 kv_cache_dtype: str,
729 q_scale: float,
730 k_scale: float,
731 + prefill_backend: MLAPrefillBackendEnum,
732 ):
733 """
734 Test that all backends produce similar outputs to a reference implementation
@@ -729,6 +765,24 @@ def test_backend_correctness(
765 if not backends_to_test:
766 pytest.skip(f"No backends support kv_cache_dtype={kv_cache_dtype}")
767
768 + # Skip prefill backends that can't satisfy capability/deps/R1 constraints.
769 + from vllm.v1.attention.backends.mla.prefill.selector import (
770 + MLAPrefillSelectorConfig,
771 + )
772 +
773 + try:
774 + prefill_invalid_reasons = prefill_backend.get_class().validate_configuration(
775 + current_platform.get_device_capability(),
776 + MLAPrefillSelectorConfig(dtype=torch.bfloat16, is_r1_compatible=True),
777 + )
778 + except ImportError:
779 + prefill_invalid_reasons = ["ImportError"]
780 + if prefill_invalid_reasons:
781 + pytest.skip(
782 + f"Prefill backend {prefill_backend.name} unavailable: "
783 + f"{prefill_invalid_reasons}"
784 + )
785 +
786 batch_spec = BATCH_SPECS[batch_spec_name]
787 is_spec_decode_test = batch_spec_name.startswith("spec_decode")
788 unique_block_sizes = sorted(set(BACKEND_BLOCK_SIZES[b] for b in backends_to_test))
@@ -799,9 +853,13 @@ def test_backend_correctness(
853 assert kv_lora_rank + qk_rope_head_dim == head_size, (
854 f"MLA dimensions don't match: {total_head_size} != {head_size}"
855 )
802 - decode_scale = 1.0 / (total_head_size**0.5)
856 qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
857 prefill_scale = qk_head_dim**-0.5
858 + # MLA reuses prefill_scale for the decode path: production sets
859 + # impl.scale = 1/sqrt(qk_head_dim) and the decode kernels apply it even
860 + # though the latent attention runs at head_size dimensions. Keeping the
861 + # reference here in sync with run_attention_backend's impl.scale.
862 + decode_scale = prefill_scale
863
864 # 2. Generate data and compute SDPA reference output for MLA
865 all_q_vllm, all_kv_c_vllm, all_k_pe_vllm = [], [], []
@@ -1092,6 +1150,7 @@ def test_backend_correctness(
1150 qk_rope_head_dim,
1151 v_head_dim,
1152 mock_kv_b_proj,
1153 + prefill_backend=prefill_backend,
1154 q_scale=q_scale,
1155 k_scale=k_scale,
1156 kv_cache_dtype=kv_cache_dtype,
vllm/model_executor/layers/attention/mla_attention.py
+1
@@ -1362,6 +1362,7 @@ def backend_supports_prefill_query_quantization() -> bool:
1362 return backend_cls.get_name() in (
1363 "FLASHINFER",
1364 "TRTLLM_RAGGED",
1365 + "TOKENSPEED_MLA",
1366 )
1367
1368
vllm/platforms/cuda.py
+4
@@ -110,6 +110,10 @@ def _get_backend_priorities(
110
111 return [
112 AttentionBackendEnum.FLASHINFER_MLA,
113 + # R1 dims + FP8 KV only; rejected by supports_combination
114 + # otherwise. Behind FLASHINFER_MLA: wins past bs≈8, regresses
115 + # at bs≤2.
116 + AttentionBackendEnum.TOKENSPEED_MLA,
117 AttentionBackendEnum.CUTLASS_MLA,
118 AttentionBackendEnum.FLASH_ATTN_MLA,
119 AttentionBackendEnum.FLASHMLA,
vllm/v1/attention/backends/mla/prefill/registry.py
+4
@@ -43,6 +43,10 @@ class MLAPrefillBackendEnum(Enum, metaclass=_MLAPrefillBackendEnumMeta):
43 "vllm.v1.attention.backends.mla.prefill.trtllm_ragged."
44 "TrtllmRaggedPrefillBackend"
45 )
46 + TOKENSPEED_MLA = (
47 + "vllm.v1.attention.backends.mla.prefill.tokenspeed_mla."
48 + "TokenspeedMLAPrefillBackend"
49 + )
50
51 def get_path(self) -> str:
52 """Get the fully qualified class path for this backend."""
vllm/v1/attention/backends/mla/prefill/selector.py
+1
@@ -67,6 +67,7 @@ def _get_mla_prefill_backend_priorities(
67 MLAPrefillBackendEnum.FLASH_ATTN,
68 MLAPrefillBackendEnum.TRTLLM_RAGGED,
69 MLAPrefillBackendEnum.FLASHINFER,
70 + MLAPrefillBackendEnum.TOKENSPEED_MLA,
71 ]
72 else: # Hopper (SM90) and older
73 return [
vllm/v1/attention/backends/mla/prefill/tokenspeed_mla.py new
+180
@@ -0,0 +1,180 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""TokenSpeed CuTe DSL backend for MLA prefill."""
4 +
5 +from typing import TYPE_CHECKING
6 +
7 +import torch
8 +
9 +from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend
10 +
11 +if TYPE_CHECKING:
12 + from vllm.config import VllmConfig
13 + from vllm.model_executor.layers.attention.mla_attention import (
14 + MLACommonPrefillMetadata,
15 + )
16 + from vllm.platforms.interface import DeviceCapability
17 +
18 +
19 +class TokenspeedMLAPrefillBackend(MLAPrefillBackend):
20 + """TokenSpeed CuTe DSL backend for MLA prefill."""
21 +
22 + requires_r1_mla_dimensions = True
23 +
24 + @staticmethod
25 + def get_name() -> str:
26 + return "TOKENSPEED_MLA"
27 +
28 + @classmethod
29 + def supports_compute_capability(cls, device_capability: "DeviceCapability") -> bool:
30 + return device_capability.major == 10
31 +
32 + _INSTALL_HINT = (
33 + "tokenspeed_mla package is not installed. "
34 + "Install it with: `uv pip install tokenspeed-mla`"
35 + )
36 +
37 + @classmethod
38 + def is_available(cls) -> bool:
39 + try:
40 + from tokenspeed_mla import (
41 + tokenspeed_mla_prefill, # noqa: F401
42 + )
43 +
44 + return True
45 + except ImportError:
46 + return False
47 +
48 + @classmethod
49 + def validate_configuration(
50 + cls,
51 + device_capability,
52 + selector_config,
53 + ) -> list[str]:
54 + # Replace the generic "required dependencies not available" message
55 + # from the base class with a specific install hint so users know
56 + # exactly which package to install when they explicitly select this
57 + # backend without having tokenspeed_mla installed.
58 + reasons = super().validate_configuration(device_capability, selector_config)
59 + return [
60 + cls._INSTALL_HINT if r == "required dependencies not available" else r
61 + for r in reasons
62 + ]
63 +
64 + def __init__(
65 + self,
66 + num_heads: int,
67 + scale: float,
68 + kv_lora_rank: int,
69 + qk_nope_head_dim: int,
70 + qk_rope_head_dim: int,
71 + v_head_dim: int,
72 + vllm_config: "VllmConfig",
73 + ) -> None:
74 + super().__init__(
75 + num_heads=num_heads,
76 + scale=scale,
77 + kv_lora_rank=kv_lora_rank,
78 + qk_nope_head_dim=qk_nope_head_dim,
79 + qk_rope_head_dim=qk_rope_head_dim,
80 + v_head_dim=v_head_dim,
81 + vllm_config=vllm_config,
82 + )
83 +
84 + # Pre-JIT BF16 and FP8 prefill kernels. Idempotent — also called from
85 + # TokenspeedMLAImpl.__init__; second call is a no-op.
86 + from tokenspeed_mla import warmup_compile_prefill
87 +
88 + for q_dtype in (torch.bfloat16, torch.float8_e4m3fn):
89 + warmup_compile_prefill(
90 + q_dtype=q_dtype,
91 + d_qk=qk_nope_head_dim + qk_rope_head_dim,
92 + d_v=v_head_dim,
93 + enable_pdl=False,
94 + )
95 +
96 + def prepare_metadata(
97 + self,
98 + prefill_metadata: "MLACommonPrefillMetadata",
99 + ) -> None:
100 + super().prepare_metadata(prefill_metadata)
101 + # Kernel signature requires `seq_lens` but the implementation never reads
102 + # it (per-batch lengths are derived from `cum_seq_lens` diffs); compute
103 + # for parity with trtllm_ragged. cuda-graph padding in
104 + # `query_start_loc` is saturated to `total_num_tokens`
105 + # (gpu_model_runner.py:1905), so trailing diffs are 0 and padded batches
106 + # are kernel no-ops — same reason trtllm passes the padded length as
107 + # batch_size directly.
108 + self._query_seq_lens = (
109 + prefill_metadata.query_start_loc[1:] - prefill_metadata.query_start_loc[:-1]
110 + )
111 +
112 + def run_prefill_new_tokens(
113 + self,
114 + q: torch.Tensor,
115 + k: torch.Tensor,
116 + v: torch.Tensor,
117 + return_softmax_lse: bool,
118 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
119 + from tokenspeed_mla import tokenspeed_mla_prefill
120 +
121 + # `v` arrives as the second half of `kv_nope.split(...)` in
122 + # mla_attention.forward_mha — a non-contiguous view of `kv_nope` along
123 + # dim=-1. The kernel does `v.reshape(1, total_kv, h_k, 1, d_v)` which
124 + # would silently copy on a non-contiguous tensor; force contiguity here
125 + # so the copy (if any) happens once outside the kernel call.
126 + v = v.contiguous()
127 +
128 + ret = tokenspeed_mla_prefill(
129 + query=q,
130 + key=k,
131 + value=v,
132 + seq_lens=self._query_seq_lens,
133 + cum_seq_lens=self._prefill_metadata.query_start_loc,
134 + max_seq_len=self._prefill_metadata.max_query_len,
135 + batch_size=self._query_seq_lens.shape[0],
136 + softmax_scale=self.scale,
137 + is_causal=True,
138 + return_lse=return_softmax_lse,
139 + enable_pdl=False,
140 + )
141 +
142 + if isinstance(ret, tuple):
143 + # Convert from (q_len, num_heads) to (num_heads, q_len)
144 + return ret[0], ret[1].transpose(0, 1).contiguous()
145 + return ret
146 +
147 + def run_prefill_context_chunk(
148 + self,
149 + chunk_idx: int,
150 + q: torch.Tensor,
151 + k: torch.Tensor,
152 + v: torch.Tensor,
153 + ) -> tuple[torch.Tensor, torch.Tensor]:
154 + from tokenspeed_mla import tokenspeed_mla_prefill
155 +
156 + assert self._prefill_metadata.chunked_context is not None
157 + chunked = self._prefill_metadata.chunked_context
158 +
159 + # See note in run_prefill_new_tokens — `v` is a split-view of `kv_nope`
160 + # in `_compute_prefill_context` and arrives non-contiguous.
161 + v = v.contiguous()
162 +
163 + attn_out, lse = tokenspeed_mla_prefill(
164 + query=q,
165 + key=k,
166 + value=v,
167 + seq_lens=chunked.seq_lens[chunk_idx],
168 + cum_seq_lens=chunked.cu_seq_lens[chunk_idx],
169 + max_seq_len=chunked.max_seq_lens[chunk_idx],
170 + batch_size=chunked.seq_lens[chunk_idx].shape[0],
171 + softmax_scale=self.scale,
172 + is_causal=False,
173 + return_lse=True,
174 + cum_seq_lens_q=self._prefill_metadata.query_start_loc,
175 + max_seq_len_q=self._prefill_metadata.max_query_len,
176 + enable_pdl=False,
177 + )
178 +
179 + # Convert from (q_len, num_heads) to (num_heads, q_len)
180 + return attn_out, lse.transpose(0, 1).contiguous()
vllm/v1/attention/backends/mla/tokenspeed_mla.py new
+277
@@ -0,0 +1,277 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""TokenSpeed CuTe DSL MLA decode backend (Blackwell, FP8 KV cache only)."""
4 +
5 +from typing import ClassVar
6 +
7 +import torch
8 +
9 +from vllm.config.cache import CacheDType
10 +from vllm.logger import init_logger
11 +from vllm.model_executor.layers.attention.mla_attention import (
12 + MLACommonBackend,
13 + MLACommonImpl,
14 + MLACommonMetadata,
15 + MLACommonMetadataBuilder,
16 + QueryLenSupport,
17 +)
18 +from vllm.platforms.interface import DeviceCapability
19 +from vllm.utils.torch_utils import is_quantized_kv_cache
20 +from vllm.v1.attention.backend import (
21 + AttentionCGSupport,
22 + AttentionLayer,
23 + AttentionType,
24 + MultipleOf,
25 +)
26 +from vllm.v1.attention.backends.utils import KVCacheLayoutType
27 +
28 +logger = init_logger(__name__)
29 +
30 +# Workspace upper bound for tokenspeed_mla_decode (per-device, lazy):
31 +# num_sms * num_heads * MAX_Q_LEN * (kv_lora_rank + 1) * sizeof(float32)
32 +# Matches the kernel's `get_workspace_size` formula. MAX_Q_LEN=8 covers up to
33 +# EAGLE3 / MTP-2 spec decoding query lengths; larger q_len fails the kernel's
34 +# own buffer check.
35 +_TOKENSPEED_MAX_Q_LEN = 8
36 +
37 +_g_workspace: dict[torch.device, torch.Tensor] = {}
38 +
39 +
40 +def _get_workspace(
41 + device: torch.device, num_heads: int, kv_lora_rank: int
42 +) -> torch.Tensor:
43 + from tokenspeed_mla import get_num_sm
44 +
45 + needed = (
46 + get_num_sm(device) * num_heads * _TOKENSPEED_MAX_Q_LEN * (kv_lora_rank + 1) * 4
47 + )
48 + existing = _g_workspace.get(device)
49 + if existing is None or existing.numel() < needed:
50 + _g_workspace[device] = torch.empty(needed, dtype=torch.int8, device=device)
51 + return _g_workspace[device]
52 +
53 +
54 +class TokenspeedMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
55 + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
56 + query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.UNIFORM
57 +
58 +
59 +class TokenspeedMLABackend(MLACommonBackend):
60 + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
61 + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
62 + "fp8",
63 + "fp8_e4m3",
64 + ]
65 +
66 + @staticmethod
67 + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
68 + return [32, 64]
69 +
70 + @staticmethod
71 + def get_name() -> str:
72 + return "TOKENSPEED_MLA"
73 +
74 + @staticmethod
75 + def get_impl_cls() -> type["TokenspeedMLAImpl"]:
76 + return TokenspeedMLAImpl
77 +
78 + @staticmethod
79 + def get_builder_cls() -> type["TokenspeedMLAMetadataBuilder"]:
80 + return TokenspeedMLAMetadataBuilder
81 +
82 + @classmethod
83 + def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
84 + return capability.major == 10
85 +
86 + @classmethod
87 + def supports_combination(
88 + cls,
89 + head_size: int,
90 + dtype: torch.dtype,
91 + kv_cache_dtype: CacheDType | None,
92 + block_size: int | None,
93 + use_mla: bool,
94 + has_sink: bool,
95 + use_sparse: bool,
96 + device_capability: DeviceCapability,
97 + ) -> str | None:
98 + # Surface a clear install hint up front rather than letting a raw
99 + # ModuleNotFoundError fire deep inside `forward_mqa` at first request.
100 + try:
101 + import tokenspeed_mla # noqa: F401
102 + except ImportError:
103 + return (
104 + "tokenspeed_mla package is not installed. "
105 + "Install it with: `uv pip install tokenspeed-mla`"
106 + )
107 +
108 + # tokenspeed_mla CuTe DSL kernel is shape-specialized for DeepSeek R1
109 + # MLA dimensions (qk_nope=128, qk_rope=64, v=128). Reject anything else.
110 + from vllm.config import get_current_vllm_config
111 +
112 + vllm_config = get_current_vllm_config()
113 + if vllm_config.model_config is not None:
114 + hf_text_config = vllm_config.model_config.hf_text_config
115 + qk_nope_head_dim = getattr(hf_text_config, "qk_nope_head_dim", 0)
116 + qk_rope_head_dim = getattr(hf_text_config, "qk_rope_head_dim", 0)
117 + v_head_dim = getattr(hf_text_config, "v_head_dim", 0)
118 + if qk_nope_head_dim != 128 or qk_rope_head_dim != 64 or v_head_dim != 128:
119 + return (
120 + "tokenspeed_mla requires DeepSeek R1 MLA dimensions "
121 + "(qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128), "
122 + f"got ({qk_nope_head_dim}, {qk_rope_head_dim}, {v_head_dim})"
123 + )
124 + return None
125 +
126 + @classmethod
127 + def get_required_kv_cache_layout(cls) -> "KVCacheLayoutType | None":
128 + return "HND"
129 +
130 +
131 +class TokenspeedMLAImpl(MLACommonImpl[MLACommonMetadata]):
132 + def __init__(
133 + self,
134 + num_heads: int,
135 + head_size: int,
136 + scale: float,
137 + num_kv_heads: int,
138 + alibi_slopes: list[float] | None,
139 + sliding_window: int | None,
140 + kv_cache_dtype: str,
141 + logits_soft_cap: float | None,
142 + attn_type: str,
143 + kv_sharing_target_layer_name: str | None,
144 + # MLA Specific Arguments
145 + **mla_args,
146 + ) -> None:
147 + super().__init__(
148 + num_heads,
149 + head_size,
150 + scale,
151 + num_kv_heads,
152 + alibi_slopes,
153 + sliding_window,
154 + kv_cache_dtype,
155 + logits_soft_cap,
156 + attn_type,
157 + kv_sharing_target_layer_name,
158 + **mla_args,
159 + )
160 +
161 + unsupported_features = [alibi_slopes, sliding_window, logits_soft_cap]
162 + if any(unsupported_features):
163 + raise NotImplementedError(
164 + "TokenspeedMLAImpl does not support one of the following: "
165 + "alibi_slopes, sliding_window, logits_soft_cap"
166 + )
167 +
168 + if attn_type != AttentionType.DECODER:
169 + raise NotImplementedError(
170 + "Encoder self-attention and "
171 + "encoder/decoder cross-attention "
172 + "are not implemented for "
173 + "TokenspeedMLAImpl"
174 + )
175 +
176 + if not is_quantized_kv_cache(self.kv_cache_dtype):
177 + raise NotImplementedError(
178 + "TokenspeedMLAImpl requires an FP8 KV cache "
179 + "(--kv-cache-dtype fp8 or fp8_e4m3); "
180 + f"got kv_cache_dtype={self.kv_cache_dtype!r}."
181 + )
182 +
183 + # Allocate (or fetch the cached) workspace lazily on first forward —
184 + # __init__ runs before the device is necessarily set on the worker;
185 + # we know it for sure at forward time when we see the input tensor.
186 + self._workspace_buffer: torch.Tensor | None = None
187 + self.softmax_scale: float | None = None
188 + self.output_scale: float | None = None
189 +
190 + # Pre-JIT BF16 and FP8 prefill kernels here too — decode impl always
191 + # runs when tokenspeed is selected, prefill backend may not (user can
192 + # pair with flash_attn / trtllm). Idempotent.
193 + from tokenspeed_mla import warmup_compile_prefill
194 +
195 + for q_dtype in (torch.bfloat16, torch.float8_e4m3fn):
196 + warmup_compile_prefill(
197 + q_dtype=q_dtype,
198 + d_qk=self.qk_nope_head_dim + self.qk_rope_head_dim,
199 + d_v=self.v_head_dim,
200 + enable_pdl=False,
201 + )
202 +
203 + def forward_mqa(
204 + self,
205 + q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
206 + kv_c_and_k_pe_cache: torch.Tensor,
207 + attn_metadata: MLACommonMetadata,
208 + layer: AttentionLayer,
209 + ) -> tuple[torch.Tensor, torch.Tensor | None]:
210 + from tokenspeed_mla import tokenspeed_mla_decode
211 +
212 + assert kv_c_and_k_pe_cache.numel() > 0
213 + assert attn_metadata.decode is not None
214 +
215 + if isinstance(q, tuple):
216 + q_nope, q_pe = q
217 + q = torch.cat([q_nope, q_pe], dim=-1)
218 +
219 + # supports_quant_query_input=True (set in MLACommonImpl) tells the
220 + # pipeline to concat+FP8-quantize Q upstream via _decode_concat_quant_fp8_op.
221 + # The kernel is shape-specialized for FP8 Q + FP8 KV, so anything else
222 + # here means the upstream quant didn't run and the kernel will produce
223 + # garbage.
224 + assert q.dtype == torch.float8_e4m3fn, (
225 + f"TokenspeedMLAImpl expected FP8 query (supports_quant_query_input=True), "
226 + f"got {q.dtype}. Pipeline isinstance(q, tuple)={isinstance(q, tuple)}, "
227 + f"q_scale={layer._q_scale_float}, k_scale={layer._k_scale_float}."
228 + )
229 +
230 + # tokenspeed_mla_decode expects query shape
231 + # (num_decodes, q_len_per_request, num_heads, head_dim).
232 + if attn_metadata.num_decode_tokens % attn_metadata.num_decodes != 0:
233 + logger.warning_once(
234 + """TokenspeedMLAImpl got a query of uneven length.
235 + This usually indicates an issue in batch reordering
236 + or incorrect setup in dummy_run."""
237 + )
238 + q = q.unsqueeze(1)
239 + else:
240 + q = q.view(attn_metadata.num_decodes, -1, q.shape[-2], q.shape[-1])
241 +
242 + if self.softmax_scale is None:
243 + # FP8 KV cache is mandatory for this backend, so q_scale/k_scale
244 + # always apply. softmax_scale is bmm1; output_scale is bmm2 — both
245 + # required to recover the correct attention output from the FP8
246 + # KV cache (V is stored as V_real/k_scale).
247 + self.softmax_scale = (
248 + self.scale * layer._q_scale_float * layer._k_scale_float
249 + )
250 + self.output_scale = layer._k_scale_float
251 +
252 + if self._workspace_buffer is None:
253 + self._workspace_buffer = _get_workspace(
254 + q.device, self.num_heads, self.kv_lora_rank
255 + )
256 +
257 + # vLLM kv_c_and_k_pe_cache is already (num_blocks, block_size, head_size).
258 + # tokenspeed_mla_decode wants 3D — pass as-is (no unsqueeze, unlike trtllm).
259 + o = tokenspeed_mla_decode(
260 + query=q,
261 + kv_cache=kv_c_and_k_pe_cache,
262 + workspace_buffer=self._workspace_buffer,
263 + kv_lora_rank=self.kv_lora_rank,
264 + qk_rope_head_dim=self.qk_rope_head_dim,
265 + block_tables=attn_metadata.decode.block_table,
266 + seq_lens=attn_metadata.decode.seq_lens,
267 + max_seq_len=attn_metadata.max_seq_len,
268 + softmax_scale=self.softmax_scale,
269 + output_scale=self.output_scale,
270 + enable_pdl=False,
271 + )
272 +
273 + # Flatten the output for consistent shape
274 + o = o.view(-1, o.shape[-2], o.shape[-1])
275 +
276 + # tokenspeed_mla_decode does not return LSE.
277 + return o, None
vllm/v1/attention/backends/registry.py
+3
@@ -63,6 +63,9 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
63 FLASHINFER_MLA = (
64 "vllm.v1.attention.backends.mla.flashinfer_mla.FlashInferMLABackend"
65 )
66 + TOKENSPEED_MLA = (
67 + "vllm.v1.attention.backends.mla.tokenspeed_mla.TokenspeedMLABackend"
68 + )
69 FLASHINFER_MLA_SPARSE = (
70 "vllm.v1.attention.backends.mla.flashinfer_mla_sparse."
71 "FlashInferMLASparseBackend"