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

Revert "[AMD] Add MoE weights and scales padding" (#21067)

Browse Source
This commit is contained in:
Alison Shao
2026-03-20 20:26:17 -07:00
committed by GitHub
parent a75e74b534
commit 048d90e165
8 changed files with 36 additions and 131 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py
+2 -2
View File
@@ -6,6 +6,7 @@
from __future__ import annotations
import functools
import os
from typing import TYPE_CHECKING, List, Optional
import torch
@@ -13,7 +14,6 @@ import torch.nn.functional as F
import triton.language as tl
from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig
from sglang.srt.layers.moe.utils import get_moe_padding_size
from sglang.srt.utils import (
cpu_has_amx_support,
get_bool_env_var,
@@ -75,7 +75,7 @@ if not _is_cuda and not _is_hip and not _is_xpu:
# Fallback: vllm not available, will use native PyTorch implementations
_has_vllm_ops = False
padding_size = get_moe_padding_size(_use_aiter)
padding_size = 128 if bool(int(os.getenv("SGLANG_MOE_PADDING", "0"))) else 0
@register_custom_op(mutates_args=["hidden_states"])
python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_kernels.py
+2 -2
View File
@@ -1,6 +1,7 @@
from __future__ import annotations
import functools
import os
from collections import OrderedDict
from typing import Any, Dict, List, Optional
@@ -9,7 +10,6 @@ import triton
import triton.language as tl
from sglang.srt.batch_invariant_ops import is_batch_invariant_mode_enabled
from sglang.srt.layers.moe.utils import get_moe_padding_size
from sglang.srt.layers.quantization.fp8_kernel import (
per_token_group_quant_fp8,
scaled_fp8_quant,
@@ -49,7 +49,7 @@ elif _is_cpu and _is_cpu_amx_available:
elif _is_hip:
pass
padding_size = get_moe_padding_size(_use_aiter)
padding_size = 128 if bool(int(os.getenv("SGLANG_MOE_PADDING", "0"))) else 0
def support_tensor_descriptor():
python/sglang/srt/layers/moe/fused_moe_triton/layer.py
+2 -14
View File
@@ -440,14 +440,8 @@ class FusedMoE(torch.nn.Module):
# Use narrow_padded_param_and_loaded_weight for:
# 1. CPU (always)
# 2. GPU with flashinfer_trtllm padding (when intermediate_size is padded to 128)
# 3. GPU with Aiter padding
# This handles the case where the loaded weights are smaller than the padded expert_data
aiter_padded = (
_use_aiter
and hasattr(self, "w2_weight")
and getattr(self.w2_weight, "weight_padded", False)
)
use_padded_loading = _is_cpu or self.use_flashinfer_trtllm_moe or aiter_padded
use_padded_loading = _is_cpu or self.use_flashinfer_trtllm_moe
if use_padded_loading:
expert_data, loaded_weight = narrow_padded_param_and_loaded_weight(
expert_data,
@@ -520,14 +514,8 @@ class FusedMoE(torch.nn.Module):
# Use narrow_padded_param_and_loaded_weight for:
# 1. CPU (always)
# 2. GPU with flashinfer_trtllm padding (when intermediate_size is padded to 128)
# 3. GPU with Aiter padding
# This handles the case where the loaded weights are smaller than the padded expert_data
aiter_padded = (
_use_aiter
and hasattr(self, "w2_weight")
and getattr(self.w2_weight, "weight_padded", False)
)
use_padded_loading = _is_cpu or self.use_flashinfer_trtllm_moe or aiter_padded
use_padded_loading = _is_cpu or self.use_flashinfer_trtllm_moe
if use_padded_loading:
expert_data, loaded_weight = narrow_padded_param_and_loaded_weight(
expert_data,
python/sglang/srt/layers/moe/utils.py
-30
View File
@@ -1,13 +1,10 @@
from __future__ import annotations
import logging
import os
from contextlib import contextmanager
from enum import Enum, IntEnum
from typing import TYPE_CHECKING, Optional
import torch
from sglang.srt.distributed.parallel_state import get_moe_expert_parallel_world_size
from sglang.srt.layers.dp_attention import (
get_attention_dp_size,
@@ -344,30 +341,3 @@ class RoutingMethodType(IntEnum):
TopK = (5,)
# Unspecified
Unspecified = 6
def get_moe_padding_size(is_aiter_moe):
if is_aiter_moe:
return 128
else:
return 128 if bool(int(os.getenv("SGLANG_MOE_PADDING", "0"))) else 0
def get_moe_weight_sizes(inter_dim, is_concat, is_packed, is_aiter_moe):
w13_up_dim = 2 * inter_dim if is_concat else inter_dim
w2_down_dim = inter_dim // 2 if is_packed else inter_dim
if is_aiter_moe:
padding_size = get_moe_padding_size(True)
align_aiter = lambda n: ((n + padding_size - 1) // padding_size) * padding_size
is_padded = (w2_down_dim % padding_size) > 0
if is_padded:
w2_down_dim = align_aiter(w2_down_dim)
# up proj + gate fusion : 2x
if is_concat:
w13_up_dim = w2_down_dim * 2
# packed
if hasattr(torch, "float4_e2m1fn_x2") and is_packed:
w13_up_dim *= 2
return (w13_up_dim, w2_down_dim, False if not is_aiter_moe else is_padded)
python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8_moe.py
+4 -18
View File
@@ -12,10 +12,7 @@ from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
FlashInferTrtllmFp8MoeQuantInfo,
)
from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo
from sglang.srt.layers.moe.utils import (
get_moe_runner_backend,
get_moe_weight_sizes,
)
from sglang.srt.layers.moe.utils import get_moe_runner_backend
from sglang.srt.layers.quantization.compressed_tensors.schemes import (
CompressedTensorsMoEScheme,
)
@@ -123,22 +120,11 @@ class CompressedTensorsW8A8Fp8MoE(CompressedTensorsMoEScheme):
f"weight quantization block_k = {block_k}."
)
w13_up_dim, w2_down_dim, weight_padded = get_moe_weight_sizes(
intermediate_size_per_partition,
is_aiter_moe=True,
is_concat=True,
is_packed=False,
)
extra_weight_attrs.update(
{"weight_padded": weight_padded},
)
# WEIGHTS
w13_weight = torch.nn.Parameter(
torch.empty(
num_experts,
w13_up_dim,
2 * intermediate_size_per_partition,
hidden_size,
dtype=params_dtype,
),
@@ -151,7 +137,7 @@ class CompressedTensorsW8A8Fp8MoE(CompressedTensorsMoEScheme):
torch.empty(
num_experts,
hidden_size,
w2_down_dim,
intermediate_size_per_partition,
dtype=params_dtype,
),
requires_grad=False,
@@ -175,7 +161,7 @@ class CompressedTensorsW8A8Fp8MoE(CompressedTensorsMoEScheme):
w13_weight_scale = torch.nn.Parameter(
torch.ones(
num_experts,
w13_up_dim,
2 * intermediate_size_per_partition,
1,
dtype=torch.float32,
),
python/sglang/srt/layers/quantization/fp8.py
+20 -37
View File
@@ -26,12 +26,7 @@ from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
FlashInferTrtllmFp8MoeQuantInfo,
)
from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo
from sglang.srt.layers.moe.utils import (
RoutingMethodType,
get_moe_padding_size,
get_moe_runner_backend,
get_moe_weight_sizes,
)
from sglang.srt.layers.moe.utils import RoutingMethodType, get_moe_runner_backend
from sglang.srt.layers.parameter import (
BlockQuantScaleParameter,
ModelWeightParameter,
@@ -783,38 +778,27 @@ class Fp8MoEMethod(FusedMoEMethodBase):
if self.quant_config.is_checkpoint_fp8_serialized:
params_dtype = torch.uint32 if _use_hip_int4 else torch.float8_e4m3fn
tp_size = get_tensor_model_parallel_world_size()
w13_up_dim, w2_up_dim, weight_padded = get_moe_weight_sizes(
intermediate_size_per_partition,
is_aiter_moe=True,
is_concat=True,
is_packed=False,
)
if self.block_quant:
block_n, block_k = (
self.quant_config.weight_block_size[0],
self.quant_config.weight_block_size[1],
)
padding_size = get_moe_padding_size(_use_aiter)
if not (_use_aiter and padding_size == block_n == block_k):
# NOTE(HandH1998): To ensure proper alignment of the block-wise quantization scales, the output_size of the weights for both the gate and up layers must be divisible by block_n.
# Required by column parallel or enabling merged weights
if intermediate_size_per_partition % block_n != 0:
# NOTE(HandH1998): To ensure proper alignment of the block-wise quantization scales, the output_size of the weights for both the gate and up layers must be divisible by block_n.
# Required by column parallel or enabling merged weights
if intermediate_size_per_partition % block_n != 0:
raise ValueError(
f"The output_size of gate's and up's weight = "
f"{intermediate_size_per_partition} is not divisible by "
f"weight quantization block_n = {block_n}."
)
if tp_size > 1:
# Required by row parallel
if intermediate_size_per_partition % block_k != 0:
raise ValueError(
f"The output_size of gate's and up's weight = "
f"The input_size of down's weight = "
f"{intermediate_size_per_partition} is not divisible by "
f"weight quantization block_n = {block_n}."
f"weight quantization block_k = {block_k}."
)
if tp_size > 1:
# Required by row parallel
if intermediate_size_per_partition % block_k != 0:
raise ValueError(
f"The input_size of down's weight = "
f"{intermediate_size_per_partition} is not divisible by "
f"weight quantization block_k = {block_k}."
)
# WEIGHTS
if _is_hip and _use_hip_int4:
@@ -841,7 +825,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
w13_weight = torch.nn.Parameter(
torch.empty(
num_experts,
w13_up_dim,
2 * intermediate_size_per_partition,
hidden_size,
dtype=params_dtype,
),
@@ -851,16 +835,12 @@ class Fp8MoEMethod(FusedMoEMethodBase):
torch.empty(
num_experts,
hidden_size,
w2_up_dim,
intermediate_size_per_partition,
dtype=params_dtype,
),
requires_grad=False,
)
extra_weight_attrs.update(
{"weight_padded": weight_padded},
)
layer.register_parameter("w13_weight", w13_weight)
set_weight_attrs(w13_weight, extra_weight_attrs)
@@ -1421,7 +1401,10 @@ class Fp8MoEMethod(FusedMoEMethodBase):
layer.w2_weight_scale1[expert_id] *= layer.w2_weight_scale[expert_id]
def process_weights_hip_scale_padding(self, layer: Module):
padding_size = get_moe_padding_size(_use_aiter)
from sglang.srt.layers.moe.fused_moe_triton.fused_moe import (
padding_size, # Avoid circular import
)
if _use_aiter:
layer.w13_weight = torch.nn.Parameter(
shuffle_weight(layer.w13_weight.data, (16, 16)),
python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4_moe.py
+5 -23
View File
@@ -8,7 +8,6 @@ from typing import TYPE_CHECKING, Any
import torch
from sglang.srt.layers.moe import MoeRunnerConfig
from sglang.srt.layers.moe.utils import get_moe_weight_sizes
from sglang.srt.layers.quantization.quark.schemes import QuarkMoEScheme
from sglang.srt.utils import (
get_bool_env_var,
@@ -74,20 +73,10 @@ class QuarkW4A4MXFp4MoE(QuarkMoEScheme):
from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported
w13_up_dim, w2_down_dim, weight_padded = get_moe_weight_sizes(
intermediate_size_per_partition,
is_aiter_moe=True,
is_concat=True,
is_packed=True,
)
# Add the quantization method used (per tensor/grouped/channel)
# to ensure the weight scales are loaded in properly
extra_weight_attrs.update(
{
"quant_method": FusedMoeWeightScaleSupported.BLOCK.value,
"weight_padded": weight_padded,
},
{"quant_method": FusedMoeWeightScaleSupported.BLOCK.value}
)
params_dtype = torch.uint8
@@ -96,7 +85,7 @@ class QuarkW4A4MXFp4MoE(QuarkMoEScheme):
w13_weight = torch.nn.Parameter(
torch.empty(
num_experts,
w13_up_dim,
2 * intermediate_size_per_partition,
hidden_size // 2,
dtype=params_dtype,
),
@@ -110,7 +99,7 @@ class QuarkW4A4MXFp4MoE(QuarkMoEScheme):
torch.empty(
num_experts,
hidden_size,
w2_down_dim,
intermediate_size_per_partition // 2,
dtype=params_dtype,
),
requires_grad=False,
@@ -123,24 +112,17 @@ class QuarkW4A4MXFp4MoE(QuarkMoEScheme):
w13_weight_scale = torch.nn.Parameter(
torch.ones(
num_experts,
w13_up_dim,
2 * intermediate_size_per_partition,
hidden_size // OCP_MX_BLOCK_SIZE,
dtype=params_dtype,
),
requires_grad=False,
)
W2_SCALE_DIVIDEND = w2_down_dim * 2
W2_SCALE_DIVISOR = intermediate_size_per_partition
scaling_up = lambda dividend, divisor: (dividend * W2_SCALE_DIVIDEND) // (
divisor * W2_SCALE_DIVISOR
)
w2_weight_scale = torch.nn.Parameter(
torch.ones(
num_experts,
hidden_size,
scaling_up(intermediate_size_per_partition, OCP_MX_BLOCK_SIZE),
intermediate_size_per_partition // OCP_MX_BLOCK_SIZE,
dtype=params_dtype,
),
requires_grad=False,
python/sglang/srt/model_executor/model_runner.py
+1 -5
View File
@@ -162,7 +162,6 @@ from sglang.srt.utils import (
empty_context,
enable_show_time_cost,
get_available_gpu_memory,
get_bool_env_var,
get_cpu_ids_by_node,
init_custom_process_group,
is_hip,
@@ -199,7 +198,6 @@ _is_hip = is_hip()
_is_npu = is_npu()
_is_cpu_amx_available = cpu_has_amx_support()
_is_cpu_arm64 = is_host_cpu_arm64()
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
if _is_npu:
from sglang.srt.hardware_backend.npu.utils import init_npu_backend
@@ -801,9 +799,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
f"moe_intermediate_size {moe_intermediate_size} must be divisible by moe_tp_size ({moe_tp_size}) which is tp_size ({self.tp_size}) divided by moe_ep_size ({self.moe_ep_size})."
)
if (
moe_intermediate_size // moe_tp_size
) % weight_block_size_n != 0 and not _use_aiter:
if (moe_intermediate_size // moe_tp_size) % weight_block_size_n != 0:
raise ValueError(
f"For quantized MoE models, please make sure ({moe_intermediate_size=} / {moe_tp_size=}) % {weight_block_size_n=} == 0 "
f"where moe_tp_size is equal to tp_size ({self.tp_size}) divided by ep_size ({self.moe_ep_size}). "