minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 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",
scripts/release/bump_flashinfer_version.py
+6 -3
View File
@@ -20,7 +20,8 @@ def read_current_flashinfer_version(repo_root: Path) -> str:
pyproject = repo_root / "python" / "pyproject.toml"
content = pyproject.read_text()
match = re.search(
r"flashinfer_python==(\d+\.\d+\.\d+(?:rc\d+|\.post\d+)?)", content
r"flashinfer_python(?:\[[^\]]+\])?==" r"(\d+\.\d+\.\d+(?:rc\d+|\.post\d+)?)",
content,
)
if not match:
raise ValueError(f"Could not find flashinfer_python version in {pyproject}")
@@ -39,8 +40,10 @@ def replace_flashinfer_version(
name = file_path.name
if name == "pyproject.toml":
new_content = new_content.replace(
f"flashinfer_python=={old_version}", f"flashinfer_python=={new_version}"
new_content = re.sub(
rf"(flashinfer_python(?:\[[^\]]+\])?==){re.escape(old_version)}",
rf"\g<1>{new_version}",
new_content,
)
new_content = new_content.replace(
f"flashinfer_cubin=={old_version}", f"flashinfer_cubin=={new_version}"
test/manual/test_dsa_alias_cli_registry_env.py
+1
View File
@@ -37,6 +37,7 @@ class TestDSAChoicesAndFields(unittest.TestCase):
def test_dsa_choices_is_canonical(self):
self.assertIn("fa3", self.DSA_CHOICES)
self.assertIn("tilelang", self.DSA_CHOICES)
self.assertIn("flashinfer_sparse_mla", self.DSA_CHOICES)
def test_nsa_choices_is_alias(self):
self.assertIs(
test/registered/unit/model_executor/test_pool_configurator.py
+43 -1
View File
@@ -78,6 +78,8 @@ def _make_model_runner(
disaggregation_mode="null",
max_running_requests=None,
disaggregation_decode_extra_slots=0,
kv_lora_rank=512,
qk_rope_head_dim=64,
):
"""Create a mock ModelRunner with the fields configurators need."""
mr = MagicMock()
@@ -96,6 +98,8 @@ def _make_model_runner(
mc = SimpleNamespace()
mc.head_dim = head_dim
mc.v_head_dim = v_head_dim
mc.kv_lora_rank = kv_lora_rank
mc.qk_rope_head_dim = qk_rope_head_dim
mc.is_hybrid_swa = is_hybrid_swa
mc.full_attention_layer_ids = (
full_attention_layer_ids
@@ -113,7 +117,6 @@ def _make_model_runner(
mc.hf_config.get_text_config = lambda: mc.hf_config
mc.linear_attn_registry_result = None
mr.model_config = mc
mr.kv_cache_dtype = "fake_bf16"
sa = SimpleNamespace()
@@ -132,6 +135,7 @@ def _make_model_runner(
sa.disaggregation_mode = disaggregation_mode
sa.max_running_requests = max_running_requests
sa.disaggregation_decode_extra_slots = disaggregation_decode_extra_slots
sa.enable_hisparse = False
sa.enable_dsa_cache_layer_split = False
sa.kv_cache_dtype = "auto"
mr.server_args = sa
@@ -230,6 +234,44 @@ class TestDefaultConfigurator(unittest.TestCase):
self.assertIsNone(config.full_max_total_num_tokens)
self.assertIsNone(config.swa_max_total_num_tokens)
@patch(
"sglang.srt.model_executor.pool_configurator.get_dsa_index_head_dim",
return_value=128,
)
@patch(
"sglang.srt.model_executor.pool_configurator.is_deepseek_dsa",
return_value=True,
)
@patch(
"sglang.srt.mem_cache.kv_cache_configurator.calculate_mla_kv_cache_dim",
side_effect=(576, 656),
)
def test_dsa_mla_cell_size_uses_backend_kv_layout(
self, mock_calculate_mla_kv_cache_dim, _mock_is_dsa, _mock_index_head_dim
):
num_layers = 2
raw = _make_model_runner(
num_layers=num_layers,
use_mla_backend=True,
)
packed = _make_model_runner(
num_layers=num_layers,
use_mla_backend=True,
)
with mock_cpu_env(kv_size=1):
from sglang.srt.model_executor.pool_configurator import (
DefaultPoolConfigurator,
)
raw_configurator = DefaultPoolConfigurator(raw)
packed_configurator = DefaultPoolConfigurator(packed)
# The DSA indexer adds 128 FP8 values and one FP32 scale (4 bytes).
self.assertEqual(raw_configurator._cell_size, (576 + 132) * num_layers)
self.assertEqual(packed_configurator._cell_size, (656 + 132) * num_layers)
self.assertEqual(mock_calculate_mla_kv_cache_dim.call_count, 2)
class TestHybridSWAConfigurator(unittest.TestCase):
"""Hybrid SWA: full/swa split, ratio, memory invariant."""
test/registered/unit/test_flashinfer_sparse_mla.py
+123
View File
@@ -0,0 +1,123 @@
import sys
import unittest
from types import ModuleType
from unittest.mock import patch
import torch
from sglang.kernels.ops.attention.flash_mla_sm120 import (
_validate_flashinfer_sparse_mla_backend,
flashinfer_sparse_mla_forward,
)
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=1, suite="base-a-test-cpu")
class TestFlashInferSparseMLAAdapter(unittest.TestCase):
def _mock_flashinfer(self, op):
flashinfer = ModuleType("flashinfer")
flashinfer.__path__ = []
mla = ModuleType("flashinfer.mla")
mla.trtllm_batch_decode_with_kv_cache_mla = op
flashinfer.mla = mla
return patch.dict(
sys.modules,
{"flashinfer": flashinfer, "flashinfer.mla": mla},
)
def test_maps_sglang_layout_to_public_flashinfer_api(self):
captured = {}
def fake_op(**kwargs):
captured.update(kwargs)
query = kwargs["query"]
return query.new_full((*query.shape[:-1], kwargs["kv_lora_rank"]), 2)
with self._mock_flashinfer(fake_op):
output = flashinfer_sparse_mla_forward(
q=torch.zeros((2, 8, 576), dtype=torch.bfloat16),
kv_cache=torch.zeros((128, 1, 656), dtype=torch.uint8),
indices=torch.tensor(
[[7, 9, -1, -1], [4, 6, 8, -1]], dtype=torch.int32
),
seq_lens=torch.tensor([2, 3], dtype=torch.int32),
workspace_buffer=torch.zeros(1024, dtype=torch.uint8),
page_size=64,
kv_cache_dim=656,
qk_nope_head_dim=192,
kv_lora_rank=512,
qk_rope_head_dim=64,
sm_scale=0.125,
skip_softmax_threshold_scale_factor=0.25,
)
self.assertEqual(tuple(captured["query"].shape), (2, 1, 8, 576))
self.assertEqual(tuple(captured["kv_cache"].shape), (2, 1, 64, 656))
self.assertEqual(tuple(captured["block_tables"].shape), (2, 1, 4))
self.assertEqual(
captured["block_tables"].tolist(),
[[[7, 9, -1, -1]], [[4, 6, 8, -1]]],
)
self.assertEqual(captured["seq_lens"].tolist(), [2, 3])
self.assertEqual(captured["max_seq_len"], 4)
self.assertEqual(captured["sparse_mla_top_k"], 4)
self.assertEqual(captured["qk_nope_head_dim"], 192)
self.assertEqual(captured["bmm1_scale"], 0.125)
self.assertEqual(captured["bmm2_scale"], 1.0)
self.assertEqual(captured["kv_scale_format"], "arbitrary_fp32")
self.assertEqual(captured["skip_softmax_threshold_scale_factor"], 0.25)
self.assertNotIn("backend", captured)
self.assertEqual(tuple(output.shape), (2, 8, 512))
self.assertTrue(torch.all(output == 2))
class TestFlashInferSparseMLABackendGate(unittest.TestCase):
def _validate(self, prefill, decode, model_arch="GlmMoeDsaForCausalLM"):
return _validate_flashinfer_sparse_mla_backend(
model_arch=model_arch,
device_sm_major=12,
kv_cache_dtype=torch.float8_e4m3fn,
prefill_impl=prefill,
decode_impl=decode,
)
def test_accepts_flashinfer_for_both_phases(self):
for model_arch in (
"GlmMoeDsaForCausalLM",
"GlmMoeDsaForCausalLMNextN",
):
with self.subTest(model_arch=model_arch):
self.assertTrue(
self._validate(
"flashinfer_sparse_mla",
"flashinfer_sparse_mla",
model_arch,
)
)
def test_rejects_other_or_mixed_backends(self):
for prefill, decode in (
("trtllm", "trtllm"),
("flashinfer_sparse_mla", "trtllm"),
):
with self.subTest(prefill=prefill, decode=decode):
with self.assertRaisesRegex(ValueError, "only flashinfer_sparse_mla"):
self._validate(prefill, decode)
def test_reports_unsupported_configuration(self):
with self.assertRaises(ValueError) as error:
self._validate(
"flashinfer_sparse_mla",
"flashinfer_sparse_mla",
"DeepseekV3ForCausalLM",
)
message = str(error.exception)
self.assertIn("model_arch='DeepseekV3ForCausalLM'", message)
self.assertIn("sm_major=12", message)
self.assertIn("kv_cache_dtype=torch.float8_e4m3fn", message)
if __name__ == "__main__":
unittest.main()
test/registered/unit/test_model_overrides.py
+14
View File
@@ -990,6 +990,20 @@ class TestGoldenModelOverrides(_IsolatedPublish):
self.assertEqual(
_dsa_split_backend_resolution(_view(arch="LlamaForCausalLM")), {}
)
with (
patch("sglang.srt.configs.model_config.is_deepseek_dsa", return_value=True),
patch.object(overrides_module, "is_npu", return_value=False),
patch.object(overrides_module, "is_xpu", return_value=False),
patch.object(overrides_module, "is_hip", return_value=False),
patch("torch.cuda.get_device_capability", return_value=(12, 0)),
):
self.assertEqual(
_dsa_split_backend_resolution(_view(arch="GlmMoeDsaForCausalLM")),
{
"dsa_prefill_backend": "flashinfer_sparse_mla",
"dsa_decode_backend": "flashinfer_sparse_mla",
},
)
with (
patch("sglang.srt.configs.model_config.is_deepseek_dsa", return_value=True),
patch.object(overrides_module, "is_npu", return_value=False),