From f2d79903cf537be975383be97f2e4da91c4264ff Mon Sep 17 00:00:00 2001 From: Jingxin Pan Date: Fri, 14 Aug 2026 11:52:44 -0700 Subject: [PATCH] Integrate with MoonEP and MXFP4 experts on the DeepGEMM runner --- python/sglang/srt/environ.py | 4 + .../srt/layers/moe/fused_moe_triton/layer.py | 14 +- .../srt/layers/moe/moe_runner/deep_gemm.py | 168 +++++++++++ .../srt/layers/moe/token_dispatcher/moonep.py | 165 ++++++++-- .../moe/token_dispatcher/moonep_weights.py | 281 ++++++++++++++++++ .../sglang/srt/layers/quantization/mxfp4.py | 132 +++++++- python/sglang/srt/server_args.py | 7 +- 7 files changed, 719 insertions(+), 52 deletions(-) create mode 100644 python/sglang/srt/layers/moe/token_dispatcher/moonep_weights.py diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index c8a91366776e..8dfa89457b36 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -980,7 +980,11 @@ class Envs: # -1 uses MoonEP's training-safe default B = E / EP. SGLANG_MOONEP_NUM_PREFETCH_SLOTS = EnvInt(-1) SGLANG_MOONEP_TOKEN_PADDING = EnvInt(128) + # Decode-phase token capacity; <= 0 derives it from max_running_requests. + SGLANG_MOONEP_DECODE_MAX_DISPATCH_TOKENS_PER_RANK = EnvInt(-1) SGLANG_MOONEP_NUM_SMS = EnvInt(32) + # MoonEP's static shapes should be capturable; off until that is shown. + SGLANG_ENABLE_MOONEP_CUDA_GRAPH = EnvBool(False) SGLANG_PPLX_NUM_MAX_DISPATCH_TOKENS_PER_RANK = EnvInt(128) SGLANG_ENABLE_MOE_DEFERRED_FINALIZE = EnvBool(True) # DeepSeek/GLM MoE (deepseek_v2.py): quantize the (dp-gathered) MoE input diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py index f781993652a9..da0417698ae4 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -397,15 +397,15 @@ def __init__( f"quant_method={type(self.quant_method).__name__})." ) - moonep_global_weight_storage = get_moe_a2a_backend().is_moonep() - if moonep_global_weight_storage: - if quant_config is not None: - raise NotImplementedError( - "MoonEP PoC supports unquantized BF16 MoE weights only." - ) + # Quantized experts instead keep the normal EP + # shard and are relocated into a symmetric VMM range after loading + moonep_global_weight_storage = ( + get_moe_a2a_backend().is_moonep() and quant_config is None + ) + if get_moe_a2a_backend().is_moonep(): if num_fused_shared_experts != 0: raise NotImplementedError( - "MoonEP PoC does not support fused shared experts yet." + "MoonEP does not support fused shared experts yet." ) self.quant_method.create_weights( diff --git a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py index 67bdb5ea8b6f..cd7d17e29ea8 100644 --- a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py +++ b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py @@ -50,6 +50,10 @@ DeepEPNormalCombineInput, DeepEPNormalDispatchOutput, ) + from sglang.srt.layers.moe.token_dispatcher.moonep import ( + MoonEPCombineInput, + MoonEPDispatchOutput, + ) from sglang.srt.layers.moe.token_dispatcher.standard import ( StandardCombineInput, StandardDispatchOutput, @@ -1240,6 +1244,170 @@ def post_permute_deep_gemm_to_deepep_normal( ) +def _moonep_m_indices( + cu_seqlens: torch.Tensor, + expert_ids: torch.Tensor, + all_tokens: int, +) -> torch.Tensor: + """Expand MoonEP's per-group segment ends into DeepGEMM's per-row group ids. + + ``cu_seqlens[g]`` is the *end* offset of group ``g`` (no leading zero), so + ``searchsorted(..., right=True)`` maps a row to the group that owns it and + naturally skips empty groups. Two kinds of rows get ``-1``, which DeepGEMM + skips entirely: rows past the last segment (MoonEP pads the receive buffer + to a static ``NvS``) and rows whose group is an unfilled prefetch slot + (``expert_ids`` already carries ``-1`` there). + + ``expert_ids`` -- not the group index -- is the value DeepGEMM wants: it + indexes the leading dimension of the expert weight tensor. + """ + num_groups = expert_ids.numel() + rows = torch.arange(all_tokens, device=cu_seqlens.device, dtype=cu_seqlens.dtype) + group = torch.searchsorted(cu_seqlens, rows, right=True) + m_indices = expert_ids[group.clamp(max=num_groups - 1)].to(torch.int32) + return torch.where(group < num_groups, m_indices, torch.full_like(m_indices, -1)) + + +@register_pre_permute("moonep", "deep_gemm") +def pre_permute_moonep_to_deep_gemm( + dispatch_output: MoonEPDispatchOutput, + quant_info: DeepGemmMoeQuantInfo, + runner_config: MoeRunnerConfig, + running_state: dict, +) -> DeepGemmRunnerInput: + """MoonEP dispatch output -> DeepGEMM m-grouped contiguous input. + + Unlike the deepep/standard pre-permutes there is no scatter here: MoonEP's + ``dispatch`` already returns rows grouped by expert and padded to + ``token_padding``, which matches DeepGEMM's + ``get_mk_alignment_for_contiguous_layout()``. All that is left is deriving + ``m_indices`` and, for quantized experts, quantizing the activations. + """ + hidden_states = dispatch_output.hidden_states + if hidden_states.ndim != 2: + raise ValueError( + f"MoonEP hidden states must be [NvS, H], got {hidden_states.shape}" + ) + + all_tokens = hidden_states.shape[0] + running_state["all_tokens"] = all_tokens + running_state["hidden_states_shape"] = hidden_states.shape + running_state["hidden_states_dtype"] = hidden_states.dtype + running_state["hidden_states_device"] = hidden_states.device + # Carried through because MoonEP's combine reconstructs from the plan and + # does not apply route weights itself -- the post-permute must. + running_state["route_weights_nvs"] = dispatch_output.route_weights_nvs + running_state["plan"] = dispatch_output.plan + running_state["num_tokens"] = dispatch_output.num_tokens + + expert_ids = dispatch_output.expert_ids + if quant_info.w13_weight.dtype != torch.bfloat16: + # Quantized experts live in MoonEP's symmetric pool, so the duplicated + # ones have to be pulled in before the GEMM reads them, and the plan's + # global expert ids have to become pool rows. This runs here rather + # than in DeepEPMoE.run_moe_core because MXFP4 on DeepGEMM sets + # deprecate_flag, which delegates past that method entirely. + from sglang.srt.layers.moe.token_dispatcher import moonep_weights + from sglang.srt.layers.moe.token_dispatcher.moonep import MoonEPBuffer + + layer_id = runner_config.layer_id + assert layer_id is not None, "MoonEP pre-permute needs runner_config.layer_id" + weight_pairs, scale_pairs = moonep_weights.prefetch_pairs(layer_id) + MoonEPBuffer.get_existing_buffer().prefetch_weight( + plan=dispatch_output.plan, + async_finish=False, + weight_pairs=weight_pairs, + scale_pairs=scale_pairs or None, + experts_to_copy=moonep_weights.expert_rows( + layer_id, + dispatch_output.plan.experts_to_copy[get_tp_group().rank_in_group], + ), + ) + # The parameters are this rank's slice of the pool, but m_indices + # addresses the whole symmetric range -- a duplicated expert's rows + # live in another rank's chunk. Point the GEMM at the full ranges, of + # which the parameters are a sub-view. + pool = moonep_weights.get_pool() + quant_info.w13_weight = pool.ranges[moonep_weights.W13_WEIGHT].view(torch.int8) + quant_info.w2_weight = pool.ranges[moonep_weights.W2_WEIGHT].view(torch.int8) + quant_info.w13_scale = pool.ranges[moonep_weights.W13_SCALE].permute(0, 2, 1) + quant_info.w2_scale = pool.ranges[moonep_weights.W2_SCALE].permute(0, 2, 1) + + expert_ids = moonep_weights.group_rows( + layer_id, expert_ids, runner_config.num_experts + ) + + m_indices = _moonep_m_indices(dispatch_output.cu_seqlens, expert_ids, all_tokens) + running_state["m_indices"] = m_indices + + if quant_info.w13_weight.dtype == torch.bfloat16: + return DeepGemmRunnerInput( + hidden_states=hidden_states, + # Unused by the BF16 contiguous GEMM, but the field is non-optional. + hidden_states_scale=torch.empty( + (all_tokens, 1), device=hidden_states.device, dtype=torch.float32 + ), + use_masked_gemm=False, + m_indices=m_indices, + ) + + from sglang.kernels.ops.quantization.fp8_kernel import ( + sglang_per_token_group_quant_fp8, + ) + + block_k = quant_info.block_shape[1] if quant_info.block_shape else 128 + running_state["mxfp8_act_gran_k"] = block_k + hidden_states_fp8, hidden_states_scale = sglang_per_token_group_quant_fp8( + hidden_states, + block_k, + column_major_scales=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, + scale_tma_aligned=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, + scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, + ) + return DeepGemmRunnerInput( + hidden_states=hidden_states_fp8, + hidden_states_scale=hidden_states_scale, + use_masked_gemm=False, + m_indices=m_indices, + ) + + +@register_post_permute("deep_gemm", "moonep") +def post_permute_deep_gemm_to_moonep( + runner_output: DeepGemmRunnerOutput, + quant_info: DeepGemmMoeQuantInfo, + runner_config: MoeRunnerConfig, + running_state: dict, +) -> MoonEPCombineInput: + """DeepGEMM output -> MoonEP combine input, still in dispatched row order. + + No gather: MoonEP's ``combine`` consumes the ``[NvS, H]`` layout directly. + + Rows DeepGEMM skipped (``m_indices == -1``) were never written, so they + still hold whatever ``torch.empty`` left behind and must be zeroed before + combine reduces them. Zeroing cannot be folded into the route-weight + multiply below, because uninitialized memory may decode to NaN and + ``0 * NaN`` is NaN. + """ + from sglang.srt.layers.moe.token_dispatcher.moonep import MoonEPCombineInput + + hidden_states = runner_output.hidden_states + hidden_states.masked_fill_((running_state["m_indices"] < 0).unsqueeze(-1), 0.0) + + route_weights_nvs = running_state["route_weights_nvs"] + if route_weights_nvs is not None: + hidden_states.mul_( + route_weights_nvs.to(dtype=hidden_states.dtype).unsqueeze(-1) + ) + + return MoonEPCombineInput( + hidden_states=hidden_states, + route_weights_nvs=route_weights_nvs, + plan=running_state["plan"], + num_tokens=running_state["num_tokens"], + ) + + def _varlen_deep_gemm_situ_mul_quant( gateup_output: torch.Tensor, masked_m: torch.Tensor, diff --git a/python/sglang/srt/layers/moe/token_dispatcher/moonep.py b/python/sglang/srt/layers/moe/token_dispatcher/moonep.py index 55ab3be63268..62b68b323c8d 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/moonep.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/moonep.py @@ -15,11 +15,12 @@ DispatchOutput, DispatchOutputFormat, ) -from sglang.srt.layers.moe.topk import TopKOutput -from sglang.srt.layers.moe.topk import TopKOutputChecker +from sglang.srt.layers.moe.topk import TopKOutput, TopKOutputChecker from sglang.srt.layers.moe.utils import DeepEPMode +_DECODE_TOKENS_PER_REQUEST_HEADROOM = 8 + _MOONEP_UNSUPPORTED_MESSAGE = ( "MoonEP MoE A2A is recognized by SGLang, but the runtime dispatcher is not " "implemented yet. MoonEP is not a drop-in DeepEP-compatible backend: it " @@ -451,9 +452,7 @@ def run_moonep_bf16_expert( f"shape {cu_seqlens.shape}" ) if route_weights_nvs is not None and route_weights_nvs.ndim != 1: - raise ValueError( - f"route_weights_nvs must be 1D, got {route_weights_nvs.shape}" - ) + raise ValueError(f"route_weights_nvs must be 1D, got {route_weights_nvs.shape}") output = torch.empty_like(hidden_states) prev = 0 @@ -496,6 +495,33 @@ def run_moonep_bf16_expert( ) +def _resolve_decode_capacity(prefill_capacity: int) -> int | None: + """Token capacity for decode batches, or None to reuse the prefill one. + + Decode is bounded by the number of running requests, which is orders of + magnitude below a prefill chunk, so giving it its own smaller buffer is + what keeps a decode step from running the MoE over a full chunk's worth of + padding. Both capacities are config-derived, so every rank picks the same + one without communicating. + """ + override = envs.SGLANG_MOONEP_DECODE_MAX_DISPATCH_TOKENS_PER_RANK.get() + if override > 0: + capacity = override + else: + from sglang.srt.server_args import get_global_server_args + + # None until the scheduler resolves it; fall back to one capacity then. + max_running = get_global_server_args().max_running_requests + if max_running is None: + return None + # Speculative decoding submits several tokens per request per step. + capacity = int(max_running) * _DECODE_TOKENS_PER_REQUEST_HEADROOM + + token_padding = envs.SGLANG_MOONEP_TOKEN_PADDING.get() + capacity = -(-capacity // token_padding) * token_padding + return None if capacity >= prefill_capacity else capacity + + class MoonEPDispatcher(BaseDispatcher): """MoonEP dispatcher for the initial BF16 inference PoC.""" @@ -527,13 +553,38 @@ def __init__( self.num_max_dispatch_tokens_per_rank = ( envs.SGLANG_MOONEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK.get() ) + self.decode_max_dispatch_tokens_per_rank = _resolve_decode_capacity( + self.num_max_dispatch_tokens_per_rank + ) self.num_prefetch_slots = None @staticmethod def _raise_unimplemented() -> NoReturn: raise NotImplementedError(_MOONEP_UNSUPPORTED_MESSAGE) - def _get_buffer(self): + def _phase_capacity(self) -> int: + """Static token capacity for the current phase. + + MoonEP's buffers are statically shaped and ``dispatch`` asserts an + exact ``S x K`` input, so every batch is padded up to the capacity. One + capacity sized for prefill makes decode pay for it: at K3 a batch of 8 + tokens would run the MoE over 16384, a ~2000x inflation. + + The two capacities come from server args rather than the batch, because + the buffer is created collectively -- picking from a runtime token + count would let ranks disagree and deadlock. The phase flag is uniform + across the group for the same reason. + """ + if self.decode_max_dispatch_tokens_per_rank is None: + return self.num_max_dispatch_tokens_per_rank + + from sglang.srt.layers.dp_attention import get_is_extend_in_batch + + if get_is_extend_in_batch(): + return self.num_max_dispatch_tokens_per_rank + return self.decode_max_dispatch_tokens_per_rank + + def _get_buffer(self, capacity: int | None = None): if self.hidden_size is None or self.num_experts is None: raise ValueError( "MoonEPDispatcher requires hidden_size and num_experts to " @@ -544,7 +595,9 @@ def _get_buffer(self): hidden_size=self.hidden_size, router_topk=self.router_topk, num_experts=self.num_experts, - num_max_dispatch_tokens_per_rank=self.num_max_dispatch_tokens_per_rank, + num_max_dispatch_tokens_per_rank=( + self._phase_capacity() if capacity is None else capacity + ), num_prefetch_slots=self.num_prefetch_slots, ) @@ -559,9 +612,9 @@ def _pad_to_capacity( hidden_states: torch.Tensor, topk_ids: torch.Tensor, topk_weights: torch.Tensor, + capacity: int, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]: num_tokens = int(hidden_states.shape[0]) - capacity = int(self.num_max_dispatch_tokens_per_rank) if num_tokens > capacity: raise ValueError( "MoonEP runtime batch has more tokens than its static buffer " @@ -587,34 +640,58 @@ def _pad_to_capacity( ) def _tokens_per_expert(self, topk_ids: torch.Tensor) -> torch.Tensor: + """Local token count per expert. + + ``torch.bincount`` would be the obvious call, but on CUDA it reads the + input's max back to the host to size its output, and a device-to-host + copy is illegal under CUDA graph capture. Scattering into a + fixed-length buffer keeps the whole thing on device. + """ assert self.num_experts is not None - return torch.bincount( - topk_ids.reshape(-1).to(dtype=torch.int64), - minlength=self.num_experts, - ).to(dtype=torch.int32) + flat = topk_ids.reshape(-1).to(dtype=torch.int64) + counts = torch.zeros( + self.num_experts, dtype=torch.int32, device=topk_ids.device + ) + counts.scatter_add_(0, flat, torch.ones_like(flat, dtype=torch.int32)) + return counts def _expert_ids_from_plan( self, cu_seqlens: torch.Tensor, plan: Any, ) -> torch.Tensor: + """Which expert each VM group carries: its own id for the first + ``num_experts`` groups, the duplicated expert's id for the prefetch + slots after them, and -1 for groups that received no tokens. + + Stays on device. The obvious loop costs one ``.item()`` per group -- + 896 host syncs per layer, ~82k per forward at K3's depth, which + dominates decode. + """ assert self.num_experts is not None + num_experts = int(self.num_experts) num_groups = int(cu_seqlens.numel()) - expert_ids = torch.full_like(cu_seqlens, -1) experts_to_copy = plan.experts_to_copy if experts_to_copy.ndim == 2: experts_to_copy = experts_to_copy[self._get_rank()] - prev = 0 - for group_id in range(num_groups): - cur = int(cu_seqlens[group_id].item()) - if cur > prev: - if group_id < self.num_experts: - expert_ids[group_id] = group_id - else: - expert_ids[group_id] = experts_to_copy[group_id - self.num_experts] - prev = cur - return expert_ids + num_slots = num_groups - num_experts + assert 0 <= num_slots <= int(experts_to_copy.numel()), ( + f"MoonEP plan has {experts_to_copy.numel()} prefetch slots but " + f"cu_seqlens describes {num_slots}" + ) + ids = torch.cat( + [ + torch.arange( + num_experts, device=cu_seqlens.device, dtype=cu_seqlens.dtype + ), + experts_to_copy[:num_slots].to(dtype=cu_seqlens.dtype), + ] + ) + # cu_seqlens holds segment *ends*, so a group is live when its end + # moved past the previous one's. + starts = torch.cat([cu_seqlens.new_zeros(1), cu_seqlens[:-1]]) + return torch.where(cu_seqlens > starts, ids, torch.full_like(ids, -1)) def dispatch( self, @@ -632,13 +709,17 @@ def dispatch( if self.num_experts is None: raise ValueError("MoonEPDispatcher requires num_experts.") + # One capacity for both the padding and the buffer: MoonEP asserts the + # dispatch input matches the buffer's static S exactly. + capacity = self._phase_capacity() hidden_states, topk_ids, topk_weights, num_tokens = self._pad_to_capacity( hidden_states, topk_output.topk_ids, topk_output.topk_weights, + capacity, ) tokens_per_expert = self._tokens_per_expert(topk_ids) - buffer = self._get_buffer() + buffer = self._get_buffer(capacity) hidden_nvsh, route_weights_nvs, cu_seqlens, plan = buffer.dispatch( hidden_states, topk_weights, @@ -694,14 +775,42 @@ def combine_b(self): def prefetch_weight( self, plan: Any, - weight_layout: MoonEPExpertWeightLayout, + weight_layout: MoonEPExpertWeightLayout | None = None, + layer_id: int | None = None, ) -> None: + """Fill this rank's prefetch slots with the duplicated experts. + + ``weight_layout`` is the BF16 PoC form: one contiguous ``[E+B]`` block + per projection, indexed by global expert id. ``layer_id`` selects the + symmetric-memory form, where the sources are VMM ranges shared by every + layer and the plan's expert ids have to be remapped to rows before the + copy can find them. + """ + assert (weight_layout is None) != (layer_id is None), ( + "MoonEPDispatcher.prefetch_weight takes exactly one of " + "weight_layout or layer_id" + ) + if weight_layout is not None: + self._get_buffer().prefetch_weight( + plan=plan, + async_finish=False, + full_gate_weight=weight_layout.full_gate_weight, + full_up_weight=weight_layout.full_up_weight, + full_down_weight=weight_layout.full_down_weight, + ) + return + + from sglang.srt.layers.moe.token_dispatcher import moonep_weights + + weight_pairs, scale_pairs = moonep_weights.prefetch_pairs(layer_id) self._get_buffer().prefetch_weight( plan=plan, async_finish=False, - full_gate_weight=weight_layout.full_gate_weight, - full_up_weight=weight_layout.full_up_weight, - full_down_weight=weight_layout.full_down_weight, + weight_pairs=weight_pairs, + scale_pairs=scale_pairs or None, + experts_to_copy=moonep_weights.expert_rows( + layer_id, plan.experts_to_copy[self._get_rank()] + ), ) def register_deepep_dispatch_hook(self, hook): diff --git a/python/sglang/srt/layers/moe/token_dispatcher/moonep_weights.py b/python/sglang/srt/layers/moe/token_dispatcher/moonep_weights.py new file mode 100644 index 000000000000..b3fae437a18a --- /dev/null +++ b/python/sglang/srt/layers/moe/token_dispatcher/moonep_weights.py @@ -0,0 +1,281 @@ +from __future__ import annotations + +import logging +import math +from typing import Optional + +import torch + +logger = logging.getLogger(__name__) + +# Sub-ranges of the pool, keyed by the attribute they end up backing. The +# scales are named apart from the parameters because what lives here is the +# post-transform runtime layout, not the checkpoint layout the loader fills. +W13_WEIGHT = "w13_weight" +W2_WEIGHT = "w2_weight" +W13_SCALE = "w13_weight_scale_runtime" +W2_SCALE = "w2_weight_scale_runtime" + +_pool: Optional[MoonEPWeightPool] = None + + +class MoonEPWeightPool: + """Process-global symmetric ranges, sliced one block per MoE layer.""" + + def __init__( + self, + num_layers: int, + num_local_experts: int, + num_prefetch_slots: int, + ep_rank: int, + ep_size: int, + group, + specs: dict[str, tuple[tuple[int, ...], torch.dtype]], + ): + from moonep.buffer import create_nvl_dist_tensor + + self.num_layers = num_layers + self.num_local_experts = num_local_experts + self.num_prefetch_slots = num_prefetch_slots + self.ep_rank = ep_rank + self.ep_size = ep_size + self.block_rows = num_local_experts + num_prefetch_slots + self.chunk_rows = _resolve_chunk_rows(specs, num_layers * self.block_rows) + self._layers: dict[int, int] = {} + + self.ranges = { + kind: create_nvl_dist_tensor( + [self.chunk_rows, *trailing], + dtype, + ep_rank, + ep_size, + group=group, + # Rotated so this rank's chunk is at the base of the mapping. + # DeepGEMM's tvm_ffi bindings resolve a tensor's device from + # its base pointer, and the unrotated base is always rank 0's + # memory -- every other rank then fails the device check. + local_first=True, + ) + for kind, (trailing, dtype) in specs.items() + } + resident = sum(t.numel() * t.element_size() for t in self.ranges.values()) + logger.info( + "MoonEP: symmetric expert pool for %d layers, %d local experts + %d " + "prefetch slots each, chunk_rows=%d (%.1f GB resident per rank)", + num_layers, + num_local_experts, + num_prefetch_slots, + self.chunk_rows, + resident / ep_size / 1e9, + ) + + def layer_offset(self, layer_id: int) -> int: + """Rows before this layer's block. Layers are numbered by the order + they ask for storage, not by ``layer_id``: with pipeline parallelism a + rank holds an arbitrary slice of the model's layer ids, and the pool is + sized for what this rank actually builds.""" + index = self._layers.get(layer_id) + if index is None: + index = len(self._layers) + if index >= self.num_layers: + raise RuntimeError( + f"MoonEP expert pool was sized for {self.num_layers} layers " + f"but layer {layer_id} is the {index + 1}th to ask for " + "storage; the layer count derived from the model config is " + "wrong for this model" + ) + self._layers[layer_id] = index + return index * self.block_rows + + def chunk_start(self, owner_rank: int) -> int: + """First row of ``owner_rank``'s chunk in this rank's mapping.""" + from moonep.buffer import local_first_chunk_index + + index = local_first_chunk_index(owner_rank, self.ep_rank, self.ep_size) + return index * self.chunk_rows + + def local_view(self, kind: str, layer_id: int) -> torch.Tensor: + """This rank's ``[num_local_experts, ...]`` slice: what the loader + writes and what the parameter is bound to.""" + start = self.chunk_start(self.ep_rank) + self.layer_offset(layer_id) + return self.ranges[kind][start : start + self.num_local_experts] + + def slot_view(self, kind: str, layer_id: int) -> torch.Tensor: + """This layer's ``[num_prefetch_slots, ...]`` copy destinations.""" + start = ( + self.chunk_start(self.ep_rank) + + self.layer_offset(layer_id) + + self.num_local_experts + ) + return self.ranges[kind][start : start + self.num_prefetch_slots] + + +def _minimum_aligned_rows(trailing_shape, dtype: torch.dtype) -> int: + """Smallest row count whose bytes land on a VMM granularity boundary.""" + from moonep.buffer import pad_dim0_for_alignment + + return pad_dim0_for_alignment([1, *trailing_shape], dtype) + + +def _resolve_chunk_rows( + specs: dict[str, tuple[tuple[int, ...], torch.dtype]], rows_in_use: int +) -> int: + """One row count that is granularity-aligned for every range at once. + + A weight and its scale must share a row indexing, so they cannot each pick + their own padding. A row count works for a tensor when + ``rows * bytes_per_row`` is a multiple of the VMM granularity, so the + common answer is the least common multiple of each one's minimum, rounded + up past the rows actually in use. + """ + step = 1 + for trailing, dtype in specs.values(): + step = math.lcm(step, _minimum_aligned_rows(trailing, dtype)) + return math.ceil(rows_in_use / step) * step + + +def _num_moe_layers() -> int: + """How many layers will ask the pool for storage. + + Counted from the config rather than tracked dynamically because the ranges + are fixed-size VMM mappings that cannot grow; ``layer_offset`` raises if + the count turns out to be too small. + """ + from sglang.srt.runtime_context import process_model_config + + config = process_model_config() + num_layers = int(config.num_hidden_layers) + first_dense = config.first_k_dense_replace or 0 + freq = getattr(config.hf_text_config, "moe_layer_freq", 1) or 1 + return sum(1 for i in range(num_layers) if i >= first_dense and i % freq == 0) + + +def get_pool() -> Optional[MoonEPWeightPool]: + return _pool + + +def alloc_expert_tensors( + layer: torch.nn.Module, + specs: dict[str, tuple[tuple[int, ...], torch.dtype]], +) -> dict[str, torch.Tensor]: + """Per-layer views of the symmetric pool, creating it on the first call. + + ``specs`` maps a sub-range name to ``(trailing_shape, dtype)`` for a single + expert row. Every layer must pass the same specs -- they share one + allocation. Collective over the EP group, so all ranks have to build their + layers in the same order. + """ + global _pool + + from sglang.srt.distributed import get_tp_group + from sglang.srt.layers.moe.token_dispatcher.moonep import ( + get_moonep_num_prefetch_slots, + ) + + if _pool is None: + group = get_tp_group().device_group + ep_size = torch.distributed.get_world_size(group) + _pool = MoonEPWeightPool( + num_layers=_num_moe_layers(), + num_local_experts=int(layer.num_local_experts), + num_prefetch_slots=get_moonep_num_prefetch_slots( + int(layer.num_experts), ep_size + ), + ep_rank=torch.distributed.get_rank(group), + ep_size=ep_size, + group=group, + specs=specs, + ) + elif set(specs) != set(_pool.ranges): + raise NotImplementedError( + "MoonEP places every layer's experts in one pool, so all layers " + f"must request the same tensors; got {sorted(specs)} after " + f"{sorted(_pool.ranges)}" + ) + + return {kind: _pool.local_view(kind, layer.layer_id) for kind in specs} + + +def expert_rows(layer_id: int, expert_ids: torch.Tensor) -> torch.Tensor: + """Global expert ids -> rows of the symmetric range. Negative ids (unused + prefetch slots) pass through so DeepGEMM still skips them.""" + assert _pool is not None, "MoonEP expert pool was never created" + epn = _pool.num_local_experts + # The mapping is local-first, so an owner's chunk index is relative to this + # rank -- row numbers differ per rank, which is fine because every consumer + # of them (m_indices, experts_to_copy) is computed locally. + owner = expert_ids // epn + chunk = (owner - _pool.ep_rank) % _pool.ep_size + rows = chunk * _pool.chunk_rows + _pool.layer_offset(layer_id) + expert_ids % epn + return torch.where(expert_ids < 0, expert_ids, rows).to(torch.int32) + + +def group_rows( + layer_id: int, expert_ids: torch.Tensor, num_global_experts: int +) -> torch.Tensor: + """MoonEP's per-VM-group expert ids -> rows of the symmetric range. + + The first ``num_global_experts`` groups are experts addressed by global id; + the groups after them are this rank's prefetch slots, whose ids name the + *source* expert but whose tokens must read the slot the copy landed in. + Empty and unfilled groups keep their -1 and stay skipped. + """ + assert _pool is not None, "MoonEP expert pool was never created" + rows = expert_rows(layer_id, expert_ids) + tail = expert_ids[num_global_experts:] + slot_base = ( + _pool.chunk_start(_pool.ep_rank) + + _pool.layer_offset(layer_id) + + _pool.num_local_experts + ) + slots = torch.arange( + slot_base, + slot_base + tail.numel(), + device=expert_ids.device, + dtype=torch.int32, + ) + rows[num_global_experts:] = torch.where(tail < 0, tail.to(torch.int32), slots) + return rows + + +def prefetch_pairs( + layer_id: int, +) -> tuple[ + list[tuple[torch.Tensor, torch.Tensor]], list[tuple[torch.Tensor, torch.Tensor]] +]: + """``(weight_pairs, scale_pairs)`` for ``Buffer.prefetch_weight``. + + Scales are re-tiled by the copy and so are kept apart from the weights. + The scale views are the *storage* orientation, which is what a byte copy + has to move -- see the MN-major note in the module docstring. + """ + assert _pool is not None, "MoonEP expert pool was never created" + weights = [ + (_pool.ranges[k], _pool.slot_view(k, layer_id)) for k in (W13_WEIGHT, W2_WEIGHT) + ] + scales = [ + (_pool.ranges[k], _pool.slot_view(k, layer_id)) + for k in (W13_SCALE, W2_SCALE) + if k in _pool.ranges + ] + return weights, scales + + +def assert_resident(layer: torch.nn.Module, kind: str, tensor: torch.Tensor) -> None: + """Fail loudly if a post-load step swapped a pooled tensor for a private one. + + SGLang's quant methods routinely rebind weights after loading, and a + rebind that lands outside the pool silently costs MoonEP its remote + readability -- prefetch would then copy from memory no peer can see. + """ + assert _pool is not None, "MoonEP expert pool was never created" + full = _pool.ranges[kind] + start = full.data_ptr() + end = start + full.numel() * full.element_size() + if not (start <= tensor.data_ptr() < end): + raise RuntimeError( + f"MoonEP: layer {layer.layer_id} {kind} left the symmetric pool " + "after loading. Some post-load step replaced the tensor instead of " + "writing into it, which would make remote prefetch read private " + "memory." + ) diff --git a/python/sglang/srt/layers/quantization/mxfp4.py b/python/sglang/srt/layers/quantization/mxfp4.py index 4ad7bfc3967c..6b864ff482ec 100644 --- a/python/sglang/srt/layers/quantization/mxfp4.py +++ b/python/sglang/srt/layers/quantization/mxfp4.py @@ -324,6 +324,52 @@ def get_scaled_act_names(self) -> List[str]: return [] +def _maybe_alloc_moonep_expert_pool( + layer, + *, + intermediate_size: int, + hidden_size: int, + weight_dtype: torch.dtype, + mxfp4_block: int, +) -> Optional[dict]: + """This layer's slice of MoonEP's symmetric expert pool, or None. + + The scale entries describe DeepGEMM's *runtime* layout in storage order: + ``transform_sf_into_required_layout`` returns ``[E, MN, K/128]`` int32 with + stride ``(.., 1, MN)``, whose contiguous bytes are the transpose. Four + e8m0 exponents pack into one int32, so the byte count matches the + ``[E, MN, K/32]`` uint8 the checkpoint carries. + """ + from sglang.srt.layers.moe.token_dispatcher import moonep_weights + from sglang.srt.layers.moe.utils import get_moe_a2a_backend + + if not get_moe_a2a_backend().is_moonep(): + return None + + scale_pack = 4 * mxfp4_block # e8m0 bytes per int32 x elements per byte + return moonep_weights.alloc_expert_tensors( + layer, + { + moonep_weights.W13_WEIGHT: ( + (2 * intermediate_size, hidden_size // 2), + weight_dtype, + ), + moonep_weights.W2_WEIGHT: ( + (hidden_size, intermediate_size // 2), + weight_dtype, + ), + moonep_weights.W13_SCALE: ( + (hidden_size // scale_pack, 2 * intermediate_size), + torch.int32, + ), + moonep_weights.W2_SCALE: ( + (intermediate_size // scale_pack, hidden_size), + torch.int32, + ), + }, + ) + + class Mxfp4MoEMethod(FusedMoEMethodBase): def __init__( @@ -334,6 +380,7 @@ def __init__( self.prefix = prefix self.topk_indices_dtype = None + self.moonep_pooled = None self.use_triton_kernels = get_moe_runner_backend().is_triton_kernels() self.with_bias = False self.use_flashinfer = get_moe_runner_backend().is_flashinfer_mxfp4() @@ -464,13 +511,31 @@ def create_weights( self.intermediate_size_per_partition = intermediate_size_per_partition_after_pad self.hidden_size = hidden_size + + from sglang.srt.layers.moe.token_dispatcher import moonep_weights + + # MoonEP needs every expert row readable by its peers, so the weights + # come from a symmetric VMM pool instead of a private allocation. The + # scales are pooled in their post-transform runtime layout, which + # process_weights_after_loading copies into rather than replacing. + self.moonep_pooled = _maybe_alloc_moonep_expert_pool( + layer, + intermediate_size=intermediate_size_per_partition_after_pad, + hidden_size=hidden_size, + weight_dtype=weight_dtype, + mxfp4_block=mxfp4_block, + ) + # Fused gate_up_proj (column parallel) w13_weight = torch.nn.Parameter( - torch.zeros( - layer.num_local_experts, - 2 * intermediate_size_per_partition_after_pad, - hidden_size // 2, - dtype=weight_dtype, + self._expert_storage( + moonep_weights.W13_WEIGHT, + ( + layer.num_local_experts, + 2 * intermediate_size_per_partition_after_pad, + hidden_size // 2, + ), + weight_dtype, ), requires_grad=False, ) @@ -508,11 +573,14 @@ def create_weights( # down_proj (row parallel) w2_weight = torch.nn.Parameter( - torch.zeros( - layer.num_local_experts, - hidden_size, - intermediate_size_per_partition_after_pad // 2, - dtype=weight_dtype, + self._expert_storage( + moonep_weights.W2_WEIGHT, + ( + layer.num_local_experts, + hidden_size, + intermediate_size_per_partition_after_pad // 2, + ), + weight_dtype, ), requires_grad=False, ) @@ -543,6 +611,21 @@ def create_weights( layer.register_parameter("w2_weight_bias", w2_weight_bias) set_weight_attrs(w2_weight_bias, extra_weight_attrs) + def _expert_storage( + self, kind: str, shape: tuple[int, ...], dtype: torch.dtype + ) -> torch.Tensor: + """Zeroed storage for an expert tensor, from MoonEP's symmetric pool + when that backend is active and a private allocation otherwise.""" + if self.moonep_pooled is None: + return torch.zeros(*shape, dtype=dtype) + + pooled = self.moonep_pooled[kind] + assert tuple(pooled.shape) == shape and pooled.dtype == dtype, ( + f"MoonEP pool gave {kind} as {tuple(pooled.shape)}/{pooled.dtype}, " + f"expected {shape}/{dtype}" + ) + return pooled.zero_() + def process_weights_after_loading(self, layer): if self.use_marlin: from sglang.srt.layers.quantization.marlin_utils import ( @@ -575,21 +658,32 @@ def process_weights_after_loading(self, layer): if self.use_deep_gemm: from deep_gemm import transform_sf_into_required_layout + from sglang.srt.layers.moe.token_dispatcher import moonep_weights + # Packed fp4 (e2m1 x2 per byte) weights: DeepGEMM expects int8. + # A re-view keeps the storage, which is what lets MoonEP's pooled + # weights survive this step. layer.w13_weight.data = layer.w13_weight.data.view(torch.int8) layer.w2_weight.data = layer.w2_weight.data.view(torch.int8) + if self.moonep_pooled is not None: + moonep_weights.assert_resident( + layer, moonep_weights.W13_WEIGHT, layer.w13_weight.data + ) + moonep_weights.assert_resident( + layer, moonep_weights.W2_WEIGHT, layer.w2_weight.data + ) # Checkpoint scales are uint8 e8m0 (biased exponents). DeepGEMM # SM100 needs them in packed-UE8M0 TMA-aligned MN-major layout. # Round-trip through fp32 is exact (values are powers of two). - for scale_name, weight in ( - ("w13_weight_scale", layer.w13_weight), - ("w2_weight_scale", layer.w2_weight), + for scale_name, weight, pool_kind in ( + ("w13_weight_scale", layer.w13_weight, moonep_weights.W13_SCALE), + ("w2_weight_scale", layer.w2_weight, moonep_weights.W2_SCALE), ): scale = getattr(layer, scale_name) num_experts, n, _ = scale.data.shape k = weight.shape[2] * 2 scale_f32 = scale.data.view(torch.float8_e8m0fnu).to(torch.float32) - scale.data = transform_sf_into_required_layout( + transformed = transform_sf_into_required_layout( scale_f32, mn=n, k=k, @@ -597,6 +691,16 @@ def process_weights_after_loading(self, layer): num_groups=num_experts, disable_ue8m0_cast=False, ) + if self.moonep_pooled is None: + scale.data = transformed + continue + # The checkpoint-layout buffer the loader filled is private + # memory; the runtime layout has to end up in the pool instead, + # so copy rather than rebind. The pooled range is stored + # transposed, matching the MN-major result's real byte order. + pooled = self.moonep_pooled[pool_kind].permute(0, 2, 1) + scale.data = pooled.copy_(transformed) + moonep_weights.assert_resident(layer, pool_kind, scale.data) if get_moe_a2a_backend().is_megamoe(): # MegaMoE consumes the same transformed sf, plus its own # interleaved/UTCCP weight layout. K3 routes EVERY batch diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 6f3414c54d10..15254d5235ac 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -6861,11 +6861,12 @@ def _handle_a2a_moe(self): self.cuda_graph_config.decode.backend = Backend.DISABLED self.cuda_graph_config.prefill.backend = Backend.DISABLED - if a2a_backend == "moonep": + if a2a_backend == "moonep" and not envs.SGLANG_ENABLE_MOONEP_CUDA_GRAPH.get(): logger.warning( - "MoonEP MoE is enabled in experimental BF16 PoC mode. " "Cuda graph is disabled while the eager MoonEP dispatch/" - "prefetch/compute/combine path is validated." + "prefetch/compute/combine path is validated. MoonEP's buffers " + "are statically shaped, so capture should be possible; set " + "SGLANG_ENABLE_MOONEP_CUDA_GRAPH=1 to try it." ) self.cuda_graph_config.decode.backend = Backend.DISABLED self.cuda_graph_config.prefill.backend = Backend.DISABLED