[2/N][DSv41] Integrate DeepGEMM MegaGate routing

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",

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

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

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

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