[HIP] [OPUS] Unify OPUS GEMM/BMM interfaces and use Torch workspaces

.gitignore

@@ -93,6 +93,6 @@ aiter/ops/opus/configs/
 *.swo
 *.swp
 
-# Debug assembly dumps from csrc/opus_gemm/gen_co/build_co.py. The .co files
-# next to them ARE tracked (the wheel needs them); these are just for analysis.
+# Debug assembly dumps produced by the OPUS offline CO builder. The adjacent
+# .co runtime images remain tracked and are included in wheels.
 csrc/opus_gemm/gen_co/*/asm/

aiter/jit/core.py

@@ -472,13 +472,8 @@ def get_config_file(self, env_name, default_file, tuned_file_name):
 gfxs = get_gfx_list()
 AITER_ASM_DIR = f"{AITER_META_DIR}/hsa/"
 os.environ["AITER_ASM_DIR"] = AITER_ASM_DIR
-# Pre-compiled opus code objects (csrc/opus_gemm/gen_co/<arch>/<symbol>.co).
-# Resolved HERE, at import time, for the same reason AITER_ASM_DIR is: the
-# module also carries -DOPUS_GEN_CO_DIR, but that macro is baked when the module
-# is COMPILED, which for a prebuilt wheel is the build tree's aiter_meta/ (a
-# directory setup.py deletes when it is done). Only an import-time value knows
-# where the files ended up after install. The C++ loader prefers this env var
-# over the macro, so an explicit user override still wins.
+# Runtime root for pre-compiled OPUS code objects. Preserve an explicit user
+# override so locally rebuilt images can be tested without rebuilding Aiter.
 OPUS_GEN_CO_DIR = f"{AITER_CSRC_DIR}/opus_gemm/gen_co"
 os.environ.setdefault("OPUS_GEN_CO_DIR", OPUS_GEN_CO_DIR)
 
@@ -945,16 +940,7 @@ def MainFunc():
             flags_hip += ["-mllvm -amdgpu-coerce-illegal-types=1"]
         if get_gfx() != "gfx942" and int(os.getenv("AITER_FP4x2", "1")) > 0:
             flags_hip += ["-D__Float4_e2m1fn_x2"]
-        # Cluster launch is a HOST-side API question (hipDrvLaunchKernelEx +
-        # HIP_LAUNCH_CONFIG appear in ROCm 7.0), not a question about which GPU
-        # this machine has -- so the arch test must accept a cross-compile for
-        # gfx1250 the way the gfx1250 flags in optCompilerConfig.json already do.
-        # get_gfx() alone reads the LAST entry of a multi-arch GPU_ARCHS, which
-        # left "gfx1250;gfx942" building gfx1250 kernels whose cluster launch
-        # path was compiled out (AiterAsmKernelFast then rejects them at launch).
-        if (
-            get_gfx() == "gfx1250" or "gfx1250" in os.environ.get("GPU_ARCHS", "")
-        ) and hip_version >= Version("7.0.0"):
+        if "gfx1250" in gfxs and hip_version >= Version("7.0.0"):
             flags_hip += ["-DAITER_ENABLE_CLUSTER_LAUNCH"]
 
         if not torch_exclude:

aiter/jit/optCompilerConfig.json

@@ -328,7 +328,7 @@
             "'-mllvm --amdgpu-mfma-vgpr-form'",
             "('-mllvm -amdgpu-expert-scheduling-mode' if (get_gfx() == 'gfx1250' or 'gfx1250' in os.environ.get('GPU_ARCHS', '')) else '')",
             "('-mllvm -enable-post-misched=1' if (get_gfx() == 'gfx1250' or 'gfx1250' in os.environ.get('GPU_ARCHS', '')) else '')",
-            "f'-DOPUS_GEN_CO_DIR=\\\\\"{AITER_CSRC_DIR}/opus_gemm/gen_co\\\\\"'"
+            "f'-DOPUS_GEN_CO_DIR=\\\\\\\"{AITER_CSRC_DIR}/opus_gemm/gen_co\\\\\\\"'"
         ],
         "flags_extra_hip_per_source": {
             "*a8w8_bs_bpreshuf*.device.cu": [

aiter/ops/batched_gemm_op_a8w8.py

@@ -18,7 +18,9 @@
 from ..jit.utils.chip_info import get_gfx_runtime as get_gfx
 from ..jit.utils.torch_guard import torch_compile_guard
 from ..utility import dtypes
-from .gemm_op_common import get_padded_m
+from .opus.policy import (
+    resolve_a8w8_mxscale_bmm_plan as _resolve_a8w8_mxscale_bmm_plan,
+)
 
 
 def gen_batched_gemm_a8w8_fake_tensors(
@@ -88,7 +90,7 @@ def get_CKBatchedGEMM_config(
             get_CKBatchedGEMM_config.has_gfx = True
         else:
             logger.warning(
-                f"{AITER_CONFIGS.AITER_CONFIG_A8W8_BATCHED_GEMM_FILE} has no 'gfx' column -- "
+                f"{AITER_CONFIGS.AITER_CONFIG_A8W8_BATCHED_GEMM_FILE} has no 'gfx' column; "
                 "falling back to cu_num-only key. Re-run the tuner or migrate the CSV."
             )
             get_CKBatchedGEMM_config.ck_batched_gemm_dict = (
@@ -150,108 +152,23 @@ def batched_gemm_a8w8_CK(
 
 
 # ---------------------------------------------------------------------------
-# Shared tuned-CSV lookup for the mxscale batched GEMM.
-#
-# Shaped like tuned_gemm.py's multi-backend lookup: this layer locates the row
-# and never interprets the kernel identifier, since that differs per backend
-# (opus names kernels with an integer kernelId, flydsl with a kernelName). The
-# row comes back whole, libtype included, so a caller can dispatch on it;
-# libtype also filters up front for CSVs that carry one row per (shape, backend)
-# rather than a single cross-backend winner per shape.
-
-# Tuner bookkeeping rather than selection inputs, so the lookup log drops them
-# and stays readable.
-_TUNED_PERF_COLUMNS = ("us", "tflops", "bw", "errRatio")
+# gfx950 MXFP8 BMM high-level caller. Tuned-row and heuristic selection live
+# in ``opus.policy``; this module owns only the hot launch cache,
+# output allocation and split-one/workspace execution choice.
 
 
 @functools.cache
-def _load_mxscale_bmm_tuned(libtype: str | None = None) -> dict:
-    """{(gfx,b,m,n,k): row} from the mxscale BMM tuned CSV; {} if it is missing."""
-    path = AITER_CONFIGS.AITER_CONFIG_BATCHED_GEMM_A8W8_BLOCKSCALE_MXSCALE_FILE
-    try:
-        df = pd.read_csv(path).drop_duplicates()
-    except FileNotFoundError:
-        logger.warning("mxscale BMM tuned CSV not found at %s", path)
-        return {}
-    if libtype is not None and "libtype" in df.columns:
-        df = df[df["libtype"] == libtype]
-    return df.set_index(["gfx", "b", "m", "n", "k"]).to_dict("index")
+def _get_mxscale_bmm_launchers():
+    """Resolve the checked split-1 launcher and workspace planner once."""
+    from .opus import opus_bmm
+    from .opus.gemm_op_a8w8 import _opus_gemm_a8w8_mxscale_bmm_launch_raw
 
+    return _opus_gemm_a8w8_mxscale_bmm_launch_raw, opus_bmm
 
-@functools.lru_cache(maxsize=1024)
-def lookup_mxscale_bmm_config(
-    b: int, m: int, n: int, k: int, *, libtype: str | None = None
-):
-    """Exact tuned row for this shape, else one at a padded M.
-
-    Same exact-then-two-granularities walk over the shared C++ getPaddedM that
-    the CK / asm / a16w16 lookups use. A bucket table built from the CSV's own M
-    values was the alternative and bought nothing: over every M up to the
-    largest tuned one, both cover the same shapes and reach the same kernel on
-    131070 of 131072 M, so this keeps the one rounding rule the repo already has.
-
-    Cached per shape like get_CKGEMM_config, and for the same reason: getPaddedM
-    is a ctypes hop into C++ at ~10us, and the padded levels run on every call
-    whose M is not itself a tuned row. DPA+MTP decode is exactly that case (M is
-    the ragged token count a rank happened to get), and paying it once per layer
-    per step cost ~1% end-to-end before this. The row is shared, so callers must
-    treat it as read-only.
-
-    Returns the row, or None when no level hits. The log prints the row whole
-    instead of named fields, so a backend gets its own kernel identifier
-    reported without this layer knowing which column holds it.
-    """
-    gfx = get_gfx()
-    path = AITER_CONFIGS.AITER_CONFIG_BATCHED_GEMM_A8W8_BLOCKSCALE_MXSCALE_FILE
-    tuned = _load_mxscale_bmm_tuned(libtype)
 
-    row, padded_m = None, m
-    for gl in (None, 0, 1):
-        padded_m = m if gl is None else get_padded_m(m, n, k, gl)
-        row = tuned.get((gfx, b, padded_m, n, k))
-        if row is not None:
-            break
-
-    if row is None:
-        logger.info(
-            f"shape is B:{b}, M:{m}, N:{n}, K:{k}, not found tuned/padded config "
-            f"in {path}, the caller will fall back!"
-        )
-        return None
-
-    if AITER_LOG_TUNED_CONFIG:
-        cfg = {c: v for c, v in row.items() if c not in _TUNED_PERF_COLUMNS}
-        if padded_m == m:
-            logger.info(
-                f"shape is B:{b}, M:{m}, N:{n}, K:{k}, is tuned on gfx = {gfx} "
-                f"in {path}, config is {cfg}!"
-            )
-        else:
-            logger.info(
-                f"shape is B:{b}, M:{m}, N:{n}, K:{k}, exact miss on gfx = {gfx}; "
-                f"using padded_M: {padded_m} config {cfg} from {path}!"
-            )
-    return row
-
-
-# ---------------------------------------------------------------------------
-# fp8 e8m0 mxscale (block-scale) batched GEMM -- public entry for the family.
-#
-# This file is the per-family (a8w8 batched) public surface, not a CK-only
-# file: like aiter/ops/gemm_op_a8w8.py hosts gemm_a8w8 (CK rowwise) +
-# gemm_a8w8_blockscale (ck/cktile/triton/asm) side by side and lazy-imports
-# backend impls, we host the mxscale batched entry here too. The concrete
-# kernels stay in their backend dirs (opus -> aiter.ops.opus.bmm_op).
-#
-# Dispatch follows tuned_gemm.mm: look the shape up once here, then let the
-# winning row's libtype pick the backend, which is why the lookup runs
-# unfiltered -- the tuner writes one winning row per shape and its libtype says
-# who won. A second backend then only has to add rows and a branch below; it
-# does not repeat the lookup.
-
-# Untuned shapes go to opus: it is the backend carrying a shape heuristic for
-# rows the CSV does not have.
-_MXSCALE_BMM_DEFAULT_LIBTYPE = "opus"
+_MXSCALE_BMM_LAUNCH_PLANS: dict[
+    tuple[int, int, int, int], tuple[object, object, int, int]
+] = {}
 
 
 def _batched_gemm_a8w8_mxscale_impl(
@@ -261,40 +178,53 @@ def _batched_gemm_a8w8_mxscale_impl(
     w_scale: Tensor,
     dtype: torch.dtype = dtypes.bf16,
 ) -> Tensor:
-    """Eager tuned-CSV lookup + libtype dispatch; returns token-major [M, G, N].
-
-    Kept unwrapped (plain Python) so tests can introspect the real dispatch
-    (which kernelId a shape resolves to) on meta tensors. The public
-    ``batched_gemm_a8w8_mxscale`` is the torch.compile-guarded custom op over
-    this; a caller that must write into its own (e.g. batch-major) output buffer
-    calls the opus backend (``aiter.ops.opus.bmm_op.bmm_a8w8_mxscale_opus``)
-    directly, which keeps the ``out=`` argument.
-    """
-    from .opus.bmm_op import bmm_a8w8_mxscale_opus
-
-    m, g, k = int(x.shape[0]), int(x.shape[1]), int(x.shape[2])
-    n = int(wo_a.shape[1])
-
-    cfg = lookup_mxscale_bmm_config(g, m, n, k)
-    libtype = cfg["libtype"] if cfg is not None else _MXSCALE_BMM_DEFAULT_LIBTYPE
-    if libtype != "opus":
-        raise NotImplementedError(
-            f"tuned row for B:{g}, M:{m}, N:{n}, K:{k} wants libtype "
-            f"{libtype!r}, which has no batched mxscale backend here yet"
+    # This body executes behind the public custom-op boundary, so real eager
+    # tensors carry concrete integer dimensions here.  Avoid four redundant
+    # Python int() conversions on every short BMM launch.
+    m, g, k = x.shape
+    n = wo_a.shape[1]
+    plan_key = (g, m, n, k)
+    try:
+        raw_launch, opus_bmm, kid, split_k = _MXSCALE_BMM_LAUNCH_PLANS[plan_key]
+    except KeyError:
+        raw_launch, opus_bmm = _get_mxscale_bmm_launchers()
+        kid, split_k = _resolve_a8w8_mxscale_bmm_plan(g, m, n, k)
+        _MXSCALE_BMM_LAUNCH_PLANS[plan_key] = (
+            raw_launch,
+            opus_bmm,
+            kid,
+            split_k,
         )
 
-    # Reading opus columns is this branch's job; whether that kernel can run
-    # this M, and what to do when it cannot, is the backend's.
-    return bmm_a8w8_mxscale_opus(
-        x,
+    Y = torch.empty((m, g, n), dtype=dtype, device=x.device)
+    if split_k <= 1:
+        # The shape resolver already returns a final canonical global kid.
+        # Enter the checked C++ launcher directly for the common no-workspace
+        # path instead of repeating the unified public routing contract.  The
+        # C++ boundary still validates dtype, shape, device, stride, arch and
+        # exact kid.  Workspace cases retain the unified Python planner below.
+        raw_launch(
+            x,
+            wo_a,
+            Y,
+            x_scale,
+            w_scale,
+            None,
+            kid,
+            max(1, split_k),
+        )
+        return Y
+    opus_bmm(
+        x.transpose(0, 1),
         wo_a,
-        x_scale,
-        w_scale,
-        None,
-        dtype=dtype,
-        kernelId=int(cfg["kernelId"]) if cfg is not None else None,
-        splitK=int(cfg["splitK"]) if cfg is not None else None,
+        Y.transpose(0, 1),
+        kid=kid,
+        layout="mxscale_bmm",
+        x_scale=x_scale.transpose(0, 1),
+        w_scale=w_scale,
+        split_k=split_k,
     )
+    return Y
 
 
 def _batched_gemm_a8w8_mxscale_fake(
@@ -304,7 +234,6 @@ def _batched_gemm_a8w8_mxscale_fake(
     w_scale: Tensor,
     dtype: torch.dtype = dtypes.bf16,
 ) -> Tensor:
-    # token-major [M, G, N]; mirrors the eager allocation in bmm_a8w8_mxscale_opus.
     return torch.empty(
         (x.shape[0], x.shape[1], wo_a.shape[1]),
         dtype=dtype,
@@ -320,31 +249,7 @@ def batched_gemm_a8w8_mxscale(
     w_scale: Tensor,
     dtype: torch.dtype = dtypes.bf16,
 ) -> Tensor:
-    """fp8 e8m0 mxscale (128x128 block-scale) batched GEMM.
-
-    mmajor DSV4 wo_a layout (matches the opus kernels + op test):
-
-    * ``x``       : [M, G, K] fp8 activation (per-token e8m0; transposed view
-                    of batch-major [G, M, K]).
-    * ``wo_a``    : [G, N, K] fp8 weight (batch-major).
-    * ``x_scale`` : [M, G, K/128] uint8 e8m0 activation scale.
-    * ``w_scale`` : [G, N/128, K/128] uint8 e8m0 weight scale.
-
-    Returns a fresh **token-major** [M, G, N] output. This entry is
-    torch.compile-guarded (registered as an ``aiter::`` custom op with a meta
-    kernel), so a framework can call it inside a compiled graph without the
-    tuned-CSV lookup / heuristic being traced. A caller that must write into its
-    own preallocated (e.g. batch-major) buffer uses
-    ``aiter.ops.opus.bmm_op.bmm_a8w8_mxscale_opus`` directly (it keeps ``out=``).
-
-    Note this is *microscaling* (e8m0) block scale -- distinct from
-    ``gemm_a8w8_blockscale`` which uses fp32 block scale. Scale type is baked
-    into the name so a future fp32-block batched variant stays separate.
-
-    The shape is looked up in the tuned CSV and the winning row's libtype picks
-    the backend. No kernel override lives on this entry: how a kernel is named is
-    backend-specific, so pin one at the backend (aiter.ops.opus.bmm_op).
-    """
+    """Run gfx950 E8M0 MXFP8 BMM and return token-major ``[M,G,N]``."""
     return _batched_gemm_a8w8_mxscale_impl(x, wo_a, x_scale, w_scale, dtype=dtype)
 
 

aiter/ops/deepgemm.py

@@ -6,7 +6,7 @@
 Hosts the CK-backed `deepgemm_ck` binding plus a thin `deepgemm()`
 wrapper. Opus entries have been extracted under `aiter.ops.opus.*`;
 see `aiter.ops.opus.gemm_a16w16_opus` for BF16 matmul and
-`aiter.ops.opus.opus_gemm_a16w16_tune` for id-based kernel selection.
+`aiter.ops.opus.opus_gemm` for exact-id kernel selection.
 
 `opus_gemm_a16w16_tune` is kept here as a deprecation shim for one
 release to ease migration from the old aggregate entry.
@@ -18,7 +18,7 @@
 from torch import Tensor
 
 from ..jit.core import compile_ops
-from .opus.gemm_op_a16w16 import opus_gemm_a16w16_tune as _opus_tune
+from . import opus as _opus
 
 
 @compile_ops("module_deepgemm", fc_name="deepgemm")
@@ -52,12 +52,22 @@ def opus_gemm_a16w16_tune(
 ) -> torch.Tensor:
     warnings.warn(
         "aiter.ops.deepgemm.opus_gemm_a16w16_tune has moved to "
-        "aiter.ops.opus.gemm_op_a16w16.opus_gemm_a16w16_tune; this "
+        "aiter.ops.opus.opus_gemm; this "
         "shim will be removed in a future release.",
         DeprecationWarning,
         stacklevel=2,
     )
-    return _opus_tune(XQ, WQ, Y, kernelId, splitK)
+    if all(tensor.dim() == 3 and tensor.size(0) == 1 for tensor in (XQ, WQ, Y)):
+        _opus.opus_gemm(
+            XQ.squeeze(0),
+            WQ.squeeze(0),
+            Y.squeeze(0),
+            kid=kernelId,
+            split_k=splitK,
+        )
+    else:
+        _opus.opus_bmm(XQ, WQ, Y, kid=kernelId, split_k=splitK)
+    return Y
 
 
 __all__ = [

aiter/ops/gemm_op_a8w8.py

@@ -1060,7 +1060,10 @@ def gemm_a8w8_blockscale_bpreshuffle(
     )
     # Triton path first: it allocates its own output, so skip the Y buffer the
     # ck/asm paths below need.
-    if (config is not None and config["libtype"] == "triton") or get_gfx() == "gfx1250":
+    if (config is not None and config["libtype"] == "triton") or (
+        get_gfx() == "gfx1250"
+        and not (config is not None and config["libtype"] == "opus")
+    ):
         # kernelName optionally carries the backend hint ("triton"/"gluon");
         # anything else -> None (auto gluon->triton detection). config=None lets
         # the triton impl load its own tuned config internally. WQ is already
@@ -1103,13 +1106,18 @@ def gemm_a8w8_blockscale_bpreshuffle(
                 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,
-            )
+            # Only an OPUS-tuned row enters the OPUS exact-kid interface.
+            kid = int(config["kernelId"])
+            from aiter.ops.opus import opus_gemm
 
-            return opus_gemm_a8w8_blockscale_bpreshuffle_tune(
-                XQ, WQ, x_scale, w_scale, Y, kernelId=kernelId
+            return opus_gemm(
+                XQ,
+                WQ,
+                Y,
+                kid=kid,
+                layout="bpreshuffle",
+                x_scale=x_scale,
+                w_scale=w_scale,
             )
         elif libtype == "flydsl" and is_flydsl_available():
             return gemm_a8w8_mxfp8_128_bpreshuffle_flydsl(