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

[BugFix] Accuracy and function Issue when run ptpc quant model (#13157)

Browse Source 
Co-authored-by: yuechguo <yuechguo@amd.com>
This commit is contained in:
Morpheus Guo
2025-11-18 01:54:58 -08:00
committed by GitHub
co-authored by yuechguo
parent cfcf2758aa
commit 7e88b9c128
1 changed files with 29 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+29 -19
View File
@@ -51,10 +51,10 @@ _is_cuda = is_cuda()
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
if _use_aiter: if _use_aiter:
from aiter import ActivationType, QuantType
from aiter.fused_moe import fused_moe
from aiter.ops.shuffle import shuffle_weight from aiter.ops.shuffle import shuffle_weight
from sglang.srt.layers.moe.rocm_moe_utils import rocm_fused_experts_tkw1
if _is_cuda: if _is_cuda:
from sgl_kernel import fused_marlin_moe from sgl_kernel import fused_marlin_moe
@@ -292,7 +292,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
max_w13_scales, requires_grad=False max_w13_scales, requires_grad=False
) )
if _use_aiter: if self.weight_quant.strategy == QuantizationStrategy.CHANNEL and _use_aiter:
with torch.no_grad(): with torch.no_grad():
# Pre-shuffle weights # Pre-shuffle weights
layer.w13_weight = torch.nn.Parameter( layer.w13_weight = torch.nn.Parameter(
@@ -325,23 +325,33 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
moe_runner_config = self.moe_runner_config moe_runner_config = self.moe_runner_config
if ( if _use_aiter and self.weight_quant.strategy == QuantizationStrategy.CHANNEL:
_use_aiter assert not moe_runner_config.no_combine, "unsupported"
and self.weight_quant.strategy == QuantizationStrategy.CHANNEL
and moe_runner_config.apply_router_weight_on_input
):
topk_weights, topk_ids, _ = topk_output topk_weights, topk_ids, _ = topk_output
output = rocm_fused_experts_tkw1( if moe_runner_config.apply_router_weight_on_input:
hidden_states=x, assert (
w1=layer.w13_weight, topk_weights.dim() == 2
w2=layer.w2_weight, ), "`topk_weights` should be in shape (num_tokens, topk)"
topk_weights=topk_weights, _, topk = topk_weights.shape
topk_ids=topk_ids, assert (
activation=moe_runner_config.activation, topk == 1
apply_router_weight_on_input=moe_runner_config.apply_router_weight_on_input, ), "Only support topk=1 when `apply_router_weight_on_input` is True"
use_fp8_w8a8=True, x = x * topk_weights.to(x.dtype)
per_channel_quant=self.weight_quant.strategy topk_weights = torch.ones_like(
== QuantizationStrategy.CHANNEL, topk_weights, dtype=torch.float32
) # topk_weights must be FP32 (float32)
output = fused_moe(
x,
layer.w13_weight,
layer.w2_weight,
topk_weights,
topk_ids,
activation=(
ActivationType.Silu
if moe_runner_config.activation == "silu"
else ActivationType.Gelu
),
quant_type=QuantType.per_Token,
w1_scale=layer.w13_weight_scale, w1_scale=layer.w13_weight_scale,
w2_scale=layer.w2_weight_scale, w2_scale=layer.w2_weight_scale,
a1_scale=layer.w13_input_scale, a1_scale=layer.w13_input_scale,