From ebf11be355dfe754e1bbf82c4ef9204efc642f48 Mon Sep 17 00:00:00 2001 From: LuminolT Date: Wed, 8 Jul 2026 12:05:46 +0800 Subject: [PATCH] fix(moe): gate fp8 megamoe on deepgemm capability Avoid entering the FP8 MegaMoE fused path when the installed DeepGEMM exposes fp8_mega_moe but cannot allocate an fp8xfp8 MegaMoE symmetric buffer. Current B300 sgl-deep-gemm asserts that MegaMoE buffer mma_type is fp8xfp4, so FP8 weights must fall back until dependency support lands. Tested: PYTHONPYCACHEPREFIX=/private/tmp/sglang_pycache python3 -m py_compile python/sglang/srt/layers/moe/mega_moe.py test/registered/unit/moe/test_glm_megamoe.py Tested: git diff --check Remote-tested: root@95.133.252.48 B300 DeepGEMM capability probe showed fp8xfp8 MegaMoE buffer fails with mma_type_str == fp8xfp4. --- python/sglang/srt/layers/moe/mega_moe.py | 13 ++++++++++++- test/registered/unit/moe/test_glm_megamoe.py | 2 +- 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/layers/moe/mega_moe.py b/python/sglang/srt/layers/moe/mega_moe.py index db7ce6871..5fb0f3879 100644 --- a/python/sglang/srt/layers/moe/mega_moe.py +++ b/python/sglang/srt/layers/moe/mega_moe.py @@ -91,6 +91,17 @@ def _get_mega_moe_symm_buffer( return buf +def _deep_gemm_supports_fp8_mega_moe() -> bool: + if not deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM: + return False + if not torch.cuda.is_available(): + return False + # Current sgl-deep-gemm exposes fp8_mega_moe, but its Blackwell MegaMoE + # symmetric buffer only accepts fp8xfp4. SM90 uses a separate buffer path + # that can back fp8_mega_moe. + return torch.cuda.get_device_capability()[0] == 9 + + def should_use_mega_moe(moe, hidden_states: torch.Tensor) -> bool: if not get_moe_a2a_backend().is_megamoe(): return False @@ -98,7 +109,7 @@ def should_use_mega_moe(moe, hidden_states: torch.Tensor) -> bool: return False weight_format = getattr(moe.experts, "_mega_moe_weight_format", None) if weight_format == "fp8": - return deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM + return _deep_gemm_supports_fp8_mega_moe() if weight_format != "fp4": return False if not envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_FP4_ACTS.get(): diff --git a/test/registered/unit/moe/test_glm_megamoe.py b/test/registered/unit/moe/test_glm_megamoe.py index 73a7a738d..e3e43c2a6 100644 --- a/test/registered/unit/moe/test_glm_megamoe.py +++ b/test/registered/unit/moe/test_glm_megamoe.py @@ -82,7 +82,7 @@ class TestGLMMegaMoE(unittest.TestCase): with patch.object( mega_moe, "get_moe_a2a_backend", return_value=_MegaBackend() ), patch.object( - mega_moe.deep_gemm_wrapper, "ENABLE_JIT_DEEPGEMM", True + mega_moe, "_deep_gemm_supports_fp8_mega_moe", return_value=True ), envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_FP4_ACTS.override(False): self.assertTrue(mega_moe.should_use_mega_moe(moe, hidden_states))