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
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),