diff --git a/python/sglang/srt/layers/moe/mega_moe.py b/python/sglang/srt/layers/moe/mega_moe.py new file mode 100644 index 000000000..5a8ce8f56 --- /dev/null +++ b/python/sglang/srt/layers/moe/mega_moe.py @@ -0,0 +1,423 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""MegaMoE forward path and expert-weight layout helpers.""" + +from __future__ import annotations + +import os +from contextlib import nullcontext +from typing import TYPE_CHECKING, Optional + +import torch + +from sglang.srt.environ import envs +from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo +from sglang.srt.layers.dp_attention import get_dp_global_num_tokens +from sglang.srt.layers.moe.topk import select_experts +from sglang.srt.layers.moe.utils import get_moe_a2a_backend +from sglang.srt.model_executor.runner import get_is_capture_mode + +if TYPE_CHECKING: + from deep_gemm import SymmBuffer + + from sglang.srt.model_executor.forward_batch_info import ForwardBatch + + +_MEGA_MOE_SYMM_BUFFER: dict = {} +_MEGA_MOE_DG_ENV_APPLIED = False + + +def _apply_mega_moe_dg_env() -> None: + """Forward SGLang's MegaMoE FP4 flags to DeepGEMM once.""" + global _MEGA_MOE_DG_ENV_APPLIED + if _MEGA_MOE_DG_ENV_APPLIED: + return + if envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_FP4_ACTS.get(): + os.environ.setdefault("DG_USE_FP4_ACTS", "1") + if envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_MXF4_KIND.get(): + os.environ.setdefault("DG_USE_MXF4_KIND", "1") + _MEGA_MOE_DG_ENV_APPLIED = True + + +def _get_mega_moe_symm_buffer( + group, + num_experts: int, + num_max_tokens_per_rank: int, + num_topk: int, + hidden: int, + intermediate_hidden: int, +) -> SymmBuffer: + import deep_gemm + + _apply_mega_moe_dg_env() + + key = ( + id(group), + num_max_tokens_per_rank, + num_experts, + num_topk, + hidden, + intermediate_hidden, + ) + buf = _MEGA_MOE_SYMM_BUFFER.get(key) + if buf is None: + buf = deep_gemm.get_symm_buffer_for_mega_moe( + group, + num_experts, + num_max_tokens_per_rank, + num_topk, + hidden, + intermediate_hidden, + use_fp8_dispatch=True, + activation="swiglu", + ) + _MEGA_MOE_SYMM_BUFFER[key] = buf + return buf + + +def should_use_mega_moe(moe, hidden_states: torch.Tensor) -> bool: + if not get_moe_a2a_backend().is_megamoe(): + return False + if not getattr(moe.experts, "_mega_moe_weights_built", False): + return False + # This fork does not carry the DSV4 JIT FP8 pre-dispatch path yet. Keep the + # fast path gated to DeepGEMM's FP4 activation pre-dispatch API. + if not envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_FP4_ACTS.get(): + return False + if get_is_capture_mode(): + return True + + global_num_tokens = get_dp_global_num_tokens() + if global_num_tokens: + max_tokens_per_rank = max(global_num_tokens) + else: + max_tokens_per_rank = hidden_states.shape[0] + cap = envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK.get() + return max_tokens_per_rank <= cap + + +def forward_mega_moe( + moe, + hidden_states: torch.Tensor, + forward_batch: Optional[ForwardBatch] = None, + input_ids_global: Optional[torch.Tensor] = None, +) -> torch.Tensor: + num_tokens = hidden_states.shape[0] + + sbo_overlap_flag = ( + getattr(moe, "alt_stream", None) is not None + and getattr(moe, "num_fused_shared_experts", 0) == 0 + and num_tokens > 0 + and get_is_capture_mode() + ) + + if sbo_overlap_flag: + current_stream = torch.cuda.current_stream() + moe.alt_stream.wait_stream(current_stream) + shared_output = moe._forward_shared_experts(hidden_states) + mega_stream_ctx = torch.cuda.stream(moe.alt_stream) + else: + shared_output = moe._forward_shared_experts(hidden_states) + mega_stream_ctx = nullcontext() + + with mega_stream_ctx: + y = _run_mega_routed( + moe, hidden_states, forward_batch, input_ids_global, num_tokens + ) + + if sbo_overlap_flag: + current_stream.wait_stream(moe.alt_stream) + + if shared_output is not None: + y.add_(shared_output) + return y + + +def _call_gate(moe, hidden_states: torch.Tensor, forward_batch: Optional[ForwardBatch]): + try: + return moe.gate(hidden_states, forward_batch=forward_batch) + except TypeError as exc: + if "forward_batch" not in str(exc): + raise + return moe.gate(hidden_states) + + +def _select_mega_topk( + moe, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + forward_batch: Optional[ForwardBatch], + input_ids_global: Optional[torch.Tensor], +): + num_token_non_padded = ( + forward_batch.num_token_non_padded if forward_batch is not None else None + ) + expert_location_dispatch_info = ExpertLocationDispatchInfo.init_new( + layer_id=getattr(moe, "layer_id", None), + ) + + if getattr(moe, "is_hash", False): + return moe.topk( + hidden_states, + router_logits, + num_token_non_padded=num_token_non_padded, + expert_location_dispatch_info=expert_location_dispatch_info, + input_ids=input_ids_global, + ) + + topk_config = getattr(moe.topk, "topk_config", None) + if topk_config is None: + return moe.topk( + hidden_states, + router_logits, + num_token_non_padded=num_token_non_padded, + expert_location_dispatch_info=expert_location_dispatch_info, + ) + + return select_experts( + hidden_states=hidden_states, + router_logits=router_logits, + topk_config=topk_config, + layer_id=getattr(moe, "layer_id", None), + num_token_non_padded=num_token_non_padded, + expert_location_dispatch_info=expert_location_dispatch_info, + ) + + +def _get_mega_top_k(moe) -> int: + topk_config = getattr(getattr(moe, "topk", None), "topk_config", None) + if topk_config is not None: + return topk_config.top_k + return getattr(moe, "top_k") + getattr(moe, "num_fused_shared_experts", 0) + + +def _run_mega_routed( + moe, + hidden_states: torch.Tensor, + forward_batch: Optional[ForwardBatch], + input_ids_global: Optional[torch.Tensor], + num_tokens: int, +) -> torch.Tensor: + import deep_gemm + + from sglang.srt.distributed.parallel_state import get_moe_ep_group + + if not envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_FP4_ACTS.get(): + raise RuntimeError( + "MegaMoE FP8 activation pre-dispatch requires DSV4 JIT kernels in this " + "fork. Set SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_FP4_ACTS=1 to use the " + "DeepGEMM FP4 activation path." + ) + + hidden_size = getattr(moe.experts, "_mega_moe_hidden_size", hidden_states.shape[-1]) + + if num_tokens > 0: + router_logits = _call_gate(moe, hidden_states, forward_batch) + topk_output = _select_mega_topk( + moe, hidden_states, router_logits, forward_batch, input_ids_global + ) + topk_ids = topk_output.topk_ids + topk_weights = topk_output.topk_weights + else: + topk_ids = None + topk_weights = None + + ep_group = get_moe_ep_group().device_group + num_experts = moe.experts.num_experts + top_k = _get_mega_top_k(moe) + intermediate_size = getattr( + moe.experts, + "_mega_moe_intermediate_size", + getattr(moe.config, "moe_intermediate_size"), + ) + num_max_tokens_per_rank = ( + envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK.get() + ) + assert num_tokens <= num_max_tokens_per_rank, ( + f"mega MoE: num_tokens={num_tokens} exceeds cap " + f"SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK=" + f"{num_max_tokens_per_rank}; raise the env var or shrink " + f"cuda_graph_max_bs / chunked_prefill_size accordingly" + ) + + buf = _get_mega_moe_symm_buffer( + ep_group, + num_experts=num_experts, + num_max_tokens_per_rank=num_max_tokens_per_rank, + num_topk=top_k, + hidden=hidden_size, + intermediate_hidden=intermediate_size, + ) + + if num_tokens > 0: + topk_ids_in = topk_ids.to(torch.int32) + topk_weights_in = topk_weights.to(torch.float32) + else: + topk_ids_in = hidden_states.new_empty((0, top_k), dtype=torch.int32) + topk_weights_in = hidden_states.new_empty((0, top_k), dtype=torch.float32) + + deep_gemm.mega_moe_pre_dispatch( + hidden_states, + topk_ids_in, + topk_weights_in, + buf.x, + buf.x_sf, + buf.topk_idx, + buf.topk_weights, + num_tokens=num_tokens, + group_size=32, + use_fp4_acts=True, + ) + + y = torch.empty( + (max(num_tokens, 1), hidden_size), + dtype=torch.bfloat16, + device=hidden_states.device, + ) + swiglu_limit = getattr(moe.config, "swiglu_limit", None) + deep_gemm.fp8_fp4_mega_moe( + y, + moe.experts.mega_l1_weights, + moe.experts.mega_l2_weights, + buf, + recipe=(1, 1, 32), + activation="swiglu", + activation_clamp=swiglu_limit, + fast_math=True, + ) + y = y[:num_tokens] + + if not getattr(moe.experts, "should_fuse_routed_scaling_factor_in_topk", False): + y.mul_(getattr(moe, "routed_scaling_factor", 1.0)) + return y + + +def _interleave_mega_moe_gate_up(t: torch.Tensor, gran: int = 8) -> torch.Tensor: + num_groups, n, *rest = t.shape + half = n // 2 + gate = t[:, :half].reshape(num_groups, half // gran, gran, *rest) + up = t[:, half:].reshape(num_groups, half // gran, gran, *rest) + result = torch.stack([gate, up], dim=2).reshape(num_groups, n, *rest) + return torch.empty_like(t).copy_(result) + + +def _interleave_mega_moe_l1_weights( + l1_weights: tuple[torch.Tensor, torch.Tensor], +) -> tuple[torch.Tensor, torch.Tensor]: + return ( + _interleave_mega_moe_gate_up(l1_weights[0]), + _interleave_mega_moe_gate_up(l1_weights[1]), + ) + + +def _transpose_mega_moe_sf_for_utccp(sf: torch.Tensor) -> torch.Tensor: + num_groups, mn, packed_sf_k = sf.shape + assert sf.dtype == torch.int and mn % 128 == 0 + result = ( + sf.reshape(num_groups, -1, 4, 32, packed_sf_k) + .transpose(2, 3) + .reshape(num_groups, mn, packed_sf_k) + ) + return torch.empty_like(sf).copy_(result) + + +def _get_mega_moe_scale_tensors(experts) -> tuple[torch.Tensor, torch.Tensor]: + if hasattr(experts, "w13_weight_scale_inv") and hasattr( + experts, "w2_weight_scale_inv" + ): + return experts.w13_weight_scale_inv.data, experts.w2_weight_scale_inv.data + if hasattr(experts, "w13_weight_scale") and hasattr(experts, "w2_weight_scale"): + return experts.w13_weight_scale.data, experts.w2_weight_scale.data + raise AttributeError( + "MegaMoE requires w13/w2 block scale tensors on the expert layer." + ) + + +def build_mega_moe_experts_weights( + experts, + *, + preserve_runner_layout: bool = True, +) -> None: + from deep_gemm import transform_sf_into_required_layout + + if getattr(experts, "_mega_moe_weights_built", False): + return + + w13 = experts.w13_weight.data + w2 = experts.w2_weight.data + w13_sf_src, w2_sf_src = _get_mega_moe_scale_tensors(experts) + + num_groups, n1, half_k1 = w13.shape + k1 = half_k1 * 2 + _, n2, half_k2 = w2.shape + k2 = half_k2 * 2 + + w13_sf = transform_sf_into_required_layout( + w13_sf_src.to(torch.float32), + mn=n1, + k=k1, + recipe=(1, 32), + num_groups=num_groups, + disable_ue8m0_cast=False, + ) + w2_sf = transform_sf_into_required_layout( + w2_sf_src.to(torch.float32), + mn=n2, + k=k2, + recipe=(1, 32), + num_groups=num_groups, + disable_ue8m0_cast=False, + ) + + if preserve_runner_layout: + w13_base = w13.clone() + w2_base = w2.clone() + w13_interleaved, w13_sf_interleaved = _interleave_mega_moe_l1_weights( + (w13_base, w13_sf) + ) + experts.mega_l1_weights = ( + w13_interleaved, + _transpose_mega_moe_sf_for_utccp(w13_sf_interleaved), + ) + experts.mega_l2_weights = (w2_base, _transpose_mega_moe_sf_for_utccp(w2_sf)) + experts._mega_moe_preserve_runner_layout = True + elif envs.SGLANG_OPT_FIX_MEGA_MOE_MEMORY.get(): + w13_interleaved, w13_sf_interleaved = _interleave_mega_moe_l1_weights( + (w13, w13_sf) + ) + w13_sf_utccp = _transpose_mega_moe_sf_for_utccp(w13_sf_interleaved) + w2_sf_utccp = _transpose_mega_moe_sf_for_utccp(w2_sf) + + experts.w13_weight.data = w13_interleaved + if hasattr(experts, "w13_weight_scale_inv"): + experts.w13_weight_scale_inv.data = w13_sf_interleaved + experts.w13_weight_scale_inv.format_ue8m0 = True + if hasattr(experts, "w2_weight_scale_inv"): + experts.w2_weight_scale_inv.data = w2_sf + experts.w2_weight_scale_inv.format_ue8m0 = True + + experts.mega_l1_weights = (experts.w13_weight.data, w13_sf_utccp) + experts.mega_l2_weights = (experts.w2_weight.data, w2_sf_utccp) + experts._mega_moe_preserve_runner_layout = False + else: + from deep_gemm import transform_weights_for_mega_moe + + l1_pair, l2_pair = transform_weights_for_mega_moe((w13, w13_sf), (w2, w2_sf)) + experts.mega_l1_weights = l1_pair + experts.mega_l2_weights = l2_pair + experts._mega_moe_preserve_runner_layout = False + + experts._mega_moe_hidden_size = k1 + experts._mega_moe_intermediate_size = k2 + experts._mega_moe_weights_built = True diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index 599d199f5..bf0636c03 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/python/sglang/srt/layers/quantization/modelopt_quant.py @@ -2111,6 +2111,13 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): weight_scale.dtype == torch.float8_e4m3fn ), f"{name} Weight Blockscale must be represented as FP8-E4M3" + if get_moe_a2a_backend().is_megamoe(): + from sglang.srt.layers.moe.mega_moe import ( + build_mega_moe_experts_weights, + ) + + build_mega_moe_experts_weights(layer, preserve_runner_layout=True) + # Weight processing based on strategy if ( self.enable_flashinfer_trtllm_moe