chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[2/N] MoE Refactor: Unify weight loader and quant methods (#8397)

Browse Source
This commit is contained in:
Cheng Wan
2025-07-27 01:00:21 -07:00
committed by GitHub
parent 36d6f0ba5b
commit bf0f448fe5
5 changed files with 221 additions and 590 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/ep_moe/layer.py
+87 -217
View File
@@ -30,13 +30,13 @@ from sglang.srt.layers.quantization.base_config import (
QuantizationConfig,
QuantizeMethodBase,
)
from sglang.srt.layers.quantization.fp8 import Fp8EPMoEMethod
from sglang.srt.layers.quantization.fp8 import Fp8Config, Fp8MoEMethod
from sglang.srt.layers.quantization.fp8_kernel import (
is_fp8_fnuz,
sglang_per_token_group_quant_fp8,
sglang_per_token_quant_fp8,
)
from sglang.srt.layers.quantization.unquant import UnquantizedEPMoEMethod
from sglang.srt.layers.quantization.unquant import UnquantizedFusedMoEMethod
from sglang.srt.layers.quantization.w4afp8 import W4AFp8Config, W4AFp8MoEMethod
from sglang.srt.managers.schedule_batch import global_server_args_dict
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
@@ -62,8 +62,6 @@ use_flashinfer_trtllm_moe = (
if not (_is_npu or _is_hip):
from sgl_kernel import silu_and_mul
from sglang.srt.layers.moe.cutlass_w4a8_moe import cutlass_w4a8_moe
if _use_aiter:
from aiter import ActivationType, QuantType
from aiter.fused_moe import fused_moe
@@ -162,7 +160,7 @@ def _get_tile_tokens_dim(num_tokens, top_k, num_experts):
return tile_tokens_dim
class EPMoE(torch.nn.Module):
class EPMoE(FusedMoE):
"""
MoE Expert Parallel Impl
@@ -184,51 +182,60 @@ class EPMoE(torch.nn.Module):
routed_scaling_factor: Optional[float] = None,
use_per_token_if_dynamic: bool = True,
):
super().__init__()
super().__init__(
num_experts=num_experts,
hidden_size=hidden_size,
intermediate_size=intermediate_size,
top_k=top_k,
layer_id=layer_id,
params_dtype=params_dtype,
quant_config=quant_config,
tp_size=tp_size,
prefix=prefix,
activation=activation,
routed_scaling_factor=routed_scaling_factor,
enable_ep_moe=True,
skip_quant=True,
)
if params_dtype is None:
params_dtype = torch.get_default_dtype()
self.tp_size = (
tp_size if tp_size is not None else get_tensor_model_parallel_world_size()
)
self.tp_rank = get_tensor_model_parallel_rank()
self.layer_id = layer_id
self.num_experts = num_experts
assert self.num_experts % self.tp_size == 0
self.num_experts_per_partition, self.expert_map = self.determine_expert_map()
self.start_expert_id = self.tp_rank * self.num_experts_per_partition
self.end_expert_id = self.start_expert_id + self.num_experts_per_partition - 1
self.num_local_experts, self.expert_map = self.determine_expert_map()
self.start_expert_id = self.ep_rank * self.num_local_experts
self.end_expert_id = self.start_expert_id + self.num_local_experts - 1
self.top_k = top_k
self.intermediate_size = intermediate_size
self.activation = activation
self.routed_scaling_factor = routed_scaling_factor
self.use_per_token_if_dynamic = use_per_token_if_dynamic
# TODO(ch-wan): move quant preparation to FusedMoE
if quant_config is None:
self.quant_method: Optional[QuantizeMethodBase] = UnquantizedEPMoEMethod()
self.quant_method: Optional[QuantizeMethodBase] = (
UnquantizedFusedMoEMethod()
)
self.use_fp8_w8a8 = False
self.use_block_quant = False
self.block_shape = None
self.activation_scheme = None
self.use_w4afp8 = False
self.w13_input_scale = None
self.w2_input_scale = None
self.w13_weight_scale = None
self.w2_weight_scale = None
elif isinstance(quant_config, W4AFp8Config):
self.quant_method: Optional[QuantizeMethodBase] = W4AFp8MoEMethod(
quant_config
)
self.use_w4afp8 = True
self.use_fp8_w8a8 = False
self.use_block_quant = False
self.fp8_dtype = torch.float8_e4m3fn
self.w13_input_scale = None
self.w2_input_scale = None
self.w13_weight_scale = None
self.w2_weight_scale = None
self.activation_scheme = quant_config.moe_activation_scheme
else:
self.quant_method: Optional[QuantizeMethodBase] = Fp8EPMoEMethod(
quant_config
)
elif isinstance(quant_config, Fp8Config):
self.quant_method: Optional[QuantizeMethodBase] = Fp8MoEMethod(quant_config)
self.use_fp8_w8a8 = True
self.use_block_quant = getattr(self.quant_method, "block_quant", False)
self.block_shape = (
@@ -238,11 +245,13 @@ class EPMoE(torch.nn.Module):
)
self.fp8_dtype = torch.float8_e4m3fn
self.activation_scheme = quant_config.activation_scheme
self.use_w4afp8 = False
else:
raise ValueError(f"Unsupported quant_config: {quant_config}")
self.quant_config = quant_config
self.quant_method.create_weights(
layer=self,
num_experts_per_partition=self.num_experts_per_partition,
num_experts=self.num_local_experts,
hidden_size=hidden_size,
intermediate_size=self.intermediate_size,
params_dtype=params_dtype,
@@ -251,19 +260,6 @@ class EPMoE(torch.nn.Module):
self.grouped_gemm_runner = None
self.w13_weight_fp8 = (
self.w13_weight,
(
self.w13_weight_scale_inv
if self.use_block_quant
else self.w13_weight_scale
),
)
self.w2_weight_fp8 = (
self.w2_weight,
self.w2_weight_scale_inv if self.use_block_quant else self.w2_weight_scale,
)
# Adapted from https://github.com/vllm-project/vllm/blob/9fb52e523abf7bdaf7e60cf2971edb5a1b13dc08/vllm/model_executor/layers/fused_moe/layer.py#L544C1-L586C43
# Modifications: use determine_expert_map as a class internal function, set 'global_num_experts' rather than '-1' for experts not assigned to the current rank.
def determine_expert_map(self) -> Tuple[int, Optional[torch.Tensor]]:
@@ -282,8 +278,8 @@ class EPMoE(torch.nn.Module):
Contains global_num_experts for experts not assigned to the current rank.
Returns None if ep_size is 1.
"""
ep_size = self.tp_size
ep_rank = self.tp_rank
ep_size = self.ep_size
ep_rank = self.ep_rank
global_num_experts = self.num_experts
assert ep_size > 0
@@ -293,7 +289,7 @@ class EPMoE(torch.nn.Module):
local_num_experts = global_num_experts // ep_size
expert_map = torch.full(
(global_num_experts,), self.num_experts, dtype=torch.int32
(global_num_experts,), global_num_experts, dtype=torch.int32
)
if ep_rank < (ep_size - 1):
expert_map[
@@ -318,6 +314,20 @@ class EPMoE(torch.nn.Module):
hidden_states: torch.Tensor,
topk_output: TopKOutput,
):
self.w13_weight_fp8 = (
self.w13_weight,
(
self.w13_weight_scale_inv
if self.use_block_quant
else self.w13_weight_scale
),
)
self.w2_weight_fp8 = (
self.w2_weight,
self.w2_weight_scale_inv if self.use_block_quant else self.w2_weight_scale,
)
assert self.quant_method is not None
assert self.activation == "silu"
hidden_states_shape = hidden_states.shape
@@ -457,7 +467,10 @@ class EPMoE(torch.nn.Module):
return output
def forward_normal(self, hidden_states: torch.Tensor, topk_output: TopKOutput):
assert self.quant_method is not None
return self.quant_method.apply(self, hidden_states, topk_output)
def run_moe(self, hidden_states: torch.Tensor, topk_output: TopKOutput):
topk_weights, topk_ids, _ = topk_output
hidden_states_shape = hidden_states.shape
@@ -470,53 +483,11 @@ class EPMoE(torch.nn.Module):
use_per_token_if_dynamic=self.use_per_token_if_dynamic,
)
if self.use_w4afp8:
local_topk_ids = topk_ids
if self.expert_map is not None:
"Translate info from expert_map to topk_ids"
local_topk_ids = torch.where(
self.expert_map[topk_ids] != self.num_experts,
self.expert_map[topk_ids],
self.num_experts,
)
output = cutlass_w4a8_moe(
self.start_expert_id,
self.end_expert_id,
self.num_experts,
hidden_states,
self.w13_weight,
self.w2_weight,
self.w13_weight_scale_inv,
self.w2_weight_scale_inv,
topk_weights,
topk_ids,
local_topk_ids,
self.quant_method.a_strides1,
self.quant_method.b_strides1,
self.quant_method.c_strides1,
self.quant_method.a_strides2,
self.quant_method.b_strides2,
self.quant_method.c_strides2,
self.quant_method.s_strides13,
self.quant_method.s_strides2,
self.quant_method.expert_offsets,
self.quant_method.problem_sizes1,
self.quant_method.problem_sizes2,
self.w13_input_scale,
self.w2_input_scale,
)
return output
if self.grouped_gemm_runner is None:
self.grouped_gemm_runner = GroupedGemmRunner(
hidden_states.device,
use_flashinfer=False, # TODO: use flashinfer
use_per_token_if_dynamic=self.use_per_token_if_dynamic,
)
num_experts = self.num_experts
reorder_topk_ids, src2dst, seg_indptr = run_moe_ep_preproess(
topk_ids, self.num_experts
topk_ids,
num_experts,
)
gateup_input = torch.empty(
@@ -524,7 +495,7 @@ class EPMoE(torch.nn.Module):
device=hidden_states.device,
dtype=(
self.fp8_dtype
if ((self.use_fp8_w8a8 or self.use_w4afp8) and not self.use_block_quant)
if self.use_fp8_w8a8 and not self.use_block_quant
else hidden_states.dtype
),
)
@@ -535,7 +506,7 @@ class EPMoE(torch.nn.Module):
else:
max_value = (
torch.max(hidden_states)
.repeat(self.num_experts_per_partition)
.repeat(self.num_local_experts)
.to(torch.float32)
)
self.w13_input_scale = max_value / torch.finfo(self.fp8_dtype).max
@@ -576,7 +547,7 @@ class EPMoE(torch.nn.Module):
seg_indptr_cur_rank = seg_indptr[self.start_expert_id : self.end_expert_id + 2]
weight_indices_cur_rank = torch.arange(
0,
self.num_experts_per_partition,
self.num_local_experts,
device=hidden_states_device,
dtype=torch.int64,
)
@@ -586,17 +557,13 @@ class EPMoE(torch.nn.Module):
b=self.w13_weight,
c=None,
c_dtype=hidden_states_dtype,
batch_size=self.num_experts_per_partition,
batch_size=self.num_local_experts,
weight_column_major=True,
seg_indptr=seg_indptr_cur_rank,
weight_indices=weight_indices_cur_rank,
use_fp8_w8a8=self.use_fp8_w8a8,
scale_a=self.w13_input_scale,
scale_b=(
self.w13_weight_scale_inv
if self.use_block_quant
else self.w13_weight_scale
),
scale_b=self.w13_weight_scale,
block_shape=self.block_shape,
)
del gateup_input
@@ -653,7 +620,7 @@ class EPMoE(torch.nn.Module):
down_input, self.w2_input_scale = sglang_per_token_quant_fp8(down_input)
else:
self.w2_input_scale = torch.ones(
self.num_experts_per_partition,
self.num_local_experts,
dtype=torch.float32,
device=hidden_states_device,
)
@@ -669,17 +636,13 @@ class EPMoE(torch.nn.Module):
a=down_input,
b=self.w2_weight,
c=down_output,
batch_size=self.num_experts_per_partition,
batch_size=self.num_local_experts,
weight_column_major=True,
seg_indptr=seg_indptr_cur_rank,
weight_indices=weight_indices_cur_rank,
use_fp8_w8a8=self.use_fp8_w8a8,
scale_a=self.w2_input_scale,
scale_b=(
self.w2_weight_scale_inv
if self.use_block_quant
else self.w2_weight_scale
),
scale_b=self.w2_weight_scale,
block_shape=self.block_shape,
)
del down_input
@@ -782,107 +745,14 @@ class EPMoE(torch.nn.Module):
return
expert_id = expert_id - self.start_expert_id
if shard_id not in ("w1", "w2", "w3"):
raise ValueError(
f"shard_id must be ['w1','w2','w3'] but " f"got {shard_id}."
)
# Special case for fp8 scales.
if "scale" in weight_name:
self._load_fp8_scale(
param.data,
loaded_weight,
weight_name,
shard_id,
expert_id,
)
return
# Flashinfer assumes w31 format for w13_weight. Same for the scales.
if use_flashinfer_trtllm_moe:
actual_shard_id = {"w1": "w3", "w3": "w1", "w2": "w2"}[shard_id]
else:
actual_shard_id = shard_id
if actual_shard_id == "w2":
param.data[expert_id] = loaded_weight
elif actual_shard_id == "w1":
param.data[expert_id][: self.intermediate_size, :] = loaded_weight
elif actual_shard_id == "w3":
param.data[expert_id][self.intermediate_size :, :] = loaded_weight
else:
raise ValueError(f"Expected shard_id w1,w2 or w3 but got {actual_shard_id}")
def _load_fp8_scale(
self,
param: torch.nn.Parameter,
loaded_weight: torch.Tensor,
weight_name: str,
shard_id: str,
expert_id: int,
) -> None:
param_data = param.data
# Input scales can be loaded directly and should be equal.
if "input_scale" in weight_name:
if self.use_w4afp8:
if shard_id == "w1":
param_data[expert_id][0] = loaded_weight
elif shard_id == "w3":
param_data[expert_id][1] = loaded_weight
else:
param_data[expert_id] = loaded_weight
return
if (
(shard_id == "w1" or shard_id == "w3")
and param_data[expert_id] != 1
and (param_data[expert_id] - loaded_weight).abs() > 1e-5
):
raise ValueError(
"input_scales of w1 and w3 of a layer "
f"must be equal. But got {param_data[expert_id]} "
f"vs. {loaded_weight}"
)
param_data[expert_id] = loaded_weight
# Weight scales
elif "weight_scale" in weight_name:
if self.use_block_quant:
if use_flashinfer_trtllm_moe:
actual_shard_id = {"w1": "w3", "w3": "w1", "w2": "w2"}[shard_id]
else:
actual_shard_id = shard_id
block_n, block_k = self.block_shape[0], self.block_shape[1]
if actual_shard_id == "w1":
param_data[expert_id][
: (self.intermediate_size + block_n - 1) // block_n, :
] = loaded_weight
elif actual_shard_id == "w3":
param_data[expert_id][
(self.intermediate_size + block_n - 1) // block_n :, :
] = loaded_weight
else: # w2
param_data[expert_id] = loaded_weight
elif self.use_w4afp8:
if shard_id == "w1":
param_data[expert_id][: self.intermediate_size, :] = loaded_weight
elif shard_id == "w3":
param_data[expert_id][self.intermediate_size :, :] = loaded_weight
else:
param_data[expert_id] = loaded_weight
# If we are in merged column case (gate_up_proj)
else:
if shard_id in ("w1", "w3"):
# We have to keep the weight scales of w1 and w3 because
# we need to re-quantize w1/w3 weights after weight loading.
idx = 0 if shard_id == "w1" else 1
param_data[expert_id][idx] = loaded_weight
# If we are in the row parallel case (down_proj)
else:
param_data[expert_id] = loaded_weight
self._weight_loader_impl(
param=param,
loaded_weight=loaded_weight,
weight_name=weight_name,
shard_id=shard_id,
expert_id=expert_id,
)
return
class DeepEPMoE(EPMoE):
@@ -932,13 +802,13 @@ class DeepEPMoE(EPMoE):
deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
), f"DeepEP {self.deepep_mode} mode requires deep_gemm"
if _use_aiter:
# expert_mask is of size (self.num_experts_per_partition + 1),
# expert_mask is of size (self.num_local_experts + 1),
# the extra 1 is for invalid rank_id (in original deepep, the invalid rank_id is -1, but aiter does not allow -1, we use a mask to make those ids invalid)
# for instance, if we have 4 experts on this rank, we would have a expert_mask like:
# self.expert_mask = [1, 1, 1, 1, 0]
# idx from 0-3 is valid and will be processed, while idx == 4 will be masked out
self.expert_mask = torch.zeros(
(self.num_experts_per_partition + 1),
(self.num_local_experts + 1),
device=torch.cuda.current_device(),
dtype=torch.int,
)
@@ -1011,13 +881,13 @@ class DeepEPMoE(EPMoE):
if self.activation_scheme == "dynamic" and not self.use_block_quant:
max_value = (
torch.max(hidden_states)
.repeat(self.num_experts_per_partition)
.repeat(self.num_local_experts)
.to(torch.float32)
)
self.w13_input_scale = max_value / torch.finfo(self.fp8_dtype).max
weight_indices_cur_rank = torch.arange(
0,
self.num_experts_per_partition,
self.num_local_experts,
device=hidden_states.device,
dtype=torch.int64,
)
@@ -1029,7 +899,7 @@ class DeepEPMoE(EPMoE):
b=self.w13_weight,
c=None,
c_dtype=hidden_states.dtype,
batch_size=self.num_experts_per_partition,
batch_size=self.num_local_experts,
weight_column_major=True,
seg_indptr=seg_indptr,
weight_indices=weight_indices_cur_rank,
@@ -1063,7 +933,7 @@ class DeepEPMoE(EPMoE):
)
if self.w2_input_scale is None and not self.use_block_quant:
self.w2_input_scale = torch.ones(
self.num_experts_per_partition,
self.num_local_experts,
dtype=torch.float32,
device=hidden_states_device,
)
@@ -1076,7 +946,7 @@ class DeepEPMoE(EPMoE):
reorder_topk_ids,
self.w2_input_scale,
0,
self.num_experts_per_partition - 1,
self.num_local_experts - 1,
BLOCK_SIZE=512,
)
else:
@@ -1096,7 +966,7 @@ class DeepEPMoE(EPMoE):
a=down_input,
b=self.w2_weight,
c=down_output,
batch_size=self.num_experts_per_partition,
batch_size=self.num_local_experts,
weight_column_major=True,
seg_indptr=seg_indptr,
weight_indices=weight_indices_cur_rank,
@@ -1121,9 +991,9 @@ class DeepEPMoE(EPMoE):
return hidden_states
# in original deepep, idx == -1 meaning invalid and will not be processed.
# aiter does not accept -1, we use a expert mask to make these idx invalid
# (idx == num_experts_per_partition) meaning not used in aiter fused_moe
# (idx == num_local_experts) meaning not used in aiter fused_moe
topk_idx_copy = topk_idx.to(torch.int32)
topk_idx_copy[topk_idx_copy == -1] = self.num_experts_per_partition
topk_idx_copy[topk_idx_copy == -1] = self.num_local_experts
return fused_moe(
hidden_states,