minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Support Flashinfer Cute-DSL MLA attention (#24737)

Browse Source 
Co-authored-by: b8zhong <b8zhong@users.noreply.github.com>
This commit is contained in:
Brayden Zhong
2026-05-28 00:21:32 -07:00
committed by GitHub
co-authored by b8zhong
parent e31ea50df8
commit 50e0b3b77f
8 changed files with 101 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/attention_registry.py
+9
View File
@@ -74,6 +74,15 @@ def create_tokenspeed_mla_backend(runner):
return TokenspeedMLABackend(runner)
@register_attention_backend("cutedsl_mla")
def create_cutedsl_mla_backend(runner):
if not runner.use_mla_backend:
raise ValueError("cutedsl_mla backend can only be used with MLA models.")
from sglang.srt.layers.attention.trtllm_mla_backend import TRTLLMMLABackend
return TRTLLMMLABackend(runner, backend="cute-dsl")
@register_attention_backend("aiter")
def create_aiter_backend(runner):
from sglang.srt.layers.attention.aiter_backend import AiterAttnBackend
python/sglang/srt/layers/attention/trtllm_mla_backend.py
+35 -9
View File
@@ -229,6 +229,11 @@ def _quantize_fp8_qkv(q, k, v, layer):
global_zero_init_workspace_buffer = None
# cute-dsl needs its own workspace: it overwrites the buffer with split-KV
# partials, which corrupts the trtllm-gen multiCtasKv counters that rely on the
# zero-init buffer (they share it under attention-backend=cutedsl_mla, where
# draft-extend falls back to trtllm-gen) and deadlocks the reduction.
global_cute_dsl_workspace_buffer = None
@dataclass
@@ -263,6 +268,7 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
skip_prefill: bool = False,
kv_indptr_buf: Optional[torch.Tensor] = None,
q_indptr_decode_buf: Optional[torch.Tensor] = None,
backend: str = "trtllm-gen",
):
super().__init__(
model_runner,
@@ -286,6 +292,7 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
self.kv_cache_dim = self.kv_lora_rank + self.qk_rope_head_dim
# Runtime parameters
self.backend = backend
self.scaling = config.scaling
self.data_type = model_runner.kv_cache_dtype
self.q_data_type = model_runner.dtype
@@ -294,14 +301,26 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
# Workspace allocation
self.workspace_size = DEFAULT_WORKSPACE_SIZE_MB * 1024 * 1024
global global_zero_init_workspace_buffer
if global_zero_init_workspace_buffer is None:
global_zero_init_workspace_buffer = torch.zeros(
self.workspace_size,
dtype=torch.uint8,
device=model_runner.device,
)
self.workspace_buffer = global_zero_init_workspace_buffer
if self.backend == "cute-dsl":
# Separate buffer from trtllm-gen (see note above); safe to share
# among cute-dsl instances.
global global_cute_dsl_workspace_buffer
if global_cute_dsl_workspace_buffer is None:
global_cute_dsl_workspace_buffer = torch.zeros(
self.workspace_size,
dtype=torch.int8,
device=model_runner.device,
)
self.workspace_buffer = global_cute_dsl_workspace_buffer
else:
global global_zero_init_workspace_buffer
if global_zero_init_workspace_buffer is None:
global_zero_init_workspace_buffer = torch.zeros(
self.workspace_size,
dtype=torch.int8,
device=model_runner.device,
)
self.workspace_buffer = global_zero_init_workspace_buffer
# CUDA graph state
self.decode_cuda_graph_metadata = {}
@@ -807,6 +826,7 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
seq_lens_i32 = (
seq_lens if seq_lens.dtype == torch.int32 else seq_lens.to(torch.int32)
)
extra_kwargs = {"backend": self.backend} if self.backend != "trtllm-gen" else {}
return flashinfer.decode.trtllm_batch_decode_with_kv_cache_mla(
query=query,
kv_cache=kv_cache,
@@ -819,6 +839,7 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
max_seq_len=max_seq_len,
bmm1_scale=bmm1_scale,
skip_softmax_threshold_scale_factor=envs.SGLANG_SKIP_SOFTMAX_DECODE_THRESHOLD_SCALE_FACTOR.get(),
**extra_kwargs,
)
def _run_prefill_kernel(
@@ -1224,7 +1245,11 @@ class TRTLLMMLAMultiStepDraftBackend(FlashInferMLAMultiStepDraftBackend):
"""Multi-step draft backend for TRT-LLM MLA used by EAGLE."""
def __init__(
self, model_runner: "ModelRunner", topk: int, speculative_num_steps: int
self,
model_runner: "ModelRunner",
topk: int,
speculative_num_steps: int,
backend: str = "trtllm-gen",
):
super().__init__(model_runner, topk, speculative_num_steps)
@@ -1234,6 +1259,7 @@ class TRTLLMMLAMultiStepDraftBackend(FlashInferMLAMultiStepDraftBackend):
skip_prefill=True,
kv_indptr_buf=self.kv_indptr[i],
q_indptr_decode_buf=self.q_indptr_decode,
backend=backend,
)
def init_forward_metadata(self, forward_batch: ForwardBatch):
python/sglang/srt/model_executor/model_runner.py
+2
View File
@@ -250,6 +250,7 @@ MLA_ATTENTION_BACKENDS = [
"fa4",
"triton",
"flashmla",
"cutedsl_mla",
"cutlass_mla",
"trtllm_mla",
"tokenspeed_mla",
@@ -264,6 +265,7 @@ CHUNKED_PREFIX_CACHE_SUPPORTED_ATTENTION_BACKENDS = [
"fa3",
"fa4",
"flashmla",
"cutedsl_mla",
"cutlass_mla",
"trtllm_mla",
"tokenspeed_mla",
python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py
+2 -1
View File
@@ -698,7 +698,8 @@ class DeepseekMLAForwardMixin:
) and get_attn_backend().kv_cache_dtype == torch.float8_e4m3fn
return (
self.current_attention_backend in ("trtllm_mla", "tokenspeed_mla")
self.current_attention_backend
in ("trtllm_mla", "tokenspeed_mla", "cutedsl_mla")
and (
forward_batch.forward_mode.is_decode_or_idle()
or forward_batch.forward_mode.is_target_verify()
python/sglang/srt/models/deepseek_common/utils.py
+1
View File
@@ -62,6 +62,7 @@ FORWARD_ABSORB_CORE_ATTENTION_BACKENDS = [
"flashinfer",
"cutlass_mla",
"trtllm_mla",
"cutedsl_mla",
"tokenspeed_mla",
"ascend",
"intel_xpu",
python/sglang/srt/server_args.py
+30
View File
@@ -166,6 +166,7 @@ ATTENTION_BACKEND_CHOICES = [
"flashinfer",
"flashmla",
"trtllm_mla",
"cutedsl_mla",
"tokenspeed_mla",
"trtllm_mha",
"dual_chunk_flash_attn",
@@ -2822,6 +2823,35 @@ class ServerArgs:
f"got {self.kv_cache_dtype}."
)
if (
self.attention_backend == "cutedsl_mla"
or self.decode_attention_backend == "cutedsl_mla"
or self.prefill_attention_backend == "cutedsl_mla"
):
assert (
self.prefill_attention_backend != "cutedsl_mla"
), "CuteDSL MLA only supports decoding for now"
if not is_sm100_supported():
raise ValueError(
"CuteDSL MLA backend is only supported on Blackwell GPUs (SM100). Please use a different backend."
)
if self.page_size not in [32, 64]:
logger.warning(
f"CuteDSL MLA only supports page_size of 32 or 64, changing page_size from {self.page_size} to 64."
)
self.page_size = 64
if self.kv_cache_dtype not in [
"fp8_e4m3",
"bf16",
"bfloat16",
"auto",
]:
raise ValueError(
"CuteDSL MLA backend only supports kv-cache-dtype of fp8_e4m3, bf16, or auto."
)
if self.prefill_attention_backend is None:
self.prefill_attention_backend = "trtllm_mla"
if (
self.attention_backend == "trtllm_mha"
or self.decode_attention_backend == "trtllm_mha"
python/sglang/srt/speculative/draft_utils.py
+11 -2
View File
@@ -53,6 +53,7 @@ class DraftBackendFactory:
"flashmla": self._create_flashmla_decode_backend,
"trtllm_mha": self._create_trtllm_mha_decode_backend,
"trtllm_mla": self._create_trtllm_mla_decode_backend,
"cutedsl_mla": self._create_cutedsl_mla_decode_backend,
"tokenspeed_mla": self._create_tokenspeed_mla_decode_backend,
"dsa": self._create_dsa_decode_backend,
"nsa": self._create_dsa_decode_backend, # Deprecated alias for "dsa"
@@ -81,6 +82,8 @@ class DraftBackendFactory:
"flashmla": self._create_flashmla_prefill_backend,
"trtllm_mha": self._create_trtllm_mha_prefill_backend,
"trtllm_mla": self._create_trtllm_mla_prefill_backend,
# cute-dsl MLA only supports decode; draft-extend falls back to trtllm-gen.
"cutedsl_mla": self._create_trtllm_mla_prefill_backend,
"tokenspeed_mla": self._create_tokenspeed_mla_prefill_backend,
"dsa": self._create_dsa_prefill_backend,
"nsa": self._create_dsa_prefill_backend, # Deprecated alias for "dsa"
@@ -188,7 +191,7 @@ class DraftBackendFactory:
self.draft_model_runner, self.topk, self.speculative_num_steps
)
def _create_trtllm_mla_decode_backend(self):
def _create_trtllm_mla_decode_backend(self, backend: str = "trtllm-gen"):
if not get_global_server_args().use_mla_backend:
raise ValueError(
"trtllm_mla backend requires MLA model (use_mla_backend=True)."
@@ -199,9 +202,15 @@ class DraftBackendFactory:
)
return TRTLLMMLAMultiStepDraftBackend(
self.draft_model_runner, self.topk, self.speculative_num_steps
self.draft_model_runner,
self.topk,
self.speculative_num_steps,
backend=backend,
)
def _create_cutedsl_mla_decode_backend(self):
return self._create_trtllm_mla_decode_backend(backend="cute-dsl")
def _create_tokenspeed_mla_decode_backend(self):
if not get_global_server_args().use_mla_backend:
raise ValueError(