Merge branch 'ROCm:main' into main

Author: iAmir97

.github/workflows/aiter-test.yaml

@@ -12,7 +12,8 @@ concurrency:
   cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}
 
 env:
-  DOCKER_IMAGE: "rocm/pytorch:latest"
+  # TODO: Revert to rocm/pytorch:latest once CK adds ROCm 7.2 support
+  DOCKER_IMAGE: "rocm/pytorch:latest@sha256:683765a52c61341e1674fe730ab3be861a444a45a36c0a8caae7653a08a0e208"
 
 jobs:
   check-signal:

3rdparty/composable_kernel

@@ -1643,11 +1643,13 @@ def fused_topk(
         M, topk, dtype=dtypes.i32, device=hidden_states.device
     )
 
-    if (
-        get_gfx() == "gfx942"
-        and (expert, topk) in [(128, 6), (128, 8), (256, 6), (256, 8)]
-        and gating_output.dtype == dtypes.fp32
-    ):
+    if (expert, topk) in [
+        (128, 4),
+        (128, 6),
+        (128, 8),
+        (256, 6),
+        (256, 8),
+    ] and gating_output.dtype in [dtypes.bf16, dtypes.fp32]:
         if topk_weights is None:
             topk_weights = torch.empty(
                 (M + 3) // 4 * 4, topk, dtype=dtypes.fp32, device=hidden_states.device

aiter/fused_moe.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: MIT
-# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
+# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
 
 import torch
 import os
@@ -500,6 +500,7 @@ def FinalFunc():
                 kernelName=kernelName1,
                 activation=activation,
                 quant_type=q_type,
+                dst_type=dtype,
             ),
             functools.partial(
                 aiter.ck_moe_stage2_fwd,

aiter/fused_moe_dp_shared_expert.py

@@ -443,7 +443,8 @@
         "verbose": "False",
         "blob_gen_cmd": [
             "f'{AITER_META_DIR}/hsa/codegen.py -m fmoe_2stages --output_dir {{}}'",
-            "f'{AITER_META_DIR}/hsa/codegen.py -m fmoe --output_dir {{}}'"
+            "f'{AITER_META_DIR}/hsa/codegen.py -m fmoe --output_dir {{}}'",
+            "f'{AITER_META_DIR}/hsa/codegen.py -m topksoftmax --output_dir {{}}'"
         ]
     },
     "module_moe_ck2stages": {
@@ -563,6 +564,21 @@
         "verbose": "False",
         "blob_gen_cmd": "f'{CK_DIR}/example/ck_tile/10_rmsnorm2d/generate.py --api fwd --gen_blobs --working_path {{}}'"
     },
+    "module_rmsnorm_quant": {
+        "srcs": [
+            "f'{AITER_CSRC_DIR}/kernels/rmsnorm_quant_kernels.cu'",
+            "f'{AITER_CSRC_DIR}/pybind/rmsnorm_quant_pybind.cu'"
+        ],
+        "flags_extra_cc": [],
+        "flags_extra_hip": ["'-ffast-math'"],
+        "extra_ldflags": "None",
+        "extra_include": [
+            "f'{AITER_CSRC_DIR}/include/ck_tile'",
+            "f'{AITER_CSRC_DIR}/include/opus'"
+        ],
+        "verbose": "False",
+        "blob_gen_cmd": "''"
+    },
     "module_smoothquant": {
         "srcs": [
             "f'{AITER_CSRC_DIR}/py_itfs_ck/smoothquant_kernels.cu'",

aiter/jit/optCompilerConfig.json

@@ -553,7 +553,7 @@ def ck_moe_stage1_fwd(
         activation.value,
         int(splitk) if splitk is not None else splitk,
         use_non_temporal_load,
-        dtype2str_dict[dst_type],
+        None if dst_type is None else dtype2str_dict[dst_type],
     )
     return out
 

aiter/ops/moe_op.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: MIT
-# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
+# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
 
 import torch
 from torch import Tensor
@@ -59,16 +59,20 @@ def rms_norm(
     ...
 
 
-@compile_ops("module_rmsnorm", gen_fake=gen_rms_norm_fake_tensor)
 def rmsnorm2d_fwd(
     input: torch.Tensor,
     weight: torch.Tensor,
     epsilon: float,
     use_model_sensitive_rmsnorm: int = 0,
-) -> Tensor: ...
+) -> Tensor:
+    out = torch.empty_like(input, dtype=input.dtype, device=input.device)
+    if use_model_sensitive_rmsnorm > 0 or input.shape[-1] > 8192:
+        rmsnorm2d_fwd_ck(out, input, weight, epsilon, use_model_sensitive_rmsnorm)
+    else:
+        rmsnorm(out, input, weight, epsilon)
+    return out
 
 
-@compile_ops("module_rmsnorm")
 def rmsnorm2d_fwd_with_add(
     out: Tensor,
     input: Tensor,
@@ -77,7 +81,19 @@ def rmsnorm2d_fwd_with_add(
     weight: Tensor,
     epsilon: float,
     use_model_sensitive_rmsnorm: int = 0,
-) -> None: ...
+) -> None:
+    if use_model_sensitive_rmsnorm > 0 or input.shape[-1] > 8192:
+        rmsnorm2d_fwd_with_add_ck(
+            out,
+            input,
+            residual_in,
+            residual_out,
+            weight,
+            epsilon,
+            use_model_sensitive_rmsnorm,
+        )
+    else:
+        add_rmsnorm(out, input, residual_in, residual_out, weight, epsilon)
 
 
 @compile_ops("module_rmsnorm")
@@ -107,18 +123,26 @@ def rmsnorm2d_fwd_with_add_smoothquant(
 ) -> None: ...
 
 
-@compile_ops("module_rmsnorm")
 def rmsnorm2d_fwd_with_dynamicquant(
     out: Tensor,
     input: Tensor,
     yscale: Tensor,
     weight: Tensor,
     epsilon: float,
     use_model_sensitive_rmsnorm: int = 0,
-) -> None: ...
+    group_size: int = 0,
+    shuffle_scale: bool = False,
+) -> None:
+    if use_model_sensitive_rmsnorm > 0 or input.shape[-1] > 8192:
+        assert group_size == 0, "group_size is not supported for ck rmsnorm"
+        assert not shuffle_scale, "shuffle_scale is not supported for ck rmsnorm"
+        rmsnorm2d_fwd_with_dynamicquant_ck(
+            out, input, yscale, weight, epsilon, use_model_sensitive_rmsnorm
+        )
+    else:
+        rmsnorm_quant(out, input, yscale, weight, epsilon, group_size, shuffle_scale)
 
 
-@compile_ops("module_rmsnorm")
 def rmsnorm2d_fwd_with_add_dynamicquant(
     out: Tensor,
     input: Tensor,
@@ -128,4 +152,124 @@ def rmsnorm2d_fwd_with_add_dynamicquant(
     weight: Tensor,
     epsilon: float,
     use_model_sensitive_rmsnorm: int = 0,
+    group_size: int = 0,
+    shuffle_scale: bool = False,
+) -> None:
+    if use_model_sensitive_rmsnorm > 0 or input.shape[-1] > 8192:
+        assert group_size == 0, "group_size is not supported for ck rmsnorm"
+        assert not shuffle_scale, "shuffle_scale is not supported for ck rmsnorm"
+        rmsnorm2d_fwd_with_add_dynamicquant_ck(
+            out,
+            input,
+            residual_in,
+            residual_out,
+            yscale,
+            weight,
+            epsilon,
+            use_model_sensitive_rmsnorm,
+        )
+    else:
+        add_rmsnorm_quant(
+            out,
+            input,
+            residual_in,
+            residual_out,
+            yscale,
+            weight,
+            epsilon,
+            group_size,
+            shuffle_scale,
+        )
+
+
+@compile_ops(
+    "module_rmsnorm", gen_fake=gen_rms_norm_fake_tensor, fc_name="rmsnorm2d_fwd"
+)
+def rmsnorm2d_fwd_ck(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    epsilon: float,
+    use_model_sensitive_rmsnorm: int = 0,
+) -> Tensor: ...
+
+
+@compile_ops("module_rmsnorm", fc_name="rmsnorm2d_fwd_with_add")
+def rmsnorm2d_fwd_with_add_ck(
+    out: Tensor,
+    input: Tensor,
+    residual_in: Tensor,
+    residual_out: Tensor,
+    weight: Tensor,
+    epsilon: float,
+    use_model_sensitive_rmsnorm: int = 0,
+) -> None: ...
+
+
+@compile_ops("module_rmsnorm", fc_name="rmsnorm2d_fwd_with_dynamicquant")
+def rmsnorm2d_fwd_with_dynamicquant_ck(
+    out: Tensor,
+    input: Tensor,
+    yscale: Tensor,
+    weight: Tensor,
+    epsilon: float,
+    use_model_sensitive_rmsnorm: int = 0,
+) -> None: ...
+
+
+@compile_ops("module_rmsnorm", fc_name="rmsnorm2d_fwd_with_add_dynamicquant")
+def rmsnorm2d_fwd_with_add_dynamicquant_ck(
+    out: Tensor,
+    input: Tensor,
+    residual_in: Tensor,
+    residual_out: Tensor,
+    yscale: Tensor,
+    weight: Tensor,
+    epsilon: float,
+    use_model_sensitive_rmsnorm: int = 0,
+) -> None: ...
+
+
+@compile_ops("module_rmsnorm_quant")
+def add_rmsnorm_quant(
+    out: Tensor,
+    input: Tensor,
+    residual_in: Tensor,
+    residual_out: Tensor,
+    scale: Tensor,
+    weight: Tensor,
+    epsilon: float,
+    group_size: int = 0,
+    shuffle_scale: bool = False,
+) -> None: ...
+
+
+@compile_ops("module_rmsnorm_quant")
+def add_rmsnorm(
+    out: Tensor,
+    input: Tensor,
+    residual_in: Tensor,
+    residual_out: Tensor,
+    weight: Tensor,
+    epsilon: float,
+) -> None: ...
+
+
+@compile_ops("module_rmsnorm_quant")
+def rmsnorm_quant(
+    out: Tensor,
+    input: Tensor,
+    scale: Tensor,
+    weight: Tensor,
+    epsilon: float,
+    group_size: int = 0,
+    shuffle_scale: bool = False,
+) -> None: ...
+
+
+@compile_ops("module_rmsnorm_quant")
+def rmsnorm(
+    out: Tensor,
+    input: Tensor,
+    weight: Tensor,
+    epsilon: float,
 ) -> None: ...

aiter/ops/rmsnorm.py

@@ -113,6 +113,7 @@
     "moe_op_silu_fused": "moe.moe_op_silu_fused",
     "moe_op": "moe.moe_op",
     "moe_routing_sigmoid_top1_fused": "moe.moe_routing_sigmoid_top1_fused",
+    "moe_routing": "moe.moe_routing",
     "quant_moe": "moe.quant_moe",
     # Normalization modules (normalization/)
     "fused_add_rmsnorm_pad": "normalization.fused_add_rmsnorm_pad",

aiter/ops/triton/__init__.py

@@ -5,7 +5,6 @@
 import torch
 import aiter
 import aiter.ops.triton.utils._triton.arch_info as arch_info
-from aiter.ops.triton.utils.types import torch_to_triton_dtype
 
 import triton
 import triton.language as tl
@@ -3365,6 +3364,8 @@ def pa_decode_gluon(
         raise RuntimeError(
             "This version triton is not support gluon jit mode, please upgrade to 3.5.0 or higher!"
         )
+    from aiter.ops.triton.utils.types import torch_to_triton_dtype
+
     cdna_version = get_cdna_version()
     assert cdna_version in [
         3,

aiter/ops/triton/gluon/pa_decode_gluon.py

@@ -1,7 +1,8 @@
 # SPDX-License-Identifier: MIT
-# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
+# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
 
 import os
+import sys
 import argparse
 import torch
 import pandas as pd
@@ -116,9 +117,10 @@ def _setup_common_arguments(self):
         )
         self.parser.add_argument(
             "--sort",
-            action="store_true",
+            type=dtypes.str2bool,
+            default=defaults.get("sort", False),
             required=False,
-            help="Arranged according to the keys",
+            help="Arranged according to the keys (True/False)",
         )
         self.parser.add_argument(
             "--errRatio",
@@ -410,6 +412,11 @@ def tune_summary(self, status):
         if not self.remain_untuned.empty:
             logger.info("untuned shapes:")
             print(self.remain_untuned)
+        if not self.remain_untuned.empty or not self.failed.empty:
+            logger.error(
+                "\033[91m[Tuning not Finished]\033[0m some shapes are not tuned or all failed, please check the result file or tune with --profile_file to get more details"
+            )
+            sys.exit(1)
 
     @abstractmethod
     def result_to_csv(self, results, file, concat=False):
@@ -480,6 +487,11 @@ def run(self, args, fast_mode=False):
 
 class GemmCommonTuner(TunerCommon):
 
+    ARG_DEFAULTS = {
+        **TunerCommon.ARG_DEFAULTS,
+        "sort": True,  # Enable sorting by default for GEMM tuners
+    }
+
     def __init__(
         self,
         name,