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

[minimax-m3] fp8 attention GEMMs on SM100 (fp8_e4m3 KV + trtllm_mha) (#30971)

Browse Source 
Co-authored-by: qiuyue <qiuyue@minimaxi.com>
Co-authored-by: xuebi <xuebi@minimaxi.com>
Co-authored-by: Xiaoyu Zhang <1182563586@qq.com>
This commit is contained in:
alumkal
2026-08-01 09:39:52 +08:00
committed by GitHub
co-authored by qiuyue xuebi Xiaoyu Zhang
parent e6a4cefc69
commit bae8eb8d6c
23 changed files with 1766 additions and 145 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/attention/unittests/dense/test_trtllm_mha.py
+148 -1
View File
@@ -6,7 +6,11 @@ import torch
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.srt.utils import is_flashinfer_available
from sglang.srt.utils.common import is_sm90_supported, is_sm120_supported
from sglang.srt.utils.common import (
is_sm90_supported,
is_sm100_supported,
is_sm120_supported,
)
from sglang.test.test_utils import CustomTestCase
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
@@ -14,6 +18,7 @@ sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kits.attention_unittest.attention_methods.dense_attention import (
DenseAttentionCase,
build_dense_attention_fixture,
run_dense_attention_case,
)
from sglang.test.kits.attention_unittest.runner_modes.cuda_graph_decode_runner import (
@@ -200,6 +205,148 @@ class TestTRTLLMMHADenseAttentionBackendCorrectness(CustomTestCase):
hidden_size=self.HIDDEN_SIZE,
)
# XQA has native page-128 kernels (any head layout). max_context_len must
# be a page multiple so the kit's per-request slot ranges stay page-aligned.
def test_page128_decode(self):
case = DenseAttentionCase(
name="trtllm_mha_xqa_decode_page128_boundary",
backend="trtllm_mha",
forward_mode=ForwardMode.DECODE,
num_heads=4,
num_kv_heads=2,
page_size=128,
prefix_lens=(127, 128, 200),
)
run_dense_attention_case(
self,
case,
head_dim=self.HEAD_DIM,
hidden_size=self.HIDDEN_SIZE,
max_context_len=512,
)
@unittest.skipIf(
not torch.cuda.is_available()
or not is_flashinfer_available()
or not is_sm100_supported(),
"CUDA + FlashInfer TRT-LLM-GEN (SM100) are required",
)
class TestTRTLLMMHAPage128TrtllmGen(CustomTestCase):
"""page_size=128 on trtllm-gen (SM100) via dynamic tokens-per-page kernels.
Those kernels only exist for GQA (q heads per kv head > 1) with equal QK/V
head dims, so every positive case here is GQA; the MHA layout must fail at
backend construction (see test_page128_mha_rejected_at_init). All cases
pass max_context_len=512: the kit's per-request slot ranges start at
``page_size + req_idx * max_context_len``, so it must be a page multiple.
"""
HEAD_DIM = 64
HIDDEN_SIZE = 256
MAX_CONTEXT_LEN = 512
DECODE_CASES = (
DenseAttentionCase(
name="trtllm_gen_gqa_decode_page128_boundary",
backend="trtllm_mha",
forward_mode=ForwardMode.DECODE,
num_heads=4,
num_kv_heads=2,
page_size=128,
prefix_lens=(127, 128, 200),
),
DenseAttentionCase(
name="trtllm_gen_gqa4_decode_page128_bsz1",
backend="trtllm_mha",
forward_mode=ForwardMode.DECODE,
num_heads=8,
num_kv_heads=2,
page_size=128,
prefix_lens=(300,),
),
)
EXTEND_CASES = (
DenseAttentionCase(
name="trtllm_gen_gqa_extend_page128",
backend="trtllm_mha",
forward_mode=ForwardMode.EXTEND,
num_heads=4,
num_kv_heads=2,
page_size=128,
prefix_lens=(0, 128),
extend_lens=(130, 5),
),
)
CUDA_GRAPH_DECODE_CASES = (
DenseAttentionCase(
name="runner_cuda_graph_trtllm_gen_gqa_decode_page128",
backend="trtllm_mha",
forward_mode=ForwardMode.DECODE,
num_heads=4,
num_kv_heads=2,
page_size=128,
prefix_lens=(127, 128, 200),
),
)
def test_page128_decode_cases(self):
for case in self.DECODE_CASES:
with self.subTest(case=case.name):
run_dense_attention_case(
self,
case,
head_dim=self.HEAD_DIM,
hidden_size=self.HIDDEN_SIZE,
max_context_len=self.MAX_CONTEXT_LEN,
)
def test_page128_extend_cases(self):
for case in self.EXTEND_CASES:
with self.subTest(case=case.name):
run_dense_attention_case(
self,
case,
head_dim=self.HEAD_DIM,
hidden_size=self.HIDDEN_SIZE,
max_context_len=self.MAX_CONTEXT_LEN,
)
def test_page128_cuda_graph_decode_cases(self):
for case in self.CUDA_GRAPH_DECODE_CASES:
with self.subTest(case=case.name):
run_dense_cuda_graph_decode_case(
self,
case,
head_dim=self.HEAD_DIM,
hidden_size=self.HIDDEN_SIZE,
max_context_len=self.MAX_CONTEXT_LEN,
)
def test_page128_mha_rejected_at_init(self):
# heads_per_kv == 1 has no page-128 trtllm-gen kernel; the backend must
# refuse at construction (not fail mid-capture with a missing-kernel
# RuntimeError from flashinfer).
case = DenseAttentionCase(
name="trtllm_gen_mha_decode_page128_rejected",
backend="trtllm_mha",
forward_mode=ForwardMode.DECODE,
num_heads=4,
num_kv_heads=4,
page_size=128,
prefix_lens=(7,),
)
with self.assertRaisesRegex(ValueError, "dynamic tokens-per-page"):
build_dense_attention_fixture(
self,
case,
head_dim=self.HEAD_DIM,
hidden_size=self.HIDDEN_SIZE,
max_context_len=self.MAX_CONTEXT_LEN,
)
if __name__ == "__main__":
unittest.main()
test/registered/unit/test_model_overrides.py
+106
View File
@@ -22,6 +22,7 @@ from sglang.srt.arg_groups.overrides import (
register_model_override,
validate_declarations,
)
from sglang.srt.environ import envs
from sglang.srt.runtime_context import (
get_context,
get_server_args,
@@ -1642,6 +1643,87 @@ class TestGoldenModelOverrides(_IsolatedPublish):
{},
)
def test_m3_fp8_attn_gemm_resolution(self):
from sglang.srt.arg_groups.overrides import _minimax_m3_overrides
from sglang.srt.server_args import m3_fp8_attn_gemm_enabled
def _args(**kw):
defaults = dict(
attention_backend="trtllm_mha",
kv_cache_dtype="fp8_e4m3",
)
defaults.update(kw)
return SimpleNamespace(**defaults)
with patch("sglang.srt.utils.common.is_sm100_supported", return_value=True):
# e4m3 + trtllm_mha + SM100: mode active
self.assertTrue(m3_fp8_attn_gemm_enabled(_args()))
# fa4 dense backend: mode inactive (no fp8-q GEMM path)
self.assertFalse(m3_fp8_attn_gemm_enabled(_args(attention_backend="fa4")))
# bf16 KV: mode inactive
self.assertFalse(m3_fp8_attn_gemm_enabled(_args(kv_cache_dtype="auto")))
# e5m2: mode inactive (fmha_sm100's variant lookup would silently
# dispatch the e4m3 kernel)
self.assertFalse(m3_fp8_attn_gemm_enabled(_args(kv_cache_dtype="fp8_e5m2")))
# SGLANG_DISABLE_M3_FP8_ATTN_GEMM kill switch wins over an
# otherwise-active config
with envs.SGLANG_DISABLE_M3_FP8_ATTN_GEMM.override(True):
self.assertFalse(m3_fp8_attn_gemm_enabled(_args()))
with patch("sglang.srt.utils.common.is_sm100_supported", return_value=False):
# non-SM100: mode inactive
self.assertFalse(m3_fp8_attn_gemm_enabled(_args()))
def _m3_args(**kw):
defaults = dict(
quantization=None,
_quantization_explicitly_unset=True,
attention_backend=None,
prefill_attention_backend=None,
decode_attention_backend=None,
page_size=None,
moe_runner_backend="auto",
kv_cache_dtype="auto",
)
defaults.update(kw)
ns = SimpleNamespace(**defaults)
ns.is_attention_backend_not_set = lambda: (
ns.attention_backend is None
and ns.prefill_attention_backend is None
and ns.decode_attention_backend is None
)
return ns
hf = SimpleNamespace()
with patch.object(overrides_module, "is_hip", return_value=False), patch.object(
overrides_module, "is_sm100_supported", return_value=True
), patch.object(overrides_module, "get_quantization_config", return_value=None):
# fp8_e4m3 KV: SM100 backend default flips to trtllm_mha (the only
# dense backend with the fp8-q GEMM path); page snaps to 128
ov = _minimax_m3_overrides(_m3_args(kv_cache_dtype="fp8_e4m3"), hf)
self.assertEqual(ov["attention_backend"], "trtllm_mha")
self.assertEqual(ov["page_size"], 128)
# auto KV: fa4 stays the SM100 default
ov = _minimax_m3_overrides(_m3_args(), hf)
self.assertEqual(ov["attention_backend"], "fa4")
self.assertEqual(ov["page_size"], 128)
# e5m2 KV: stays on fa4 + the widening Triton path, and warns
with self.assertLogs(
"sglang.srt.arg_groups.overrides", level="WARNING"
) as logs:
ov = _minimax_m3_overrides(_m3_args(kv_cache_dtype="fp8_e5m2"), hf)
self.assertEqual(ov["attention_backend"], "fa4")
self.assertIn("fp8_e5m2", "\n".join(logs.output))
# explicit backend choice is never overridden
ov = _minimax_m3_overrides(
_m3_args(kv_cache_dtype="fp8_e4m3", attention_backend="fa4"), hf
)
self.assertNotIn("attention_backend", ov)
# kill switch also reverts the SM100 backend default to fa4
with envs.SGLANG_DISABLE_M3_FP8_ATTN_GEMM.override(True):
ov = _minimax_m3_overrides(_m3_args(kv_cache_dtype="fp8_e4m3"), hf)
self.assertEqual(ov["attention_backend"], "fa4")
self.assertEqual(ov["page_size"], 128)
def test_page_constraint_passes_at_callable_level(self):
from sglang.srt.arg_groups.overrides import (
ResolvedView,
@@ -1682,6 +1764,30 @@ class TestGoldenModelOverrides(_IsolatedPublish):
),
{"page_size": 64},
)
# trtllm_mha accepts 128 (trtllm-gen dynamic tokens-per-page kernels)
self.assertEqual(
_mla_backend_page_constraints(
_view(attention_backend="trtllm_mha", page_size=128)
),
{},
)
# trtllm_mha with an unsupported page still snaps to 64
self.assertEqual(
_mla_backend_page_constraints(
_view(attention_backend="trtllm_mha", page_size=256)
),
{"page_size": 64},
)
# chained: cutlass_mla decode -> 128, then trtllm_mha prefill keeps 128
self.assertEqual(
_mla_backend_page_constraints(
_view(
decode_attention_backend="cutlass_mla",
prefill_attention_backend="trtllm_mha",
)
),
{"page_size": 128},
)
# no matching backend: nothing declared
self.assertEqual(_mla_backend_page_constraints(_view()), {})