fix

Author: zhiding512

vllm/envs.py

@@ -167,6 +167,7 @@
     VLLM_USE_FLASHINFER_MOE_MXFP4_BF16: bool = False
     VLLM_TUNED_CONFIG_FOLDER: Optional[str] = None
     VLLM_ROCM_USE_AITER_TRITON_FUSED_ROPE_ZEROS_KV_CACHE: bool = False
+    VLLM_ROCM_USE_AITER_FUSED_MOE_A16W4: bool = False
 
 
 def get_default_cache_root():
@@ -1007,6 +1008,9 @@ def get_vllm_port() -> Optional[int]:
     # BF16 (activation) x MXFP4 (weight) MoE backend.
     "VLLM_USE_FLASHINFER_MOE_MXFP4_BF16":
     lambda: bool(int(os.getenv("VLLM_USE_FLASHINFER_MOE_MXFP4_BF16", "0"))),
+    
+    "VLLM_ROCM_USE_AITER_FUSED_MOE_A16W4":
+    lambda: bool(int(os.getenv("VLLM_ROCM_USE_AITER_FUSED_MOE_A16W4", "0"))),
 
     # Control the cache sized used by the xgrammar compiler. The default
     # of 512 MB should be enough for roughly 1000 JSON schemas.

vllm/model_executor/layers/quantization/mxfp4.py

@@ -33,7 +33,7 @@
     from flashinfer import (mxfp8_quantize, shuffle_matrix_a,
                             shuffle_matrix_sf_a, trtllm_fp4_block_scale_moe)
 
-if current_platform.is_rocm():
+if current_platform.is_rocm() and envs.VLLM_ROCM_USE_AITER_FUSED_MOE_A16W4:
     import aiter
     from aiter.fused_moe import fused_topk, moe_sorting
     from aiter.ops.shuffle import shuffle_mxfp4_weight, shuffle_mxfp4_scale
@@ -374,7 +374,7 @@ def swap_every_two_rows(x, axis=-1):
             else:
                 num_warps = 8
 
-            if current_platform.is_rocm():
+            if current_platform.is_rocm() and envs.VLLM_ROCM_USE_AITER_FUSED_MOE_A16W4:
                 w13_aiter_weight = layer.w13_weight.contiguous()
                 w13_aiter_scale = layer.w13_weight_scale.contiguous()
                 w2_aiter_weight = layer.w2_weight.contiguous()
@@ -542,7 +542,7 @@ def apply(
             )[0]
             return trtllm_gen_output
         else:
-            if current_platform.is_rocm():
+            if current_platform.is_rocm() and envs.VLLM_ROCM_USE_AITER_FUSED_MOE_A16W4:
                 token_num = x.shape[0]
                 BLOCKM = 16 if token_num < 2048 else 32
                 topk_weights, topk_ids = fused_topk(x, router_logits, top_k, True)
@@ -566,8 +566,8 @@ def apply(
                     sorted_expert_ids,
                     num_valid_ids,
                     top_k,
-                    0, # n_pad_zeros
-                    0, # k_pad_zeros
+                    192, # n_pad_zeros
+                    128, # k_pad_zeros
                     None, # sorted_weights
                     None,
                     self.w13_scale_aiter_tensor,
@@ -582,8 +582,8 @@ def apply(
                     sorted_expert_ids,
                     num_valid_ids,
                     top_k,
-                    0, # n_pad_zeros
-                    0, # k_pad_zeros
+                    192, # n_pad_zeros
+                    128, # k_pad_zeros
                     sorted_weights, # sorted_weights
                     None,
                     self.w2_scale_aiter_tensor,