diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 1ab737a1fab0..c84c24e2b862 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -1070,6 +1070,9 @@ class Envs: # DeepGEMM Mega MoE # =================================================================== SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK = EnvInt(8192) + # Blackwell MegaMoE uses a whole-grid software barrier. Keep a small + # residency margin so every cluster can launch beside other streams. + SGLANG_OPT_DEEPGEMM_MEGA_MOE_RESERVED_SMS = EnvInt(2) # =================================================================== # Top-k kernels diff --git a/python/sglang/srt/layers/moe/mega_moe.py b/python/sglang/srt/layers/moe/mega_moe.py index 5709b93425c9..3485fd841a95 100644 --- a/python/sglang/srt/layers/moe/mega_moe.py +++ b/python/sglang/srt/layers/moe/mega_moe.py @@ -15,8 +15,9 @@ from __future__ import annotations +import functools import os -from contextlib import nullcontext +from contextlib import contextmanager, nullcontext from typing import TYPE_CHECKING, Optional import torch @@ -44,6 +45,43 @@ _MEGA_MOE_SYMM_BUFFER: dict = {} +@functools.lru_cache(maxsize=1) +def _mega_moe_max_num_sms() -> Optional[int]: + if _device_sm < 100: + # The SM90 MegaMoE implementation does not use the whole-grid clustered + # launch that needs a residency margin. + return None + + # Physical count, not deep_gemm.get_num_sms(): two-batch overlap and the DSA + # indexer reconfigure that process-wide, so reserving on top would compound. + num_sms = torch.cuda.get_device_properties(device="cuda").multi_processor_count + reserved_num_sms = max(envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_RESERVED_SMS.get(), 0) + return max(2, num_sms - reserved_num_sms) + + +@contextmanager +def _configure_mega_moe_deep_gemm_num_sms(deep_gemm): + max_num_sms = _mega_moe_max_num_sms() + if max_num_sms is None: + yield + return + + current_num_sms = deep_gemm.get_num_sms() + # Stay under an outer context's budget instead of claiming SMs back from it. + target_num_sms = min(max_num_sms, current_num_sms) + # Round down: the clustered launch needs an even CTA count. + target_num_sms -= target_num_sms % 2 + if target_num_sms == current_num_sms: + yield + return + + deep_gemm.set_num_sms(target_num_sms) + try: + yield + finally: + deep_gemm.set_num_sms(current_num_sms) + + def _get_mega_moe_symm_buffer( group, num_experts: int, @@ -247,16 +285,17 @@ def _run_mega_routed( device=hidden_states.device, ) swiglu_limit = getattr(moe.config, "swiglu_limit", None) - deep_gemm.fp8_fp4_mega_moe( - y, - moe.experts.mega_l1_weights, - moe.experts.mega_l2_weights, - buf, - recipe=(1, 1, 32), - activation="swiglu", - activation_clamp=swiglu_limit, - fast_math=True, - ) + with _configure_mega_moe_deep_gemm_num_sms(deep_gemm): + deep_gemm.fp8_fp4_mega_moe( + y, + moe.experts.mega_l1_weights, + moe.experts.mega_l2_weights, + buf, + recipe=(1, 1, 32), + activation="swiglu", + activation_clamp=swiglu_limit, + fast_math=True, + ) y = y[:num_tokens] if not moe.experts.should_fuse_routed_scaling_factor_in_topk: diff --git a/python/sglang/srt/models/kimi_k3.py b/python/sglang/srt/models/kimi_k3.py index c164b4c0a8bd..5d9997b11c7e 100644 --- a/python/sglang/srt/models/kimi_k3.py +++ b/python/sglang/srt/models/kimi_k3.py @@ -780,7 +780,10 @@ def _forward_mega_experts( from sglang.kernels.ops.attention.dsv4 import mega_moe_pre_dispatch from sglang.srt.distributed.parallel_state import get_moe_ep_group from sglang.srt.environ import envs - from sglang.srt.layers.moe.mega_moe import _get_mega_moe_symm_buffer + from sglang.srt.layers.moe.mega_moe import ( + _configure_mega_moe_deep_gemm_num_sms, + _get_mega_moe_symm_buffer, + ) # In SP-MoE mode (KimiK3DecoderLayer reduce-scatters the o_proj # output) the incoming rows are already this rank's token shard, so @@ -833,15 +836,16 @@ def _forward_mega_experts( dtype=torch.bfloat16, device=routed_input.device, ) - deep_gemm.fp8_fp4_mega_moe( - y, - self.experts.mega_l1_weights, - self.experts.mega_l2_weights, - buf, - recipe=(1, 1, 32), - activation="situ", - fast_math=True, - ) + with _configure_mega_moe_deep_gemm_num_sms(deep_gemm): + deep_gemm.fp8_fp4_mega_moe( + y, + self.experts.mega_l1_weights, + self.experts.mega_l2_weights, + buf, + recipe=(1, 1, 32), + activation="situ", + fast_math=True, + ) y = y[:num_tokens] if not self.experts.should_fuse_routed_scaling_factor_in_topk: if (