From 90d6d871d6700dd229a76b788bf0dca22bab6c70 Mon Sep 17 00:00:00 2001 From: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Date: Sat, 17 Oct 2026 01:21:55 -0400 Subject: [PATCH] [Attention] Add CuTe DSL backend selection for MLA context prefill (#62069) Signed-off-by: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Signed-off-by: wzhao18 --- tests/v1/attention/test_mla_backends.py | 38 +++++++++++++++++++ vllm/config/attention.py | 5 +++ .../layers/attention/mla_attention.py | 20 ++++++++++ .../backends/mla/prefill/trtllm_ragged.py | 23 +++++++++++ 4 files changed, 86 insertions(+) diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index 9d4be102d9..09a8158ff1 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -10,6 +10,7 @@ Known Issues: import sys from types import MethodType, SimpleNamespace +from typing import Literal import pytest import torch @@ -47,6 +48,9 @@ from vllm.v1.attention.backends.mla.prefill.base import MLADimensions from vllm.v1.attention.backends.mla.prefill.selector import ( MLAPrefillSelectorConfig, ) +from vllm.v1.attention.backends.mla.prefill.trtllm_ragged import ( + is_cute_dsl_mla_prefill_supported, +) from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.attention.ops.flashmla import is_flashmla_dense_supported from vllm.v1.kv_cache_interface import ( @@ -1483,6 +1487,7 @@ def run_attention_backend( # Force the prefill backend selection (None means auto-select). vllm_config.attention_config.mla_prefill_backend = prefill_backend + mla_attention_module.backend_supports_prefill_query_quantization.cache_clear() # Set the current vllm config so that get_current_vllm_config() works # in the backend implementations @@ -1641,6 +1646,8 @@ def _run_backend_correctness( v_head_dim: int, chunked_prefill_workspace_size: int | None = None, adaptive_verification: bool = False, + trtllm_ragged_backend: Literal["trtllm-gen", "cute-dsl"] = "trtllm-gen", + use_prefill_query_quantization: bool = False, ): """Test that all backends produce similar outputs to a reference implementation using torch.nn.functional.scaled_dot_product_attention. @@ -1720,6 +1727,10 @@ def _run_backend_correctness( hf_config_override["num_attention_heads"] ) vllm_config.cache_config.cache_dtype = kv_cache_dtype + vllm_config.attention_config.trtllm_ragged_backend = trtllm_ragged_backend + vllm_config.attention_config.use_prefill_query_quantization = ( + use_prefill_query_quantization + ) # For spec decode tests, add a speculative_config to set the reorder_batch_threshold if is_spec_decode_test: @@ -1910,6 +1921,11 @@ def _run_backend_correctness( q_mha = torch.cat([q_nope, q_pe], dim=-1) # [q_len, num_heads, total_dim] k_pe_full_expanded = k_pe_prefill.expand(-1, num_q_heads, -1) k_full = torch.cat([k_nope_full, k_pe_full_expanded], dim=-1) + if use_prefill_query_quantization: + fp8_dtype = current_platform.fp8_dtype() + q_mha = q_mha.to(fp8_dtype).to(dtype) + k_full = k_full.to(fp8_dtype).to(dtype) + v_full = v_full.to(fp8_dtype).to(dtype) # Create custom attention mask: # - Query tokens can attend to all context tokens @@ -2184,6 +2200,8 @@ def test_backend_correctness( _prefill_backend_dimension_params(), ) @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) +@pytest.mark.parametrize("trtllm_ragged_backend", ["trtllm-gen", "cute-dsl"]) +@pytest.mark.parametrize("use_prefill_query_quantization", [False, True]) def test_chunked_context_backend_correctness( default_vllm_config, dist_init, @@ -2192,8 +2210,26 @@ def test_chunked_context_backend_correctness( qk_nope_head_dim: int, v_head_dim: int, kv_cache_dtype: str, + trtllm_ragged_backend: Literal["trtllm-gen", "cute-dsl"], + use_prefill_query_quantization: bool, ): """Split, packed, and context-free requests match the SDPA reference.""" + if prefill_backend != MLAPrefillBackendEnum.TRTLLM_RAGGED and ( + trtllm_ragged_backend == "cute-dsl" or use_prefill_query_quantization + ): + pytest.skip("TRTLLM ragged-specific test configuration") + if use_prefill_query_quantization and kv_cache_dtype == "auto": + pytest.skip("FP8 prefill queries require an FP8 KV cache") + if trtllm_ragged_backend == "cute-dsl": + query_dtype = ( + current_platform.fp8_dtype() + if use_prefill_query_quantization + else torch.bfloat16 + ) + if not is_cute_dsl_mla_prefill_supported( + qk_nope_head_dim + 64, v_head_dim, query_dtype + ): + pytest.skip("Unsupported CuTe DSL prefill dimensions or query dtype") _run_backend_correctness( default_vllm_config, dist_init, @@ -2208,6 +2244,8 @@ def test_chunked_context_backend_correctness( qk_nope_head_dim=qk_nope_head_dim, v_head_dim=v_head_dim, chunked_prefill_workspace_size=1024, + trtllm_ragged_backend=trtllm_ragged_backend, + use_prefill_query_quantization=use_prefill_query_quantization, ) diff --git a/vllm/config/attention.py b/vllm/config/attention.py index fc3ecb0dfc..1b7c996a64 100644 --- a/vllm/config/attention.py +++ b/vllm/config/attention.py @@ -85,6 +85,11 @@ class AttentionConfig: use_prefill_query_quantization: bool = False """If set, quantize query for attention in prefill.""" + trtllm_ragged_backend: Literal["trtllm-gen", "cute-dsl"] = "trtllm-gen" + """Kernel backend for cached-context TRTLLM_RAGGED MLA prefill. + CuTe DSL requires FP8 queries for head dimension 192 and does not support + head dimension 256. New-token prefill always uses TRTLLM.""" + indexer_kv_dtype: IndexerKVDType = "auto" """Data type for the sparse-attention indexer K cache. "auto" picks the model's default (bf16 for MiniMax M3, fp8 for the DeepSeek sparse diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index caf60b009f..865913b782 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -2024,6 +2024,7 @@ def build_mla_chunked_context_metadata( dcp_local_block_size: int, dcp_virtual_block_size: int, dcp_manager: MLADCPManager | None = None, + align_indptr: bool = False, ) -> "MLACommonPrefillMetadata.ChunkedContextMetadata | None": """Build chunked-context metadata for an MLA prefill. @@ -2043,6 +2044,7 @@ def build_mla_chunked_context_metadata( dcp_local_block_size: Per-rank interleave block size for DCP. dcp_virtual_block_size: ``dcp_local_block_size * dcp_world_size``. dcp_manager: Shared MLA DCP collective manager. + align_indptr: Align packed int32 indptr slices to 16 bytes for CuTe DSL. Returns: The chunked-context metadata, or None when no prefill has any context. @@ -2123,6 +2125,15 @@ def build_mla_chunked_context_metadata( num_tokens = sum(plan.seq_lens) num_local_tokens = num_tokens + if align_indptr: + # CuTe DSL requires 16-byte-aligned int32 indptrs. + padding = -boundary_offset % 4 + cu_seq_lens_flat.extend([0] * padding) + cu_seqlens_q_flat.extend([0] * padding) + if use_dcp: + padded_local_cu_seq_lens_flat.extend([0] * padding) + boundary_offset += padding + seq_lens_flat.extend(plan.seq_lens) cu_seq_lens_flat.extend(itertools.accumulate(plan.seq_lens, initial=0)) query_base = query_start_loc[plan.request_start] @@ -2616,6 +2627,14 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]): query_start_loc_cpu[reqs_start:] - query_start_loc_cpu[reqs_start] ) + align_indptr = False + if ( + self._prefill_backend.get_name() == "TRTLLM_RAGGED" + and self.vllm_config.attention_config.trtllm_ragged_backend + == "cute-dsl" + ): + align_indptr = True + chunked_context_metadata = build_mla_chunked_context_metadata( context_lens_cpu=context_lens_cpu, prefill_query_start_loc_cpu=prefill_query_start_loc_cpu, @@ -2628,6 +2647,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]): dcp_local_block_size=self.dcp_local_block_size, dcp_virtual_block_size=self.dcp_virtual_block_size, dcp_manager=self.dcp_manager, + align_indptr=align_indptr, ) prefill_metadata = MLACommonPrefillMetadata( diff --git a/vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py b/vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py index 49c9dfa21c..a1f21cb0d1 100644 --- a/vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py +++ b/vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py @@ -22,6 +22,17 @@ if TYPE_CHECKING: from vllm.platforms.interface import DeviceCapability +def is_cute_dsl_mla_prefill_supported( + qk_head_dim: int, v_head_dim: int, query_dtype: torch.dtype +) -> bool: + # FlashInfer has no public compatibility API; replace these checks when available. + return ( + query_dtype in (torch.float16, torch.bfloat16, torch.float8_e4m3fn) + and (qk_head_dim, v_head_dim) in ((128, 128), (192, 128)) + and (qk_head_dim != 192 or query_dtype == torch.float8_e4m3fn) + ) + + class TrtllmRaggedPrefillBackend(MLAPrefillBackend): """TRT-LLM Ragged backend for MLA prefill.""" @@ -79,6 +90,7 @@ class TrtllmRaggedPrefillBackend(MLAPrefillBackend): v_head_dim=v_head_dim, vllm_config=vllm_config, ) + self._context_backend = vllm_config.attention_config.trtllm_ragged_backend self._workspace_spec = ( (envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE,), torch.uint8, @@ -198,6 +210,16 @@ class TrtllmRaggedPrefillBackend(MLAPrefillBackend): "TRTLLM ragged context prefill contains an empty query or KV row" ) + if self._context_backend == "cute-dsl": + if not is_cute_dsl_mla_prefill_supported(q.shape[-1], v.shape[-1], q.dtype): + raise ValueError( + "CuTe DSL context prefill does not support " + f"Q/K head dimension {q.shape[-1]}, " + f"V head dimension {v.shape[-1]}, and query dtype {q.dtype}." + ) + # Generic MLA supplies V as a projection slice; CuTe requires packed V. + v = v.contiguous() + if out is None: out = torch.empty( q.shape[0], @@ -225,6 +247,7 @@ class TrtllmRaggedPrefillBackend(MLAPrefillBackend): cum_seq_lens_kv=chunk.cu_seq_lens, enable_pdl=False, is_causal=False, + backend=self._context_backend, return_lse=True, out=out, skip_all_rows_active_check=True,