From aca8787c2aac639cee6106676f3c5cfc21cb549a Mon Sep 17 00:00:00 2001 From: kpham-sgl Date: Fri, 23 Oct 2026 14:38:59 -0700 Subject: [PATCH] [Perf] dsv41: Let Mega-mHC write the MegaMoE input under MegaGate With MegaGate routing in MegaMoE's buffer, the Mega-mHC kernel emits the next MoE layer's FP8 rows and its routed and shared-expert scales straight into that buffer, so the per-layer pre-dispatch quantize is skipped. Co-Authored-By: Claude Fable 5.1 --- .../sglang/kernels/ops/layernorm/mhc_mega.py | 9 ++ python/sglang/srt/layers/moe/mega_moe.py | 98 ++++++++++++++++--- .../sglang/srt/models/deepseek_v4_boundary.py | 5 + python/sglang/srt/models/deepseek_v4_mhc.py | 24 ++++- 4 files changed, 119 insertions(+), 17 deletions(-) diff --git a/python/sglang/kernels/ops/layernorm/mhc_mega.py b/python/sglang/kernels/ops/layernorm/mhc_mega.py index 58a42e9e52c9..28bf44de9993 100644 --- a/python/sglang/kernels/ops/layernorm/mhc_mega.py +++ b/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) diff --git a/python/sglang/srt/layers/moe/mega_moe.py b/python/sglang/srt/layers/moe/mega_moe.py index 66bd8276caf6..312b57f72107 100644 --- a/python/sglang/srt/layers/moe/mega_moe.py +++ b/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(). diff --git a/python/sglang/srt/models/deepseek_v4_boundary.py b/python/sglang/srt/models/deepseek_v4_boundary.py index 6828bdc22d49..426720e72026 100644 --- a/python/sglang/srt/models/deepseek_v4_boundary.py +++ b/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): diff --git a/python/sglang/srt/models/deepseek_v4_mhc.py b/python/sglang/srt/models/deepseek_v4_mhc.py index 1482fa965b14..38a68c80e4b6 100644 --- a/python/sglang/srt/models/deepseek_v4_mhc.py +++ b/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