From 06a7a595e4b62d21156b556352c8a90f2f6eb44a Mon Sep 17 00:00:00 2001 From: Tyler Michael Smith Date: Wed, 21 Jan 2026 15:56:51 -0500 Subject: [PATCH] [AMD][WideEP] Integrate aiter batched deepgemm Signed-off-by: Tyler Michael Smith --- .../moe/test_aiter_batched_deepgemm.py | 229 +++++++++++++++ vllm/_aiter_ops.py | 81 ++++++ vllm/envs.py | 6 + .../layers/fused_moe/oracle/fp8.py | 17 +- .../fused_moe/rocm_aiter_batched_moe.py | 269 ++++++++++++++++++ 5 files changed, 601 insertions(+), 1 deletion(-) create mode 100644 tests/kernels/moe/test_aiter_batched_deepgemm.py create mode 100644 vllm/model_executor/layers/fused_moe/rocm_aiter_batched_moe.py diff --git a/tests/kernels/moe/test_aiter_batched_deepgemm.py b/tests/kernels/moe/test_aiter_batched_deepgemm.py new file mode 100644 index 00000000000..7168470ba33 --- /dev/null +++ b/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) diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 3e232d61924..7bc8b36263f 100644 --- a/vllm/_aiter_ops.py +++ b/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, diff --git a/vllm/envs.py b/vllm/envs.py index 3c837228336..717c38d3009 100755 --- a/vllm/envs.py +++ b/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") diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index cdf2d291b5d..11eac89a214 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/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 ) diff --git a/vllm/model_executor/layers/fused_moe/rocm_aiter_batched_moe.py b/vllm/model_executor/layers/fused_moe/rocm_aiter_batched_moe.py new file mode 100644 index 00000000000..4d925a48dcd --- /dev/null +++ b/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, + )