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