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

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"])

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():

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_padded_loading = _is_cpu or self.use_flashinfer_trtllm_moe
-            _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
         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_padded_loading = _is_cpu or self.use_flashinfer_trtllm_moe
-            _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
         if use_padded_loading:
             expert_data, loaded_weight = narrow_padded_param_and_loaded_weight(
                 expert_data,

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)

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 (
+from sglang.srt.layers.moe.utils import get_moe_runner_backend
-    get_moe_runner_backend,
-    get_moe_weight_sizes,
-)
 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

@@ -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 (
+from sglang.srt.layers.moe.utils import RoutingMethodType, get_moe_runner_backend
-    RoutingMethodType,
-    get_moe_padding_size,
-    get_moe_runner_backend,
-    get_moe_weight_sizes,
-)
 from sglang.srt.layers.parameter import (
     BlockQuantScaleParameter,
     ModelWeightParameter,
@@ -783,22 +778,11 @@ 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:
@@ -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

@@ -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}
-                "quant_method": FusedMoeWeightScaleSupported.BLOCK.value,
-                "weight_padded": weight_padded,
-            },
         )
 
         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

@@ -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 (
+            if (moe_intermediate_size // moe_tp_size) % weight_block_size_n != 0:
-                moe_intermediate_size // moe_tp_size
-            ) % weight_block_size_n != 0 and not _use_aiter:
                 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}). "