Add DCP + Eagle support for Tokenspeed MLA backends (#48180)
Signed-off-by: Pavani Majety <pmajety@nvidia.com> Signed-off-by: Jingyi Yang <girasoleyang@gmail.com> Co-authored-by: Jingyi Yang <girasoleyang@gmail.com>
Pavani Majety committed
Jul 13, 2026 at 11:46 UTC
7fc97042c30eda7b45889162aaf64fa06d4055c4
8 files changed
+313
-23
docs/design/attention_backends.md
+1
-1
@@ -232,7 +232,7 @@ MLA decode backends are selected using the standard
232
| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %1 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
233
| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 1, 64 | Any | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
234
| `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
235
-| `TOKENSPEED_MLA` | fp16, bf16 | `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
235
+| `TOKENSPEED_MLA` | fp16, bf16 | `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
236
| `TRITON_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
237
| `XPU_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | Any | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | Any |
238
requirements/cuda.txt
+1
-1
@@ -26,7 +26,7 @@ nvidia-cutlass-dsl[cu13]==4.5.2
26
quack-kernels>=0.3.3
27
28
# Tokenspeed_MLA for faster mla with spec decode
29
-tokenspeed-mla==0.1.2; platform_system == "Linux"
29
+tokenspeed-mla==0.1.8; platform_system == "Linux"
30
31
# Humming kernels for quantization gemm
32
humming-kernels[cu13]==0.1.10
tests/kernels/attention/test_use_trtllm_attention.py
+11
@@ -154,6 +154,17 @@ def test_use_dcp_fallback(_mock):
154
assert _call(dcp_world_size=2) is False
155
156
157
+@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
158
+def test_use_dcp_fallback_prefill(_mock):
159
+ # TRTLLM prefill has no cross-rank LSE combine for DCP-sharded KV.
160
+ assert _call(dcp_world_size=2, is_prefill=True) is False
161
+
162
+
163
+@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
164
+def test_use_dcp_fallback_prefill_force_on_still_false(_mock):
165
+ assert _call(dcp_world_size=2, is_prefill=True, force_use_trtllm=True) is False
166
+
167
+
168
@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=False)
169
def test_use_platform_unsupported(_mock):
170
assert _call() is False
tests/v1/attention/test_flashinfer_dcp_spec_reorder.py
new
+73
@@ -0,0 +1,73 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""FlashInfer GQA builder: reorder threshold under DCP with spec decode."""
4
+
5
+import pytest
6
+import torch
7
+
8
+from tests.v1.attention.utils import create_vllm_config
9
+from vllm.config import SpeculativeConfig, set_current_vllm_config
10
+from vllm.platforms import current_platform
11
+from vllm.v1.attention.backends import flashinfer as flashinfer_backend
12
+from vllm.v1.attention.backends.flashinfer import (
13
+ FlashInferDecodeKernel,
14
+ FlashInferMetadataBuilder,
15
+)
16
+from vllm.v1.attention.backends.utils import PerLayerParameters
17
+from vllm.v1.kv_cache_interface import FullAttentionSpec
18
+
19
+if not current_platform.is_cuda():
20
+ pytest.skip("FlashInfer backend requires a CUDA platform.", allow_module_level=True)
21
+
22
+
23
+def test_flashinfer_gqa_dcp_spec_decode_clamps_reorder_threshold(monkeypatch):
24
+ """trtllm-gen decode receives no cp_rank/global-seq-len information, so its
25
+ end-aligned causal mask is wrong for q_len > 1 over the DCP-interleaved
26
+ local KV shard. The builder must keep reorder_batch_threshold at 1 under
27
+ DCP so spec queries take the (DCP-aware) prefill path instead.
28
+ """
29
+ vllm_config = create_vllm_config(max_model_len=1024)
30
+ vllm_config.parallel_config.decode_context_parallel_size = 2
31
+ vllm_config.speculative_config = SpeculativeConfig(
32
+ method="ngram", num_speculative_tokens=3
33
+ )
34
+
35
+ monkeypatch.setattr(
36
+ flashinfer_backend, "can_use_trtllm_attention", lambda *args, **kwargs: True
37
+ )
38
+ monkeypatch.setattr(
39
+ FlashInferMetadataBuilder,
40
+ "_get_flashinfer_trtllm_api_decode_kernel",
41
+ staticmethod(lambda: FlashInferDecodeKernel.TRTLLM_GEN),
42
+ )
43
+ monkeypatch.setattr(
44
+ flashinfer_backend,
45
+ "get_per_layer_parameters",
46
+ lambda *args, **kwargs: {
47
+ "layer.0": PerLayerParameters(
48
+ window_left=-1, logits_soft_cap=None, sm_scale=0.1, has_sinks=False
49
+ )
50
+ },
51
+ )
52
+
53
+ kv_cache_spec = FullAttentionSpec(
54
+ block_size=16,
55
+ num_kv_heads=vllm_config.model_config.get_num_kv_heads(
56
+ vllm_config.parallel_config
57
+ ),
58
+ head_size=vllm_config.model_config.get_head_size(),
59
+ dtype=vllm_config.model_config.dtype,
60
+ )
61
+ with set_current_vllm_config(vllm_config):
62
+ builder = FlashInferMetadataBuilder(
63
+ kv_cache_spec,
64
+ ["layer.0"],
65
+ vllm_config,
66
+ torch.device("cpu"),
67
+ )
68
+
69
+ # Guard against passing vacuously with the kernel disabled.
70
+ assert (
71
+ builder.flashinfer_trtllm_api_decode_kernel == FlashInferDecodeKernel.TRTLLM_GEN
72
+ )
73
+ assert builder.reorder_batch_threshold == 1
tests/v1/attention/test_mla_backends.py
+98
@@ -8,6 +8,9 @@ Known Issues:
8
test_backend_correctness[small_prefill], but passes when run alone.
9
"""
10
11
+import sys
12
+from types import SimpleNamespace
13
+
14
import pytest
15
import torch
16
@@ -32,6 +35,7 @@ from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
35
from vllm.v1.attention.backend import CommonAttentionMetadata
36
from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla
37
from vllm.v1.attention.backends.mla import flashmla as flashmla_module
38
+from vllm.v1.attention.backends.mla import tokenspeed_mla as tokenspeed_mla_module
39
from vllm.v1.attention.backends.mla.prefill import (
40
MLAPrefillBackendEnum,
41
get_mla_prefill_backend,
@@ -693,6 +697,100 @@ def test_mock_mla_dcp_fp8_decode_gathers_quantized_query(
697
)
698
699
700
+def test_tokenspeed_mla_dcp_single_token_decode_contract(monkeypatch):
701
+ decode_call = None
702
+ num_decodes = 2
703
+ tokens_per_decode = 1
704
+ num_decode_tokens = num_decodes * tokens_per_decode
705
+ dcp_world_size = 2
706
+ dcp_rank = 1
707
+ num_heads = 128
708
+ kv_lora_rank = 512
709
+ qk_rope_head_dim = 64
710
+ head_size = kv_lora_rank + qk_rope_head_dim
711
+ block_size = 64
712
+ num_blocks = 4
713
+ max_seq_len = 24
714
+
715
+ def fake_decode(**kwargs):
716
+ nonlocal decode_call
717
+ decode_call = kwargs
718
+ q = kwargs["query"]
719
+ out = torch.empty(
720
+ q.shape[0],
721
+ q.shape[1],
722
+ q.shape[2],
723
+ kv_lora_rank,
724
+ dtype=torch.bfloat16,
725
+ )
726
+ lse = torch.empty(q.shape[0], q.shape[1], q.shape[2], dtype=torch.float32)
727
+ return out, lse
728
+
729
+ monkeypatch.setitem(
730
+ sys.modules,
731
+ "tokenspeed_mla",
732
+ SimpleNamespace(tokenspeed_mla_decode=fake_decode),
733
+ )
734
+
735
+ impl = object.__new__(tokenspeed_mla_module.TokenspeedMLAImpl)
736
+ impl.dcp_world_size = dcp_world_size
737
+ impl.dcp_rank = dcp_rank
738
+ impl.cp_kv_cache_interleave_size = 1
739
+ impl.need_to_return_lse_for_decode = True
740
+ impl.kv_lora_rank = kv_lora_rank
741
+ impl.qk_rope_head_dim = qk_rope_head_dim
742
+ impl.num_heads = num_heads
743
+ impl.scale = 1.0
744
+ impl.softmax_scale = 1.0
745
+ impl.output_scale = 1.0
746
+ impl._workspace_buffer = torch.empty(1, dtype=torch.int8)
747
+
748
+ metadata = SimpleNamespace(
749
+ num_decodes=num_decodes,
750
+ num_decode_tokens=num_decode_tokens,
751
+ max_seq_len=max_seq_len,
752
+ decode=SimpleNamespace(
753
+ block_table=torch.empty((num_decodes, 1), dtype=torch.int32),
754
+ seq_lens=torch.tensor([16, max_seq_len], dtype=torch.int32),
755
+ dcp_tot_seq_lens=torch.tensor([16, max_seq_len], dtype=torch.int32),
756
+ ),
757
+ )
758
+ q = torch.empty(num_decode_tokens, num_heads, head_size, dtype=torch.float8_e4m3fn)
759
+ kv_cache = torch.empty(
760
+ num_blocks,
761
+ block_size,
762
+ head_size,
763
+ dtype=torch.float8_e4m3fn,
764
+ )
765
+
766
+ out, lse = impl.forward_mqa(
767
+ q,
768
+ kv_cache,
769
+ metadata,
770
+ SimpleNamespace(_q_scale_float=2.0, _k_scale_float=3.0),
771
+ )
772
+
773
+ assert out.shape == (num_decode_tokens, num_heads, kv_lora_rank)
774
+ assert lse is not None
775
+ assert lse.shape == (num_decode_tokens, num_heads)
776
+
777
+ assert decode_call is not None
778
+ assert decode_call["query"].shape == (
779
+ num_decodes,
780
+ tokens_per_decode,
781
+ num_heads,
782
+ head_size,
783
+ )
784
+ torch.testing.assert_close(decode_call["seq_lens"], metadata.decode.seq_lens)
785
+ torch.testing.assert_close(decode_call["block_tables"], metadata.decode.block_table)
786
+ torch.testing.assert_close(
787
+ decode_call["causal_seqs"], metadata.decode.dcp_tot_seq_lens
788
+ )
789
+ assert decode_call["return_lse"] is True
790
+ assert decode_call["cp_world"] == dcp_world_size
791
+ assert decode_call["cp_rank"] == dcp_rank
792
+
793
+
794
@pytest.mark.parametrize("is_fp8_kvcache", [False, True], ids=["bf16", "fp8"])
795
def test_flashmla_dcp_decode_metadata_uses_gathered_query_heads(
796
monkeypatch, is_fp8_kvcache
vllm/utils/flashinfer.py
+5
-3
@@ -439,11 +439,13 @@ def use_trtllm_attention(
439
if force_use_trtllm is not None and not force_use_trtllm:
440
return False
441
442
- # Decode context parallel is not supported
442
+ # TRTLLM prefill attends only the DCP-local KV shard and has no
443
+ # cross-rank LSE combine, so it cannot be used with DCP; fall back to
444
+ # FlashInfer's DCP prefill path. TRTLLM decode under DCP is selected
445
+ # separately (all-gathered query heads + LSE combine in forward).
446
if dcp_world_size > 1:
447
logger.warning_once(
445
- "Trtllm does not support returning LSE and as a result "
446
- "does not support DCP, reverting to FlashInfer"
448
+ "TRTLLM prefill does not support DCP, reverting to FlashInfer"
449
)
450
return False
451
vllm/v1/attention/backends/flashinfer.py
+76
-8
@@ -752,12 +752,32 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
752
if can_use_xqa_or_trtllm_gen_decode
753
else None
754
)
755
+ if (
756
+ self.use_dcp
757
+ and self.flashinfer_trtllm_api_decode_kernel == FlashInferDecodeKernel.XQA
758
+ ):
759
+ logger.warning_once(
760
+ "FlashInfer XQA decode does not support returning LSE and "
761
+ "therefore does not support DCP, reverting to native FlashInfer "
762
+ "decode."
763
+ )
764
+ self.use_trtllm_decode_attention = False
765
+ self.flashinfer_trtllm_api_decode_kernel = None
766
supports_spec_as_decode = (
767
self.flashinfer_trtllm_api_decode_kernel
768
== FlashInferDecodeKernel.TRTLLM_GEN
769
)
770
self._init_reorder_batch_threshold(
760
- 1, supports_spec_as_decode=supports_spec_as_decode
771
+ 1,
772
+ supports_spec_as_decode=supports_spec_as_decode,
773
+ # trtllm-gen decode receives no cp_rank/global-seq-len information,
774
+ # so its end-aligned causal mask is wrong for q_len > 1 over the
775
+ # DCP-interleaved local KV shard (spec token i misses up to
776
+ # (dcp_world_size - 1) * (q_len - 1 - i) KV entries, including its
777
+ # own). Keep the threshold at 1 under DCP so spec queries take the
778
+ # DCP-aware prefill path, until the kernel is CP-aware (compare
779
+ # flash_attn_varlen_func's cp_world_size/cp_rank/cp_tot_seqused_k).
780
+ supports_dcp_with_varlen=False,
781
)
782
783
self._cascade_wrapper = None # Wrapper for cascade attention
@@ -1127,9 +1147,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1147
has_sinks=self.has_sinks,
1148
has_spec=uses_spec_reorder,
1149
)
1130
- decode_with_flashinfer_trtllm_api = (
1131
- causal and self.use_trtllm_decode_attention and self.dcp_world_size <= 1
1132
- )
1150
+ decode_with_flashinfer_trtllm_api = causal and self.use_trtllm_decode_attention
1151
1152
if not causal and self.use_dcp:
1153
raise NotImplementedError(
@@ -1230,7 +1248,13 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1248
# Compute paged_kv_indices if necessary
1249
# paged_kv_indices is only needed for FlashInfer native paths;
1250
# XQA/trtllm-gen paths use block_tables directly on GPU.
1233
- needs_paged_kv_indices = use_cascade or not all_uses_trtllm
1251
+ needs_native_paged_prefill = num_prefills > 0 and not prefill_use_trtllm
1252
+ needs_native_paged_decode = (
1253
+ num_decodes > 0 and not decode_with_flashinfer_trtllm_api
1254
+ )
1255
+ needs_paged_kv_indices = (
1256
+ use_cascade or needs_native_paged_prefill or needs_native_paged_decode
1257
+ )
1258
if needs_paged_kv_indices:
1259
assert num_blocks_np is not None
1260
assert seq_lens_np is not None
@@ -1306,6 +1330,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1330
assert qo_indptr_prefill_cpu.shape[0] == num_prefills + 1
1331
1332
if prefill_use_trtllm:
1333
+ # TRTLLM prefill has no cross-rank combine for DCP-sharded KV;
1334
+ # use_trtllm_attention never selects it when DCP is enabled.
1335
+ assert not self.use_dcp
1336
# Create GPU versions
1337
qo_indptr_prefill_gpu = (
1338
qo_indptr[prefill_start:] - qo_indptr[prefill_start]
@@ -1332,7 +1359,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1359
max_q_len_prefill = int(query_lens_prefill_cpu.max().item())
1360
attn_metadata.prefill = TRTLLMPrefill(
1361
block_tables=block_table_tensor[prefill_start:],
1335
- seq_lens=seq_lens[prefill_start:],
1362
+ seq_lens=prefill_seq_lens,
1363
cum_seq_lens_q=qo_indptr_prefill_gpu,
1364
cum_seq_lens_kv=paged_kv_indptr_prefill_gpu,
1365
max_q_len=max_q_len_prefill,
@@ -1409,10 +1436,16 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1436
f"Got {num_decode_tokens=} and {num_decodes=}."
1437
)
1438
assert self.flashinfer_trtllm_api_decode_kernel is not None
1439
+ seq_lens_decode = seq_lens[:num_decodes]
1440
+ if self.use_dcp:
1441
+ assert common_attn_metadata.dcp_local_seq_lens is not None
1442
+ seq_lens_decode = common_attn_metadata.dcp_local_seq_lens[
1443
+ :num_decodes
1444
+ ]
1445
attn_metadata.decode = FlashInferTrtllmAPIDecode(
1446
kernel=self.flashinfer_trtllm_api_decode_kernel,
1447
block_tables=block_table_tensor[:num_decodes],
1415
- seq_lens=seq_lens[:num_decodes],
1448
+ seq_lens=seq_lens_decode,
1449
max_seq_len=max_seq_len,
1450
)
1451
else:
@@ -2086,6 +2119,18 @@ class FlashInferImpl(AttentionImpl):
2119
f"contiguous, got strides {kv_strides}"
2120
)
2121
2122
+ if use_dcp:
2123
+ assert decode_with_trtllm_gen
2124
+ if output.dtype == FP4_DTYPE:
2125
+ raise NotImplementedError(
2126
+ "DCP decode with FlashInfer trtllm-gen does not support "
2127
+ "FP4 attention output yet."
2128
+ )
2129
+ decode_query = get_dcp_group().all_gather(
2130
+ decode_query.contiguous(), dim=-2
2131
+ )
2132
+ decode_query = canonicalize_singleton_dim_strides(decode_query)
2133
+
2134
if output.dtype == FP4_DTYPE:
2135
assert self.o_sf_scale is not None
2136
out = FP4Tensor(
@@ -2124,6 +2169,19 @@ class FlashInferImpl(AttentionImpl):
2169
else self.bmm1_scale
2170
)
2171
2172
+ lse = None
2173
+ if use_dcp:
2174
+ out = torch.empty(
2175
+ decode_query.shape,
2176
+ dtype=output.dtype,
2177
+ device=decode_query.device,
2178
+ )
2179
+ lse = torch.empty(
2180
+ (decode_query.size(0), decode_query.size(1)),
2181
+ dtype=torch.float32,
2182
+ device=decode_query.device,
2183
+ )
2184
+
2185
trtllm_batch_decode_with_kv_cache(
2186
query=decode_query,
2187
kv_cache=(
@@ -2145,9 +2203,19 @@ class FlashInferImpl(AttentionImpl):
2203
kv_cache_sf=(
2204
nvfp4_kv_block_scales if self.is_kvcache_nvfp4 else None
2205
),
2206
+ lse=lse,
2207
+ return_lse=self.need_to_return_lse_for_decode,
2208
)
2209
2150
- if needs_fp8_out:
2210
+ if use_dcp:
2211
+ assert isinstance(out, torch.Tensor)
2212
+ assert lse is not None
2213
+ output[:num_decode_tokens] = self.dcp_combine(
2214
+ out,
2215
+ lse,
2216
+ get_dcp_group(),
2217
+ )
2218
+ elif needs_fp8_out:
2219
output[:num_decode_tokens].copy_(out.to(output.dtype))
2220
return output_padded
2221
vllm/v1/attention/backends/mla/tokenspeed_mla.py
+48
-10
@@ -2,7 +2,7 @@
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
5
+from typing import TYPE_CHECKING, ClassVar
6
7
import torch
8
@@ -25,6 +25,10 @@ from vllm.v1.attention.backend import (
25
)
26
from vllm.v1.attention.backends.utils import KVCacheLayoutType
27
28
+if TYPE_CHECKING:
29
+ from vllm.config import VllmConfig
30
+ from vllm.v1.kv_cache_interface import AttentionSpec
31
+
32
logger = init_logger(__name__)
33
34
# Workspace upper bound for tokenspeed_mla_decode (per-device, lazy):
@@ -55,6 +59,22 @@ class TokenspeedMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
59
_cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
60
query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.UNIFORM
61
62
+ def __init__(
63
+ self,
64
+ kv_cache_spec: "AttentionSpec",
65
+ layer_names: list[str],
66
+ vllm_config: "VllmConfig",
67
+ device: torch.device,
68
+ ) -> None:
69
+ super().__init__(
70
+ kv_cache_spec,
71
+ layer_names,
72
+ vllm_config,
73
+ device,
74
+ MLACommonMetadata,
75
+ supports_dcp_with_varlen=True,
76
+ )
77
+
78
79
class TokenspeedMLABackend(MLACommonBackend):
80
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
@@ -130,6 +150,11 @@ class TokenspeedMLABackend(MLACommonBackend):
150
151
152
class TokenspeedMLAImpl(MLACommonImpl[MLACommonMetadata]):
153
+ can_return_lse_for_decode: bool = True
154
+ # tokenspeed_mla_decode returns LSE in log2 units; its own DCP test merges
155
+ # partial outputs with exp2(lse).
156
+ lse_base_on_e: bool = False
157
+
158
def __init__(
159
self,
160
num_heads: int,
@@ -228,9 +253,15 @@ class TokenspeedMLAImpl(MLACommonImpl[MLACommonMetadata]):
253
f"q_scale={layer._q_scale_float}, k_scale={layer._k_scale_float}."
254
)
255
256
+ num_decodes = attn_metadata.num_decodes
257
+ num_decode_tokens = attn_metadata.num_decode_tokens
258
+ block_tables = attn_metadata.decode.block_table
259
+ seq_lens = attn_metadata.decode.seq_lens
260
+ causal_seqs = attn_metadata.decode.dcp_tot_seq_lens
261
+
262
# tokenspeed_mla_decode expects query shape
263
# (num_decodes, q_len_per_request, num_heads, head_dim).
233
- if attn_metadata.num_decode_tokens % attn_metadata.num_decodes != 0:
264
+ if num_decode_tokens % num_decodes != 0:
265
logger.warning_once(
266
"""TokenspeedMLAImpl got a query of uneven length.
267
This usually indicates an issue in batch reordering
@@ -238,7 +269,7 @@ class TokenspeedMLAImpl(MLACommonImpl[MLACommonMetadata]):
269
)
270
q = q.unsqueeze(1)
271
else:
241
- q = q.view(attn_metadata.num_decodes, -1, q.shape[-2], q.shape[-1])
272
+ q = q.view(num_decodes, -1, q.shape[-2], q.shape[-1])
273
274
if self.softmax_scale is None:
275
# FP8 KV cache is mandatory for this backend, so q_scale/k_scale
@@ -257,22 +288,29 @@ class TokenspeedMLAImpl(MLACommonImpl[MLACommonMetadata]):
288
289
# vLLM kv_c_and_k_pe_cache is already (num_blocks, block_size, head_size).
290
# tokenspeed_mla_decode wants 3D — pass as-is (no unsqueeze, unlike trtllm).
260
- o = tokenspeed_mla_decode(
291
+ return_lse = self.need_to_return_lse_for_decode
292
+ kernel_out = tokenspeed_mla_decode(
293
query=q,
294
kv_cache=kv_c_and_k_pe_cache,
295
workspace_buffer=self._workspace_buffer,
296
kv_lora_rank=self.kv_lora_rank,
297
qk_rope_head_dim=self.qk_rope_head_dim,
266
- block_tables=attn_metadata.decode.block_table,
267
- seq_lens=attn_metadata.decode.seq_lens,
298
+ block_tables=block_tables,
299
+ seq_lens=seq_lens,
300
max_seq_len=attn_metadata.max_seq_len,
301
softmax_scale=self.softmax_scale,
302
output_scale=self.output_scale,
303
enable_pdl=False,
304
+ return_lse=return_lse,
305
+ causal_seqs=causal_seqs if self.dcp_world_size > 1 else None,
306
+ cp_world=self.dcp_world_size,
307
+ cp_rank=self.dcp_rank,
308
)
309
+ if return_lse:
310
+ o, lse = kernel_out
311
+ lse = lse.view(-1, lse.shape[-1])
312
+ else:
313
+ o, lse = kernel_out, None
314
274
- # Flatten the output for consistent shape
315
o = o.view(-1, o.shape[-2], o.shape[-1])
276
-
277
- # tokenspeed_mla_decode does not return LSE.
278
- return o, None
316
+ return o, lse