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

feat: SM120 (Blackwell Desktop) support for GLM-5.1 inference (#26928)

Browse Source
This commit is contained in:
Void
2026-07-28 14:52:34 -07:00
committed by GitHub
parent 9c0dbf508f
commit 7f438a6031
9 changed files with 371 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/attention/flash_mla_sm120.py
+82
View File
@@ -19,8 +19,15 @@ import triton
import triton.language as tl
from sglang.srt.environ import envs
from sglang.srt.utils import is_hip
logger = logging.getLogger(__name__)
_is_hip = is_hip()
_GLM_DSA_MODEL_ARCHS = (
"GlmMoeDsaForCausalLM",
"GlmMoeDsaForCausalLMNextN",
)
# Page layout constants for DSv4-Flash (MODEL1):
# nope_dim = 448, rope_dim = 64, quantize_block_size = 64
@@ -474,3 +481,78 @@ def _flash_mla_flashinfer(
)
return (output.unsqueeze(1), None)
def _validate_flashinfer_sparse_mla_backend(
*,
model_arch: str,
device_sm_major: int,
kv_cache_dtype: torch.dtype,
prefill_impl: str,
decode_impl: str,
) -> bool:
selected = {prefill_impl, decode_impl}
uses_flashinfer_sparse_mla = "flashinfer_sparse_mla" in selected
is_glm_sm12_fp8 = (
model_arch in _GLM_DSA_MODEL_ARCHS
and device_sm_major == 12
and kv_cache_dtype == torch.float8_e4m3fn
and not _is_hip
)
if uses_flashinfer_sparse_mla and not is_glm_sm12_fp8:
raise ValueError(
"flashinfer_sparse_mla supports only GLM DSA with FP8 KV cache "
"on NVIDIA SM120/SM121; "
f"got model_arch={model_arch!r}, sm_major={device_sm_major}, "
f"kv_cache_dtype={kv_cache_dtype}, prefill_impl={prefill_impl!r}, "
f"decode_impl={decode_impl!r}."
)
if is_glm_sm12_fp8:
unsupported = selected - {"flashinfer_sparse_mla"}
if unsupported:
raise ValueError(
"GLM DSA with FP8 KV cache on NVIDIA SM120/SM121 supports "
"only flashinfer_sparse_mla, "
f"but got {sorted(unsupported)}."
)
return uses_flashinfer_sparse_mla
def flashinfer_sparse_mla_forward(
q: torch.Tensor,
kv_cache: torch.Tensor,
indices: torch.Tensor,
seq_lens: torch.Tensor,
workspace_buffer: torch.Tensor,
*,
page_size: int,
kv_cache_dim: int,
qk_nope_head_dim: int,
kv_lora_rank: int,
qk_rope_head_dim: int,
sm_scale: float,
skip_softmax_threshold_scale_factor: float | None,
) -> torch.Tensor:
"""Run FlashInfer's SM120 sparse MLA kernel on SGLang's packed DSA cache."""
from flashinfer.mla import trtllm_batch_decode_with_kv_cache_mla
topk = indices.shape[1]
result = trtllm_batch_decode_with_kv_cache_mla(
query=q.unsqueeze(1),
kv_cache=kv_cache.view(torch.uint8)
.view(-1, page_size, kv_cache_dim)
.unsqueeze(1),
workspace_buffer=workspace_buffer,
qk_nope_head_dim=qk_nope_head_dim,
kv_lora_rank=kv_lora_rank,
qk_rope_head_dim=qk_rope_head_dim,
block_tables=indices.unsqueeze(1),
seq_lens=seq_lens,
max_seq_len=topk,
sparse_mla_top_k=topk,
bmm1_scale=float(sm_scale),
bmm2_scale=1.0,
kv_scale_format="arbitrary_fp32",
skip_softmax_threshold_scale_factor=skip_softmax_threshold_scale_factor,
)
return result.squeeze(1)
python/sglang/srt/arg_groups/overrides.py
+19
View File
@@ -1293,6 +1293,25 @@ def _dsa_split_backend_resolution(view: Any) -> dict:
user_set_prefill = view.dsa_prefill_backend is not None
user_set_decode = view.dsa_decode_backend is not None
declared: Dict[str, Any] = {}
model_arch = hf_config.architectures[0]
is_glm_sm12_fp8 = (
model_arch == "GlmMoeDsaForCausalLM"
and major == 12
and kv_cache_dtype == "fp8_e4m3"
and not is_hip()
)
if is_glm_sm12_fp8:
backend = "flashinfer_sparse_mla"
if not user_set_prefill:
declared["dsa_prefill_backend"] = backend
if not user_set_decode:
declared["dsa_decode_backend"] = backend
logger.warning(
"Set DSA backends for GLM FP8 KV Cache on SM120/SM121: "
f"prefill={backend}, decode={backend}."
)
return declared
if view.enable_hisparse:
from sglang.srt.arg_groups.hisparse_hook import _hisparse_default_backend
python/sglang/srt/layers/attention/dsa_backend.py
+82 -2
View File
@@ -331,7 +331,13 @@ class DSAIndexerMetadata(BaseIndexerMetadata):
_DSA_IMPL_T: TypeAlias = Literal[
"flashmla_sparse", "flashmla_sparse_q8", "flashmla_kv", "fa3", "tilelang", "trtllm"
"flashmla_sparse",
"flashmla_sparse_q8",
"flashmla_kv",
"flashinfer_sparse_mla",
"fa3",
"tilelang",
"trtllm",
]
@@ -490,8 +496,29 @@ class DeepseekSparseAttnBackend(
self._q8kv8_identity_scale: Optional[torch.Tensor] = None
self._q8kv8_qpad_buf: Optional[torch.Tensor] = None
from sglang.kernels.ops.attention.flash_mla_sm120 import (
_validate_flashinfer_sparse_mla_backend,
)
uses_flashinfer_sparse_mla = _validate_flashinfer_sparse_mla_backend(
model_arch=model_runner.model_config.hf_config.architectures[0],
device_sm_major=self.device_sm_major,
kv_cache_dtype=self.kv_cache_dtype,
prefill_impl=self.dsa_prefill_impl,
decode_impl=self.dsa_decode_impl,
)
if uses_flashinfer_sparse_mla:
self.workspace_buffer = get_buffer(
"dsa_flashinfer_sparse_mla_workspace",
lambda: torch.zeros(
envs.SGLANG_FLASHINFER_WORKSPACE_SIZE.get(),
dtype=torch.uint8,
device=model_runner.device,
),
)
# Allocate global workspace buffer for TRT-LLM kernels (ragged attention on SM100/B200, or trtllm decode)
if self.device_sm_major >= 10 or self.dsa_decode_impl == "trtllm":
elif self.device_sm_major >= 10 or self.dsa_decode_impl == "trtllm":
self.workspace_buffer = get_buffer(
"dsa_trtllm_workspace",
lambda: torch.empty(
@@ -2090,6 +2117,19 @@ class DeepseekSparseAttnBackend(
sm_scale=layer.scaling,
v_head_dim=layer.v_head_dim,
)
elif dsa_impl == "flashinfer_sparse_mla":
if q_rope is not None:
q_all = concat_mla_absorb_q_general(q_nope, q_rope)
if topk_transform_method == TopkTransformMethod.RAGGED:
page_table_1 = topk_indices
return self._forward_flashinfer_sparse_mla(
q_all=q_all,
kv_cache=kv_cache,
page_table_1=page_table_1,
seq_lens=metadata.dsa_cache_seqlens_int32,
sm_scale=layer.scaling,
skip_softmax_threshold_scale_factor=envs.SGLANG_SKIP_SOFTMAX_PREFILL_THRESHOLD_SCALE_FACTOR.get(),
)
elif dsa_impl == "flashmla_kv":
if q_rope is not None:
q_all = concat_mla_absorb_q_general(q_nope, q_rope)
@@ -2233,6 +2273,17 @@ class DeepseekSparseAttnBackend(
sm_scale=layer.scaling,
v_head_dim=layer.v_head_dim,
)
elif self.dsa_decode_impl == "flashinfer_sparse_mla":
if q_all is None:
q_all = concat_mla_absorb_q_general(q_nope, q_rope)
return self._forward_flashinfer_sparse_mla(
q_all=q_all,
kv_cache=kv_cache,
page_table_1=page_table_1,
seq_lens=metadata.dsa_cache_seqlens_int32,
sm_scale=layer.scaling,
skip_softmax_threshold_scale_factor=envs.SGLANG_SKIP_SOFTMAX_DECODE_THRESHOLD_SCALE_FACTOR.get(),
)
elif self.dsa_decode_impl == "flashmla_kv":
if q_rope is not None:
q_all = concat_mla_absorb_q_general(q_nope, q_rope)
@@ -2502,6 +2553,35 @@ class DeepseekSparseAttnBackend(
o = o[:, :num_heads, :]
return o
def _forward_flashinfer_sparse_mla(
self,
q_all: torch.Tensor,
kv_cache: torch.Tensor,
page_table_1: torch.Tensor,
seq_lens: torch.Tensor,
sm_scale: float,
skip_softmax_threshold_scale_factor: float | None,
) -> torch.Tensor:
from sglang.kernels.ops.attention.flash_mla_sm120 import (
flashinfer_sparse_mla_forward,
)
assert self.workspace_buffer is not None
return flashinfer_sparse_mla_forward(
q=q_all,
kv_cache=kv_cache,
indices=page_table_1,
seq_lens=seq_lens,
workspace_buffer=self.workspace_buffer,
page_size=self.real_page_size,
kv_cache_dim=self.kv_cache_dim,
qk_nope_head_dim=self.qk_nope_head_dim,
kv_lora_rank=self.kv_lora_rank,
qk_rope_head_dim=self.qk_rope_head_dim,
sm_scale=sm_scale,
skip_softmax_threshold_scale_factor=skip_softmax_threshold_scale_factor,
)
def _forward_flashmla_kv(
self,
q_all: torch.Tensor,
python/sglang/srt/server_args.py
+1
View File
@@ -324,6 +324,7 @@ DSA_CHOICES = [
"flashmla_sparse_q8",
"flashmla_kv",
"flashmla_auto",
"flashinfer_sparse_mla",
"fa3",
"tilelang",
"aiter",