From 7fc97042c30eda7b45889162aaf64fa06d4055c4 Mon Sep 17 00:00:00 2001 From: Pavani Majety Date: Mon, 13 Jul 2026 11:46:02 -0700 Subject: [PATCH] Add DCP + Eagle support for Tokenspeed MLA backends (#48180) Signed-off-by: Pavani Majety Signed-off-by: Jingyi Yang Co-authored-by: Jingyi Yang --- docs/design/attention_backends.md | 2 +- requirements/cuda.txt | 2 +- .../attention/test_use_trtllm_attention.py | 11 +++ .../test_flashinfer_dcp_spec_reorder.py | 73 ++++++++++++++ tests/v1/attention/test_mla_backends.py | 98 +++++++++++++++++++ vllm/utils/flashinfer.py | 8 +- vllm/v1/attention/backends/flashinfer.py | 84 ++++++++++++++-- .../attention/backends/mla/tokenspeed_mla.py | 58 +++++++++-- 8 files changed, 313 insertions(+), 23 deletions(-) create mode 100644 tests/v1/attention/test_flashinfer_dcp_spec_reorder.py diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index 23df00672f..4bf3822f3a 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -232,7 +232,7 @@ MLA decode backends are selected using the standard | `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %1 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A | | `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 1, 64 | Any | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | N/A | | `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A | -| `TOKENSPEED_MLA` | fp16, bf16 | `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x | +| `TOKENSPEED_MLA` | fp16, bf16 | `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x | | `TRITON_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | Any | | `XPU_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | Any | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | Any | diff --git a/requirements/cuda.txt b/requirements/cuda.txt index 116aca3c7c..91a5799768 100644 --- a/requirements/cuda.txt +++ b/requirements/cuda.txt @@ -26,7 +26,7 @@ nvidia-cutlass-dsl[cu13]==4.5.2 quack-kernels>=0.3.3 # Tokenspeed_MLA for faster mla with spec decode -tokenspeed-mla==0.1.2; platform_system == "Linux" +tokenspeed-mla==0.1.8; platform_system == "Linux" # Humming kernels for quantization gemm humming-kernels[cu13]==0.1.10 diff --git a/tests/kernels/attention/test_use_trtllm_attention.py b/tests/kernels/attention/test_use_trtllm_attention.py index 136774bc74..8ae4ac7a32 100644 --- a/tests/kernels/attention/test_use_trtllm_attention.py +++ b/tests/kernels/attention/test_use_trtllm_attention.py @@ -154,6 +154,17 @@ def test_use_dcp_fallback(_mock): assert _call(dcp_world_size=2) is False +@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True) +def test_use_dcp_fallback_prefill(_mock): + # TRTLLM prefill has no cross-rank LSE combine for DCP-sharded KV. + assert _call(dcp_world_size=2, is_prefill=True) is False + + +@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True) +def test_use_dcp_fallback_prefill_force_on_still_false(_mock): + assert _call(dcp_world_size=2, is_prefill=True, force_use_trtllm=True) is False + + @patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=False) def test_use_platform_unsupported(_mock): assert _call() is False diff --git a/tests/v1/attention/test_flashinfer_dcp_spec_reorder.py b/tests/v1/attention/test_flashinfer_dcp_spec_reorder.py new file mode 100644 index 0000000000..cfdbc9211c --- /dev/null +++ b/tests/v1/attention/test_flashinfer_dcp_spec_reorder.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""FlashInfer GQA builder: reorder threshold under DCP with spec decode.""" + +import pytest +import torch + +from tests.v1.attention.utils import create_vllm_config +from vllm.config import SpeculativeConfig, set_current_vllm_config +from vllm.platforms import current_platform +from vllm.v1.attention.backends import flashinfer as flashinfer_backend +from vllm.v1.attention.backends.flashinfer import ( + FlashInferDecodeKernel, + FlashInferMetadataBuilder, +) +from vllm.v1.attention.backends.utils import PerLayerParameters +from vllm.v1.kv_cache_interface import FullAttentionSpec + +if not current_platform.is_cuda(): + pytest.skip("FlashInfer backend requires a CUDA platform.", allow_module_level=True) + + +def test_flashinfer_gqa_dcp_spec_decode_clamps_reorder_threshold(monkeypatch): + """trtllm-gen decode receives no cp_rank/global-seq-len information, so its + end-aligned causal mask is wrong for q_len > 1 over the DCP-interleaved + local KV shard. The builder must keep reorder_batch_threshold at 1 under + DCP so spec queries take the (DCP-aware) prefill path instead. + """ + vllm_config = create_vllm_config(max_model_len=1024) + vllm_config.parallel_config.decode_context_parallel_size = 2 + vllm_config.speculative_config = SpeculativeConfig( + method="ngram", num_speculative_tokens=3 + ) + + monkeypatch.setattr( + flashinfer_backend, "can_use_trtllm_attention", lambda *args, **kwargs: True + ) + monkeypatch.setattr( + FlashInferMetadataBuilder, + "_get_flashinfer_trtllm_api_decode_kernel", + staticmethod(lambda: FlashInferDecodeKernel.TRTLLM_GEN), + ) + monkeypatch.setattr( + flashinfer_backend, + "get_per_layer_parameters", + lambda *args, **kwargs: { + "layer.0": PerLayerParameters( + window_left=-1, logits_soft_cap=None, sm_scale=0.1, has_sinks=False + ) + }, + ) + + kv_cache_spec = FullAttentionSpec( + block_size=16, + num_kv_heads=vllm_config.model_config.get_num_kv_heads( + vllm_config.parallel_config + ), + head_size=vllm_config.model_config.get_head_size(), + dtype=vllm_config.model_config.dtype, + ) + with set_current_vllm_config(vllm_config): + builder = FlashInferMetadataBuilder( + kv_cache_spec, + ["layer.0"], + vllm_config, + torch.device("cpu"), + ) + + # Guard against passing vacuously with the kernel disabled. + assert ( + builder.flashinfer_trtllm_api_decode_kernel == FlashInferDecodeKernel.TRTLLM_GEN + ) + assert builder.reorder_batch_threshold == 1 diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index 315c77de39..18d1d309c1 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -8,6 +8,9 @@ Known Issues: test_backend_correctness[small_prefill], but passes when run alone. """ +import sys +from types import SimpleNamespace + import pytest import torch @@ -32,6 +35,7 @@ from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE from vllm.v1.attention.backend import CommonAttentionMetadata from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla from vllm.v1.attention.backends.mla import flashmla as flashmla_module +from vllm.v1.attention.backends.mla import tokenspeed_mla as tokenspeed_mla_module from vllm.v1.attention.backends.mla.prefill import ( MLAPrefillBackendEnum, get_mla_prefill_backend, @@ -693,6 +697,100 @@ def test_mock_mla_dcp_fp8_decode_gathers_quantized_query( ) +def test_tokenspeed_mla_dcp_single_token_decode_contract(monkeypatch): + decode_call = None + num_decodes = 2 + tokens_per_decode = 1 + num_decode_tokens = num_decodes * tokens_per_decode + dcp_world_size = 2 + dcp_rank = 1 + num_heads = 128 + kv_lora_rank = 512 + qk_rope_head_dim = 64 + head_size = kv_lora_rank + qk_rope_head_dim + block_size = 64 + num_blocks = 4 + max_seq_len = 24 + + def fake_decode(**kwargs): + nonlocal decode_call + decode_call = kwargs + q = kwargs["query"] + out = torch.empty( + q.shape[0], + q.shape[1], + q.shape[2], + kv_lora_rank, + dtype=torch.bfloat16, + ) + lse = torch.empty(q.shape[0], q.shape[1], q.shape[2], dtype=torch.float32) + return out, lse + + monkeypatch.setitem( + sys.modules, + "tokenspeed_mla", + SimpleNamespace(tokenspeed_mla_decode=fake_decode), + ) + + 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.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 = 1.0 + impl.softmax_scale = 1.0 + impl.output_scale = 1.0 + impl._workspace_buffer = torch.empty(1, dtype=torch.int8) + + metadata = SimpleNamespace( + num_decodes=num_decodes, + num_decode_tokens=num_decode_tokens, + max_seq_len=max_seq_len, + decode=SimpleNamespace( + block_table=torch.empty((num_decodes, 1), dtype=torch.int32), + seq_lens=torch.tensor([16, max_seq_len], dtype=torch.int32), + dcp_tot_seq_lens=torch.tensor([16, max_seq_len], dtype=torch.int32), + ), + ) + q = torch.empty(num_decode_tokens, num_heads, head_size, dtype=torch.float8_e4m3fn) + kv_cache = torch.empty( + num_blocks, + block_size, + head_size, + dtype=torch.float8_e4m3fn, + ) + + out, lse = impl.forward_mqa( + q, + kv_cache, + metadata, + SimpleNamespace(_q_scale_float=2.0, _k_scale_float=3.0), + ) + + assert out.shape == (num_decode_tokens, num_heads, kv_lora_rank) + assert lse is not None + assert lse.shape == (num_decode_tokens, num_heads) + + assert decode_call is not None + assert decode_call["query"].shape == ( + num_decodes, + tokens_per_decode, + num_heads, + head_size, + ) + torch.testing.assert_close(decode_call["seq_lens"], metadata.decode.seq_lens) + torch.testing.assert_close(decode_call["block_tables"], metadata.decode.block_table) + torch.testing.assert_close( + decode_call["causal_seqs"], metadata.decode.dcp_tot_seq_lens + ) + assert decode_call["return_lse"] is True + assert decode_call["cp_world"] == dcp_world_size + assert decode_call["cp_rank"] == dcp_rank + + @pytest.mark.parametrize("is_fp8_kvcache", [False, True], ids=["bf16", "fp8"]) def test_flashmla_dcp_decode_metadata_uses_gathered_query_heads( monkeypatch, is_fp8_kvcache diff --git a/vllm/utils/flashinfer.py b/vllm/utils/flashinfer.py index 998c3b63df..1334d110b4 100644 --- a/vllm/utils/flashinfer.py +++ b/vllm/utils/flashinfer.py @@ -439,11 +439,13 @@ def use_trtllm_attention( if force_use_trtllm is not None and not force_use_trtllm: return False - # Decode context parallel is not supported + # TRTLLM prefill attends only the DCP-local KV shard and has no + # cross-rank LSE combine, so it cannot be used with DCP; fall back to + # FlashInfer's DCP prefill path. TRTLLM decode under DCP is selected + # separately (all-gathered query heads + LSE combine in forward). if dcp_world_size > 1: logger.warning_once( - "Trtllm does not support returning LSE and as a result " - "does not support DCP, reverting to FlashInfer" + "TRTLLM prefill does not support DCP, reverting to FlashInfer" ) return False diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index 08ea5efdeb..7342823943 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -752,12 +752,32 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]): if can_use_xqa_or_trtllm_gen_decode else None ) + if ( + self.use_dcp + and self.flashinfer_trtllm_api_decode_kernel == FlashInferDecodeKernel.XQA + ): + logger.warning_once( + "FlashInfer XQA decode does not support returning LSE and " + "therefore does not support DCP, reverting to native FlashInfer " + "decode." + ) + self.use_trtllm_decode_attention = False + self.flashinfer_trtllm_api_decode_kernel = None supports_spec_as_decode = ( self.flashinfer_trtllm_api_decode_kernel == FlashInferDecodeKernel.TRTLLM_GEN ) self._init_reorder_batch_threshold( - 1, supports_spec_as_decode=supports_spec_as_decode + 1, + supports_spec_as_decode=supports_spec_as_decode, + # trtllm-gen decode receives no cp_rank/global-seq-len information, + # so its end-aligned causal mask is wrong for q_len > 1 over the + # DCP-interleaved local KV shard (spec token i misses up to + # (dcp_world_size - 1) * (q_len - 1 - i) KV entries, including its + # own). Keep the threshold at 1 under DCP so spec queries take the + # DCP-aware prefill path, until the kernel is CP-aware (compare + # flash_attn_varlen_func's cp_world_size/cp_rank/cp_tot_seqused_k). + supports_dcp_with_varlen=False, ) self._cascade_wrapper = None # Wrapper for cascade attention @@ -1127,9 +1147,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]): has_sinks=self.has_sinks, has_spec=uses_spec_reorder, ) - decode_with_flashinfer_trtllm_api = ( - causal and self.use_trtllm_decode_attention and self.dcp_world_size <= 1 - ) + decode_with_flashinfer_trtllm_api = causal and self.use_trtllm_decode_attention if not causal and self.use_dcp: raise NotImplementedError( @@ -1230,7 +1248,13 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]): # Compute paged_kv_indices if necessary # paged_kv_indices is only needed for FlashInfer native paths; # XQA/trtllm-gen paths use block_tables directly on GPU. - needs_paged_kv_indices = use_cascade or not all_uses_trtllm + needs_native_paged_prefill = num_prefills > 0 and not prefill_use_trtllm + needs_native_paged_decode = ( + num_decodes > 0 and not decode_with_flashinfer_trtllm_api + ) + needs_paged_kv_indices = ( + use_cascade or needs_native_paged_prefill or needs_native_paged_decode + ) if needs_paged_kv_indices: assert num_blocks_np is not None assert seq_lens_np is not None @@ -1306,6 +1330,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]): assert qo_indptr_prefill_cpu.shape[0] == num_prefills + 1 if prefill_use_trtllm: + # TRTLLM prefill has no cross-rank combine for DCP-sharded KV; + # use_trtllm_attention never selects it when DCP is enabled. + assert not self.use_dcp # Create GPU versions qo_indptr_prefill_gpu = ( qo_indptr[prefill_start:] - qo_indptr[prefill_start] @@ -1332,7 +1359,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]): max_q_len_prefill = int(query_lens_prefill_cpu.max().item()) attn_metadata.prefill = TRTLLMPrefill( block_tables=block_table_tensor[prefill_start:], - seq_lens=seq_lens[prefill_start:], + seq_lens=prefill_seq_lens, cum_seq_lens_q=qo_indptr_prefill_gpu, cum_seq_lens_kv=paged_kv_indptr_prefill_gpu, max_q_len=max_q_len_prefill, @@ -1409,10 +1436,16 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]): f"Got {num_decode_tokens=} and {num_decodes=}." ) assert self.flashinfer_trtllm_api_decode_kernel is not None + seq_lens_decode = seq_lens[:num_decodes] + if self.use_dcp: + assert common_attn_metadata.dcp_local_seq_lens is not None + seq_lens_decode = common_attn_metadata.dcp_local_seq_lens[ + :num_decodes + ] attn_metadata.decode = FlashInferTrtllmAPIDecode( kernel=self.flashinfer_trtllm_api_decode_kernel, block_tables=block_table_tensor[:num_decodes], - seq_lens=seq_lens[:num_decodes], + seq_lens=seq_lens_decode, max_seq_len=max_seq_len, ) else: @@ -2086,6 +2119,18 @@ class FlashInferImpl(AttentionImpl): f"contiguous, got strides {kv_strides}" ) + if use_dcp: + assert decode_with_trtllm_gen + if output.dtype == FP4_DTYPE: + raise NotImplementedError( + "DCP decode with FlashInfer trtllm-gen does not support " + "FP4 attention output yet." + ) + decode_query = get_dcp_group().all_gather( + decode_query.contiguous(), dim=-2 + ) + decode_query = canonicalize_singleton_dim_strides(decode_query) + if output.dtype == FP4_DTYPE: assert self.o_sf_scale is not None out = FP4Tensor( @@ -2124,6 +2169,19 @@ class FlashInferImpl(AttentionImpl): else self.bmm1_scale ) + lse = None + if use_dcp: + out = torch.empty( + decode_query.shape, + dtype=output.dtype, + device=decode_query.device, + ) + lse = torch.empty( + (decode_query.size(0), decode_query.size(1)), + dtype=torch.float32, + device=decode_query.device, + ) + trtllm_batch_decode_with_kv_cache( query=decode_query, kv_cache=( @@ -2145,9 +2203,19 @@ class FlashInferImpl(AttentionImpl): kv_cache_sf=( nvfp4_kv_block_scales if self.is_kvcache_nvfp4 else None ), + lse=lse, + return_lse=self.need_to_return_lse_for_decode, ) - if needs_fp8_out: + if use_dcp: + assert isinstance(out, torch.Tensor) + assert lse is not None + output[:num_decode_tokens] = self.dcp_combine( + out, + lse, + get_dcp_group(), + ) + elif needs_fp8_out: output[:num_decode_tokens].copy_(out.to(output.dtype)) return output_padded diff --git a/vllm/v1/attention/backends/mla/tokenspeed_mla.py b/vllm/v1/attention/backends/mla/tokenspeed_mla.py index 0f819fe8ce..2daf85ccb3 100644 --- a/vllm/v1/attention/backends/mla/tokenspeed_mla.py +++ b/vllm/v1/attention/backends/mla/tokenspeed_mla.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """TokenSpeed CuTe DSL MLA decode backend (Blackwell, FP8 KV cache only).""" -from typing import ClassVar +from typing import TYPE_CHECKING, ClassVar import torch @@ -25,6 +25,10 @@ from vllm.v1.attention.backend import ( ) from vllm.v1.attention.backends.utils import KVCacheLayoutType +if TYPE_CHECKING: + from vllm.config import VllmConfig + from vllm.v1.kv_cache_interface import AttentionSpec + logger = init_logger(__name__) # Workspace upper bound for tokenspeed_mla_decode (per-device, lazy): @@ -55,6 +59,22 @@ class TokenspeedMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]): _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.UNIFORM + def __init__( + self, + kv_cache_spec: "AttentionSpec", + layer_names: list[str], + vllm_config: "VllmConfig", + device: torch.device, + ) -> None: + super().__init__( + kv_cache_spec, + layer_names, + vllm_config, + device, + MLACommonMetadata, + supports_dcp_with_varlen=True, + ) + class TokenspeedMLABackend(MLACommonBackend): supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16] @@ -130,6 +150,11 @@ class TokenspeedMLABackend(MLACommonBackend): class TokenspeedMLAImpl(MLACommonImpl[MLACommonMetadata]): + can_return_lse_for_decode: 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 + def __init__( self, num_heads: int, @@ -228,9 +253,15 @@ class TokenspeedMLAImpl(MLACommonImpl[MLACommonMetadata]): f"q_scale={layer._q_scale_float}, k_scale={layer._k_scale_float}." ) + num_decodes = attn_metadata.num_decodes + num_decode_tokens = attn_metadata.num_decode_tokens + block_tables = attn_metadata.decode.block_table + seq_lens = attn_metadata.decode.seq_lens + causal_seqs = attn_metadata.decode.dcp_tot_seq_lens + # tokenspeed_mla_decode expects query shape # (num_decodes, q_len_per_request, num_heads, head_dim). - if attn_metadata.num_decode_tokens % attn_metadata.num_decodes != 0: + if num_decode_tokens % num_decodes != 0: logger.warning_once( """TokenspeedMLAImpl got a query of uneven length. This usually indicates an issue in batch reordering @@ -238,7 +269,7 @@ class TokenspeedMLAImpl(MLACommonImpl[MLACommonMetadata]): ) q = q.unsqueeze(1) else: - q = q.view(attn_metadata.num_decodes, -1, q.shape[-2], q.shape[-1]) + q = q.view(num_decodes, -1, q.shape[-2], q.shape[-1]) if self.softmax_scale is None: # FP8 KV cache is mandatory for this backend, so q_scale/k_scale @@ -257,22 +288,29 @@ class TokenspeedMLAImpl(MLACommonImpl[MLACommonMetadata]): # vLLM kv_c_and_k_pe_cache is already (num_blocks, block_size, head_size). # tokenspeed_mla_decode wants 3D — pass as-is (no unsqueeze, unlike trtllm). - o = tokenspeed_mla_decode( + return_lse = self.need_to_return_lse_for_decode + kernel_out = tokenspeed_mla_decode( query=q, kv_cache=kv_c_and_k_pe_cache, workspace_buffer=self._workspace_buffer, kv_lora_rank=self.kv_lora_rank, qk_rope_head_dim=self.qk_rope_head_dim, - block_tables=attn_metadata.decode.block_table, - seq_lens=attn_metadata.decode.seq_lens, + block_tables=block_tables, + seq_lens=seq_lens, max_seq_len=attn_metadata.max_seq_len, softmax_scale=self.softmax_scale, output_scale=self.output_scale, enable_pdl=False, + return_lse=return_lse, + causal_seqs=causal_seqs if self.dcp_world_size > 1 else None, + cp_world=self.dcp_world_size, + cp_rank=self.dcp_rank, ) + if return_lse: + o, lse = kernel_out + lse = lse.view(-1, lse.shape[-1]) + else: + o, lse = kernel_out, None - # Flatten the output for consistent shape o = o.view(-1, o.shape[-2], o.shape[-1]) - - # tokenspeed_mla_decode does not return LSE. - return o, None + return o, lse