[Perf] dsv41: Let Mega-mHC write the MegaMoE input under MegaGate

python/sglang/kernels/ops/layernorm/mhc_mega.py

@@ -28,6 +28,7 @@ def mhc_mega_boundary(
     sinkhorn_iters,
     y_fp8=None,
     y_gemm_sf=None,
+    moe_input=None,
 ):
     """Update the residual, collapse with carried pre, and compute next stats.
 
@@ -44,6 +45,11 @@ def mhc_mega_boundary(
     next_pre = torch.empty((rows, hc, 1), device=x.device, dtype=torch.float32)
     next_post = torch.empty_like(next_pre)
     next_comb = torch.empty_like(comb)
+    y_routed_sf = y_shared_sf = None
+    shared_sf_block_m = 0
+    if moe_input is not None:
+        assert y_fp8 is None
+        y_fp8, y_routed_sf, y_shared_sf, shared_sf_block_m = moe_input
     deep_gemm.mega_mhc(
         x=x,
         residual=residual,
@@ -69,5 +75,8 @@ def mhc_mega_boundary(
         y_bf16=normalized,
         y_fp8=y_fp8,
         y_gemm_sf=y_gemm_sf,
+        y_routed_sf=y_routed_sf,
+        y_shared_sf=y_shared_sf,
+        shared_sf_block_m=shared_sf_block_m,
     )
     return updated, normalized, (next_pre.squeeze(-1), next_post.squeeze(-1), next_comb)

python/sglang/srt/layers/moe/mega_moe.py

@@ -17,7 +17,7 @@
 
 import functools
 from contextlib import contextmanager, nullcontext
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, NamedTuple, Optional
 
 import torch
 
@@ -270,6 +270,68 @@ def _mega_moe_buffer(moe: DeepseekV2MoE):
     )
 
 
+MEGA_MOE_INPUT_TAG = "_mega_moe_input_in_buffer"
+
+
+class MegaMoeInput(NamedTuple):
+    buf: SymmBuffer
+    num_tokens: int
+    x: torch.Tensor
+    routed_sf: torch.Tensor
+    shared_sf: torch.Tensor
+    shared_block_m: int
+
+    def views(self):
+        return (self.x, self.routed_sf, self.shared_sf, self.shared_block_m)
+
+    def mark(self, rows: torch.Tensor) -> None:
+        setattr(rows, MEGA_MOE_INPUT_TAG, (self.buf, self.num_tokens))
+
+
+def _mega_moe_input_in_buffer(rows: torch.Tensor, buf, num_tokens: int) -> bool:
+    tag = getattr(rows, MEGA_MOE_INPUT_TAG, None)
+    return tag is not None and tag[0] is buf and tag[1] == num_tokens
+
+
+def mega_moe_input_target(
+    moe: DeepseekV2MoE, hidden_states: torch.Tensor, forward_batch
+) -> Optional[MegaMoeInput]:
+    if not (
+        should_use_mega_gate(moe, hidden_states, forward_batch)
+        and get_moe_a2a_backend().is_megamoe()
+        and is_mega_moe_experts_ready(moe.experts)
+        and moe.mega_shared_l1_weights is not None
+        and _mega_moe_mma_type(moe.experts) == "fp8xfp4"
+    ):
+        return None
+    import deep_gemm
+
+    buf = _mega_moe_buffer(moe)
+    num_tokens = hidden_states.shape[0]
+    x = buf.x[:num_tokens]
+    if x.dtype != torch.float8_e4m3fn:
+        x = x.view(torch.float8_e4m3fn)
+    routed_sf = buf.x_sf[:num_tokens]
+    if not routed_sf.is_contiguous():
+        return None
+    shared_block_m = deep_gemm.get_block_m_for_mega_moe(
+        num_ranks=buf.group.size(),
+        num_experts=buf.num_experts,
+        num_max_tokens_per_rank=buf.num_max_tokens_per_rank,
+        num_tokens=num_tokens,
+        num_topk=buf.num_topk,
+        mma_type=buf.mma_type,
+    )
+    return MegaMoeInput(
+        buf,
+        num_tokens,
+        x,
+        routed_sf,
+        buf.shared_l1_acts_sf[:num_tokens],
+        shared_block_m,
+    )
+
+
 def _run_mega_routed(
     moe: DeepseekV2MoE,
     hidden_states: torch.Tensor,
@@ -474,21 +536,25 @@ def run_mega_routed_experts(
                 num_topk=buf.num_topk,
                 mma_type=buf.mma_type,
             )
-        mega_moe_pre_dispatch(
-            hidden_states,
-            buf.topk_idx if routing_in_buffer else topk_ids_in,
-            buf.topk_weights if routing_in_buffer else topk_weights_in,
-            buf.x,
-            buf.x_sf,
-            buf.topk_idx,
-            buf.topk_weights,
-            quant_group_size=32,
-            shared_x_sf=(
-                buf.shared_l1_acts_sf if shared_l1_weights is not None else None
-            ),
-            shared_block_m=shared_block_m,
-            copy_topk=not routing_in_buffer,
-        )
+        if not (
+            routing_in_buffer
+            and _mega_moe_input_in_buffer(hidden_states, buf, num_tokens)
+        ):
+            mega_moe_pre_dispatch(
+                hidden_states,
+                buf.topk_idx if routing_in_buffer else topk_ids_in,
+                buf.topk_weights if routing_in_buffer else topk_weights_in,
+                buf.x,
+                buf.x_sf,
+                buf.topk_idx,
+                buf.topk_weights,
+                quant_group_size=32,
+                shared_x_sf=(
+                    buf.shared_l1_acts_sf if shared_l1_weights is not None else None
+                ),
+                shared_block_m=shared_block_m,
+                copy_topk=not routing_in_buffer,
+            )
 
     # Allocate at least one row so y has a non-null CUDA data_ptr;
     # the DeepGEMM tvm-ffi binding rejects nullptr in convert_to_torch_tensor().

python/sglang/srt/models/deepseek_v4_boundary.py

@@ -188,6 +188,11 @@ def next_boundary(self):
             hc=hc,
             accepts_mxfp8_deepgemm=self.attention
             and self.layer.self_attn.accepts_deepgemm_mxfp8_input(),
+            mega_moe_input=(
+                None
+                if self.attention
+                else mhc.mega_moe_input_of(self.layer, self.batch)
+            ),
         )
 
     def _after_post(self, residual, norm):

python/sglang/srt/models/deepseek_v4_mhc.py

@@ -22,6 +22,7 @@
 
 from __future__ import annotations
 
+from functools import partial
 from typing import Any, Callable, NamedTuple, Optional, Tuple, TypeAlias, Union
 
 import msgspec
@@ -239,6 +240,7 @@ class HcNextBoundary(NamedTuple):
     norm_fusable: bool
     hc: Optional[HcSubLayer] = None
     accepts_mxfp8_deepgemm: bool = False
+    mega_moe_input: Optional[Callable[[torch.Tensor], Any]] = None
 
 
 def is_deferred_finalize(routed) -> bool:
@@ -250,12 +252,26 @@ def is_deferred_finalize(routed) -> bool:
     return isinstance(routed, FlashInferTrtllmDeferredFinalizeOutput)
 
 
+def mega_moe_input_of(
+    layer, forward_batch: Optional[ForwardBatch]
+) -> Optional[Callable[[torch.Tensor], Any]]:
+    if not envs.SGLANG_OPT_DEEPGEMM_MEGA_GATE.get():
+        return None
+    mlp = getattr(layer, "mlp", None)
+    if not hasattr(mlp, "mega_shared_l1_weights"):
+        return None
+    from sglang.srt.layers.moe.mega_moe import mega_moe_input_target
+
+    return partial(mega_moe_input_target, mlp, forward_batch=forward_batch)
+
+
 def make_boundary(
     norm: RMSNorm,
     *,
     accepts_mxfp8: bool,
     hc: Optional[HcSubLayer] = None,
     accepts_mxfp8_deepgemm: bool = False,
+    mega_moe_input: Optional[Callable[[torch.Tensor], Any]] = None,
 ) -> HcNextBoundary:
     """Can this hyper-connection's norm fold into the previous post, and does its
     sublayer take a pre-quantized input."""
@@ -272,6 +288,7 @@ def make_boundary(
         ),
         hc=hc,
         accepts_mxfp8_deepgemm=accepts_mxfp8_deepgemm,
+        mega_moe_input=mega_moe_input,
     )
 
 
@@ -658,9 +675,11 @@ def _post_fusion(
         and y.shape[0] > 0
     ):
         next_hc = next.hc
-        y_fp8 = y_sf = None
+        y_fp8 = y_sf = moe_input = None
         if next.accepts_mxfp8_deepgemm and fused_mxfp8_rows(y.shape[0]):
             y_fp8, y_sf = empty_deepgemm_mxfp8_pair(*y.shape, y.device)
+        elif next.mega_moe_input is not None:
+            moe_input = next.mega_moe_input(y)
         updated, normalized, stats = mhc_mega_boundary(
             y,
             residual,
@@ -677,9 +696,12 @@ def _post_fusion(
             next_hc.cfg.sinkhorn_iters,
             y_fp8=y_fp8,
             y_gemm_sf=y_sf,
+            moe_input=None if moe_input is None else moe_input.views(),
         )
         if y_fp8 is not None:
             return HcState(updated, pre, HcQuantized(normalized, (y_fp8, y_sf)), stats)
+        if moe_input is not None:
+            moe_input.mark(normalized)
         return HcState(updated, pre, HcNormed(normalized), stats)
     if (
         next is not None