[AMD][WideEP] Integrate aiter batched deepgemm

Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>

tests/kernels/moe/test_aiter_batched_deepgemm.py

@@ -0,0 +1,229 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Tests for AITER batched deepgemm MoE implementation.
+
+These tests compare AiterBatchedExperts against BatchedTritonExperts
+to verify numerical correctness of the AITER deepgemm kernel integration.
+"""
+
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+
+# Skip all tests if not on ROCm
+if not current_platform.is_rocm():
+    pytest.skip("AITER tests require ROCm", allow_module_level=True)
+
+from vllm._aiter_ops import rocm_aiter_ops
+from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config
+from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
+    BatchedPrepareAndFinalize,
+    BatchedTritonExperts,
+)
+from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel
+from vllm.model_executor.layers.fused_moe.rocm_aiter_batched_moe import (
+    AiterBatchedExperts,
+)
+
+from .utils import make_dummy_moe_config
+
+BLOCK_SIZE = [128, 128]
+
+
+def is_aiter_deepgemm_supported() -> bool:
+    """Check if AITER deepgemm is supported on current device."""
+    if not current_platform.is_rocm():
+        return False
+    if not rocm_aiter_ops.is_deepgemm_enabled():
+        return False
+    try:
+        gpu_arch = torch.cuda.get_device_properties("cuda").gcnArchName
+        return "gfx942" in gpu_arch
+    except Exception:
+        return False
+
+
+def make_block_quant_fp8_weights(
+    E: int, N: int, K: int, block_size: list[int]
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Create FP8 quantized weights with block scales for testing.
+
+    Args:
+        E: Number of experts
+        N: Intermediate dimension (2*N for gate+up)
+        K: Hidden dimension
+        block_size: Block size for quantization [block_k, block_n]
+
+    Returns:
+        (w1, w2, w1_scale, w2_scale) tuple of quantized weights and scales
+    """
+    device = "cuda"
+    fp8_dtype = torch.float8_e4m3fn
+    fp8_info = torch.finfo(fp8_dtype)
+
+    block_k, block_n = block_size
+    num_k_blocks = (K + block_k - 1) // block_k
+    num_n_blocks_w1 = (2 * N + block_n - 1) // block_n
+    num_n_blocks_w2 = (K + block_n - 1) // block_n
+
+    # Create random weights and clamp to FP8 range
+    w1 = torch.randn(E, 2 * N, K, device=device, dtype=torch.bfloat16) / 10.0
+    w2 = torch.randn(E, K, N, device=device, dtype=torch.bfloat16) / 10.0
+
+    w1.clamp_(fp8_info.min, fp8_info.max)
+    w2.clamp_(fp8_info.min, fp8_info.max)
+
+    # Quantize to FP8
+    w1_fp8 = w1.to(fp8_dtype)
+    w2_fp8 = w2.to(fp8_dtype)
+
+    # Create per-block scales
+    w1_scale = torch.ones(E, num_n_blocks_w1, num_k_blocks, device=device, dtype=torch.float32)
+    w2_scale = torch.ones(E, num_n_blocks_w2, num_k_blocks, device=device, dtype=torch.float32)
+
+    return w1_fp8, w2_fp8, w1_scale, w2_scale
+
+
+def calc_diff(a: torch.Tensor, b: torch.Tensor) -> float:
+    """Calculate relative difference between two tensors."""
+    a_flat = a.flatten().float()
+    b_flat = b.flatten().float()
+    diff = torch.abs(a_flat - b_flat)
+    max_val = torch.max(torch.abs(a_flat), torch.abs(b_flat))
+    # Avoid division by zero
+    max_val = torch.clamp(max_val, min=1e-6)
+    rel_diff = diff / max_val
+    return rel_diff.mean().item()
+
+
+@pytest.mark.skipif(
+    not is_aiter_deepgemm_supported(),
+    reason="Requires AITER deepgemm on MI300X (gfx942)"
+)
+@pytest.mark.parametrize("E", [8, 16])  # number of experts
+@pytest.mark.parametrize("T", [128, 256])  # tokens per expert
+@pytest.mark.parametrize("K", [128, 256])  # hidden dim
+@pytest.mark.parametrize("N", [256, 512])  # intermediate dim per expert
+@pytest.mark.parametrize("topk", [1, 2])
+def test_aiter_batched_deepgemm_vs_triton(
+    E: int, T: int, K: int, N: int, topk: int
+):
+    """Compare AiterBatchedExperts to BatchedTritonExperts."""
+
+    device = "cuda"
+    w1, w2, w1_s, w2_s = make_block_quant_fp8_weights(E, N, K, BLOCK_SIZE)
+
+    M = E * T  # total tokens
+    a = torch.randn(M, K, device=device, dtype=torch.bfloat16) / 10.0
+    fp8_info = torch.finfo(torch.float8_e4m3fn)
+    a.clamp_(fp8_info.min, fp8_info.max)
+
+    # random router outputs → top-k indices / weights
+    router_logits = torch.randn(M, E, device=device, dtype=torch.float32)
+    topk_weights, topk_ids = torch.topk(router_logits, k=topk, dim=-1)
+    topk_weights = torch.nn.functional.softmax(topk_weights, dim=-1)
+
+    # token number for each expert
+    cnt = torch.bincount(topk_ids.flatten(), minlength=E)
+    max_cnt = int(cnt.max().item())
+    # next power of 2 for max token number
+    max_num_tokens = 1 << (max_cnt - 1).bit_length()
+
+    prep_finalize = BatchedPrepareAndFinalize(
+        max_num_tokens=max_num_tokens,
+        num_local_experts=E,
+        num_dispatchers=1,
+        rank=0,
+    )
+
+    quant_config = fp8_w8a8_moe_quant_config(
+        w1_scale=w1_s,
+        w2_scale=w2_s,
+        per_act_token_quant=False,
+        block_shape=BLOCK_SIZE,
+    )
+
+    moe_config = make_dummy_moe_config()
+
+    # triton (reference)
+    triton_experts = BatchedTritonExperts(
+        max_num_tokens=max_num_tokens,
+        num_dispatchers=1,
+        quant_config=quant_config,
+        moe_config=moe_config,
+    )
+    mk_triton = FusedMoEModularKernel(prep_finalize, triton_experts)
+
+    out_triton = mk_triton(
+        hidden_states=a,
+        w1=w1,
+        w2=w2,
+        topk_weights=topk_weights,
+        topk_ids=topk_ids,
+        inplace=False,
+        global_num_experts=E,
+    )
+
+    # AITER batched deepgemm
+    aiter_experts = AiterBatchedExperts(
+        max_num_tokens=max_num_tokens,
+        num_dispatchers=1,
+        quant_config=quant_config,
+        moe_config=moe_config,
+    )
+    mk_aiter = FusedMoEModularKernel(prep_finalize, aiter_experts)
+
+    out_aiter = mk_aiter(
+        hidden_states=a,
+        w1=w1,
+        w2=w2,
+        topk_weights=topk_weights,
+        topk_ids=topk_ids,
+        inplace=False,
+        global_num_experts=E,
+    )
+
+    diff = calc_diff(out_aiter, out_triton)
+    assert diff < 1e-2, f"Output diff too large: {diff}"
+
+
+@pytest.mark.skipif(
+    not is_aiter_deepgemm_supported(),
+    reason="Requires AITER deepgemm on MI300X (gfx942)"
+)
+def test_aiter_batched_deepgemm_supports_device():
+    """Test that _supports_current_device correctly identifies MI300X."""
+    assert AiterBatchedExperts._supports_current_device()
+
+
+@pytest.mark.skipif(
+    not current_platform.is_rocm(),
+    reason="Requires ROCm platform"
+)
+def test_aiter_batched_deepgemm_activation_format():
+    """Test that AiterBatchedExperts uses BatchedExperts format."""
+    from vllm.model_executor.layers.fused_moe.modular_kernel import (
+        FusedMoEActivationFormat,
+    )
+    assert AiterBatchedExperts.activation_format() == FusedMoEActivationFormat.BatchedExperts
+
+
+@pytest.mark.skipif(
+    not current_platform.is_rocm(),
+    reason="Requires ROCm platform"
+)
+def test_aiter_batched_deepgemm_quant_scheme():
+    """Test that AiterBatchedExperts supports correct quant schemes."""
+    from vllm.model_executor.layers.quantization.utils.quant_utils import (
+        kFp8Dynamic128Sym,
+        kFp8Static128BlockSym,
+    )
+    # Should support FP8 128-block symmetric
+    assert AiterBatchedExperts._supports_quant_scheme(
+        kFp8Static128BlockSym, kFp8Dynamic128Sym
+    )
+    # Should not support unquantized
+    assert not AiterBatchedExperts._supports_quant_scheme(None, None)

vllm/_aiter_ops.py

@@ -444,6 +444,45 @@ def _rocm_aiter_gemm_a8w8_blockscale_fake(
     return Y
 
 
+def _rocm_aiter_deepgemm_impl(
+    XQ: torch.Tensor,
+    WQ: torch.Tensor,
+    Y: torch.Tensor,
+    group_layout: torch.Tensor,
+    x_scale: torch.Tensor | None = None,
+    w_scale: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """
+    AITER DeepGemm masked-M grouped GEMM kernel.
+
+    Args:
+        XQ: Quantized input activations (E, T, K)
+        WQ: Quantized weights (E, N, K) - must be shuffled with shuffle_weight
+        Y: Output tensor (E, T, N)
+        group_layout: Per-expert token counts (E,) - the masked_m tensor
+        x_scale: Optional input scales for FP8 quantization
+        w_scale: Optional weight scales for FP8 quantization
+
+    Returns:
+        Y tensor with results written in-place
+    """
+    from aiter.ops.deepgemm import deepgemm
+
+    return deepgemm(XQ, WQ, Y, group_layout, x_scale, w_scale)
+
+
+def _rocm_aiter_deepgemm_fake(
+    XQ: torch.Tensor,
+    WQ: torch.Tensor,
+    Y: torch.Tensor,
+    group_layout: torch.Tensor,
+    x_scale: torch.Tensor | None = None,
+    w_scale: torch.Tensor | None = None,
+) -> torch.Tensor:
+    # Y is modified in place, return it as-is for symbolic execution
+    return Y
+
+
 def _rocm_aiter_rms_norm_impl(
     x: torch.Tensor, weight: torch.Tensor, variance_epsilon: float
 ) -> torch.Tensor:
@@ -849,6 +888,7 @@ class rocm_aiter_ops:
     _MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS
     # TODO: Consolidate under _LINEAR_ENABLED
     _TRITON_UNQUANT_GEMM = envs.VLLM_ROCM_USE_AITER_TRITON_GEMM
+    _DEEPGEMM_ENABLED = envs.VLLM_ROCM_USE_AITER_DEEPGEMM
 
     @classmethod
     def refresh_env_variables(cls):
@@ -873,6 +913,7 @@ class rocm_aiter_ops:
         cls._TRITON_ROTARY_EMBED = envs.VLLM_ROCM_USE_AITER_TRITON_ROPE
         cls._MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS
         cls._TRITON_UNQUANT_GEMM = envs.VLLM_ROCM_USE_AITER_TRITON_GEMM
+        cls._DEEPGEMM_ENABLED = envs.VLLM_ROCM_USE_AITER_DEEPGEMM
 
     @classmethod
     @if_aiter_supported
@@ -949,6 +990,11 @@ class rocm_aiter_ops:
     def is_triton_gemm_enabled(cls) -> bool:
         return cls._AITER_ENABLED and cls._TRITON_UNQUANT_GEMM
 
+    @classmethod
+    @if_aiter_supported
+    def is_deepgemm_enabled(cls) -> bool:
+        return cls._AITER_ENABLED and cls._DEEPGEMM_ENABLED
+
     @staticmethod
     @if_aiter_supported
     def register_ops_once() -> None:
@@ -1029,6 +1075,14 @@ class rocm_aiter_ops:
                 fake_impl=_rocm_aiter_gemm_a8w8_blockscale_fake,
             )
 
+            direct_register_custom_op(
+                op_name="rocm_aiter_deepgemm",
+                op_func=_rocm_aiter_deepgemm_impl,
+                mutates_args=["Y"],
+                fake_impl=_rocm_aiter_deepgemm_fake,
+                dispatch_key=current_platform.dispatch_key,
+            )
+
             direct_register_custom_op(
                 op_name="rocm_aiter_rms_norm",
                 op_func=_rocm_aiter_rms_norm_impl,
@@ -1195,6 +1249,33 @@ class rocm_aiter_ops:
             A, B, As, Bs, output_dtype
         )
 
+    @staticmethod
+    def deepgemm(
+        XQ: torch.Tensor,
+        WQ: torch.Tensor,
+        Y: torch.Tensor,
+        group_layout: torch.Tensor,
+        x_scale: torch.Tensor | None = None,
+        w_scale: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """
+        AITER DeepGemm masked-M grouped GEMM kernel.
+
+        Args:
+            XQ: Quantized input activations (E, T, K)
+            WQ: Quantized weights (E, N, K) - must be shuffled with shuffle_weight
+            Y: Output tensor (E, T, N)
+            group_layout: Per-expert token counts (E,) - the masked_m tensor
+            x_scale: Optional input scales for FP8 quantization
+            w_scale: Optional weight scales for FP8 quantization
+
+        Returns:
+            Y tensor with results
+        """
+        return torch.ops.vllm.rocm_aiter_deepgemm(
+            XQ, WQ, Y, group_layout, x_scale, w_scale
+        )
+
     @staticmethod
     def fused_moe(
         hidden_states: torch.Tensor,

vllm/envs.py

@@ -128,6 +128,7 @@ if TYPE_CHECKING:
     VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION: bool = False
     VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: bool = False
     VLLM_ROCM_USE_AITER_TRITON_GEMM: bool = True
+    VLLM_ROCM_USE_AITER_DEEPGEMM: bool = True
     VLLM_ROCM_USE_SKINNY_GEMM: bool = True
     VLLM_ROCM_FP8_PADDING: bool = True
     VLLM_ROCM_MOE_PADDING: bool = True
@@ -1055,6 +1056,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
     "VLLM_ROCM_USE_AITER_TRITON_GEMM": lambda: (
         os.getenv("VLLM_ROCM_USE_AITER_TRITON_GEMM", "True").lower() in ("true", "1")
     ),
+    # Whether to use AITER's DeepGemm masked-M grouped GEMM kernel.
+    # By default is enabled for batched MoE on ROCm.
+    "VLLM_ROCM_USE_AITER_DEEPGEMM": lambda: (
+        os.getenv("VLLM_ROCM_USE_AITER_DEEPGEMM", "True").lower() in ("true", "1")
+    ),
     # use rocm skinny gemms
     "VLLM_ROCM_USE_SKINNY_GEMM": lambda: (
         os.getenv("VLLM_ROCM_USE_SKINNY_GEMM", "True").lower() in ("true", "1")

vllm/model_executor/layers/fused_moe/oracle/fp8.py

@@ -49,6 +49,7 @@ class Fp8MoeBackend(Enum):
     TRITON = "TRITON"
     BATCHED_TRITON = "BATCHED_TRITON"
     AITER = "AITER"
+    AITER_BATCHED_DEEPGEMM = "AITER_BATCHED_DEEPGEMM"
     VLLM_CUTLASS = "VLLM_CUTLASS"
     BATCHED_VLLM_CUTLASS = "BATCHED_VLLM_CUTLASS"
 
@@ -108,6 +109,13 @@ def backend_to_kernel_cls(
 
         return AiterExperts
 
+    elif backend == Fp8MoeBackend.AITER_BATCHED_DEEPGEMM:
+        from vllm.model_executor.layers.fused_moe.rocm_aiter_batched_moe import (
+            AiterBatchedExperts,
+        )
+
+        return AiterBatchedExperts
+
     elif backend == Fp8MoeBackend.VLLM_CUTLASS:
         from vllm.model_executor.layers.fused_moe.triton_cutlass_moe import (
             TritonOrCutlassExperts,
@@ -144,6 +152,7 @@ def select_fp8_moe_backend(
     # NOTE: the kernels are selected in the following order.
     AVAILABLE_BACKENDS = [
         Fp8MoeBackend.AITER,
+        Fp8MoeBackend.AITER_BATCHED_DEEPGEMM,
         Fp8MoeBackend.FLASHINFER_TRTLLM,
         Fp8MoeBackend.FLASHINFER_CUTLASS,
         Fp8MoeBackend.DEEPGEMM,
@@ -294,8 +303,14 @@ def select_fp8_moe_backend(
     if envs.is_set("VLLM_ROCM_USE_AITER") or envs.is_set("VLLM_ROCM_USE_AITER_MOE"):
         if not envs.VLLM_ROCM_USE_AITER or not envs.VLLM_ROCM_USE_AITER_MOE:
             AVAILABLE_BACKENDS.remove(Fp8MoeBackend.AITER)
+            if Fp8MoeBackend.AITER_BATCHED_DEEPGEMM in AVAILABLE_BACKENDS:
+                AVAILABLE_BACKENDS.remove(Fp8MoeBackend.AITER_BATCHED_DEEPGEMM)
         else:
-            backend = Fp8MoeBackend.AITER
+            # Select batched or standard AITER backend based on activation format
+            if activation_format == mk.FusedMoEActivationFormat.BatchedExperts:
+                backend = Fp8MoeBackend.AITER_BATCHED_DEEPGEMM
+            else:
+                backend = Fp8MoeBackend.AITER
             return _return_or_raise(
                 backend, config, weight_key, activation_key, activation_format
             )

vllm/model_executor/layers/fused_moe/rocm_aiter_batched_moe.py

@@ -0,0 +1,269 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+AITER-based batched MoE implementation using DeepGemm masked-M grouped GEMM.
+
+This module provides an optimized batched experts implementation for AMD ROCm
+GPUs using AITER's deepgemm kernel. The kernel operates on the masked-M
+batched format where:
+  - Activations are shaped (E, T, K) where E=num_experts, T=max_tokens, K=hidden
+  - Each expert processes a variable number of tokens tracked by tokens_per_expert
+  - Weights must be shuffled using AITER's shuffle_weight for optimal performance
+
+NOTE: Currently only supported on MI300X (gfx942). Support for MI325X (gfx950)
+is pending in AITER.
+"""
+
+import torch
+
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+from vllm._aiter_ops import rocm_aiter_ops
+from vllm.platforms import current_platform
+from vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe import (
+    persistent_masked_m_silu_mul_quant,
+)
+from vllm.model_executor.layers.fused_moe.config import (
+    FusedMoEConfig,
+    FusedMoEParallelConfig,
+    FusedMoEQuantConfig,
+)
+from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
+    TopKWeightAndReduceDelegate,
+)
+from vllm.model_executor.layers.fused_moe.utils import _resize_cache
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    QuantKey,
+    kFp8Dynamic128Sym,
+    kFp8Static128BlockSym,
+)
+from vllm.utils.deep_gemm import DeepGemmQuantScaleFMT
+
+
+class AiterBatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):
+    """
+    AITER-based batched experts using DeepGemm masked-M grouped GEMM.
+
+    This implementation uses AITER's deepgemm kernel for the GEMMs and
+    a Triton kernel for the fused SiLU+Mul+Quantize activation.
+
+    The workflow for each forward pass is:
+    1. GEMM1: input @ w1 -> intermediate (gate, up projections)
+    2. Activation: SiLU(gate) * up, then quantize to FP8
+    3. GEMM2: activated @ w2 -> output (down projection)
+    """
+
+    def __init__(
+        self,
+        moe_config: FusedMoEConfig,
+        quant_config: FusedMoEQuantConfig,
+        max_num_tokens: int,
+        num_dispatchers: int,
+    ):
+        """
+        Initialize AITER batched experts.
+
+        Args:
+            moe_config: MoE configuration
+            quant_config: Quantization configuration
+            max_num_tokens: Maximum number of tokens from a DP Rank
+            num_dispatchers: The number of DP dispatchers
+        """
+        super().__init__(
+            moe_config=moe_config,
+            quant_config=quant_config,
+            max_num_tokens=max_num_tokens,
+            num_dispatchers=num_dispatchers,
+        )
+        # AITER deepgemm requires 128x128 block quantization
+        assert self.block_shape == [128, 128], (
+            f"AITER DeepGemm requires block_shape=[128, 128], got {self.block_shape}"
+        )
+        assert self.quant_config.use_fp8_w8a8, (
+            "AITER DeepGemm requires FP8 W8A8 quantization"
+        )
+
+        # Track if weights have been shuffled
+        self._weights_shuffled = False
+
+    @staticmethod
+    def activation_format() -> mk.FusedMoEActivationFormat:
+        return mk.FusedMoEActivationFormat.BatchedExperts
+
+    @staticmethod
+    def _supports_current_device() -> bool:
+        # AITER deepgemm only supports MI300X (gfx942) currently
+        # TODO: Add gfx950 support when AITER adds it
+        if not rocm_aiter_ops.is_deepgemm_enabled():
+            return False
+        if not current_platform.is_rocm():
+            return False
+        from vllm.platforms.rocm import on_mi3xx
+        # Check for gfx942 specifically (MI300X)
+        # on_mi3xx() returns True for gfx942 and gfx950, but deepgemm
+        # only supports gfx942 currently
+        if not on_mi3xx():
+            return False
+        try:
+            gpu_arch = torch.cuda.get_device_properties("cuda").gcnArchName
+            return "gfx942" in gpu_arch
+        except Exception:
+            return False
+
+    @staticmethod
+    def _supports_no_act_and_mul() -> bool:
+        return False
+
+    @staticmethod
+    def _supports_quant_scheme(
+        weight_key: QuantKey | None,
+        activation_key: QuantKey | None,
+    ) -> bool:
+        SUPPORTED_W_A = [(kFp8Static128BlockSym, kFp8Dynamic128Sym)]
+        return (weight_key, activation_key) in SUPPORTED_W_A
+
+    @staticmethod
+    def _supports_activation(activation: str) -> bool:
+        return activation in ["silu"]
+
+    @staticmethod
+    def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
+        return True
+
+    def supports_chunking(self) -> bool:
+        return False
+
+    def supports_expert_map(self) -> bool:
+        return False
+
+    def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
+        return TopKWeightAndReduceDelegate()
+
+    def workspace_shapes(
+        self,
+        M: int,
+        N: int,
+        K: int,
+        topk: int,
+        global_num_experts: int,
+        local_num_experts: int,
+        expert_tokens_meta: mk.ExpertTokensMetadata | None,
+        activation: str,
+    ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
+        assert self.num_dispatchers is not None
+        assert self.max_num_tokens is not None
+        num_dispatchers = self.num_dispatchers
+        num_experts = local_num_experts
+        max_num_tokens = M if self.max_num_tokens is None else self.max_num_tokens
+        activation_out_dim = self.adjust_N_for_activation(N, activation)
+        workspace13 = (num_experts, max_num_tokens * num_dispatchers, max(K, N))
+        workspace2 = (num_experts, max_num_tokens * num_dispatchers, activation_out_dim)
+        output = (num_experts, max_num_tokens * num_dispatchers, K)
+        return (workspace13, workspace2, output)
+
+    def _shuffle_weights_if_needed(
+        self,
+        w1: torch.Tensor,
+        w2: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Shuffle weights for AITER deepgemm if not already shuffled.
+
+        AITER's deepgemm kernel requires weights to be in a specific
+        16x16 tiled layout for optimal memory access patterns.
+        """
+        if self._weights_shuffled:
+            return w1, w2
+
+        # Check if weights are already marked as shuffled
+        if hasattr(w1, "is_shuffled") and w1.is_shuffled:
+            self._weights_shuffled = True
+            return w1, w2
+
+        # Shuffle weights using AITER's shuffle_weight
+        w1_shuffled, w2_shuffled = rocm_aiter_ops.shuffle_weights(
+            w1, w2, layout=(16, 16)
+        )
+        self._weights_shuffled = True
+        return w1_shuffled, w2_shuffled
+
+    def apply(
+        self,
+        output: torch.Tensor,
+        hidden_states: torch.Tensor,
+        w1: torch.Tensor,
+        w2: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        activation: str,
+        global_num_experts: int,
+        expert_map: torch.Tensor | None,
+        a1q_scale: torch.Tensor | None,
+        a2_scale: torch.Tensor | None,
+        workspace13: torch.Tensor,
+        workspace2: torch.Tensor,
+        expert_tokens_meta: mk.ExpertTokensMetadata | None,
+        apply_router_weight_on_input: bool,
+    ):
+        """
+        Execute batched experts using AITER deepgemm.
+
+        The forward pass consists of:
+        1. GEMM1: hidden_states @ w1 -> workspace1 (gate, up projections)
+        2. Activation: SiLU(gate) * up, quantize -> a2q, a2q_scale
+        3. GEMM2: a2q @ w2 -> output (down projection)
+        """
+        assert expert_tokens_meta is not None
+        expert_num_tokens = expert_tokens_meta.expert_num_tokens
+
+        assert hidden_states.ndim == 3, (
+            f"Expected 3D hidden_states (E, T, K), got {hidden_states.ndim}D"
+        )
+        assert self.block_shape is not None
+
+        a1q = hidden_states
+        E, _, K = hidden_states.size()
+        _, N, _ = w1.size()  # w1 is (E, N, K) for AITER
+
+        E, max_num_tokens, N, K, _ = self.moe_problem_size(
+            hidden_states, w1, w2, topk_ids
+        )
+
+        workspace1 = _resize_cache(workspace13, (E, max_num_tokens, N))
+
+        # Allocate output tensor for GEMM1
+        # Note: AITER deepgemm writes to Y in-place
+        workspace1.zero_()
+
+        # GEMM1: a1q @ w1 -> workspace1
+        # Shapes: a1q (E, T, K), w1 (E, N, K), workspace1 (E, T, N)
+        rocm_aiter_ops.deepgemm(
+            XQ=a1q,
+            WQ=w1,
+            Y=workspace1,
+            group_layout=expert_num_tokens.to(torch.int32),
+            x_scale=a1q_scale,
+            w_scale=self.w1_scale,
+        )
+
+        # Activation: SiLU(gate) * up, then quantize to FP8
+        # Use the same Triton kernel as BatchedDeepGemmExperts
+        quant_scale_fmt = DeepGemmQuantScaleFMT.FLOAT32
+        a2q, a2q_scale = persistent_masked_m_silu_mul_quant(
+            workspace1,
+            expert_num_tokens,
+            quant_scale_fmt=quant_scale_fmt,
+        )
+
+        # Zero output before GEMM2
+        output.zero_()
+
+        # GEMM2: a2q @ w2 -> output
+        # Shapes: a2q (E, T, H), w2 (E, K, H), output (E, T, K)
+        rocm_aiter_ops.deepgemm(
+            XQ=a2q,
+            WQ=w2,
+            Y=output,
+            group_layout=expert_num_tokens.to(torch.int32),
+            x_scale=a2q_scale,
+            w_scale=self.w2_scale,
+        )