EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 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

85
python/sglang/srt/layers/quantization/w4afp8.py
View File

@@ -1,7 +1,7 @@
from __future__ import annotations
import logging
from typing import Any, Dict, List, Optional
from typing import TYPE_CHECKING, Any, Dict, List, Optional
import torch
from torch.nn import Module
@@ -17,6 +17,9 @@ from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod
from sglang.srt.layers.quantization.utils import is_layer_skipped
from sglang.srt.utils import set_weight_attrs
if TYPE_CHECKING:
from sglang.srt.layers.moe.ep_moe.layer import EPMoE, TopKOutput
ACTIVATION_SCHEMES = ["static", "dynamic"]
logger = logging.getLogger(__name__)
@@ -84,13 +87,14 @@ class W4AFp8Config(QuantizationConfig):
self, layer: torch.nn.Module, prefix: str
) -> Optional[QuantizeMethodBase]:
from sglang.srt.layers.linear import LinearBase
from sglang.srt.layers.moe.ep_moe.layer import EPMoE
from sglang.srt.layers.moe.fused_moe_triton import FusedMoE
if isinstance(layer, LinearBase):
if is_layer_skipped(prefix, self.ignored_layers):
return UnquantizedLinearMethod()
return Fp8LinearMethod(self)
elif isinstance(layer, FusedMoE):
elif isinstance(layer, EPMoE):
return W4AFp8MoEMethod(self)
return None
@@ -105,8 +109,8 @@ class W4AFp8MoEMethod(FusedMoEMethodBase):
def create_weights(
self,
layer: Module,
num_experts_per_partition: int,
layer: EPMoE,
num_experts: int,
hidden_size: int,
intermediate_size: int,
params_dtype: torch.dtype,
@@ -117,7 +121,7 @@ class W4AFp8MoEMethod(FusedMoEMethodBase):
# Fused gate_up_proj (column parallel)
w13_weight = torch.nn.Parameter(
torch.empty(
num_experts_per_partition,
num_experts,
intermediate_size * 2,
hidden_size // 2,
dtype=torch.int8,
@@ -130,7 +134,7 @@ class W4AFp8MoEMethod(FusedMoEMethodBase):
# down_proj (row parallel)
w2_weight = torch.nn.Parameter(
torch.empty(
num_experts_per_partition,
num_experts,
hidden_size,
intermediate_size // 2,
dtype=torch.int8,
@@ -142,7 +146,7 @@ class W4AFp8MoEMethod(FusedMoEMethodBase):
w13_weight_scale = torch.nn.Parameter(
torch.zeros(
num_experts_per_partition,
num_experts,
2 * intermediate_size,
hidden_size // self.quant_config.group_size,
dtype=torch.float32,
@@ -154,7 +158,7 @@ class W4AFp8MoEMethod(FusedMoEMethodBase):
w2_weight_scale = torch.nn.Parameter(
torch.zeros(
num_experts_per_partition,
num_experts,
hidden_size,
intermediate_size // self.quant_config.group_size,
dtype=torch.float32,
@@ -166,14 +170,14 @@ class W4AFp8MoEMethod(FusedMoEMethodBase):
# Input scales
w13_input_scale = torch.nn.Parameter(
torch.ones((num_experts_per_partition, 2), dtype=torch.bfloat16),
torch.ones((num_experts, 2), dtype=torch.bfloat16),
requires_grad=False,
)
layer.register_parameter("w13_input_scale", w13_input_scale)
set_weight_attrs(w13_input_scale, extra_weight_attrs)
w2_input_scale = torch.nn.Parameter(
torch.ones(num_experts_per_partition, dtype=torch.bfloat16),
torch.ones(num_experts, dtype=torch.bfloat16),
requires_grad=False,
)
layer.register_parameter("w2_input_scale", w2_input_scale)
@@ -183,25 +187,25 @@ class W4AFp8MoEMethod(FusedMoEMethodBase):
device = layer.w13_weight.device
self.a_strides1 = torch.full(
(num_experts_per_partition, 3),
(num_experts, 3),
hidden_size,
device=device,
dtype=torch.int64,
)
self.c_strides1 = torch.full(
(num_experts_per_partition, 3),
(num_experts, 3),
2 * intermediate_size,
device=device,
dtype=torch.int64,
)
self.a_strides2 = torch.full(
(num_experts_per_partition, 3),
(num_experts, 3),
intermediate_size,
device=device,
dtype=torch.int64,
)
self.c_strides2 = torch.full(
(num_experts_per_partition, 3),
(num_experts, 3),
hidden_size,
device=device,
dtype=torch.int64,
@@ -212,13 +216,13 @@ class W4AFp8MoEMethod(FusedMoEMethodBase):
self.s_strides2 = self.c_strides2
self.expert_offsets = torch.empty(
(num_experts_per_partition + 1), dtype=torch.int32, device=device
(num_experts + 1), dtype=torch.int32, device=device
)
self.problem_sizes1 = torch.empty(
(num_experts_per_partition, 3), dtype=torch.int32, device=device
(num_experts, 3), dtype=torch.int32, device=device
)
self.problem_sizes2 = torch.empty(
(num_experts_per_partition, 3), dtype=torch.int32, device=device
(num_experts, 3), dtype=torch.int32, device=device
)
return
@@ -266,3 +270,50 @@ class W4AFp8MoEMethod(FusedMoEMethodBase):
[w2_input_scale_max], dtype=dtype, device=device
)
layer.w2_input_scale = Parameter(new_w2_input_scale, requires_grad=False)
def apply(
self,
layer: EPMoE,
hidden_states: torch.Tensor,
topk_output: TopKOutput,
) -> torch.Tensor:
# TODO(ch-wan): move it out of this class
from sglang.srt.layers.moe.cutlass_w4a8_moe import cutlass_w4a8_moe
topk_ids, topk_weights, _ = topk_output
local_topk_ids = topk_ids
if layer.expert_map is not None:
"Translate info from expert_map to topk_ids"
local_topk_ids = torch.where(
layer.expert_map[topk_ids] != layer.num_experts,
layer.expert_map[topk_ids],
layer.num_experts,
)
return cutlass_w4a8_moe(
layer.start_expert_id,
layer.end_expert_id,
layer.num_experts,
hidden_states,
layer.w13_weight,
layer.w2_weight,
layer.w13_weight_scale_inv,
layer.w2_weight_scale_inv,
topk_weights,
topk_ids,
local_topk_ids,
self.a_strides1,
self.b_strides1,
self.c_strides1,
self.a_strides2,
self.b_strides2,
self.c_strides2,
self.s_strides13,
self.s_strides2,
self.expert_offsets,
self.problem_sizes1,
self.problem_sizes2,
layer.w13_input_scale,
layer.w2_input_scale,
)