Repository: vllm-project/vllm
Pull request #17110: [FEAT] [ROCm]: Add AITER CK 2 Stages MoE support
Author: login: tjtanaa
Labels: name: rocm
description: Related to AMD ROCm
name: ready
description: ONLY add when PR is ready to merge/full CI is needed
Created: 2025-04-24T13:34:40Z | Merged: 2025-05-14T10:03:12Z

# Description

This PR introduces the use of `ck_moe_2stages` from `ROCm/aiter`.


## Performance gain

### meta-llama/Llama-4-Maverick-17B-128E-Instruct

V1 Engine
Input: Output = 1000:1000

| Metric | No AITER | AITER 1 Stage | AITER 2 Stage | % Gain (2 Stage vs Baseline) |
|--------|----------|---------------|---------------|------------------------------|
| Benchmark Duration (s) | 67.39 | 64.84 | 59.52 | 11.7% |
| Request Throughput (req/s) | 7.42 | 7.71 | 8.40 | 13.2% |
| Request Goodput (req/s) | 0.61 | 0.71 | 1.01 | 65.6% |
| Output Token Throughput (tok/s) | 6,135.14 | 6,437.38 | 6,898.15 | 12.4% |
| Total Token Throughput (tok/s) | 13,469.17 | 14,059.20 | 15,201.54 | 12.9% |
| Mean TTFT (ms) | 7,285.70 | 7,133.57 | 6,187.27 | 15.1% |
| Median TTFT (ms) | 7,231.37 | 6,962.17 | 6,006.79 | 16.9% |
| Mean TPOT (ms) | 52.90 | 50.15 | 47.16 | 10.9% |
| Median ITL (ms) | 34.30 | 31.46 | 29.33 | 14.5% |


### amd/Mixtral-8x7B-Instruct-v0.1-FP8-KV

V1 Engine
Input: Output = 1000:1000

| Metric | No AIter | AIter 2Stage | Improvement |
|--------|----------|--------------|-------------|
| Benchmark duration (s) | 80.38 | 76.48 | 4.85% faster |
| Total generated tokens | 436,389 | 482,066 | 10.47% more |
| Request throughput (req/s) | 6.22 | 6.54 | 5.14% higher |
| Output token throughput (tok/s) | 5,429.41 | 6,303.20 | 16.09% higher |
| Total token throughput (tok/s) | 11,578.38 | 12,765.36 | 10.25% higher |
| Median TTFT (ms) | 9,477.47 | 9,574.82 | 1.03% slower |
| Median TPOT (ms) | 50.20 | 46.73 | 6.91% faster |
| Median ITL (ms) | 39.45 | 36.85 | 6.59% faster |

### mistralai/Mixtral-8x7B-v0.1

V1 Engine
Input: Output = 1000:1000

| Metric | No AITER | AITER 1-Stage | AITER 2-Stage | % Gain (2-Stage vs No AITER) | % Gain (2-Stage vs 1-Stage) |
|--------|----------|---------------|---------------|------------------------------|------------------------------|
| Benchmark Duration (s) | 108.38 | 122.76 | 100.61 | 7.2% | 18.0% |
| Request Throughput (req/s) | 4.61 | 4.07 | 4.97 | 7.8% | 22.1% |
| Request Goodput (req/s) | 0.02 | 0.01 | 0.02 | 0.0% | 100.0% |
| Output Token Throughput (tok/s) | 4445.47 | 3907.71 | 4773.95 | 7.4% | 22.2% |
| Total Token Throughput (tok/s) | 9005.44 | 7933.76 | 9686.44 | 7.6% | 22.1% |
| Mean TTFT (ms) | 13133.12 | 20749.40 | 12677.77 | 3.5% | 38.9% |
| Median TTFT (ms) | 12911.29 | 20420.47 | 12429.12 | 3.7% | 39.1% |
| P99 TTFT (ms) | 27070.03 | 41767.87 | 25574.30 | 5.5% | 38.8% |
| Mean TPOT (ms) | 83.69 | 104.94 | 78.76 | 5.9% | 24.9% |
| Median TPOT (ms) | 64.66 | 73.77 | 61.07 | 5.6% | 17.2% |
| P99 TPOT (ms) | 399.80 | 642.01 | 393.91 | 1.5% | 38.6% |
| Mean ITL (ms) | 62.21 | 70.75 | 58.46 | 6.0% | 17.4% |
| Median ITL (ms) | 52.15 | 55.18 | 48.90 | 6.2% | 11.4% |
| P99 ITL (ms) | 403.70 | 645.63 | 396.17 | 1.9% | 38.6% |


# Lm_eval 
Dataset: GSM8K
Engine: V1

### AITER 1 stage amd/Mixtral-8x7B-Instruct-v0.1-FP8-KV fallback to 2 stage moe

vllm (pretrained=amd/Mixtral-8x7B-Instruct-v0.1-FP8-KV,tensor_parallel_size=2,max_model_len=10000,trust_remote_code=True), gen_kwargs: (None), limit: None, num_fewshot: 5, batch_size: auto
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.6528|±  |0.0131|
|     |       |strict-match    |     5|exact_match|↑  |0.6490|±  |0.0131|

### AITER 2 stage amd/Mixtral-8x7B-Instruct-v0.1-FP8-KV

[2025-04-30 02:18:36] INFO evaluation_tracker.py:272: Output path not provided, skipping saving results aggregated
vllm (pretrained=amd/Mixtral-8x7B-Instruct-v0.1-FP8-KV,tensor_parallel_size=2,max_model_len=10000,trust_remote_code=True), gen_kwargs: (None), limit: None, num_fewshot: 5, batch_size: auto
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.6422|±  |0.0132|
|     |       |strict-match    |     5|exact_match|↑  |0.6391|±  |0.0132|


### AITER 1 stage meta-llama/Llama-4-Scout-17B-16E-Instruct

vllm (pretrained=meta-llama/Llama-4-Scout-17B-16E-Instruct,tensor_parallel_size=8,max_model_len=10000,trust_remote_code=True), gen_kwargs: (None), limit: None, num_fewshot: 5, batch_size: auto
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.9143|±  |0.0077|
|     |       |strict-match    |     5|exact_match|↑  |0.8961|±  |0.0084|


### AITER 2 stage meta-llama/Llama-4-Scout-17B-16E-Instruct

vllm (pretrained=meta-llama/Llama-4-Scout-17B-16E-Instruct,tensor_parallel_size=8,max_model_len=10000,trust_remote_code=True), gen_kwargs: (None), limit: None, num_fewshot: 5, batch_size: auto
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.9227|±  |0.0074|
|     |       |strict-match    |     5|exact_match|↑  |0.9060|±  |0.0080|



### AITER 1 stage mistralai/Mixtral-8x7B-Instruct-v0.1 FP16

[2025-05-05 09:29:26] INFO evaluation_tracker.py:272: Output path not provided, skipping saving results aggregated
vllm (pretrained=mistralai/Mixtral-8x7B-Instruct-v0.1,tensor_parallel_size=2,max_model_len=10000,trust_remote_code=True), gen_kwargs: (None), limit: None, num_fewshot: 5, batch_size: auto
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.6429|±  |0.0132|
|     |       |strict-match    |     5|exact_match|↑  |0.6399|±  |0.0132|

### AITER 2 stage mistralai/Mixtral-8x7B-Instruct-v0.1FP16

vllm (pretrained=mistralai/Mixtral-8x7B-Instruct-v0.1,tensor_parallel_size=2,max_model_len=10000,trust_remote_code=True), gen_kwargs: (None), limit: None, num_fewshot: 5, batch_size: auto
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.6513|±  |0.0131|
|     |       |strict-match    |     5|exact_match|↑  |0.6475|±  |0.0132|

### AITER 1 stage mistralai/Mixtral-8x7B-Instruct-v0.1 Dynamic FP8

vllm (pretrained=mistralai/Mixtral-8x7B-Instruct-v0.1,tensor_parallel_size=2,max_model_len=10000,quantization=fp8,trust_remote_code=True), gen_kwargs: (None), limit: None, num_fewshot: 5, batch_size: auto
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.5777|±  |0.0136|
|     |       |strict-match    |     5|exact_match|↑  |0.5701|±  |0.0136|

### AITER 2 stage mistralai/Mixtral-8x7B-Instruct-v0.1 Dynamic FP8 (fall back to 1 stage)

vllm (pretrained=mistralai/Mixtral-8x7B-Instruct-v0.1,tensor_parallel_size=2,max_model_len=10000,quantization=fp8,trust_remote_code=True), gen_kwargs: (None), limit: None, num_fewshot: 5, batch_size: auto
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.5792|±  |0.0136|
|     |       |strict-match    |     5|exact_match|↑  |0.5701|±  |0.0136|

## Updated unit tests
Given that the `rocm_aiter_fused_experts` has been removed from the function `dispatch_fused_experts_func` in `vllm/model_executor/layers/fused_moe/fused_moe.py`.



## Future TODO:
```
# TODO: remove @cache to allow the use case
# where there are multiple instances of LLM
# running on the same process/program with different
# configurations. Blockage is that dispatch_topk_func()
# in vllm/model_executor/layers/fused_moe/fused_moe.py
# uses this function in critical path.
@cache
def is_rocm_aiter_moe_enabled() -> bool:
    return current_platform.is_rocm() \
        and envs.VLLM_ROCM_USE_AITER_MOE \
        and envs.VLLM_ROCM_USE_AITER
```


<!--- pyml disable-next-line no-emphasis-as-heading -->



Diff:
diff --git a/tests/kernels/moe/test_moe.py b/tests/kernels/moe/test_moe.py
index c1d0940f26cb..96b090136e3c 100644
--- a/tests/kernels/moe/test_moe.py
+++ b/tests/kernels/moe/test_moe.py
@@ -224,9 +224,16 @@ def test_mixtral_moe(dtype: torch.dtype, padding: bool, use_rocm_aiter: bool,
     """Make sure our Mixtral MoE implementation agrees with the one from
     huggingface."""
 
+    # clear the cache before every test
+    from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
+        is_rocm_aiter_moe_enabled)
+    is_rocm_aiter_moe_enabled.cache_clear()
     if use_rocm_aiter:
         monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
 
+        if dtype == torch.float32:
+            pytest.skip("AITER ROCm test skip for float32")
+
     # Instantiate our and huggingface's MoE blocks
     config = MixtralConfig()
     hf_moe = MixtralSparseMoeBlock(config).to(dtype).to("cuda")
diff --git a/tests/model_executor/test_enabled_custom_ops.py b/tests/model_executor/test_enabled_custom_ops.py
index 93453ddb657c..e957db5b3f16 100644
--- a/tests/model_executor/test_enabled_custom_ops.py
+++ b/tests/model_executor/test_enabled_custom_ops.py
@@ -8,10 +8,8 @@
 from vllm.model_executor.layers.activation import (GeluAndMul,
                                                    ReLUSquaredActivation,
                                                    SiluAndMul)
-from vllm.model_executor.layers.fused_moe.fused_moe import (
-    dispatch_fused_experts_func, dispatch_topk_func,
-    torch_vllm_inplace_fused_experts, torch_vllm_outplace_fused_experts,
-    vllm_topk_softmax)
+from vllm.model_executor.layers.fused_moe.fused_moe import (dispatch_topk_func,
+                                                            vllm_topk_softmax)
 from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
     is_rocm_aiter_moe_enabled)
 from vllm.model_executor.layers.layernorm import (
@@ -142,24 +140,6 @@ def test_topk_dispatch(use_rocm_aiter: str, monkeypatch):
         assert topk_func == vllm_topk_softmax
 
 
-@pytest.mark.parametrize("use_rocm_aiter", ["0", "1"])
-@pytest.mark.parametrize("inplace", [True, False])
-def test_fused_experts_dispatch(use_rocm_aiter: str, inplace: bool,
-                                monkeypatch):
-
-    monkeypatch.setenv("VLLM_ROCM_USE_AITER", use_rocm_aiter)
-    is_rocm_aiter_moe_enabled.cache_clear()
-    fused_experts_func = dispatch_fused_experts_func(inplace)
-    if current_platform.is_rocm() and int(use_rocm_aiter):
-        from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
-            rocm_aiter_fused_experts)
-        assert fused_experts_func == rocm_aiter_fused_experts
-    elif inplace:
-        assert fused_experts_func == torch_vllm_inplace_fused_experts
-    else:
-        assert fused_experts_func == torch_vllm_outplace_fused_experts
-
-
 @pytest.mark.parametrize("add_residual", [True, False])
 @pytest.mark.parametrize("use_rocm_aiter", ["0", "1"])
 @pytest.mark.parametrize("use_rocm_aiter_norm", ["0", "1"])
diff --git a/vllm/model_executor/layers/fused_moe/fused_moe.py b/vllm/model_executor/layers/fused_moe/fused_moe.py
index 8c28cedbcd77..7bf4243305ac 100644
--- a/vllm/model_executor/layers/fused_moe/fused_moe.py
+++ b/vllm/model_executor/layers/fused_moe/fused_moe.py
@@ -1100,9 +1100,6 @@ def torch_vllm_outplace_fused_experts(**kwargs) -> torch.Tensor:
 
 
 def dispatch_fused_experts_func(inplace: bool) -> Callable[..., torch.Tensor]:
-    if is_rocm_aiter_moe_enabled():
-        from .rocm_aiter_fused_moe import rocm_aiter_fused_experts
-        return rocm_aiter_fused_experts
     if inplace:
         return torch_vllm_inplace_fused_experts
     return torch_vllm_outplace_fused_experts
diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py
index f74e38bde6b8..14f360e3bbf3 100644
--- a/vllm/model_executor/layers/fused_moe/layer.py
+++ b/vllm/model_executor/layers/fused_moe/layer.py
@@ -84,6 +84,16 @@ def apply(
 class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
     """MoE method without quantization."""
 
+    def __init__(self):
+        super().__init__()
+
+        self.rocm_aiter_moe_enabled = is_rocm_aiter_moe_enabled()
+        if self.rocm_aiter_moe_enabled:
+            from .rocm_aiter_fused_moe import rocm_aiter_fused_experts
+            self.rocm_aiter_fused_experts = rocm_aiter_fused_experts
+        else:
+            self.rocm_aiter_fused_experts = None  # type: ignore
+
     def create_weights(self, layer: torch.nn.Module, num_experts: int,
                        hidden_size: int, intermediate_size_per_partition: int,
                        params_dtype: torch.dtype, **extra_weight_attrs):
@@ -126,11 +136,13 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
         layer.w2_weight.data = self._maybe_pad_weight(layer.w2_weight.data)
         # Lazy import to avoid importing triton.
         from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
-            is_rocm_aiter_moe_enabled, shuffle_weights)
-        if is_rocm_aiter_moe_enabled():
-            # reshaping weights is required for aiter moe kernel.
-            shuffled_w13, shuffled_w2 = shuffle_weights(
-                layer.w13_weight.data, layer.w2_weight.data)
+            shuffle_weights)
+
+        if self.rocm_aiter_moe_enabled:
+            # use 2stage ck moe layout
+            shuffled_w13, shuffled_w2 = shuffle_weights(layer.w13_weight.data,
+                                                        layer.w2_weight.data,
+                                                        layout=(32, 32))
 
             layer.w13_weight.data = shuffled_w13
             layer.w2_weight.data = shuffled_w2
@@ -211,6 +223,16 @@ def forward_cuda(
             scoring_func=scoring_func,
             e_score_correction_bias=e_score_correction_bias)
 
+        if self.rocm_aiter_moe_enabled:
+            return self.rocm_aiter_fused_experts(
+                hidden_states=x,
+                w1=layer.w13_weight,
+                w2=layer.w2_weight,
+                topk_weights=topk_weights,
+                topk_ids=topk_ids,
+                activation=activation,
+                apply_router_weight_on_input=apply_router_weight_on_input)
+
         return fused_experts(
             hidden_states=x,
             w1=layer.w13_weight,
diff --git a/vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py b/vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py
index 02e75402f4c8..1bf12395afee 100644
--- a/vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py
+++ b/vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py
@@ -20,7 +20,7 @@ def rocm_aiter_asm_moe_tkw1_impl(
         hidden_states: torch.Tensor,
         w1: torch.Tensor,
         w2: torch.Tensor,
-        topk_weight: torch.Tensor,
+        topk_weights: torch.Tensor,
         topk_ids: torch.Tensor,
         fc1_scale: Optional[torch.Tensor] = None,
         fc2_scale: Optional[torch.Tensor] = None,
@@ -40,7 +40,7 @@ def rocm_aiter_asm_moe_tkw1_impl(
     return asm_moe_tkw1(hidden_states,
                         w1,
                         w2,
-                        topk_weight,
+                        topk_weights,
                         topk_ids,
                         fc1_scale=fc1_scale,
                         fc2_scale=fc2_scale,
@@ -56,7 +56,7 @@ def rocm_aiter_asm_moe_tkw1_fake(
         hidden_states: torch.Tensor,
         w1: torch.Tensor,
         w2: torch.Tensor,
-        topk_weight: torch.Tensor,
+        topk_weights: torch.Tensor,
         topk_ids: torch.Tensor,
         fc1_scale: Optional[torch.Tensor] = None,
         fc2_scale: Optional[torch.Tensor] = None,
@@ -69,23 +69,6 @@ def rocm_aiter_asm_moe_tkw1_fake(
     return torch.empty_like(hidden_states)
 
 
-def rocm_aiter_ck_moe_impl(hidden_states: torch.Tensor, w1: torch.Tensor,
-                           w2: torch.Tensor, topk_weights: torch.Tensor,
-                           topk_ids: torch.Tensor) -> torch.Tensor:
-    from aiter import ck_moe
-    return ck_moe(hidden_states=hidden_states,
-                  w1=w1,
-                  w2=w2,
-                  topk_weights=topk_weights,
-                  topk_ids=topk_ids)
-
-
-def rocm_aiter_ck_moe_fake(hidden_states: torch.Tensor, w1: torch.Tensor,
-                           w2: torch.Tensor, topk_weights: torch.Tensor,
-                           topk_ids: torch.Tensor) -> torch.Tensor:
-    return torch.empty_like(hidden_states)
-
-
 def rocm_aiter_fmoe_fp8_blockscale_g1u1_impl(
         topk_ids: torch.Tensor,
         topk_weights: torch.Tensor,
@@ -151,7 +134,7 @@ def rocm_aiter_fmoe_fp8_blockscale_g1u1_fake(
 def rocm_aiter_asm_moe_impl(hidden_states: torch.Tensor,
                             w1: torch.Tensor,
                             w2: torch.Tensor,
-                            topk_weight: torch.Tensor,
+                            topk_weights: torch.Tensor,
                             topk_ids: torch.Tensor,
                             fc1_scale: Optional[torch.Tensor] = None,
                             fc2_scale: Optional[torch.Tensor] = None,
@@ -174,7 +157,7 @@ def rocm_aiter_asm_moe_impl(hidden_states: torch.Tensor,
     return rocm_aiter_asm_fmoe.asm_moe(hidden_states=hidden_states,
                                        w1=w1,
                                        w2=w2,
-                                       topk_weight=topk_weight,
+                                       topk_weight=topk_weights,
                                        topk_ids=topk_ids,
                                        fc1_scale=fc1_scale,
                                        fc2_scale=fc2_scale,
@@ -187,7 +170,7 @@ def rocm_aiter_asm_moe_impl(hidden_states: torch.Tensor,
 def rocm_aiter_asm_moe_fake(hidden_states: torch.Tensor,
                             w1: torch.Tensor,
                             w2: torch.Tensor,
-                            topk_weight: torch.Tensor,
+                            topk_weights: torch.Tensor,
                             topk_ids: torch.Tensor,
                             fc1_scale: Optional[torch.Tensor] = None,
                             fc2_scale: Optional[torch.Tensor] = None,
@@ -198,6 +181,49 @@ def rocm_aiter_asm_moe_fake(hidden_states: torch.Tensor,
     return torch.empty_like(hidden_states)
 
 
+def rocm_aiter_ck_moe_2stages_impl(
+    hidden_states: torch.Tensor,
+    w1: torch.Tensor,
+    w2: torch.Tensor,
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    fc1_scale: Optional[torch.Tensor] = None,
+    fc2_scale: Optional[torch.Tensor] = None,
+    a1_scale: Optional[torch.Tensor] = None,
+    a2_scale: Optional[torch.Tensor] = None,
+    block_size: Optional[list[int]] = None,
+    expert_mask: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    from aiter.fused_moe_bf16_asm import ck_moe_2stages
+    return ck_moe_2stages(a1=hidden_states,
+                          w1=w1,
+                          w2=w2,
+                          topk_weight=topk_weights,
+                          topk_ids=topk_ids,
+                          fc1_scale=fc1_scale,
+                          fc2_scale=fc2_scale,
+                          a1_scale=a1_scale,
+                          a2_scale=a2_scale,
+                          block_size=block_size,
+                          expert_mask=expert_mask)
+
+
+def rocm_aiter_ck_moe_2stages_fake(
+    hidden_states: torch.Tensor,
+    w1: torch.Tensor,
+    w2: torch.Tensor,
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    fc1_scale: Optional[torch.Tensor] = None,
+    fc2_scale: Optional[torch.Tensor] = None,
+    a1_scale: Optional[torch.Tensor] = None,
+    a2_scale: Optional[torch.Tensor] = None,
+    block_size: Optional[list[int]] = None,
+    expert_mask: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    return torch.empty_like(hidden_states)
+
+
 def rocm_aiter_topk_softmax_impl(topk_weights: torch.Tensor,
                                  topk_indices: torch.Tensor,
                                  token_expert_indices: torch.Tensor,
@@ -257,14 +283,6 @@ def rocm_aiter_biased_grouped_topk_fake(
         dispatch_key=current_platform.dispatch_key,
     )
 
-    direct_register_custom_op(
-        op_name="rocm_aiter_ck_moe",
-        op_func=rocm_aiter_ck_moe_impl,
-        mutates_args=[],
-        fake_impl=rocm_aiter_ck_moe_fake,
-        dispatch_key=current_platform.dispatch_key,
-    )
-
     direct_register_custom_op(
         op_name="rocm_aiter_fmoe_fp8_blockscale_g1u1",
         op_func=rocm_aiter_fmoe_fp8_blockscale_g1u1_impl,
@@ -281,6 +299,14 @@ def rocm_aiter_biased_grouped_topk_fake(
         dispatch_key=current_platform.dispatch_key,
     )
 
+    direct_register_custom_op(
+        op_name="rocm_aiter_ck_moe_2stages",
+        op_func=rocm_aiter_ck_moe_2stages_impl,
+        mutates_args=[],
+        fake_impl=rocm_aiter_ck_moe_2stages_fake,
+        dispatch_key=current_platform.dispatch_key,
+    )
+
     direct_register_custom_op(
         op_name="rocm_aiter_topk_softmax",
         op_func=rocm_aiter_topk_softmax_impl,
@@ -330,29 +356,21 @@ def rocm_aiter_biased_group_topk(
     return topk_weights, topk_ids
 
 
-def rocm_aiter_fused_experts(hidden_states: torch.Tensor,
-                             w1: torch.Tensor,
-                             w2: torch.Tensor,
-                             topk_weights: torch.Tensor,
-                             topk_ids: torch.Tensor,
-                             inplace: bool = False,
-                             activation: str = "silu",
-                             apply_router_weight_on_input: bool = False,
-                             use_fp8_w8a8: bool = False,
-                             use_int8_w8a8: bool = False,
-                             use_int8_w8a16: bool = False,
-                             use_int4_w4a16: bool = False,
-                             per_channel_quant: bool = False,
-                             global_num_experts: int = -1,
-                             expert_map: Optional[torch.Tensor] = None,
-                             w1_scale: Optional[torch.Tensor] = None,
-                             w2_scale: Optional[torch.Tensor] = None,
-                             w1_zp: Optional[torch.Tensor] = None,
-                             w2_zp: Optional[torch.Tensor] = None,
-                             a1_scale: Optional[torch.Tensor] = None,
-                             a2_scale: Optional[torch.Tensor] = None,
-                             block_shape: Optional[list[int]] = None,
-                             allow_deep_gemm: bool = False) -> torch.Tensor:
+def rocm_aiter_fused_experts(
+        hidden_states: torch.Tensor,
+        w1: torch.Tensor,
+        w2: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        activation: str = "silu",
+        apply_router_weight_on_input: bool = False,
+        use_fp8_w8a8: bool = False,
+        per_channel_quant: bool = False,
+        w1_scale: Optional[torch.Tensor] = None,
+        w2_scale: Optional[torch.Tensor] = None,
+        a1_scale: Optional[torch.Tensor] = None,
+        a2_scale: Optional[torch.Tensor] = None,
+        block_shape: Optional[list[int]] = None) -> torch.Tensor:
 
     from vllm.model_executor.layers.quantization.utils.fp8_utils import (
         per_token_group_quant_fp8)
@@ -375,8 +393,8 @@ def rocm_aiter_fused_experts(hidden_states: torch.Tensor,
         a1, a1_scale = per_token_group_quant_fp8(hidden_states, block_shape[1])
 
         return torch.ops.vllm.rocm_aiter_fmoe_fp8_blockscale_g1u1(
-            topk_ids, topk_weights, hidden_states.dtype, expert_map, a1, w1,
-            w2, w1_scale, w2_scale, a1_scale, block_shape, None)
+            topk_ids, topk_weights, hidden_states.dtype, None, a1, w1, w2,
+            w1_scale, w2_scale, a1_scale, block_shape, None)
 
     # w8a8 per-channel quantization
     elif per_channel_quant and apply_router_weight_on_input and use_fp8_w8a8:
@@ -401,17 +419,36 @@ def rocm_aiter_fused_experts(hidden_states: torch.Tensor,
             fc2_smooth_scale=None,
             a16=False,
             per_tensor_quant_scale=None,
-            expert_mask=expert_map,
+            expert_mask=None,
             activation_str=activation)
 
     # w8a8 per-tensor activation per-tensor weight
     elif use_fp8_w8a8:
         assert not apply_router_weight_on_input, (
             "apply_router_weight_on_input is not supported for fp8_w8a8")
+
+        # - faster static per-tensor-activation static per-tensor-weight
+        #   fp8 quantization w8a8
+        if a1_scale is not None and a2_scale is not None:
+            return torch.ops.vllm.rocm_aiter_ck_moe_2stages(
+                hidden_states=hidden_states,
+                w1=w1,
+                w2=w2,
+                topk_weights=topk_weights,
+                topk_ids=topk_ids,
+                fc1_scale=w1_scale,
+                fc2_scale=w2_scale,
+                a1_scale=a1_scale,
+                a2_scale=a2_scale)
+
+        # - fallback static per-tensor-activation static per-tensor-weight
+        #   fp8 quantization w8a8
+        # - dynamic per-tensor activation static per-tensor-weight
+        #   fp8 quantization w8a8
         return torch.ops.vllm.rocm_aiter_asm_moe(hidden_states=hidden_states,
                                                  w1=w1,
                                                  w2=w2,
-                                                 topk_weight=topk_weights,
+                                                 topk_weights=topk_weights,
                                                  topk_ids=topk_ids,
                                                  fc1_scale=w1_scale,
                                                  fc2_scale=w2_scale,
@@ -431,12 +468,12 @@ def rocm_aiter_fused_experts(hidden_states: torch.Tensor,
         topk_ids = topk_ids.to(torch.int32)
         topk_weights = torch.ones_like(topk_weights, dtype=torch.float32)
 
-    # w16a16 fallback to rocm_aiter_ck_moe w16a16
-    return torch.ops.vllm.rocm_aiter_ck_moe(hidden_states=hidden_states,
-                                            w1=w1,
-                                            w2=w2,
-                                            topk_weights=topk_weights,
-                                            topk_ids=topk_ids)
+    return torch.ops.vllm.rocm_aiter_ck_moe_2stages(
+        hidden_states=hidden_states,
+        w1=w1,
+        w2=w2,
+        topk_weights=topk_weights,
+        topk_ids=topk_ids)
 
 
 def rocm_aiter_topk_softmax(topk_weights: torch.Tensor,
@@ -450,7 +487,8 @@ def rocm_aiter_topk_softmax(topk_weights: torch.Tensor,
     return topk_weights, topk_indices
 
 
-def shuffle_weights(*tensors: torch.Tensor) -> tuple[torch.Tensor, ...]:
+def shuffle_weights(*tensors: torch.Tensor,
+                    layout: tuple[int, int]) -> tuple[torch.Tensor, ...]:
     """
     Applies shuffle_weight function from AITER to each 
     input tensor and returns them.
@@ -462,7 +500,8 @@ def shuffle_weights(*tensors: torch.Tensor) -> tuple[torch.Tensor, ...]:
     A Tuple of shuffled tensors.
     """
     from aiter.ops.shuffle import shuffle_weight
-    return tuple(shuffle_weight(tensor) for tensor in tensors)
+
+    return tuple(shuffle_weight(tensor, layout=layout) for tensor in tensors)
 
 
 def expand_weights(*tensors: torch.Tensor,
@@ -484,4 +523,4 @@ def expand_weights(*tensors: torch.Tensor,
 
     return tuple(
         tensor.unsqueeze(-1).unsqueeze(-1).expand((-1, dim, -1))
-        for tensor, dim in zip(tensors, expansion_dims))
+        for tensor, dim in zip(tensors, expansion_dims))
\ No newline at end of file
diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
index a74f1f7233af..fa0067c44802 100644
--- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
@@ -125,6 +125,10 @@ def __init__(
         # Disable marlin for rocm
         if current_platform.is_rocm():
             self.use_marlin = False
+        from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
+            is_rocm_aiter_moe_enabled)
+
+        self.rocm_aiter_moe_enabled = is_rocm_aiter_moe_enabled()
 
     def create_weights(self, layer: torch.nn.Module, num_experts: int,
                        hidden_size: int, intermediate_size_per_partition: int,
@@ -276,24 +280,22 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
             layer.w13_weight_scale = torch.nn.Parameter(max_w13_scales,
                                                         requires_grad=False)
 
-        from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
-            is_rocm_aiter_moe_enabled)
-
         # Property to determine if AITER is used
-        if is_rocm_aiter_moe_enabled():
+        if self.rocm_aiter_moe_enabled:
             from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (  # noqa E501
                 rocm_aiter_fused_experts, shuffle_weights)
 
             # reshaping weights is required for aiter moe kernel.
-            shuffled_w13, shuffled_w2 = shuffle_weights(
-                layer.w13_weight.data, layer.w2_weight.data)
+            shuffled_w13, shuffled_w2 = shuffle_weights(layer.w13_weight.data,
+                                                        layer.w2_weight.data,
+                                                        layout=(16, 16))
 
             layer.w13_weight = torch.nn.Parameter(shuffled_w13,
                                                   requires_grad=False)
             layer.w2_weight = torch.nn.Parameter(shuffled_w2,
                                                  requires_grad=False)
 
-            self.fused_experts_func = rocm_aiter_fused_experts
+            self.rocm_aiter_fused_experts_func = rocm_aiter_fused_experts
         else:
             from vllm.model_executor.layers.fused_moe import fused_experts
             self.fused_experts_func = fused_experts
@@ -335,6 +337,22 @@ def apply(
             scoring_func=scoring_func,
             e_score_correction_bias=e_score_correction_bias)
 
+        if self.rocm_aiter_moe_enabled:
+            return self.rocm_aiter_fused_experts_func(
+                hidden_states=x,
+                w1=layer.w13_weight,
+                w2=layer.w2_weight,
+                topk_weights=topk_weights,
+                topk_ids=topk_ids,
+                activation=activation,
+                apply_router_weight_on_input=apply_router_weight_on_input,
+                use_fp8_w8a8=True,
+                per_channel_quant=self.weight_quant.strategy ==
+                QuantizationStrategy.CHANNEL,
+                w1_scale=layer.w13_weight_scale,
+                w2_scale=layer.w2_weight_scale,
+                a1_scale=layer.w13_input_scale,
+                a2_scale=layer.w2_input_scale)
         if self.use_marlin:
             assert activation == "silu", (
                 f"{activation} not supported for Marlin MoE.")
diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py
index 589ca7bed329..cfd398c07fb9 100644
--- a/vllm/model_executor/layers/quantization/fp8.py
+++ b/vllm/model_executor/layers/quantization/fp8.py
@@ -591,6 +591,8 @@ def process_weights_after_loading(self, layer: Module) -> None:
         from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
             expand_weights, is_rocm_aiter_moe_enabled, shuffle_weights)
 
+        self.rocm_aiter_moe_enabled = is_rocm_aiter_moe_enabled()
+
         # TODO (rob): refactor block quant into separate class.
         if self.block_quant:
             assert self.quant_config.activation_scheme == "dynamic"
@@ -616,10 +618,12 @@ def process_weights_after_loading(self, layer: Module) -> None:
             layer.w2_weight = Parameter(w2_weight, requires_grad=False)
             layer.w2_weight_scale_inv = Parameter(w2_weight_scale_inv,
                                                   requires_grad=False)
-            if is_rocm_aiter_moe_enabled():
+            if self.rocm_aiter_moe_enabled:
                 # reshaping weights is required for aiter moe kernel.
                 shuffled_w13, shuffled_w2 = shuffle_weights(
-                    layer.w13_weight.data, layer.w2_weight.data)
+                    layer.w13_weight.data,
+                    layer.w2_weight.data,
+                    layout=(16, 16))
 
                 layer.w13_weight = torch.nn.Parameter(shuffled_w13,
                                                       requires_grad=False)
@@ -663,7 +667,7 @@ def process_weights_after_loading(self, layer: Module) -> None:
                                                   requires_grad=False)
             layer.w2_weight = torch.nn.Parameter(w2_weight,
                                                  requires_grad=False)
-            if is_rocm_aiter_moe_enabled():
+            if self.rocm_aiter_moe_enabled:
                 # reshaping weights is required for aiter moe kernel.
                 w13_scales, w2_scales = expand_weights(
                     layer.w13_weight_scale.data,
@@ -676,8 +680,9 @@ def process_weights_after_loading(self, layer: Module) -> None:
                 layer.w2_weight_scale = torch.nn.Parameter(
                     w2_scales.contiguous(), requires_grad=False)
 
-                shuffled_w13, shuffled_w2 = shuffle_weights(
-                    layer.w13_weight, layer.w2_weight)
+                shuffled_w13, shuffled_w2 = shuffle_weights(layer.w13_weight,
+                                                            layer.w2_weight,
+                                                            layout=(16, 16))
 
                 layer.w13_weight = torch.nn.Parameter(shuffled_w13,
                                                       requires_grad=False)
@@ -748,7 +753,7 @@ def process_weights_after_loading(self, layer: Module) -> None:
                             dq_weight, max_w13_scales[expert_id])
                     start += shard_size
 
-            if is_rocm_aiter_moe_enabled():
+            if self.rocm_aiter_moe_enabled:
                 # reshaping weights is required for aiter moe kernel.
                 expansion_dims = [
                     layer.w13_weight.shape[1], layer.w2_weight.shape[1]
@@ -760,8 +765,9 @@ def process_weights_after_loading(self, layer: Module) -> None:
                 layer.w2_weight_scale = torch.nn.Parameter(
                     w2_scales.contiguous(), requires_grad=False)
 
-                shuffled_w13, shuffled_w2 = shuffle_weights(
-                    layer.w13_weight, layer.w2_weight)
+                shuffled_w13, shuffled_w2 = shuffle_weights(layer.w13_weight,
+                                                            layer.w2_weight,
+                                                            layout=(32, 32))
 
                 layer.w13_weight = torch.nn.Parameter(shuffled_w13,
                                                       requires_grad=False)
@@ -796,6 +802,8 @@ def apply(
         activation: str = "silu",
     ) -> torch.Tensor:
         from vllm.model_executor.layers.fused_moe import fused_experts
+        from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
+            rocm_aiter_fused_experts)
 
         topk_weights, topk_ids = FusedMoE.select_experts(
             hidden_states=x,
@@ -810,6 +818,24 @@ def apply(
             e_score_correction_bias=e_score_correction_bias,
         )
 
+        if self.rocm_aiter_moe_enabled:
+            return rocm_aiter_fused_experts(
+                x,
+                layer.w13_weight,
+                layer.w2_weight,
+                topk_weights=topk_weights,
+                topk_ids=topk_ids,
+                activation=activation,
+                use_fp8_w8a8=True,
+                apply_router_weight_on_input=apply_router_weight_on_input,
+                w1_scale=(layer.w13_weight_scale_inv
+                          if self.block_quant else layer.w13_weight_scale),
+                w2_scale=(layer.w2_weight_scale_inv
+                          if self.block_quant else layer.w2_weight_scale),
+                a1_scale=layer.w13_input_scale,
+                a2_scale=layer.w2_input_scale,
+                block_shape=self.quant_config.weight_block_size)
+
         if self.use_marlin:
             assert activation == "silu", (
                 f"{activation} not supported for Marlin MoE.")


Reviews:
user: login: SageMoore
body: In general I would like us to contain all of the 1 stage vs 2 stage dispatching logic in the `rocm_aiter_fused_experts`  function if posible
state: CHANGES_REQUESTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-08T16:16:21Z
user: login: tjtanaa
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-08T16:34:18Z
user: login: hongxiayang
body: Verified with llama4 bf16 128e model.
LGTM in general. Agreed the environment variable part can be simplified.
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-05-08T22:10:04Z
user: login: hongxiayang
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-08T22:10:35Z
user: login: hongxiayang
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-08T22:15:04Z
user: login: tjtanaa
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-09T07:15:48Z
user: login: tjtanaa
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-10T06:34:39Z
user: login: tjtanaa
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-10T06:35:57Z
user: login: tjtanaa
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-10T11:53:35Z
user: login: SageMoore
body: I'd like to propose an alternative implementation. Instead of passing this new environment variable around, let's just keep the existing `VLLM_ROCM_USE_AITER_MOE` variable and add heuristics to the code that decide when to use the one stage vs the two stage implementation.

It looks like the two stage kernel is generally faster but doesn't support dynamic quantization? Let's rework this PR so that we always use the two stage implementation when it's supported and fallback to the one stage implementation when it isn't. We don't need any logging when this fallback occurs.

If we make the change that I suggested, it looks like you can simplify the changes in fp8.py down to just some layout changes for the weights. fused_moe.py should get a bit simpler as well.

What do you think?
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-12T20:27:37Z
user: login: SageMoore
body: Looks reasonable. Thanks for cleaning up the dispatching logic!
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-13T13:32:22Z

Review comments:
diff_hunk: @@ -581,7 +581,11 @@ def create_weights(self, layer: Module, num_experts: int, hidden_size: int,
     def process_weights_after_loading(self, layer: Module) -> None:
         # Lazy import to avoid importing triton too early.
         from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
-            expand_weights, is_rocm_aiter_moe_enabled, shuffle_weights)
+            expand_weights, is_rocm_aiter_2stage_moe_enabled,
+            is_rocm_aiter_moe_enabled, shuffle_weights)
+
+        self.rocm_aiter_moe_enabled = is_rocm_aiter_moe_enabled()
+        self.rocm_aiter_2stage_moe_enabled = is_rocm_aiter_2stage_moe_enabled()
path: vllm/model_executor/layers/quantization/fp8.py
user: login: SageMoore
body: It looks like this variable is only used to warn that we are falling back to a different implementation. Let's remove it to simplify the code.
created_at: 2025-05-08T15:59:22Z
author_association: CONTRIBUTOR
diff_hunk: @@ -259,29 +318,32 @@ def rocm_aiter_topk_softmax_fake(topk_weights: torch.Tensor,
     )
 
 
-def rocm_aiter_fused_experts(hidden_states: torch.Tensor,
-                             w1: torch.Tensor,
-                             w2: torch.Tensor,
-                             topk_weights: torch.Tensor,
-                             topk_ids: torch.Tensor,
-                             inplace: bool = False,
-                             activation: str = "silu",
-                             apply_router_weight_on_input: bool = False,
-                             use_fp8_w8a8: bool = False,
-                             use_int8_w8a8: bool = False,
-                             use_int8_w8a16: bool = False,
-                             use_int4_w4a16: bool = False,
-                             per_channel_quant: bool = False,
-                             global_num_experts: int = -1,
-                             expert_map: Optional[torch.Tensor] = None,
-                             w1_scale: Optional[torch.Tensor] = None,
-                             w2_scale: Optional[torch.Tensor] = None,
-                             w1_zp: Optional[torch.Tensor] = None,
-                             w2_zp: Optional[torch.Tensor] = None,
-                             a1_scale: Optional[torch.Tensor] = None,
-                             a2_scale: Optional[torch.Tensor] = None,
-                             block_shape: Optional[List[int]] = None,
-                             allow_deep_gemm: bool = False) -> torch.Tensor:
+def rocm_aiter_fused_experts(
+    hidden_states: torch.Tensor,
+    w1: torch.Tensor,
+    w2: torch.Tensor,
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    inplace: bool = False,
+    activation: str = "silu",
+    apply_router_weight_on_input: bool = False,
+    use_fp8_w8a8: bool = False,
+    use_int8_w8a8: bool = False,
+    use_int8_w8a16: bool = False,
+    use_int4_w4a16: bool = False,
+    per_channel_quant: bool = False,
+    global_num_experts: int = -1,
+    expert_map: Optional[torch.Tensor] = None,
+    w1_scale: Optional[torch.Tensor] = None,
+    w2_scale: Optional[torch.Tensor] = None,
+    w1_zp: Optional[torch.Tensor] = None,
+    w2_zp: Optional[torch.Tensor] = None,
+    a1_scale: Optional[torch.Tensor] = None,
+    a2_scale: Optional[torch.Tensor] = None,
+    block_shape: Optional[List[int]] = None,
+    allow_deep_gemm: bool = False,
+    use_ck_moe_2stages: bool = False,
path: vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py
user: login: SageMoore
body: Instead of passing this boolean around, can you check the environment variable here?
created_at: 2025-05-08T16:01:27Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1129,7 +1129,8 @@ def fused_experts(hidden_states: torch.Tensor,
                   a1_scale: Optional[torch.Tensor] = None,
                   a2_scale: Optional[torch.Tensor] = None,
                   block_shape: Optional[List[int]] = None,
-                  allow_deep_gemm: bool = False) -> torch.Tensor:
+                  allow_deep_gemm: bool = False,
+                  use_ck_moe_2stages: bool = False) -> torch.Tensor:
path: vllm/model_executor/layers/fused_moe/fused_moe.py
user: login: SageMoore
body: Nit: I doesn't look like you need this?
created_at: 2025-05-08T16:15:04Z
author_association: CONTRIBUTOR
diff_hunk: @@ -259,29 +318,32 @@ def rocm_aiter_topk_softmax_fake(topk_weights: torch.Tensor,
     )
 
 
-def rocm_aiter_fused_experts(hidden_states: torch.Tensor,
-                             w1: torch.Tensor,
-                             w2: torch.Tensor,
-                             topk_weights: torch.Tensor,
-                             topk_ids: torch.Tensor,
-                             inplace: bool = False,
-                             activation: str = "silu",
-                             apply_router_weight_on_input: bool = False,
-                             use_fp8_w8a8: bool = False,
-                             use_int8_w8a8: bool = False,
-                             use_int8_w8a16: bool = False,
-                             use_int4_w4a16: bool = False,
-                             per_channel_quant: bool = False,
-                             global_num_experts: int = -1,
-                             expert_map: Optional[torch.Tensor] = None,
-                             w1_scale: Optional[torch.Tensor] = None,
-                             w2_scale: Optional[torch.Tensor] = None,
-                             w1_zp: Optional[torch.Tensor] = None,
-                             w2_zp: Optional[torch.Tensor] = None,
-                             a1_scale: Optional[torch.Tensor] = None,
-                             a2_scale: Optional[torch.Tensor] = None,
-                             block_shape: Optional[List[int]] = None,
-                             allow_deep_gemm: bool = False) -> torch.Tensor:
+def rocm_aiter_fused_experts(
+    hidden_states: torch.Tensor,
+    w1: torch.Tensor,
+    w2: torch.Tensor,
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    inplace: bool = False,
+    activation: str = "silu",
+    apply_router_weight_on_input: bool = False,
+    use_fp8_w8a8: bool = False,
+    use_int8_w8a8: bool = False,
+    use_int8_w8a16: bool = False,
+    use_int4_w4a16: bool = False,
+    per_channel_quant: bool = False,
+    global_num_experts: int = -1,
+    expert_map: Optional[torch.Tensor] = None,
+    w1_scale: Optional[torch.Tensor] = None,
+    w2_scale: Optional[torch.Tensor] = None,
+    w1_zp: Optional[torch.Tensor] = None,
+    w2_zp: Optional[torch.Tensor] = None,
+    a1_scale: Optional[torch.Tensor] = None,
+    a2_scale: Optional[torch.Tensor] = None,
+    block_shape: Optional[List[int]] = None,
+    allow_deep_gemm: bool = False,
+    use_ck_moe_2stages: bool = False,
path: vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py
user: login: tjtanaa
body: @SageMoore There is an RFC that highlights Accessing `envs.ENV` is very costly. RFC Issue https://github.com/vllm-project/vllm/issues/17067 .
 Thus, all the env are only invoked and stored as a property of the class during initialization stage.
created_at: 2025-05-08T16:34:18Z
author_association: COLLABORATOR
diff_hunk: @@ -1129,7 +1129,8 @@ def fused_experts(hidden_states: torch.Tensor,
                   a1_scale: Optional[torch.Tensor] = None,
                   a2_scale: Optional[torch.Tensor] = None,
                   block_shape: Optional[List[int]] = None,
-                  allow_deep_gemm: bool = False) -> torch.Tensor:
+                  allow_deep_gemm: bool = False,
+                  use_ck_moe_2stages: bool = False) -> torch.Tensor:
path: vllm/model_executor/layers/fused_moe/fused_moe.py
user: login: hongxiayang
body: agreed. this is not needed
created_at: 2025-05-08T22:10:34Z
author_association: COLLABORATOR
diff_hunk: @@ -16,11 +16,18 @@ def is_rocm_aiter_moe_enabled() -> bool:
         and envs.VLLM_ROCM_USE_AITER
 
 
+def is_rocm_aiter_2stage_moe_enabled() -> bool:
+    return current_platform.is_rocm() \
+        and envs.VLLM_ROCM_USE_AITER_2STAGE_MOE \
+        and envs.VLLM_ROCM_USE_AITER_MOE \
+        and envs.VLLM_ROCM_USE_AITER
path: vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py
user: login: hongxiayang
body: This seems too much to check 3 environment variables. envs.VLLM_ROCM_USE_AITER_2STAGE_MOE is enough as it is only used when the other two are already true. 
created_at: 2025-05-08T22:15:04Z
author_association: COLLABORATOR
diff_hunk: @@ -581,7 +581,11 @@ def create_weights(self, layer: Module, num_experts: int, hidden_size: int,
     def process_weights_after_loading(self, layer: Module) -> None:
         # Lazy import to avoid importing triton too early.
         from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
-            expand_weights, is_rocm_aiter_moe_enabled, shuffle_weights)
+            expand_weights, is_rocm_aiter_2stage_moe_enabled,
+            is_rocm_aiter_moe_enabled, shuffle_weights)
+
+        self.rocm_aiter_moe_enabled = is_rocm_aiter_moe_enabled()
+        self.rocm_aiter_2stage_moe_enabled = is_rocm_aiter_2stage_moe_enabled()
path: vllm/model_executor/layers/quantization/fp8.py
user: login: tjtanaa
body: Same comment as below
https://github.com/vllm-project/vllm/pull/17110#discussion_r2080076439

There is an RFC that highlights Accessing envs.ENV is very costly. RFC Issue https://github.com/vllm-project/vllm/issues/17067 .
Thus, all the env are only invoked and stored as a property of the class during initialization stage.
created_at: 2025-05-09T07:15:48Z
author_association: COLLABORATOR
diff_hunk: @@ -1129,7 +1129,8 @@ def fused_experts(hidden_states: torch.Tensor,
                   a1_scale: Optional[torch.Tensor] = None,
                   a2_scale: Optional[torch.Tensor] = None,
                   block_shape: Optional[List[int]] = None,
-                  allow_deep_gemm: bool = False) -> torch.Tensor:
+                  allow_deep_gemm: bool = False,
+                  use_ck_moe_2stages: bool = False) -> torch.Tensor:
path: vllm/model_executor/layers/fused_moe/fused_moe.py
user: login: tjtanaa
body: I have removed the argument `allow_deep_gemm`. `use_ck_moe_2stages` is kept as there is an RFC that highlights Accessing envs.ENV is very costly. RFC Issue https://github.com/vllm-project/vllm/issues/17067 .
Thus, all the env are only invoked and stored as a property of the class during initialization stage.
created_at: 2025-05-10T06:34:39Z
author_association: COLLABORATOR
diff_hunk: @@ -16,11 +16,18 @@ def is_rocm_aiter_moe_enabled() -> bool:
         and envs.VLLM_ROCM_USE_AITER
 
 
+def is_rocm_aiter_2stage_moe_enabled() -> bool:
+    return current_platform.is_rocm() \
+        and envs.VLLM_ROCM_USE_AITER_2STAGE_MOE \
+        and envs.VLLM_ROCM_USE_AITER_MOE \
+        and envs.VLLM_ROCM_USE_AITER
path: vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py
user: login: tjtanaa
body: The `is_rocm_aiter_2stage_moe_enabled()` has been removed since `envs.VLLM_ROCM_USE_AITER_2STAGE_MOE` is being called in the layer class during initialization only, not in the forward pass.
created_at: 2025-05-10T06:35:57Z
author_association: COLLABORATOR
diff_hunk: @@ -1129,7 +1129,8 @@ def fused_experts(hidden_states: torch.Tensor,
                   a1_scale: Optional[torch.Tensor] = None,
                   a2_scale: Optional[torch.Tensor] = None,
                   block_shape: Optional[List[int]] = None,
-                  allow_deep_gemm: bool = False) -> torch.Tensor:
+                  allow_deep_gemm: bool = False,
+                  use_ck_moe_2stages: bool = False) -> torch.Tensor:
path: vllm/model_executor/layers/fused_moe/fused_moe.py
user: login: tjtanaa
body: To pass the pre-commit tests of file `vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py` , we have adjusted the logic of assignment of fused_experts function to `fused_experts_func` to become `rocm_aiter_fused_experts_func`, following the approach in `vllm/attention/backends/rocm_flash_attn.py`, where the attention functions are assigned to different property name: `self.fa_attn_func` , `self.sdpa_attn_func` and `self.triton_attn_func` 

This also allows us to clean up the unused arguments of the function in `rocm_aiter_fused_experts` (`vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py`). Note that the `expert_map` has been removed as a bugfix, `expert_map` (integer ID of the experts) in vLLM and `expert_mask` (boolean mask of active experts on current GPU) in the AITER ops are different. The current `rocm_aiter_fused_experts` has removed the `expert_map` argument and encourage to add it back when enabling EP using AITER in future PR.


created_at: 2025-05-10T11:53:34Z
author_association: COLLABORATOR

PR comments:
user: login: github-actions[bot]
created_at: 2025-04-24T13:34:49Z
body: 👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run `fastcheck` CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your `fastcheck` build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping `simon-mo` or `khluu` to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add `ready` label to the PR or enable auto-merge.

🚀
author_association: NONE
user: login: mergify[bot]
created_at: 2025-05-05T16:41:03Z
body: This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @tjtanaa.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork
author_association: CONTRIBUTOR
user: login: hongxiayang
created_at: 2025-05-12T22:37:12Z
body: > I'd like to propose an alternative implementation. Instead of passing this new environment variable around, let's just keep the existing `VLLM_ROCM_USE_AITER_MOE` variable and add heuristics to the code that decide when to use the one stage vs the two stage implementation.
> 
> It looks like the two stage kernel is generally faster but doesn't support dynamic quantization? Let's rework this PR so that we always use the two stage implementation when it's supported and fallback to the one stage implementation when it isn't. We don't need any logging when this fallback occurs.
> 
> If we make the change that I suggested, it looks like you can simplify the changes in fp8.py down to just some layout changes for the weights. fused_moe.py should get a bit simpler as well.
> 
> What do you think?

@SageMoore Can we get this merged and do the rework as a follow up PR? 

> I'd like to propose an alternative implementation. Instead of passing this new environment variable around, let's just keep the existing `VLLM_ROCM_USE_AITER_MOE` variable and add heuristics to the code that decide when to use the one stage vs the two stage implementation.
> 
> It looks like the two stage kernel is generally faster but doesn't support dynamic quantization? Let's rework this PR so that we always use the two stage implementation when it's supported and fallback to the one stage implementation when it isn't. We don't need any logging when this fallback occurs.
> 
> If we make the change that I suggested, it looks like you can simplify the changes in fp8.py down to just some layout changes for the weights. fused_moe.py should get a bit simpler as well.
> 
> What do you think?

@SageMoore It would be better to do a different implementation in a follow-up PR. We would like to have this PR merged as this improves performance for a bunch of models.  
author_association: COLLABORATOR
user: login: mergify[bot]
created_at: 2025-05-13T11:19:40Z
body: This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @tjtanaa.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork
author_association: CONTRIBUTOR
user: login: tjtanaa
created_at: 2025-05-13T13:00:05Z
body: > I'd like to propose an alternative implementation. Instead of passing this new environment variable around, let's just keep the existing `VLLM_ROCM_USE_AITER_MOE` variable and add heuristics to the code that decide when to use the one stage vs the two stage implementation.
> 
> It looks like the two stage kernel is generally faster but doesn't support dynamic quantization? Let's rework this PR so that we always use the two stage implementation when it's supported and fallback to the one stage implementation when it isn't. We don't need any logging when this fallback occurs.
> 
> If we make the change that I suggested, it looks like you can simplify the changes in fp8.py down to just some layout changes for the weights. fused_moe.py should get a bit simpler as well.
> 
> What do you think?

@SageMoore Thank you for the suggestion, we had already cleared up this env variable and made simplifications to the fused moe functions.
author_association: COLLABORATOR
user: login: mergify[bot]
created_at: 2025-05-14T02:49:58Z
body: This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @tjtanaa.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork
author_association: CONTRIBUTOR
user: login: mergify[bot]
created_at: 2025-05-14T05:05:16Z
body: This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @tjtanaa.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork
author_association: CONTRIBUTOR