iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Attention] Add CuTe DSL backend selection for MLA context prefill (#62069)

Browse Source 
Signed-off-by: Wei Zhao <51183510+wzhao18@users.noreply.github.com>
Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>
This commit is contained in:
Wei Zhao
2026-10-17 01:21:55 -04:00
committed by GitHub
parent d74f4180cd
commit 90d6d871d6
4 changed files with 86 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/attention/test_mla_backends.py
+38
View File
@@ -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,
)
vllm/config/attention.py
+5
View File
@@ -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
vllm/model_executor/layers/attention/mla_attention.py
+20
View File
@@ -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(
vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py
+23
View File
@@ -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,