[Spec Decode] Enable fused non-causal TokenSpeed MLA for DSpark (#50911)
Signed-off-by: Shreyas Misra <shreyasm@nvidia.com> Co-authored-by: OpenAI Codex <codex@openai.com>
NVShreyas committed
Aug 4, 2026 at 11:05 UTC
8ae8337ffa6eed0823b0b827618090937512d503
2 files changed
+32
-7
tests/v1/attention/test_mla_backends.py
+24
-7
@@ -827,13 +827,25 @@ def test_mock_mla_dcp_fp8_decode_gathers_quantized_query(
827
)
828
829
830
-def test_tokenspeed_mla_dcp_single_token_decode_contract(monkeypatch):
830
+def test_tokenspeed_mla_noncausal_capability():
831
+ builder = tokenspeed_mla_module.TokenspeedMLAMetadataBuilder
832
+ assert builder.supports_non_causal_multi_token_decode
833
+ assert tokenspeed_mla_module.TokenspeedMLABackend.supports_non_causal()
834
+
835
+
836
+@pytest.mark.parametrize(
837
+ ("causal", "tokens_per_decode", "dcp_world_size", "dcp_rank"),
838
+ [
839
+ pytest.param(True, 1, 2, 1, id="causal-dcp"),
840
+ pytest.param(False, 3, 1, 0, id="noncausal-multi-token"),
841
+ ],
842
+)
843
+def test_tokenspeed_mla_decode_contract(
844
+ monkeypatch, causal, tokens_per_decode, dcp_world_size, dcp_rank
845
+):
846
decode_call = None
847
num_decodes = 2
833
- tokens_per_decode = 1
848
num_decode_tokens = num_decodes * tokens_per_decode
835
- dcp_world_size = 2
836
- dcp_rank = 1
849
num_heads = 128
850
kv_lora_rank = 512
851
qk_rope_head_dim = 64
@@ -879,6 +891,7 @@ def test_tokenspeed_mla_dcp_single_token_decode_contract(monkeypatch):
891
num_decodes=num_decodes,
892
num_decode_tokens=num_decode_tokens,
893
max_seq_len=max_seq_len,
894
+ causal=causal,
895
decode=SimpleNamespace(
896
block_table=torch.empty((num_decodes, 1), dtype=torch.int32),
897
seq_lens=torch.tensor([16, max_seq_len], dtype=torch.int32),
@@ -913,9 +926,13 @@ def test_tokenspeed_mla_dcp_single_token_decode_contract(monkeypatch):
926
)
927
torch.testing.assert_close(decode_call["seq_lens"], metadata.decode.seq_lens)
928
torch.testing.assert_close(decode_call["block_tables"], metadata.decode.block_table)
916
- torch.testing.assert_close(
917
- decode_call["causal_seqs"], metadata.decode.dcp_tot_seq_lens
918
- )
929
+ if dcp_world_size > 1:
930
+ torch.testing.assert_close(
931
+ decode_call["causal_seqs"], metadata.decode.dcp_tot_seq_lens
932
+ )
933
+ else:
934
+ assert decode_call["causal_seqs"] is None
935
+ assert decode_call["causal_mask"] is causal
936
assert decode_call["return_lse"] is True
937
assert decode_call["cp_world"] == dcp_world_size
938
assert decode_call["cp_rank"] == dcp_rank
vllm/v1/attention/backends/mla/tokenspeed_mla.py
+8
@@ -58,6 +58,9 @@ def _get_workspace(
58
class TokenspeedMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
59
_cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
60
query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.UNIFORM
61
+ # The kernel accepts an explicit causal mask, so a non-causal DSpark
62
+ # block can remain fused instead of being flattened to single tokens.
63
+ supports_non_causal_multi_token_decode: ClassVar[bool] = True
64
65
def __init__(
66
self,
@@ -103,6 +106,10 @@ class TokenspeedMLABackend(MLACommonBackend):
106
def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
107
return capability.major == 10
108
109
+ @classmethod
110
+ def supports_non_causal(cls) -> bool:
111
+ return True
112
+
113
@classmethod
114
def supports_combination(
115
cls,
@@ -302,6 +309,7 @@ class TokenspeedMLAImpl(MLACommonImpl[MLACommonMetadata]):
309
output_scale=self.output_scale,
310
enable_pdl=False,
311
return_lse=return_lse,
312
+ causal_mask=attn_metadata.causal,
313
causal_seqs=causal_seqs if self.dcp_world_size > 1 else None,
314
cp_world=self.dcp_world_size,
315
cp_rank=self.dcp_rank,