diff --git a/requirements/cuda.txt b/requirements/cuda.txt index fbc80b28e1..e0db38baad 100644 --- a/requirements/cuda.txt +++ b/requirements/cuda.txt @@ -30,7 +30,7 @@ nvidia-cutlass-dsl[cu13]==4.8.0 quack-kernels==0.6.5 # Required by MSA # Tokenspeed_MLA for faster mla with spec decode -tokenspeed-mla==0.1.8; platform_system == "Linux" +tokenspeed-mla==0.2.15; platform_system == "Linux" # Humming kernels for quantization gemm humming-kernels[cu13]==0.1.16 diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index 9ae0fc216a..1c192f3f1c 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -1098,8 +1098,14 @@ def test_flashinfer_mla_dspark_dcp_supports_target_and_draft(monkeypatch): pytest.param(False, 3, 1, 0, id="noncausal-multi-token"), ], ) +@pytest.mark.parametrize("cp_interleave_size", [1, 16, 896]) def test_tokenspeed_mla_decode_contract( - monkeypatch, causal, tokens_per_decode, dcp_world_size, dcp_rank + monkeypatch, + causal, + tokens_per_decode, + dcp_world_size, + dcp_rank, + cp_interleave_size, ): decode_call = None num_decodes = 2 @@ -1135,7 +1141,7 @@ def test_tokenspeed_mla_decode_contract( impl = object.__new__(tokenspeed_mla_module.TokenspeedMLAImpl) impl.dcp_world_size = dcp_world_size impl.dcp_rank = dcp_rank - impl.cp_kv_cache_interleave_size = 1 + impl._parallel_config = SimpleNamespace(cp_kv_cache_interleave_size=1) impl.need_to_return_lse_for_decode = True impl.kv_lora_rank = kv_lora_rank impl.qk_rope_head_dim = qk_rope_head_dim @@ -1164,6 +1170,8 @@ def test_tokenspeed_mla_decode_contract( dtype=torch.float8_e4m3fn, ) + # NIXL may resolve the interleave size after the implementation is created. + impl._parallel_config.cp_kv_cache_interleave_size = cp_interleave_size out, lse = impl.forward_mqa( q, kv_cache, @@ -1194,6 +1202,142 @@ def test_tokenspeed_mla_decode_contract( assert decode_call["return_lse"] is True assert decode_call["cp_world"] == dcp_world_size assert decode_call["cp_rank"] == dcp_rank + assert decode_call.get("cp_interleave_size", 1) == cp_interleave_size + + +@pytest.mark.skipif( + AttentionBackendEnum.TOKENSPEED_MLA not in BACKENDS_TO_TEST, + reason="TokenSpeed MLA requires an SM100-family GPU and tokenspeed-mla", +) +@pytest.mark.parametrize( + ("cp_interleave_size", "seq_len"), + [ + pytest.param(1, 5, id="token-interleave-empty-local-ranks"), + pytest.param(1, 8 + 17, id="token-interleave-full-cycle"), + pytest.param(896, 5, id="block-interleave-empty-local-ranks"), + pytest.param(896, 8 * 896 + 17, id="block-interleave-full-cycle"), + ], +) +def test_tokenspeed_mla_dcp_matches_unsharded_decode(cp_interleave_size, seq_len): + from tokenspeed_mla import tokenspeed_mla_decode + + from vllm.v1.attention.ops.dcp import ( + _lse_weighted_combine, + mask_dcp_empty_shards_, + ) + + torch.manual_seed(7) + device = torch.device("cuda") + dcp_world_size = 8 + kernel_block_size = 64 + num_heads = 128 + kv_lora_rank = 512 + qk_rope_head_dim = 64 + head_size = kv_lora_rank + qk_rope_head_dim + + def make_paged_cache(tokens: torch.Tensor): + # Empty DCP ranks still have the null block in vLLM's physical cache. + num_pages = max(1, cdiv(tokens.shape[0], kernel_block_size)) + cache = torch.zeros( + num_pages, + kernel_block_size, + head_size, + dtype=tokens.dtype, + device=device, + ) + cache.view(-1, head_size)[: tokens.shape[0]].copy_(tokens) + block_table = torch.arange( + num_pages, dtype=torch.int32, device=device + ).unsqueeze(0) + return cache, block_table + + # Four query positions exercise the MTP/DSpark causal-mask path. + query = (torch.randn(1, 4, num_heads, head_size, device=device) * 0.1).to( + torch.float8_e4m3fn + ) + global_tokens = (torch.randn(seq_len, head_size, device=device) * 0.1).to( + torch.float8_e4m3fn + ) + workspace = tokenspeed_mla_module._get_workspace(device, num_heads, kv_lora_rank) + global_cache, global_block_table = make_paged_cache(global_tokens) + global_seq_len = torch.tensor([seq_len], dtype=torch.int32, device=device) + reference, _ = tokenspeed_mla_decode( + query=query, + kv_cache=global_cache, + workspace_buffer=workspace, + kv_lora_rank=kv_lora_rank, + qk_rope_head_dim=qk_rope_head_dim, + block_tables=global_block_table, + seq_lens=global_seq_len, + max_seq_len=seq_len, + softmax_scale=head_size**-0.5, + return_lse=True, + ) + + positions = torch.arange(seq_len, device=device) + owners = positions.div(cp_interleave_size, rounding_mode="floor").remainder( + dcp_world_size + ) + max_local_seq_len = ( + cdiv(seq_len, dcp_world_size * cp_interleave_size) * cp_interleave_size + ) + partial_outputs = [] + partial_lses = [] + for dcp_rank in range(dcp_world_size): + local_tokens = global_tokens[owners == dcp_rank] + local_cache, local_block_table = make_paged_cache(local_tokens) + local_seq_len = torch.tensor( + [local_tokens.shape[0]], dtype=torch.int32, device=device + ) + impl = object.__new__(tokenspeed_mla_module.TokenspeedMLAImpl) + impl.dcp_world_size = dcp_world_size + impl.dcp_rank = dcp_rank + impl._parallel_config = SimpleNamespace( + cp_kv_cache_interleave_size=cp_interleave_size + ) + impl.need_to_return_lse_for_decode = True + impl.kv_lora_rank = kv_lora_rank + impl.qk_rope_head_dim = qk_rope_head_dim + impl.num_heads = num_heads + impl.scale = head_size**-0.5 + impl.softmax_scale = None + impl.output_scale = None + impl._workspace_buffer = workspace + metadata = SimpleNamespace( + num_decodes=1, + num_decode_tokens=query.shape[1], + max_seq_len=max_local_seq_len, + causal=True, + decode=SimpleNamespace( + block_table=local_block_table, + seq_lens=local_seq_len, + dcp_tot_seq_lens=global_seq_len, + ), + ) + output, lse = impl.forward_mqa( + query.view(-1, num_heads, head_size), + local_cache, + metadata, + SimpleNamespace(_q_scale_float=1.0, _k_scale_float=1.0), + ) + # Empty-shard normalization belongs to the downstream DCP combine path. + mask_dcp_empty_shards_( + lse, + local_seq_len, + torch.tensor([0, query.shape[1]], dtype=torch.int32, device=device), + ) + partial_outputs.append( + output.view(query.shape[1], num_heads, kv_lora_rank).float() + ) + partial_lses.append(lse.view(query.shape[1], num_heads).float()) + + # TokenSpeed returns base-2 LSE. Merge the per-rank partial attention states + # the same way MLADCPManager does before comparing to the unsharded result. + merged = _lse_weighted_combine( + torch.stack(partial_outputs), torch.stack(partial_lses), is_lse_base_on_e=False + ).unsqueeze(0) + # FP8 attention rounds differently when KV is partitioned across DCP ranks. + torch.testing.assert_close(merged, reference.float(), atol=1e-2, rtol=1e-2) @pytest.mark.parametrize("is_fp8_kvcache", [False, True], ids=["bf16", "fp8"]) diff --git a/vllm/v1/attention/backends/mla/tokenspeed_mla.py b/vllm/v1/attention/backends/mla/tokenspeed_mla.py index b3d982eac3..96859d7f98 100644 --- a/vllm/v1/attention/backends/mla/tokenspeed_mla.py +++ b/vllm/v1/attention/backends/mla/tokenspeed_mla.py @@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, ClassVar import torch +from vllm.config import get_current_vllm_config from vllm.config.cache import CacheDType from vllm.logger import init_logger from vllm.model_executor.layers.attention.mla_attention import ( @@ -155,6 +156,7 @@ class TokenspeedMLABackend(MLACommonBackend): class TokenspeedMLAImpl(MLACommonImpl[MLACommonMetadata]): can_return_lse_for_decode: bool = True supports_dcp: bool = True + supports_mtp_with_cp_non_trivial_interleave_size: bool = True # tokenspeed_mla_decode returns LSE in log2 units; its own DCP test merges # partial outputs with exp2(lse). lse_base_on_e: bool = False @@ -216,6 +218,9 @@ class TokenspeedMLAImpl(MLACommonImpl[MLACommonMetadata]): self._workspace_buffer: torch.Tensor | None = None self.softmax_scale: float | None = None self.output_scale: float | None = None + # NIXL resolves interleaving after model construction; retain the config + # rather than caching its initial interleave size. + self._parallel_config = get_current_vllm_config().parallel_config # Pre-JIT BF16 and FP8 prefill kernels here too — decode impl always # runs when tokenspeed is selected, prefill backend may not (user can @@ -311,9 +316,13 @@ class TokenspeedMLAImpl(MLACommonImpl[MLACommonMetadata]): causal_seqs=causal_seqs if self.dcp_world_size > 1 else None, cp_world=self.dcp_world_size, cp_rank=self.dcp_rank, + cp_interleave_size=self._parallel_config.cp_kv_cache_interleave_size, ) if return_lse: o, lse = kernel_out + # DCP ranks with no local KV may return undefined outputs (e.g., + # NaNs). The downstream DCP combine path must mask empty + # shards and ignore their outputs when merging partial attention. lse = lse.view(-1, lse.shape[-1]) else: o, lse = kernel_out, None