diff --git a/python/sglang/kernels/ops/moe/__init__.py b/python/sglang/kernels/ops/moe/__init__.py index 9ae48045b053..98b73131045c 100644 --- a/python/sglang/kernels/ops/moe/__init__.py +++ b/python/sglang/kernels/ops/moe/__init__.py @@ -19,6 +19,21 @@ _CUDA = frozenset({CapabilityRequirement.CUDA}) _HIP = frozenset({CapabilityRequirement.HIP}) +register_kernel( + KernelSpec( + op="moe.bf16_mega_gate", + backend=KernelBackend.DEEPGEMM, + target="sglang.kernels.ops.moe.mega_gate:bf16_mega_gate", + capabilities=frozenset( + {CapabilityRequirement.cuda(min_sm=(10, 0), max_sm=(10, 9))} + ), + format_signature=FormatSignature( + description="BF16 gate GEMM + normalized sqrtsoftplus top-k; FP32 weights, int64 logical IDs", + ), + description="DeepGEMM MegaGate routing on datacenter Blackwell.", + ) +) + register_kernel( KernelSpec( op="moe.moe_align_block_size", diff --git a/python/sglang/kernels/ops/moe/mega_gate.py b/python/sglang/kernels/ops/moe/mega_gate.py new file mode 100644 index 000000000000..49129680ed5f --- /dev/null +++ b/python/sglang/kernels/ops/moe/mega_gate.py @@ -0,0 +1,68 @@ +"""DeepGEMM's fused BF16 gate projection and expert selection (SM10x).""" + +from __future__ import annotations + +from typing import Optional + +import torch + + +def is_mega_gate_available() -> bool: + try: + import deep_gemm + except ImportError: + return False + return callable(getattr(deep_gemm, "bf16_mega_gate", None)) + + +def bf16_mega_gate( + hidden_states: torch.Tensor, + weight: torch.Tensor, + top_k: int, + *, + routed_scaling_factor: float, + ep_rank: int, + bias: Optional[torch.Tensor] = None, + image_bias: Optional[torch.Tensor] = None, + image_token_mask: Optional[torch.Tensor] = None, + valid_token_mask: Optional[torch.Tensor] = None, + hash_topk_ids: Optional[torch.Tensor] = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Return normalized sqrtsoftplus weights and logical routed expert IDs. + + Shared-slot insertion, EPLB mapping and recording belong to the runtime. + Outputs are PyTorch-owned, including during CUDA graph capture. The pinned + DeepGEMM build checks every GEMM score for finiteness before applying its + token mask, so clear padded input rows without changing the caller's tensor. + """ + import deep_gemm + + if valid_token_mask is not None: + hidden_states = torch.where(valid_token_mask[:, None], hidden_states, 0) + + ids = torch.empty( + (hidden_states.shape[0], top_k), dtype=torch.int64, device=hidden_states.device + ) + weights = torch.empty_like(ids, dtype=torch.float32) + deep_gemm.bf16_mega_gate( + hidden_states, + weight, + top_k, + use_shared_as_routed=False, + num_shared_experts=0, + routed_scaling_factor=routed_scaling_factor, + ep_rank=ep_rank, + scoring_func="sqrtsoftplus", + mask=valid_token_mask, + bias=bias, + image_bias=image_bias, + image_token_mask=image_token_mask, + fix_routing_mask=( + torch.ones_like(ids[:, 0], dtype=torch.bool) + if hash_topk_ids is not None + else None + ), + unmapped_topk_idx=hash_topk_ids, + out=(ids, weights), + ) + return weights, ids diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index db80063997a3..8d4901684835 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -1314,6 +1314,9 @@ class Envs: # residency margin so every cluster can launch beside other streams. SGLANG_OPT_DEEPGEMM_MEGA_MOE_RESERVED_SMS = EnvInt(2) SGLANG_OPT_DEEPGEMM_MEGA_MOE_FUSE_SHARED_EXPERTS = EnvBool(True) + # Opt-in until routing and end-to-end speed are validated against SGLang's + # router; small local batches retain gate GEMM + top-k. + SGLANG_OPT_DEEPGEMM_MEGA_GATE = EnvBool(False) # =================================================================== # Top-k kernels diff --git a/python/sglang/srt/layers/moe/mega_gate.py b/python/sglang/srt/layers/moe/mega_gate.py new file mode 100644 index 000000000000..0084886974f2 --- /dev/null +++ b/python/sglang/srt/layers/moe/mega_gate.py @@ -0,0 +1,172 @@ +"""DeepSeek-V4.1 MegaGate dispatch and SGLang routing bookkeeping.""" + +from __future__ import annotations + +import torch + +from sglang.kernels.ops.moe.mega_gate import bf16_mega_gate, is_mega_gate_available +from sglang.srt.environ import envs +from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder +from sglang.srt.layers.moe.topk import ( + TopKConfig, + _mask_topk_ids_padded_region, + _post_process_topk_ids, + _zero_topk_weights_padded_region, +) +from sglang.srt.runtime_context import get_exec, get_parallel + + +def should_use_mega_gate(moe, hidden_states, device_sm, dispatch_info) -> bool: + if ( + not envs.SGLANG_OPT_DEEPGEMM_MEGA_GATE.get() + or device_sm // 10 != 10 + or getattr(moe.config, "model_type", None) != "deepseek_v41" + # vLLM's Flash measurements favor the unfused router at <=16 rows. + # Keep the conservative threshold for target and draft expert counts. + or not 16 < hidden_states.shape[0] <= (1 << 20) + or hidden_states.dtype != torch.bfloat16 + or not hidden_states.is_contiguous() + or getattr(moe.topk, "enable_waterfill", False) + or get_exec().deterministic.enable_deterministic_inference + or envs.SGLANG_SIMULATE_UNIFORM_EXPERTS.get() + or envs.SGLANG_SIMULATE_ROUND_ROBIN_EXPERTS.get() + ): + return False + weight = moe.gate.weight + num_experts, hidden = weight.shape + if ( + weight.dtype != torch.bfloat16 + or not weight.is_contiguous() + or hidden != hidden_states.shape[1] + or hidden % 256 != 0 + or not 0 < num_experts <= 512 + or num_experts % 4 != 0 + or not 0 < moe.config.num_experts_per_tok <= min(32, num_experts) + ): + return False + for bias in ( + moe.gate.e_score_correction_bias, + moe.gate.e_score_correction_bias_vl, + ): + if bias is not None and ( + bias.dtype != torch.float32 or not bias.is_contiguous() + ): + return False + if dispatch_info is not None and dispatch_info.ep_dispatch_algorithm in ( + "lp", + "fake", + ): + # Preserve the LP solver's collectives and fake routing's logit override. + return False + if moe.is_hash: + if moe.topk.score_func != "sqrtsoftplus": + return False + else: + config = moe.topk.topk_config + if ( + config.scoring_func != "sqrtsoftplus" + or not config.renormalize + or config.top_k <= 1 + or config.use_grouped_topk + or config.custom_routing_function is not None + or config.torch_native + ): + return False + return is_mega_gate_available() + + +def run_mega_gate(moe, hidden_states, input_ids, num_token_non_padded, dispatch_info): + if moe.is_hash: + config = TopKConfig( + top_k=moe.topk.topk, + num_fused_shared_experts=moe.topk.num_fused_shared_experts, + routed_scaling_factor=moe.topk.routed_scaling_factor, + apply_routed_scaling_factor_on_output=moe.topk.apply_routed_scaling_factor_on_output, + # HashTopK never captures into the target's routed-expert buffer. + allow_routed_experts_capture=False, + ) + else: + config = moe.topk.topk_config + num_tokens = hidden_states.shape[0] + valid_mask = None + if num_token_non_padded is not None: + valid_mask = ( + torch.arange(num_tokens, device=hidden_states.device) < num_token_non_padded + ) + + hash_ids = None + image_mask = None + image_bias = None + if moe.is_hash: + # Match HashTopK: image IDs still select their checkpoint hash routes. + # Padding IDs may be outside the vocabulary, so sanitize before lookup. + if input_ids is None: + raise ValueError("DeepSeek-V4.1 hash routing requires input_ids.") + safe_ids = ( + torch.where(valid_mask, input_ids, 0) + if valid_mask is not None + else input_ids + ) + hash_ids = moe.topk.tid2eid[safe_ids].to(torch.int64) + elif moe.gate.e_score_correction_bias_vl is not None and input_ids is not None: + image_bias = moe.gate.e_score_correction_bias_vl + image_mask = (input_ids == moe.config.image_token_id).contiguous() + + weights, ids = bf16_mega_gate( + hidden_states, + moe.gate.weight, + moe.config.num_experts_per_tok, + routed_scaling_factor=( + config.routed_scaling_factor + if config.apply_routed_scaling_factor_on_output + else 1.0 + ), + ep_rank=get_parallel().moe_ep_rank, + bias=moe.gate.e_score_correction_bias, + image_bias=image_bias, + image_token_mask=image_mask, + valid_token_mask=valid_mask, + hash_topk_ids=hash_ids, + ) + num_shared = config.num_fused_shared_experts + if num_shared: + shared_ids = moe.gate.weight.shape[0] + torch.arange( + num_shared, dtype=ids.dtype, device=ids.device + ) + ids = torch.cat([ids, shared_ids.expand(num_tokens, -1)], dim=-1) + shared_weight = ( + 1.0 + if config.apply_routed_scaling_factor_on_output + else 1.0 / config.routed_scaling_factor + ) + weights = torch.cat( + [weights, weights.new_full((num_tokens, num_shared), shared_weight)], dim=-1 + ) + + # MegaGate emits logical routed IDs. Reuse SGLang's EPLB, capture and + # home-rank shared-slot mapping; only the expert count is needed from logits. + logits_shape = hidden_states.new_empty((0, moe.gate.weight.shape[0])) + ids, weights, recorder_ids = _post_process_topk_ids( + ids, + weights, + config, + logits_shape, + moe.layer_id, + num_token_non_padded=num_token_non_padded, + expert_location_dispatch_info=dispatch_info, + padded_rows_masked=True, + ) + if num_shared and config.apply_routed_scaling_factor_on_output: + weights[:, -num_shared:] = 1.0 + # Shared-slot remapping may overwrite weights in padded rows. Mask last, + # before recording, including when every row on this DP rank is padding. + if num_token_non_padded is not None: + _mask_topk_ids_padded_region(ids, num_token_non_padded) + _zero_topk_weights_padded_region(weights, num_token_non_padded) + if recorder_ids is not None: + _mask_topk_ids_padded_region(recorder_ids, num_token_non_padded) + if recorder_ids is not None: + get_global_expert_distribution_recorder().on_select_experts( + topk_ids=recorder_ids + ) + return weights, ids diff --git a/python/sglang/srt/layers/moe/mega_moe.py b/python/sglang/srt/layers/moe/mega_moe.py index 20d03d8cc7f7..ac211b6993a3 100644 --- a/python/sglang/srt/layers/moe/mega_moe.py +++ b/python/sglang/srt/layers/moe/mega_moe.py @@ -26,6 +26,7 @@ from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp from sglang.srt.layers.dp_attention import get_dp_global_num_tokens +from sglang.srt.layers.moe.mega_gate import run_mega_gate, should_use_mega_gate from sglang.srt.layers.moe.mega_moe_sm90 import ( is_sm90_fp8_mega_moe_available, run_sm90_mega_routed, @@ -260,35 +261,49 @@ def _run_mega_routed( hidden_size = moe.config.hidden_size if num_tokens > 0: - router_logits = moe.gate(hidden_states, forward_batch=forward_batch) - topk_kwargs = {"input_ids": input_ids_global} if moe.is_hash else {} num_token_non_padded = ( forward_batch.moe_num_token_non_padded() if forward_batch is not None else None ) dispatch_info = ExpertLocationDispatchInfo.init_new(layer_id=moe.layer_id) - if ( - not moe.is_hash - and getattr(moe.gate, "e_score_correction_bias_vl", None) is not None - ): - topk_output = vision_topk( - moe, - router_logits, - input_ids_global, - num_token_non_padded=num_token_non_padded, - expert_location_dispatch_info=dispatch_info, - ) + if should_use_mega_gate(moe, hidden_states, _device_sm, dispatch_info): + import deep_gemm + + # Keep MegaGate under the same SM budget as MegaMoE when the + # alternate stream overlaps with shared-expert work. + with _configure_mega_moe_deep_gemm_num_sms(deep_gemm): + topk_weights, topk_ids = run_mega_gate( + moe, + hidden_states, + input_ids_global, + num_token_non_padded, + dispatch_info, + ) else: - topk_output = moe.topk( - hidden_states, - router_logits, - num_token_non_padded=num_token_non_padded, - expert_location_dispatch_info=dispatch_info, - **topk_kwargs, - ) - topk_ids = topk_output.topk_ids - topk_weights = topk_output.topk_weights + router_logits = moe.gate(hidden_states, forward_batch=forward_batch) + topk_kwargs = {"input_ids": input_ids_global} if moe.is_hash else {} + if ( + not moe.is_hash + and getattr(moe.gate, "e_score_correction_bias_vl", None) is not None + ): + topk_output = vision_topk( + moe, + router_logits, + input_ids_global, + num_token_non_padded=num_token_non_padded, + expert_location_dispatch_info=dispatch_info, + ) + else: + topk_output = moe.topk( + hidden_states, + router_logits, + num_token_non_padded=num_token_non_padded, + expert_location_dispatch_info=dispatch_info, + **topk_kwargs, + ) + topk_ids = topk_output.topk_ids + topk_weights = topk_output.topk_weights else: topk_ids = None topk_weights = None diff --git a/test/registered/kernels/ops/moe/test_mega_gate.py b/test/registered/kernels/ops/moe/test_mega_gate.py new file mode 100644 index 000000000000..dcf12d368a2b --- /dev/null +++ b/test/registered/kernels/ops/moe/test_mega_gate.py @@ -0,0 +1,153 @@ +"""DeepGEMM MegaGate parity and graph replay with nonfinite padding.""" + +import unittest + +import torch + +from sglang.kernels.ops.gemm.bf16_fp32 import linear_bf16_fp32 +from sglang.kernels.ops.moe.mega_gate import bf16_mega_gate, is_mega_gate_available +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="4-gpu-b200") + + +class TestMegaGate(CustomTestCase): + @classmethod + def setUpClass(cls): + super().setUpClass() + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10: + raise unittest.SkipTest("MegaGate requires SM10x") + if not is_mega_gate_available(): + raise unittest.SkipTest("DeepGEMM does not expose bf16_mega_gate") + + @staticmethod + def _reference( + x, weight, top_k, bias, image_bias, image_mask, valid_mask, hash_ids + ): + x = torch.where(valid_mask[:, None], x, 0) + scores = torch.log1p(torch.exp(linear_bf16_fp32(x, weight))).sqrt() + if hash_ids is not None: + ids = hash_ids + else: + row_bias = ( + torch.where(image_mask[:, None], image_bias, bias) + if image_mask is not None + else bias + ) + ids = (scores + row_bias).topk(top_k, dim=-1).indices + weights = scores.gather(1, ids) + weights = weights / (weights.sum(-1, keepdim=True) + 1e-20) * 1.5 + return weights.masked_fill(~valid_mask[:, None], 0), ids.masked_fill( + ~valid_mask[:, None], -1 + ) + + def _assert_routes_close(self, actual, expected): + actual_weights, actual_ids = actual + expected_weights, expected_ids = expected + actual_ids, order = actual_ids.sort(dim=-1) + actual_weights = actual_weights.gather(1, order) + expected_ids, order = expected_ids.sort(dim=-1) + expected_weights = expected_weights.gather(1, order) + torch.testing.assert_close(actual_ids, expected_ids) + torch.testing.assert_close( + actual_weights, expected_weights, rtol=1e-3, atol=1e-4 + ) + + def test_routing_and_cuda_graph_replay(self): + torch.manual_seed(42) + # Target Flash and reduced expert-count draft gates, at exact H=5120. + for num_experts, top_k in ((384, 6), (128, 6)): + for tokens in (17, 32, 128, 512): + for routing in ("text", "image", "hash"): + with self.subTest( + num_experts=num_experts, tokens=tokens, routing=routing + ): + x = torch.randn( + (tokens, 5120), device="cuda", dtype=torch.bfloat16 + ) + weight = ( + torch.randn( + (num_experts, 5120), device="cuda", dtype=torch.bfloat16 + ) + * 0.001 + ) + # Separate selection boundaries from BF16 GEMM roundoff. + bias = torch.arange(num_experts, device="cuda").float() * 0.25 + image_bias = bias.flip(0) if routing == "image" else None + image_mask = ( + torch.zeros(tokens, device="cuda", dtype=torch.bool) + if routing == "image" + else None + ) + hash_ids = ( + torch.arange(top_k, device="cuda") + .expand(tokens, -1) + .clone() + if routing == "hash" + else None + ) + valid_mask = torch.ones(tokens, device="cuda", dtype=torch.bool) + + def run(): + return bf16_mega_gate( + x, + weight, + top_k, + routed_scaling_factor=1.5, + ep_rank=0, + bias=bias, + image_bias=image_bias, + image_token_mask=image_mask, + valid_token_mask=valid_mask, + hash_topk_ids=hash_ids, + ) + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.stream(stream): + # DeepGEMM owns per-stream barriers: initialize them + # on the capture stream before entering capture. + for _ in range(3): + run() + with torch.cuda.graph(graph, stream=stream): + actual = run() + torch.cuda.current_stream().wait_stream(stream) + for valid_rows, padding in ( + (tokens, float("nan")), + (tokens // 2, float("nan")), + (0, float("inf")), + (tokens, float("nan")), + ): + x.normal_() + valid_mask.copy_( + torch.arange(tokens, device="cuda") < valid_rows + ) + x[valid_rows:] = padding + if image_mask is not None: + image_mask.logical_not_() + if hash_ids is not None: + # The external kernel also writes this buffer. + # Refresh routes after a replay masks padded IDs. + hash_ids.copy_( + torch.arange(top_k, device="cuda") + .expand(tokens, -1) + .roll(1, dims=-1) + ) + expected = self._reference( + x, + weight, + top_k, + bias, + image_bias, + image_mask, + valid_mask, + hash_ids, + ) + graph.replay() + self._assert_routes_close(actual, expected) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/layers/moe/test_mega_moe_deepgemm_api.py b/test/registered/unit/layers/moe/test_mega_moe_deepgemm_api.py index 5965ea59ba38..126e669ed473 100644 --- a/test/registered/unit/layers/moe/test_mega_moe_deepgemm_api.py +++ b/test/registered/unit/layers/moe/test_mega_moe_deepgemm_api.py @@ -8,18 +8,27 @@ import torch +from sglang.srt.environ import envs from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel maybe_stub_sgl_kernel() -from sglang.srt.layers.moe import MoeA2ABackend, MoeRunnerBackend, mega_moe +from sglang.srt.layers.moe import MoeA2ABackend, MoeRunnerBackend +from sglang.srt.layers.moe import mega_gate as mega_gate_runtime +from sglang.srt.layers.moe import mega_moe from sglang.srt.layers.moe import topk as topk_module from sglang.srt.layers.moe.fused_moe_triton import layer as fused_moe_layer_module from sglang.srt.layers.moe.topk import TopKConfig from sglang.srt.layers.moe.utils import draft_model_build_scope from sglang.srt.layers.quantization.unquant import UnquantizedFusedMoEMethod -from sglang.srt.runtime_context import get_context, get_exec, get_flags, get_parallel +from sglang.srt.runtime_context import ( + get_context, + get_exec, + get_flags, + get_parallel, + get_resources, +) register_cpu_ci(est_time=8, suite="base-a-test-cpu") @@ -432,6 +441,298 @@ def run_experts(_experts, x, ids, weights, **kwargs): self.assertEqual(out.shape, (0, 4)) moe.gate.assert_not_called() + @staticmethod + def _mega_gate_reference(x, weight, top_k, **kwargs): + """Numerical substitute at the external DeepGEMM boundary only.""" + if not torch.isfinite(x).all(): + raise ValueError("The pinned MegaGate rejects nonfinite GEMM inputs") + scores = torch.log1p(torch.exp(x.float() @ weight.float().T)).sqrt() + if kwargs["unmapped_topk_idx"] is not None: + ids = kwargs["unmapped_topk_idx"] + else: + bias = kwargs["bias"] + if kwargs["image_bias"] is not None: + bias = torch.where( + kwargs["image_token_mask"][:, None], kwargs["image_bias"], bias + ) + ids = (scores + bias).topk(top_k, dim=-1).indices + weights = scores.gather(1, ids) + weights = weights / (weights.sum(-1, keepdim=True) + 1e-20) + weights = weights * kwargs["routed_scaling_factor"] + if kwargs["mask"] is not None: + weights = weights.masked_fill(~kwargs["mask"][:, None], 0) + ids = ids.masked_fill(~kwargs["mask"][:, None], -1) + kwargs["out"][0].copy_(ids) + kwargs["out"][1].copy_(weights) + + @classmethod + def _mega_gate_moe(cls, fused_shared, scale_in_topk, hash_routing=False): + moe = cls._v41_moe(fused_shared, scale_in_topk) + moe.config.model_type = "deepseek_v41" + moe.config.hidden_size = 256 + moe.gate.weight = torch.zeros((8, 256), dtype=torch.bfloat16) + moe.gate.e_score_correction_bias = torch.tensor( + [4.0, 3.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] + ) + moe.gate.e_score_correction_bias_vl = torch.tensor( + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 3.0, 4.0] + ) + moe.topk.topk_config.scoring_func = "sqrtsoftplus" + if hash_routing: + moe.is_hash = True + moe.gate.e_score_correction_bias = None + moe.topk = SimpleNamespace( + topk=2 + fused_shared, + num_fused_shared_experts=fused_shared, + routed_scaling_factor=2.0, + apply_routed_scaling_factor_on_output=scale_in_topk, + score_func="sqrtsoftplus", + tid2eid=torch.tensor([[6, 7]] * 100, dtype=torch.int32), + ) + return moe + + def test_mega_gate_preserves_vision_hash_shared_scaling_and_padding(self): + """A fused gate must preserve expert contributions across EP layouts. + + Nonfinite padding and out-of-vocabulary padded IDs must reach the + expert boundary with zero weight; routed scaling must apply once and + the home-rank shared expert must contribute exactly one identity. + """ + deep_gemm = self.deep_gemm + deep_gemm.bf16_mega_gate = self._mega_gate_reference + for is_hash in (False, True): + for fused_shared in (0, 1): + for scale_in_topk in (False, True): + for valid_rows in (0, 2, 32): + with self.subTest( + is_hash=is_hash, + fused_shared=fused_shared, + scale_in_topk=scale_in_topk, + valid_rows=valid_rows, + ): + moe = self._mega_gate_moe( + fused_shared, scale_in_topk, is_hash + ) + hidden = torch.ones((32, 256), dtype=torch.bfloat16) + hidden[valid_rows:] = float("nan") + input_ids = torch.tensor([7, 99] * 16) + input_ids[valid_rows:] = 1000 + batch = SimpleNamespace( + moe_num_token_non_padded=lambda: torch.tensor( + valid_rows + ) + ) + + def run_experts(_experts, x, ids, weights, **kwargs): + self.assertTrue((ids[valid_rows:] == -1).all()) + self.assertTrue((weights[valid_rows:] == 0).all()) + factors = (ids + 1).float() + factors[ids == 9] = 1.0 # rank-1 shared slot + return (factors * weights).sum( + -1, keepdim=True + ) * kwargs["routed_scaling_factor"] + + with ( + get_context().override_server_args( + model_path="dummy", + tp_size=2, + ep_size=2, + enable_deterministic_inference=False, + ), + envs.SGLANG_OPT_DEEPGEMM_MEGA_GATE.override(True), + patch.dict(sys.modules, {"deep_gemm": deep_gemm}), + patch.object(mega_moe, "_device_sm", 100), + get_flags().moe.override( + a2a_backend=MoeA2ABackend.MEGAMOE + ), + get_parallel().override(moe_ep_size=2, moe_ep_rank=1), + patch.object( + mega_moe.ExpertLocationDispatchInfo, + "init_new", + return_value=None, + ), + patch.object( + mega_moe, "run_mega_routed_experts", run_experts + ), + patch.object(topk_module, "_is_cuda", True), + patch.object( + topk_module, + "_can_fuse_padded_region", + return_value=False, + ), + ): + out = mega_moe._run_mega_routed( + moe, hidden, batch, input_ids, num_tokens=32 + ) + image_value = 15.0 + 3 * fused_shared + text_value = image_value if is_hash else 3.0 + fused_shared + expected = torch.tensor([text_value, image_value] * 16) + expected[valid_rows:] = 0.0 + torch.testing.assert_close(out[:, 0], expected) + self.assertTrue(torch.isnan(hidden[valid_rows:]).all()) + + def test_mega_gate_preserves_target_only_routed_experts_capture(self): + """Draft and hash routing must leave the target's capture buffer intact. + + Reusing normal TopK postprocessing for HashTopK otherwise opts hash + layers into capture even though HashTopK never writes this buffer. + """ + deep_gemm = ModuleType("deep_gemm") + deep_gemm.bf16_mega_gate = self._mega_gate_reference + hidden = torch.ones((32, 256), dtype=torch.bfloat16) + capture_buffer = torch.full((32, 2), -17, dtype=torch.int64) + capturer = SimpleNamespace( + capture=lambda layer_id, topk_indices: capture_buffer.copy_(topk_indices) + ) + with ( + get_resources().override(experts_capturer=capturer), + get_parallel().override(moe_ep_rank=0), + patch.dict(sys.modules, {"deep_gemm": deep_gemm}), + patch.object(topk_module, "_is_cuda", True), + ): + for routing in ("target", "draft", "hash"): + with self.subTest(routing=routing): + capture_buffer.fill_(-17) + moe = self._mega_gate_moe(0, False, routing == "hash") + if routing == "draft": + moe.topk.topk_config.allow_routed_experts_capture = False + mega_gate_runtime.run_mega_gate( + moe, hidden, torch.full((32,), 7), None, None + ) + expected = ( + torch.tensor([[0, 1]]).expand(32, -1) + if routing == "target" + else torch.full((32, 2), -17) + ) + torch.testing.assert_close(capture_buffer, expected) + + def test_mega_gate_admission_preserves_unsupported_router_fallbacks(self): + """Do not silently change routing when the fused API cannot express it.""" + moe = self._mega_gate_moe(0, False) + hidden = torch.zeros((32, 256), dtype=torch.bfloat16) + deep_gemm = ModuleType("deep_gemm") + deep_gemm.bf16_mega_gate = self._mega_gate_reference + with ( + get_context().override_server_args( + model_path="dummy", enable_deterministic_inference=False + ), + envs.SGLANG_OPT_DEEPGEMM_MEGA_GATE.override(True), + patch.dict(sys.modules, {"deep_gemm": deep_gemm}), + ): + self.assertTrue( + mega_gate_runtime.should_use_mega_gate(moe, hidden, 100, None) + ) + for tokens in (0, 1, 16): + self.assertFalse( + mega_gate_runtime.should_use_mega_gate( + moe, hidden[:tokens], 100, None + ) + ) + self.assertTrue( + mega_gate_runtime.should_use_mega_gate(moe, hidden[:17], 103, None) + ) + for sm in (90, 120): + self.assertFalse( + mega_gate_runtime.should_use_mega_gate(moe, hidden, sm, None) + ) + for name, value in ( + ("renormalize", False), + ("use_grouped_topk", True), + ("scoring_func", "sigmoid"), + ("custom_routing_function", object()), + ("torch_native", True), + ): + with patch.object(moe.topk.topk_config, name, value): + self.assertFalse( + mega_gate_runtime.should_use_mega_gate(moe, hidden, 100, None) + ) + for algorithm in ("lp", "fake"): + self.assertFalse( + mega_gate_runtime.should_use_mega_gate( + moe, + hidden, + 100, + SimpleNamespace(ep_dispatch_algorithm=algorithm), + ) + ) + with patch.object(moe.topk, "enable_waterfill", True, create=True): + self.assertFalse( + mega_gate_runtime.should_use_mega_gate(moe, hidden, 100, None) + ) + with envs.SGLANG_SIMULATE_UNIFORM_EXPERTS.override(True): + self.assertFalse( + mega_gate_runtime.should_use_mega_gate(moe, hidden, 100, None) + ) + with get_exec().deterministic.override(enable_deterministic_inference=True): + self.assertFalse( + mega_gate_runtime.should_use_mega_gate(moe, hidden, 100, None) + ) + with envs.SGLANG_OPT_DEEPGEMM_MEGA_GATE.override(False): + self.assertFalse( + mega_gate_runtime.should_use_mega_gate(moe, hidden, 100, None) + ) + del deep_gemm.bf16_mega_gate + self.assertFalse( + mega_gate_runtime.should_use_mega_gate(moe, hidden, 100, None) + ) + + def test_mega_gate_maps_eplb_before_shared_slots_and_records_routed_ids(self): + """Redundant-expert maps index routed IDs before per-rank slot insertion. + + Shared IDs must not index the EPLB map or enter expert statistics, and + padding must not be counted as the map's last physical expert. + """ + moe = self._mega_gate_moe(1, True) + hidden = torch.ones((32, 256), dtype=torch.bfloat16) + hidden[2:] = float("nan") + deep_gemm = ModuleType("deep_gemm") + deep_gemm.bf16_mega_gate = self._mega_gate_reference + # Reverse routed placement. Physical expert count includes redundancy. + dispatch_info = mega_moe.ExpertLocationDispatchInfo( + ep_dispatch_algorithm="static", + partial_logical_to_rank_dispatch_physical_map=torch.tensor( + [9, 8, 7, 6, 5, 4, 3, 2], dtype=torch.int32 + ), + partial_logical_to_all_physical_map=None, + partial_logical_to_all_physical_map_num_valid=None, + num_physical_experts=10, + ) + recorded = [] + recorder = SimpleNamespace( + on_select_experts=lambda **kw: recorded.append(kw["topk_ids"].clone()) + ) + with ( + get_context().override_server_args( + model_path="dummy", tp_size=2, ep_size=2 + ), + get_flags().moe.override(a2a_backend=MoeA2ABackend.MEGAMOE), + get_parallel().override(moe_ep_size=2, moe_ep_rank=1), + patch.dict(sys.modules, {"deep_gemm": deep_gemm}), + patch.object(topk_module, "_is_cuda", True), + patch.object(topk_module, "_can_fuse_padded_region", return_value=False), + patch.object( + mega_gate_runtime, + "get_global_expert_distribution_recorder", + return_value=recorder, + ), + ): + weights, ids = mega_gate_runtime.run_mega_gate( + moe, + hidden, + torch.tensor([7, 99] + [1000] * 30), + torch.tensor(2), + dispatch_info, + ) + # Rank 1's redundant-routed layout has shared ID 11. Logical 0/1 + # maps to physical 9/8, then shared-slot insertion shifts to 10/9. + torch.testing.assert_close(ids[:2], torch.tensor([[10, 9, 11], [2, 3, 11]])) + torch.testing.assert_close(weights[:2], torch.ones((2, 3))) + self.assertTrue((ids[2:] == -1).all()) + self.assertTrue((weights[2:] == 0).all()) + self.assertEqual(len(recorded), 1) + torch.testing.assert_close(recorded[0][:2], torch.tensor([[9, 8], [2, 3]])) + self.assertTrue((recorded[0][2:] == -1).all()) + @staticmethod def _v41_moe(fused_shared, scale_in_topk): return SimpleNamespace(