From 048d90e1651a3842253ce5291162a4840e9ee5cc Mon Sep 17 00:00:00 2001 From: Alison Shao <54658187+alisonshao@users.noreply.github.com> Date: Fri, 20 Mar 2026 20:26:17 -0700 Subject: [PATCH] Revert "[AMD] Add MoE weights and scales padding" (#21067) --- .../layers/moe/fused_moe_triton/fused_moe.py | 4 +- .../fused_moe_triton_kernels.py | 4 +- .../srt/layers/moe/fused_moe_triton/layer.py | 16 +----- python/sglang/srt/layers/moe/utils.py | 30 ---------- .../compressed_tensors_w8a8_fp8_moe.py | 22 ++----- python/sglang/srt/layers/quantization/fp8.py | 57 +++++++------------ .../quark/schemes/quark_w4a4_mxfp4_moe.py | 28 ++------- .../sglang/srt/model_executor/model_runner.py | 6 +- 8 files changed, 36 insertions(+), 131 deletions(-) diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py b/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py index 7a7d9fd0b..f2621473c 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py @@ -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"]) diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_kernels.py b/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_kernels.py index 6a7a69e27..b80e1827b 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_kernels.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_kernels.py @@ -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(): diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py index e472cc4c5..d60fb92c5 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -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, diff --git a/python/sglang/srt/layers/moe/utils.py b/python/sglang/srt/layers/moe/utils.py index 119a127a2..65da83c0e 100644 --- a/python/sglang/srt/layers/moe/utils.py +++ b/python/sglang/srt/layers/moe/utils.py @@ -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) diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8_moe.py b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8_moe.py index c76e6902e..fdd3c9622 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8_moe.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8_moe.py @@ -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, ), diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index f26099501..c7874830d 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -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)), diff --git a/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4_moe.py b/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4_moe.py index 7a6147169..df6650be8 100644 --- a/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4_moe.py +++ b/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4_moe.py @@ -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, diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 7115b90d8..62079dacf 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -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}). "