[Dsv4.1] Mega-mhc integration

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

@@ -0,0 +1,65 @@
+"""SGLang DSv4.1 prefill integration of Mega mHC.
+
+Credit: DeepSeek's DeepGEMM project provides the fused Mega mHC kernel.
+https://github.com/deepseek-ai/DeepGEMM
+"""
+
+import torch
+
+
+def mhc_mega_boundary(
+    x,
+    residual,
+    pre,
+    post,
+    comb,
+    fn,
+    scale,
+    base,
+    weight,
+    rms_eps,
+    hc_eps,
+    norm_eps,
+    sinkhorn_iters,
+):
+    """Update the residual, collapse with carried pre, and compute next stats.
+
+    The collapse uses ``pre`` from the preceding sublayer, not the new pre
+    produced from the updated residual. Preserve this DSv4.1 shifted ordering.
+    Used by the opt-in SM10x prefill path. Before CUDA graph capture, warm up
+    on the capture stream to initialize DeepGEMM's per-stream barriers.
+    """
+    import deep_gemm
+
+    rows, hidden = x.shape
+    hc = residual.shape[1]
+    updated = torch.empty_like(residual)
+    normalized = torch.empty_like(x)
+    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)
+    deep_gemm.mega_mhc(
+        x=x,
+        residual=residual,
+        shifted_prev_mix=pre.unsqueeze(-1),
+        post_mix=post.unsqueeze(-1),
+        comb_res_mix=comb,
+        fn=fn,
+        mix_scales=scale,
+        mix_bases=base,
+        hc_mult=hc,
+        hc_norm_eps=rms_eps,
+        hc_pre_eps=hc_eps,
+        hc_post_scale=2.0,
+        sinkhorn_eps=hc_eps,
+        num_sinkhorn_iters=sinkhorn_iters,
+        rmsnorm_weight=weight,
+        rmsnorm_eps=norm_eps,
+        rmsnorm_scale=1.0,
+        new_residual=updated,
+        new_prev_mix=next_pre,
+        new_post_mix=next_post,
+        new_comb_res_mix=next_comb,
+        y_bf16=normalized,
+    )
+    return updated, normalized, (next_pre.squeeze(-1), next_post.squeeze(-1), next_comb)

python/sglang/srt/environ.py

@@ -1599,6 +1599,7 @@ class Envs:
     # DeepSeek-V4.1 engram host table: keep the tables in host memory (layout
     # below) and gather rows from the GPU instead of sharding them over HBM.
     SGLANG_ENABLE_DSV41_ENGRAM_HOST_TABLE = EnvBool(False)
+    SGLANG_OPT_DSV41_MEGA_MHC_PREFILL = EnvBool(False)
     # "shared" is one buffer for the whole TP group, mapped by every rank, with no
     # lookup all-reduce (the ranks must share a PID namespace); "per_rank" is one
     # anonymous mapping per rank holding only its rows, gathered with the

python/sglang/srt/models/deepseek_v4.py

@@ -2842,12 +2842,14 @@ def _init_boundaries(self) -> None:
         self.local_boundary = mhc.make_boundary(
             self.post_attention_layernorm,
             accepts_mxfp8=False,
+            hc=self.ffn_hc,
         )
         # MOE[i] -> ATTN[i+1]
         if self._next_layer is not None:
             self.next_boundary = mhc.make_boundary(
                 self._next_layer.input_layernorm,
                 accepts_mxfp8=self._next_layer.self_attn.accepts_mxfp8_swizzled_input(),
+                hc=self._next_layer.attn_hc,
             )
 
     def refresh_mhc_norm_weight_cache(self):
@@ -3302,6 +3304,11 @@ def forward_hc_pre_from_prev(
         pre-mix. ``seam_open`` is False when the late-layer tail narrows the rows after
         this layer, so nothing precomputed for the next one would still describe it."""
         self._init_boundaries()
+        mega_mhc_prefill = (
+            self.config.model_type == "deepseek_v41"
+            and seam_open
+            and mhc.can_use_mega_mhc_prefill(self.hc_cfg, state.residual, forward_batch)
+        )
         stats_stream = None
         if mhc.use_stats_stream(self.hc_cfg, forward_batch, state.residual):
             stats_stream = self.hc_stats_stream
@@ -3312,6 +3319,7 @@ def run_attn_hc(state: mhc.HcState) -> mhc.HcState:
             quantized = [] if self.self_attn.accepts_mxfp8_swizzled_input() else None
             mhc.fork_stats_stream(stats_stream)
             x = mhc.combine(self.attn_hc, state, quantized)
+            precomputed = state.stats
             state.release()
             del state
             world_size = self.self_attn.attn_tp_size
@@ -3334,13 +3342,16 @@ def run_attn_hc(state: mhc.HcState) -> mhc.HcState:
                 stats_stream=stats_stream,
                 next=self.local_boundary,
                 world_size=world_size,
+                precomputed=precomputed,
+                mega_mhc=mega_mhc_prefill,
             )
 
         def run_ffn_hc(state: mhc.HcState) -> mhc.HcState:
             assert self.ffn_hc is not None
             residual = state.residual
             mhc.fork_stats_stream(stats_stream)
             x = mhc.combine(self.ffn_hc, state)
+            precomputed = state.stats
             state.release()
             del state
             nxt = self.next_boundary if seam_open else None
@@ -3362,6 +3373,8 @@ def run_ffn_hc(state: mhc.HcState) -> mhc.HcState:
                 stats_stream=stats_stream,
                 next=nxt,
                 world_size=self.mlp.tp_size,
+                precomputed=precomputed,
+                mega_mhc=mega_mhc_prefill,
             )
 
         return run_ffn_hc(run_attn_hc(state))

python/sglang/srt/models/deepseek_v4_mhc.py

@@ -28,15 +28,17 @@
 import torch
 import torch.nn.functional as F
 
+from sglang.kernels.ops.layernorm.mhc_mega import mhc_mega_boundary
 from sglang.kernels.ops.layernorm.mhc_post_split_h import mhc_post_split_h
+from sglang.srt.batch_invariant_ops import is_batch_invariant_mode_enabled
 from sglang.srt.environ import envs
 from sglang.srt.hardware_backend.npu.utils import is_npu_arch35
 from sglang.srt.layers.layernorm import RMSNorm
 from sglang.srt.layers.quantization.mxfp8_input import Mxfp8SwizzledInput
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_executor.runner import compile_in_capture_mode
 from sglang.srt.models.deepseek_v2 import MoEOutput, _is_hip, _is_npu, _is_xpu
-from sglang.srt.runtime_context import get_parallel, get_platform
+from sglang.srt.runtime_context import get_forward, get_parallel, get_platform
 from sglang.srt.utils import is_gfx95_supported
 
 _is_gfx95_supported = is_gfx95_supported()
@@ -129,6 +131,7 @@ class HcState(msgspec.Struct):
     streams: Union[torch.Tensor, HcPending, None]
     pre: Optional[torch.Tensor] = None
     input: Optional[HcPreOutput] = None
+    stats: Optional[HcTriplet] = None
 
     def release(self) -> None:
         """Empty the state at its last reader, the combine; the residual the post
@@ -138,6 +141,7 @@ def release(self) -> None:
         self.streams = None
         self.pre = None
         self.input = None
+        self.stats = None
 
     @property
     def residual(self) -> torch.Tensor:
@@ -179,6 +183,7 @@ class HcNextBoundary(NamedTuple):
     norm: RMSNorm
     accepts_mxfp8: bool
     norm_fusable: bool
+    hc: Optional[HcSubLayer] = None
 
 
 def is_deferred_finalize(routed) -> bool:
@@ -190,7 +195,9 @@ def is_deferred_finalize(routed) -> bool:
     return isinstance(routed, FlashInferTrtllmDeferredFinalizeOutput)
 
 
-def make_boundary(norm: RMSNorm, *, accepts_mxfp8: bool) -> HcNextBoundary:
+def make_boundary(
+    norm: RMSNorm, *, accepts_mxfp8: bool, hc: Optional[HcSubLayer] = None
+) -> HcNextBoundary:
     """Can this hyper-connection's norm fold into the previous post, and does its
     sublayer take a pre-quantized input."""
     return HcNextBoundary(
@@ -204,6 +211,7 @@ def make_boundary(norm: RMSNorm, *, accepts_mxfp8: bool) -> HcNextBoundary:
             and norm.weight.shape == (5120,)
             and norm.weight.is_contiguous()
         ),
+        hc=hc,
     )
 
 
@@ -550,27 +558,69 @@ def _compute_triplet(
     hc: HcSubLayer,
     residual: torch.Tensor,
     stats_stream: Optional[torch.cuda.Stream],
+    precomputed: Optional[HcTriplet] = None,
 ) -> HcTriplet:
     """Issue the triplet (on the side stream, beside the sublayer) and join before
     the post reads it."""
+    if precomputed is not None:
+        return precomputed
     coefficients = mix_stats(hc, residual, stats_stream)
     if stats_stream is not None:
         torch.cuda.current_stream().wait_stream(stats_stream)
     return coefficients
 
 
+def can_use_mega_mhc_prefill(
+    cfg: HcConfig, residual: torch.Tensor, forward_batch: ForwardBatch
+) -> bool:
+    return (
+        envs.SGLANG_OPT_DSV41_MEGA_MHC_PREFILL.get()
+        and can_fuse_post(cfg)
+        # DeepGEMM Mega mHC supports SM10x only.
+        and get_platform().is_sm100
+        and residual.is_cuda
+        and not _is_hip
+        and forward_batch.forward_mode.is_extend_without_speculative()
+        and 4096 <= residual.shape[0] <= 65536
+        and residual.shape[1:] == (4, 5120)
+        and residual.dtype == torch.bfloat16
+        and residual.is_contiguous()
+        and not get_forward().sp_active
+        and not is_batch_invariant_mode_enabled()
+    )
+
+
 def _post_fusion(
     hc: HcSubLayer,
     y: torch.Tensor,
     residual: torch.Tensor,
     coefficients: HcTriplet,
     next: Optional[HcNextBoundary],
+    mega_mhc: bool = False,
 ) -> HcState:
     """Step 4 without a collective: the wide-tile kernel folds the next combine +
     norm in where it serves this seam, otherwise the pure post runs and the next
     combine computes itself."""
     cfg = hc.cfg
     pre, post_mix, comb = coefficients
+    if mega_mhc and next is not None and next.norm_fusable and next.hc is not None:
+        nxt = next.hc
+        updated, normalized, stats = mhc_mega_boundary(
+            y,
+            residual,
+            pre,
+            post_mix,
+            comb,
+            nxt.fn,
+            nxt.scale,
+            nxt.base,
+            nxt.norm.weight,
+            nxt.cfg.rms_eps,
+            nxt.cfg.eps,
+            nxt.norm.variance_epsilon,
+            nxt.cfg.sinkhorn_iters,
+        )
+        return HcState(updated, pre, HcNormed(normalized), stats)
     if (
         next is not None
         and next.norm_fusable
@@ -607,6 +657,8 @@ def run_attn_post(
     stats_stream: Optional[torch.cuda.Stream],
     next: Optional[HcNextBoundary],
     world_size: int,
+    precomputed: Optional[HcTriplet] = None,
+    mega_mhc: bool = False,
 ) -> HcState:
     """The attention post. An `AttnOutput` rides the collective kernel (which also
     folds ``next``'s norm); the attention may decline the handover even when asked,
@@ -617,7 +669,7 @@ def run_attn_post(
         )
 
         assert next is not None
-        pre, post_mix, comb = _compute_triplet(hc, residual, stats_stream)
+        pre, post_mix, comb = _compute_triplet(hc, residual, stats_stream, precomputed)
         _, updated, normalized = all_reduce_mhc_post_combine_norm(
             out.partial,
             residual,
@@ -629,8 +681,8 @@ def run_attn_post(
             world_size=world_size,
         )
         return HcState(updated, pre, HcNormed(normalized))
-    coefficients = _compute_triplet(hc, residual, stats_stream)
-    return _post_fusion(hc, out, residual, coefficients, next)
+    coefficients = _compute_triplet(hc, residual, stats_stream, precomputed)
+    return _post_fusion(hc, out, residual, coefficients, next, mega_mhc)
 
 
 def run_moe_post(
@@ -641,6 +693,8 @@ def run_moe_post(
     stats_stream: Optional[torch.cuda.Stream],
     next: Optional[HcNextBoundary],
     world_size: int,
+    precomputed: Optional[HcTriplet] = None,
+    mega_mhc: bool = False,
 ) -> HcState:
     """The MoE post. A deferred finalize the push plane can carry rides the
     collective kernel (finalize + shared add + all-reduce, quantizing ``next``'s
@@ -657,7 +711,7 @@ def run_moe_post(
         # [T x top_k (padded)] view and overstates the plane load by ~6x.
         and can_fuse_all_reduce(out.routed.expert_weights.shape[0], hc.cfg.hidden)
     ):
-        pre, post_mix, comb = _compute_triplet(hc, residual, stats_stream)
+        pre, post_mix, comb = _compute_triplet(hc, residual, stats_stream, precomputed)
         args = (
             out.routed.gemm2_out,
             out.routed.expanded_idx_to_permuted_idx,
@@ -711,5 +765,5 @@ def run_moe_post(
             and should_add_replicated_moe_output()
         ):
             out += pieces.shared
-    coefficients = _compute_triplet(hc, residual, stats_stream)
-    return _post_fusion(hc, out, residual, coefficients, next)
+    coefficients = _compute_triplet(hc, residual, stats_stream, precomputed)
+    return _post_fusion(hc, out, residual, coefficients, next, mega_mhc)