add aiter moe a16w4

vllm/envs.py

@@ -182,6 +182,8 @@
     VLLM_ROCM_USE_AITER_TRITON_FUSED_ADD_RMSNORM_PAD: bool = True
     VLLM_ROCM_USE_AITER_TRITON_BF16_GEMM: bool = True
     ROCM_TRITON_MOE_PRESHUFFLE_SCALES: bool = True
+    VLLM_ROCM_USE_AITER_FUSED_MOE_A16W4: bool = False
+
 def get_default_cache_root():
     return os.getenv(
         "XDG_CACHE_HOME",
@@ -1028,6 +1030,9 @@
     # 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.
@@ -1232,29 +1237,29 @@
     # Use AITER Triton fused RMSNORM + Quantization
     "VLLM_ROCM_USE_AITER_TRITON_FUSED_RMSNORM_FP8_QUANT":
     lambda: bool(int(os.getenv("VLLM_ROCM_USE_AITER_TRITON_FUSED_RMSNORM_FP8_QUANT", "1"))),
 
     # Use AITER Triton fused elementwise multiply + elementwise addtion
     "VLLM_ROCM_USE_AITER_TRITON_FUSED_MUL_ADD":
     lambda: bool(int(os.getenv("VLLM_ROCM_USE_AITER_TRITON_FUSED_MUL_ADD", "1"))),
 
     # Use AITER Triton fused rope + zeros + reshape_and_cache
     "VLLM_ROCM_USE_AITER_TRITON_FUSED_ROPE_ZEROS_KV_CACHE":
     lambda: bool(int(os.getenv("VLLM_ROCM_USE_AITER_TRITON_FUSED_ROPE_ZEROS_KV_CACHE", "1"))),
 
     # Use AITER Triton fused FP8 per-token group quant + FP8 batched GEMM
     "VLLM_ROCM_USE_AITER_TRITON_FP8_BMM":
     lambda: bool(int(os.getenv("VLLM_ROCM_USE_AITER_TRITON_FP8_BMM", "1"))),
     "VLLM_ROCM_USE_AITER_TRITON_FP8_BMM_MAX_BATCH_SIZE":
     lambda: int(os.getenv("VLLM_ROCM_USE_AITER_TRITON_FP8_BMM_MAX_BATCH_SIZE", 256)),
 
     # Use AITER Triton fused SiLU Mul + FP4 per-token group quant
     "VLLM_ROCM_USE_AITER_TRITON_SILU_MUL_FP4_QUANT":
     lambda: bool(int(os.getenv("VLLM_ROCM_USE_AITER_TRITON_SILU_MUL_FP4_QUANT", "0"))),
 
     # Use AITER Triton fused SiLU Mul + FP8 per-token group quant
     "VLLM_ROCM_USE_AITER_TRITON_SILU_MUL_FP8_QUANT":
     lambda: bool(int(os.getenv("VLLM_ROCM_USE_AITER_TRITON_SILU_MUL_FP8_QUANT", "1"))),
 
     # Use AITER Triton fused add + rmsnorm + padding
     "VLLM_ROCM_USE_AITER_TRITON_FUSED_ADD_RMSNORM_PAD":
     lambda: bool(int(os.getenv("VLLM_ROCM_USE_AITER_TRITON_FUSED_ADD_RMSNORM_PAD", "1"))),

vllm/model_executor/layers/quantization/mxfp4.py

@@ -61,7 +61,11 @@ def should_use_flashinfer_mxfp4():
     return (_should_use_flashinfer_mxfp4_mxfp8()
             or _should_use_flashinfer_mxfp4_bf16())
 
-
+if current_platform.is_rocm() and envs.VLLM_ROCM_USE_AITER 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
+
 class Mxfp4Config(QuantizationConfig):
 
     def __init__(self, ignored_layers: Optional[list[str]] = None):
@@ -389,7 +393,6 @@ def swap_every_two_rows(x, axis=-1):
                                       requires_grad=False)
         else:
             from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
-
             w13_bias = layer.w13_bias.to(torch.float32)
             w2_bias = layer.w2_bias.to(torch.float32)
 
@@ -403,18 +406,34 @@ def swap_every_two_rows(x, axis=-1):
             else:
                 num_warps = 8
 
-            w13_weight, w13_flex, w13_scale = _swizzle_mxfp4(
-                layer.w13_weight, layer.w13_weight_scale, num_warps)
-            w2_weight, w2_flex, w2_scale = _swizzle_mxfp4(
-                layer.w2_weight, layer.w2_weight_scale, num_warps)
-
-            self.w13_precision_config = PrecisionConfig(
-                weight_scale=w13_scale, flex_ctx=FlexCtx(rhs_data=w13_flex))
-            self.w2_precision_config = PrecisionConfig(
-                weight_scale=w2_scale, flex_ctx=FlexCtx(rhs_data=w2_flex))
-
-            self.w13_weight_triton_tensor = w13_weight
-            self.w2_weight_triton_tensor = w2_weight
+            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()
+                w2_aiter_scale = layer.w2_weight_scale.contiguous()
+
+                e, n, k = w13_aiter_weight.shape
+                w13_aiter_weight = w13_aiter_weight.view(e, n // 2, 2, k).permute(0, 2, 1, 3).contiguous().view(e, n, k)
+                w13_aiter_scale = w13_aiter_scale.view(e, n // 2, 2, -1).permute(0, 2, 1, 3).contiguous().view(e, n, -1)
+
+                self.w13_weight_aiter_tensor = shuffle_mxfp4_weight(w13_aiter_weight, 16, True)
+                self.w13_scale_aiter_tensor = shuffle_mxfp4_scale(w13_aiter_scale, True)
+                self.w2_weight_aiter_tensor = shuffle_mxfp4_weight(w2_aiter_weight, 16, False)
+                self.w2_scale_aiter_tensor = shuffle_mxfp4_scale(w2_aiter_scale, False)
+                self.w13_bias_aiter_tensor = layer.w13_bias.view(-1, n // 2, 2).permute(0, 2, 1).contiguous().view(-1, n)
+            else: 
+                w13_weight, w13_flex, w13_scale = _swizzle_mxfp4(
+                    layer.w13_weight, layer.w13_weight_scale, num_warps)
+                w2_weight, w2_flex, w2_scale = _swizzle_mxfp4(
+                    layer.w2_weight, layer.w2_weight_scale, num_warps)
+
+                self.w13_precision_config = PrecisionConfig(
+                    weight_scale=w13_scale, flex_ctx=FlexCtx(rhs_data=w13_flex))
+                self.w2_precision_config = PrecisionConfig(
+                    weight_scale=w2_scale, flex_ctx=FlexCtx(rhs_data=w2_flex))
+
+                self.w13_weight_triton_tensor = w13_weight
+                self.w2_weight_triton_tensor = w2_weight
 
             # need to delete the original weights to save memory on single GPU
             del layer.w13_weight
@@ -665,6 +684,56 @@ def apply(
             )[0]
             return trtllm_gen_output
         else:
+            if current_platform.is_rocm() and envs.VLLM_ROCM_USE_AITER 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)
+                sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, moe_out = moe_sorting(
+                    topk_ids,
+                    topk_weights,
+                    self.num_experts,
+                    x.shape[1],
+                    torch.bfloat16,
+                    BLOCKM
+                )
+                _, n1, k1 = self.w13_weight_aiter_tensor.shape
+                _, k2, n2 = self.w2_weight_aiter_tensor.shape
+                D = n2 if k2 == k1 else n2*2
+                cktile_moe_out1 = torch.empty((token_num, top_k, D), dtype=torch.bfloat16, device=x.device)
+                aiter.moe_cktile2stages_gemm1(
+                    x,
+                    self.w13_weight_aiter_tensor,
+                    cktile_moe_out1,
+                    sorted_ids,
+                    sorted_expert_ids,
+                    num_valid_ids,
+                    top_k,
+                    192, # n_pad_zeros
+                    128, # k_pad_zeros
+                    None, # sorted_weights
+                    None,
+                    self.w13_scale_aiter_tensor,
+                    self.w13_bias_aiter_tensor,
+                    BLOCKM, # block_size
+                )
+                aiter.moe_cktile2stages_gemm2(
+                    cktile_moe_out1,
+                    self.w2_weight_aiter_tensor,
+                    moe_out,
+                    sorted_ids,
+                    sorted_expert_ids,
+                    num_valid_ids,
+                    top_k,
+                    192, # n_pad_zeros
+                    128, # k_pad_zeros
+                    sorted_weights, # sorted_weights
+                    None,
+                    self.w2_scale_aiter_tensor,
+                    layer.w2_bias,
+                    BLOCKM, # block_size
+                )
+                return moe_out
+
             from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (  # noqa: E501
                 triton_kernel_moe_forward)
             return triton_kernel_moe_forward(
@@ -681,4 +750,4 @@ def apply(
                 w1_precision=self.w13_precision_config,
                 w2_precision=self.w2_precision_config,
                 apply_router_weight_on_input=apply_router_weight_on_input,
-            )
+            )