diff --git a/python/sglang/kernels/ops/layernorm/mhc_mega.py b/python/sglang/kernels/ops/layernorm/mhc_mega.py new file mode 100644 index 000000000000..dd16d4287c24 --- /dev/null +++ b/python/sglang/kernels/ops/layernorm/mhc_mega.py @@ -0,0 +1,68 @@ +"""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 + +from sglang.srt.utils import is_cuda + +if is_cuda(): + import deep_gemm + + +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. + """ + 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) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 5befef1ff8d8..049e743ae0ad 100644 --- a/python/sglang/srt/environ.py +++ b/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 diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 5b82f536a5eb..35673b4b356d 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/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): @@ -3297,11 +3299,16 @@ def forward_hc_pre_from_prev( forward_batch: ForwardBatch, input_ids_global: torch.Tensor, seam_open: bool = True, + mega_mhc_seam_open: bool = True, ) -> mhc.HcState: """The layer's two hyper-connections, each collapsing with the previous one's - 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.""" + pre-mix. ``seam_open`` preserves the existing tail-mode boundary policy; + ``mega_mhc_seam_open`` disables Mega mHC for both sublayers in the layer + immediately before row selection.""" self._init_boundaries() + mega_mhc_prefill = mhc.can_use_mega_mhc_prefill( + self.hc_cfg, seam_open=mega_mhc_seam_open + ) 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,6 +3342,8 @@ 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: @@ -3341,6 +3351,7 @@ def run_ffn_hc(state: mhc.HcState) -> mhc.HcState: 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_next=self.next_boundary if mega_mhc_prefill else None, ) return run_ffn_hc(run_attn_hc(state)) @@ -4117,6 +4130,7 @@ def _forward_layers_hc_pre_from_prev( forward_batch=forward_batch, input_ids_global=input_ids_global, seam_open=tail is None, + mega_mhc_seam_open=i + 1 != self.late_layer_start, ) state = state.materialized(self.layers[self.end_layer - 1].hc_cfg) if saved_full is not None: diff --git a/python/sglang/srt/models/deepseek_v4_mhc.py b/python/sglang/srt/models/deepseek_v4_mhc.py index a0a66dc77275..769c9e60c9cb 100644 --- a/python/sglang/srt/models/deepseek_v4_mhc.py +++ b/python/sglang/srt/models/deepseek_v4_mhc.py @@ -28,6 +28,7 @@ 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.environ import envs from sglang.srt.hardware_backend.npu.utils import is_npu_arch35 @@ -129,6 +130,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 +140,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 +182,7 @@ class HcNextBoundary(NamedTuple): norm: RMSNorm accepts_mxfp8: bool norm_fusable: bool + hc: Optional[HcSubLayer] = None def is_deferred_finalize(routed) -> bool: @@ -190,7 +194,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 +210,7 @@ def make_boundary(norm: RMSNorm, *, accepts_mxfp8: bool) -> HcNextBoundary: and norm.weight.shape == (5120,) and norm.weight.is_contiguous() ), + hc=hc, ) @@ -550,27 +557,59 @@ 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, *, seam_open: bool = True) -> bool: + return ( + envs.SGLANG_OPT_DSV41_MEGA_MHC_PREFILL.get() + and get_platform().is_sm100 + and not _is_hip + and not cfg.cp_prefill + and seam_open + ) + + def _post_fusion( hc: HcSubLayer, y: torch.Tensor, residual: torch.Tensor, coefficients: HcTriplet, next: Optional[HcNextBoundary], + mega_next: Optional[HcNextBoundary] = None, ) -> 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_next is not None and mega_next.norm_fusable and mega_next.hc is not None: + nxt = mega_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 +646,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 +658,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 +670,10 @@ 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, next if mega_mhc else None + ) def run_moe_post( @@ -641,6 +684,8 @@ def run_moe_post( stats_stream: Optional[torch.cuda.Stream], next: Optional[HcNextBoundary], world_size: int, + precomputed: Optional[HcTriplet] = None, + mega_next: Optional[HcNextBoundary] = None, ) -> 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 +702,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 +756,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_next)