feat(moe): add megamoe forward and weight layout
Add a shared MegaMoE forward path and expert weight builder that can be used by GLM MoE without DeepSeek-only type assumptions. Build ModelOpt NVFP4 MegaMoE sidecar weights before FlashInfer TRTLLM alignment so the original runner layout remains available as fallback. Keep the runtime fast path gated to DeepGEMM's FP4 activation pre-dispatch because this fork does not carry the upstream DSV4 JIT FP8 pre-dispatch kernels yet. Constraint: preserve none+flashinfer_trtllm runner layout as fallback. Feature-flag: --moe-a2a-backend=megamoe and SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_FP4_ACTS. Conflict-hotspots: python/sglang/srt/layers/moe/mega_moe.py, python/sglang/srt/layers/quantization/modelopt_quant.py. Scope-risk: DeepGEMM runtime APIs and GPU e2e are not available on this local machine. Tested: PYTHONPYCACHEPREFIX=/private/tmp/sglang_pycache python3 -m py_compile python/sglang/srt/layers/moe/mega_moe.py python/sglang/srt/layers/quantization/modelopt_quant.py. Tested: git diff --check. Not-tested: GLM 5.2 MegaMoE GPU runtime; local environment lacks deep_gemm and target GPU.
This commit is contained in:
LuminolT
423
python/sglang/srt/layers/moe/mega_moe.py
Normal file
423
python/sglang/srt/layers/moe/mega_moe.py
Normal file
@@ -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
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user