[gfx942][opus] Support A8W8 blockscale BpreShuffle GEMM

aiter/jit/optCompilerConfig.json

@@ -329,6 +329,10 @@
             "('-mllvm -enable-post-misched=1' if (get_gfx() == 'gfx1250' or 'gfx1250' in os.environ.get('GPU_ARCHS', '')) else '')"
         ],
         "flags_extra_hip_per_source": {
+            "*a8w8_bs_bpreshuf*.device.cu": [
+                "-D__HIPCC_RTC__",
+                "-mllvm --disable-machine-licm"
+            ],
             "*.device.cu": [
                 "-D__HIPCC_RTC__"
             ]
@@ -341,7 +345,7 @@
         "verbose": "False",
         "is_python_module": "True",
         "is_standalone": "False",
-        "blob_gen_cmd": "f'{AITER_CSRC_DIR}/opus_gemm/gen_instances.py --working_path {{0}} --tune_files {AITER_ROOT_DIR}/aiter/configs/bf16_tuned_gemm.csv:{AITER_ROOT_DIR}/aiter/configs/model_configs/*_bf16_tuned_gemm.csv --compiled_kids_sidecar={{0}}/../../compiled_kids_opus.json'"
+        "blob_gen_cmd": "f'{AITER_CSRC_DIR}/opus_gemm/gen_instances.py --working_path {{0}} --tune_files {AITER_ROOT_DIR}/aiter/configs/bf16_tuned_gemm.csv:{AITER_ROOT_DIR}/aiter/configs/model_configs/*_bf16_tuned_gemm.csv:{AITER_CONFIGS.AITER_CONFIG_GEMM_A8W8_BLOCKSCALE_BPRESHUFFLE_FILE} --compiled_kids_sidecar={{0}}/../../compiled_kids_opus.json'"
     },
     "module_gemm_a8w8_bpreshuffle_cktile": {
         "srcs": [
@@ -1739,4 +1743,4 @@
         "verbose": "False",
         "blob_gen_cmd": "''"
     }
-}
+}

aiter/ops/gemm_op_a8w8.py

@@ -919,6 +919,15 @@ def gemm_a8w8_blockscale_bpreshuffle(
             return gemm_a8w8_blockscale_bpreshuffle_asm(
                 XQ, WQ, Y, x_scale, w_scale, splitK=splitK, kernelName=kernelName
             )
+        elif libtype == "opus":
+            kernelId = int(config["kernelId"])
+            from aiter.ops.opus.gemm_op_a8w8 import (
+                opus_gemm_a8w8_blockscale_bpreshuffle_tune,
+            )
+
+            return opus_gemm_a8w8_blockscale_bpreshuffle_tune(
+                XQ, WQ, x_scale, w_scale, Y, kernelId=kernelId
+            )
     try:
         return gemm_a8w8_blockscale_bpreshuffle_ck(XQ, WQ, x_scale, w_scale, Y)
     except RuntimeError as e:

aiter/ops/opus/__init__.py

@@ -1,17 +1,16 @@
 # SPDX-License-Identifier: MIT
 # Copyright (C) 2025-2026, Advanced Micro Devices, Inc. All rights reserved.
-"""opus kernel Python user-facing API; a16w16 today, a8w8 in follow-ups.
+"""opus kernel Python user-facing API.
 
 Public API: `gemm_a16w16_opus` (CSV lookup + C++ heuristic) and
-`opus_gemm_a16w16_tune` (id-based binding). On unsupported arch the
-two callables become stubs that raise RuntimeError on invocation, so
-`import aiter` keeps working alongside its 30+ other ops.
+`opus_gemm_a16w16_tune` (id-based binding). The gfx942 A8W8 blockscale
+bpreshuffle entry is an explicit tune API.
 """
 
 from ._arch import _detect_arch
 
 _SUPPORTED = {"gfx950", "gfx942", "gfx1250"}
-_FEATURE = "aiter.ops.opus (a16w16)"
+_FEATURE = "aiter.ops.opus"
 _HINT = (
     "opus_gemm supports gfx950 (MFMA 16x16x32 / ds_read_b64_tr / 160 KiB "
     "LDS) and gfx942 (MFMA 16x16x16 / ds_read_b128 / 64 KiB LDS). Set "
@@ -43,16 +42,28 @@ def _stub(*_args, **_kwargs):
         gemm_a16w16_opus,
         opus_gemm_workspace_init,
     )
+
+    def opus_gemm_a8w8_blockscale_bpreshuffle_tune(*args, **kwargs):
+        from .gemm_op_a8w8 import (
+            opus_gemm_a8w8_blockscale_bpreshuffle_tune as _impl,
+        )
+
+        return _impl(*args, **kwargs)
+
 else:
     # Don't raise ImportError -- aiter/__init__.py's star-import would catch
     # it and silently disable the 30+ subsequent op imports.
     gemm_a16w16_opus = _make_unsupported_arch_stub("gemm_a16w16_opus")
     opus_gemm_a16w16_tune = _make_unsupported_arch_stub("opus_gemm_a16w16_tune")
+    opus_gemm_a8w8_blockscale_bpreshuffle_tune = _make_unsupported_arch_stub(
+        "opus_gemm_a8w8_blockscale_bpreshuffle_tune"
+    )
     opus_gemm_workspace_init = _make_unsupported_arch_stub("opus_gemm_workspace_init")
 
 
 __all__ = [
     "opus_gemm_a16w16_tune",
+    "opus_gemm_a8w8_blockscale_bpreshuffle_tune",
     "gemm_a16w16_opus",
     "opus_gemm_workspace_init",
 ]

aiter/ops/opus/gemm_op_a8w8.py

@@ -0,0 +1,57 @@
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2025-2026, Advanced Micro Devices, Inc. All rights reserved.
+"""Low-level Opus gfx942 A8W8 blockscale bpreshuffle entry points."""
+
+import torch
+from torch import Tensor
+
+from ...jit.core import compile_ops
+
+
+def _gen_opus_a8w8_blockscale_bpreshuffle_fake_tensors(
+    XQ: Tensor,
+    WQ: Tensor,
+    x_scale: Tensor,
+    w_scale: Tensor,
+    Y: Tensor,
+    kernelId: int,
+) -> Tensor:
+    return Y
+
+
+@compile_ops(
+    "module_deepgemm_opus",
+    fc_name="opus_gemm_a8w8_blockscale_bpreshuffle_tune",
+    gen_fake=_gen_opus_a8w8_blockscale_bpreshuffle_fake_tensors,
+    develop=True,
+)
+def _opus_gemm_a8w8_blockscale_bpreshuffle_tune_raw(
+    XQ: Tensor,
+    WQ: Tensor,
+    x_scale: Tensor,
+    w_scale: Tensor,
+    Y: Tensor,
+    kernelId: int,
+) -> Tensor: ...
+
+
+def opus_gemm_a8w8_blockscale_bpreshuffle_tune(
+    XQ: Tensor,
+    WQ: Tensor,
+    x_scale: Tensor,
+    w_scale: Tensor,
+    Y: Tensor | None = None,
+    kernelId: int = 11000,
+) -> Tensor:
+    """Run one gfx942 Opus A8W8 blockscale bpreshuffle kernel by explicit id."""
+    if Y is None:
+        Y = torch.empty(
+            (XQ.shape[-2], WQ.shape[-2]), device=XQ.device, dtype=torch.bfloat16
+        )
+    _opus_gemm_a8w8_blockscale_bpreshuffle_tune_raw(
+        XQ, WQ, x_scale, w_scale, Y, kernelId
+    )
+    return Y
+
+
+__all__ = ["opus_gemm_a8w8_blockscale_bpreshuffle_tune"]

csrc/ck_gemm_a8w8_blockscale/gemm_a8w8_blockscale_tune.py

@@ -20,12 +20,16 @@
 from aiter.utility.mp_tuner import mp_tuner
 from aiter.ops.shuffle import shuffle_weight
 from aiter.jit.utils.chip_info import get_gfx_runtime as get_gfx
+from aiter.ops.opus.gemm_op_a8w8 import (
+    opus_gemm_a8w8_blockscale_bpreshuffle_tune,
+)
 
 sys.path.insert(0, str(Path(__file__).parent.parent))
 from ck_gemm_a8w8_blockscale_bpreshuffle.gemm_a8w8_blockscale_bpreshuffle_common import (
     kernels_list as candidate_kernels_bpreshuffle_dict,
 )
 from gemm_a8w8_blockscale_instance import candidate_kernels_dict
+from opus_gemm.opus_gemm_common import gfx942_a8w8_kernels_list
 
 # cktile
 from gemm_a8w8_blockscale_cktile_instance import (
@@ -148,6 +152,22 @@ def run_gemm_a8w8_blockscale_asm(
     )
 
 
+def run_gemm_a8w8_blockscale_opus(
+    x,
+    weight,
+    x_scale,
+    w_scale,
+    out,
+    kernel_id,
+):
+    """
+    Run gfx942 Opus a8w8 blockscale bpreshuffle tuned kernel.
+    """
+    return opus_gemm_a8w8_blockscale_bpreshuffle_tune(
+        x, weight, x_scale, w_scale, out, kernelId=kernel_id
+    )
+
+
 def generate_data(m, n, k, seed, device="cuda"):
     """
     Generate random data for testing the gemm a8w8 blockscale kernel.
@@ -221,9 +241,9 @@ def _setup_specific_arguments(self):
             "--libtype",
             type=str,
             default="all",
-            choices=["ck", "cktile", "asm", "all", "both"],
+            choices=["ck", "cktile", "asm", "opus", "all", "both"],
             required=False,
-            help="CK gemm a8w8 blockscale type to tune: ck, cktile, asm, both or all (covers all supported backends across standard/preshuffleB modes)",
+            help="CK gemm a8w8 blockscale type to tune: ck, cktile, asm, opus, both or all (covers all supported backends across standard/preshuffleB modes)",
         )
 
         self.parser.add_argument(
@@ -438,6 +458,49 @@ def get_gemm_a8w8_blockscale_tune_task(
                 )
         return tasks_ck
 
+    def get_gemm_a8w8_blockscale_opus_tune_task(
+        self,
+        info_keys,
+        seed,
+        preshuffleB,
+        run_kwargs,
+    ):
+        gfx, _, M, N, K = info_keys
+        if not preshuffleB or gfx != "gfx942":
+            return []
+
+        gemm_keys = ["x", "weight_shuffle", "x_scale_t", "w_scale", "out"]
+        ref_keys = ["x", "weight", "x_scale", "w_scale"]
+        ref_args = (ref_keys, None, dtypes.bf16)
+        tasks_opus = []
+        for kernel_id, kernel in gfx942_a8w8_kernels_list.items():
+            if N % kernel.B_N != 0 or K % kernel.B_K != 0:
+                continue
+            if not kernel.has_oob and M % kernel.B_M != 0:
+                continue
+            info = (info_keys, kernel_id, 0, kernel.name, "opus", preshuffleB)
+            gemm_args = (gemm_keys, kernel_id)
+            tasks_opus.append(
+                (
+                    info,
+                    generate_data,
+                    (M, N, K, seed),
+                    run_gemm_a8w8_blockscale_opus,
+                    gemm_args,
+                    dict(run_kwargs),
+                    run_torch,
+                    ref_args,
+                    {},
+                    None,
+                    1e-2,
+                    0.01,
+                    None,
+                    None,
+                    ("out",),
+                )
+            )
+        return tasks_opus
+
     def run_config(self, args):
         from aiter.ops.gemm_op_a8w8 import (
             gemm_a8w8_blockscale,
@@ -639,6 +702,15 @@ def tune(
                         run_kwargs,
                     )
                 )
+            if lib in ("opus", "all"):
+                task.extend(
+                    self.get_gemm_a8w8_blockscale_opus_tune_task(
+                        info_keys,
+                        seed,
+                        isPreshuffleB,
+                        run_kwargs,
+                    )
+                )
             shape_kernel_nums = len(task) - prev_task_count
 
             tasks_data.append((shape_kernel_nums, ()))

csrc/include/rocm_ops.hpp

@@ -293,6 +293,17 @@ namespace py = pybind11;
           py::arg("kernelId") = 0,            \
           py::arg("splitK")   = 0);
 
+#define OPUS_GEMM_A8W8_BLOCKSCALE_BPRESHUFFLE_TUNE_PYBIND \
+    m.def("opus_gemm_a8w8_blockscale_bpreshuffle_tune",   \
+          &opus_gemm_a8w8_blockscale_bpreshuffle_tune,    \
+          "opus_gemm_a8w8_blockscale_bpreshuffle_tune",   \
+          py::arg("XQ"),                                  \
+          py::arg("WQ"),                                  \
+          py::arg("x_scale"),                             \
+          py::arg("w_scale"),                             \
+          py::arg("Y"),                                   \
+          py::arg("kernelId"));
+
 #define OPUS_GEMM_WORKSPACE_INIT_PYBIND                              \
     m.def("opus_gemm_workspace_init",                                \
           &opus_gemm_workspace_init,                                 \

csrc/opus_gemm/codegen/common.py

@@ -15,11 +15,11 @@
 W3_KERNEL_PAIRS = {
     "a16w16_kbuf2v": "a16w16_kbuf2v_sk",
     "a16w16_kbuf2v_bk128": "a16w16_kbuf2v_bk128_sk",
-    "a16w16_kbuf1": "a16w16_kbuf1_sk",
     "a16w16_quad_mfma32_kbuf1": "a16w16_quad_mfma32_kbuf1_sk",
 }
 _NOSPLIT = tuple(W3_KERNEL_PAIRS.keys())
-_SPLITK = tuple(W3_KERNEL_PAIRS.values())
+_GFX942_SPLITK_ONLY = ("a16w16_kbuf1_sk",)
+_SPLITK = tuple(W3_KERNEL_PAIRS.values()) + _GFX942_SPLITK_ONLY
 _GFX942_A16W16_TAGS = (
     _SPLITK
     + (