diff --git a/aiter/ops/batched_gemm_op_a8w8.py b/aiter/ops/batched_gemm_op_a8w8.py index 6bb62131eb..b5281e41e8 100644 --- a/aiter/ops/batched_gemm_op_a8w8.py +++ b/aiter/ops/batched_gemm_op_a8w8.py @@ -19,6 +19,9 @@ 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 +91,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,17 +153,9 @@ 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. +# 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. _TUNED_PERF_COLUMNS = ("us", "tflops", "bw", "errRatio") @@ -173,6 +168,15 @@ def _mxscale_bmm_tuned_path(bpreshuffle: bool) -> str: ) +@functools.cache +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.cache def _load_mxscale_bmm_tuned( libtype: str | None = None, bpreshuffle: bool = False @@ -251,24 +255,14 @@ def lookup_mxscale_bmm_config( 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" +@functools.lru_cache(maxsize=1024) +def _get_mxscale_bmm_launch_plan( + g: int, + m: int, + n: int, + k: int, +) -> tuple[int, int]: + return _resolve_a8w8_mxscale_bmm_plan(g, m, n, k) def _batched_gemm_a8w8_mxscale_impl( @@ -278,41 +272,43 @@ 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 does not take a raw [G, N, K] weight; " - f"{libtype!r} rows are served by batched_gemm_a8w8_mxscale_bpreshuffle" + # 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] + raw_launch, opus_bmm = _get_mxscale_bmm_launchers() + kid, split_k = _get_mxscale_bmm_launch_plan(g, m, n, k) + + 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), ) - - # 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, + 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( @@ -322,7 +318,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, @@ -338,31 +333,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) diff --git a/aiter/ops/gemm_op_a8w8.py b/aiter/ops/gemm_op_a8w8.py index f1addcb3d5..dcad6b4f7c 100644 --- a/aiter/ops/gemm_op_a8w8.py +++ b/aiter/ops/gemm_op_a8w8.py @@ -1095,12 +1095,16 @@ def gemm_a8w8_blockscale_bpreshuffle( ) elif libtype == "opus": kernelId = int(config["kernelId"]) - from aiter.ops.opus.gemm_op_a8w8 import ( - opus_gemm_a8w8_blockscale_bpreshuffle_tune, - ) + 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=kernelId, + layout="bpreshuffle", + x_scale=x_scale, + w_scale=w_scale, ) elif libtype == "flydsl": return gemm_a8w8_mxfp8_128_bpreshuffle_flydsl( diff --git a/aiter/ops/opus/README.md b/aiter/ops/opus/README.md index 41960a9e91..6a7bf7a940 100644 --- a/aiter/ops/opus/README.md +++ b/aiter/ops/opus/README.md @@ -1,1108 +1,479 @@ -# Opus a16w16 GEMM +# OPUS GEMM and BMM Python interfaces -BF16 × BF16 → BF16/FP32 matmul backed by the opus kernel family (AMD -gfx950 / MI300X class). Provides a shape-driven Python API, a runtime -dispatcher with CSV-baked lookup + heuristic fallback, and a tuning -pipeline that populates the lookup. +OPUS exposes strict exact-kid functions for logical 2D GEMM and batch-first 3D +BMM, plus the retained shape-driven `gemm_a16w16_opus` compatibility entry. +The exact functions never select a kernel from the shape; the compatibility +entry resolves an A16W16 kid before entering the same exact path. -Underlying JIT module: `module_deepgemm_opus` -(see `aiter/jit/optCompilerConfig.json`). - ---- - -## 1. Quick Start +## Public API ```python import torch -from aiter.ops.opus import gemm_a16w16_opus - -A = torch.randn(1024, 4096, device="cuda", dtype=torch.bfloat16) -B = torch.randn(2048, 4096, device="cuda", dtype=torch.bfloat16) # [N, K] - -Y = gemm_a16w16_opus(A, B) # bf16 output, tune table + C++ heuristic -Y = gemm_a16w16_opus(A, B, dtype=torch.float32) # fp32 output -Y = gemm_a16w16_opus(A, B, out=preallocated_Y) # reuse buffer -``` - -First call triggers a JIT build of `module_deepgemm_opus` (~11s on -the dev container; see [§7.6](#76-compile-time-techniques)). -Subsequent Python processes reuse the compiled `.so`. - -**Inputs**: `A` is `[M, K]` or `[batch, M, K]` bf16. `B` is bf16 -(plain layout, not pre-shuffled) in one of two shapes: - -- `[N, K]` — only when `batch == 1`. -- `[batch, N, K]` — must be contiguous (strides `(N*K, K, 1)`); broadcast - views like `B.unsqueeze(0).expand(batch, -1, -1)` are rejected because - the opus launcher hardcodes `stride_b_batch == N*K`. Use - `B.expand(batch, -1, -1).contiguous()` (or pass a real per-batch - weight) when you need to broadcast. - -**Output**: `[M, N]` or `[batch, M, N]`. bf16 and fp32 both supported on -all bias-aware kid families (split-barrier 4..9 and a16w16_flatmm_splitk -200..299) — the splitk reduce kernel templated on `D_OUT` selects the -right path at launch time. - -**Optional bias** (per-row, broadcast across N): pass via `bias=` with one -of two shapes: - -- `[M]` — broadcast across batch; requires `batch == 1`. -- `[batch, M]` — per-batch row vector. - -bias dtype must equal `dtype` (match-output convention). Bias is folded -into the fp32 accumulator before cast → output. Both the CSV-tuned and -heuristic-fallback paths carry bias through; see [§2 Dispatch](#2-how-dispatch-works) -for the routing rules. - -### Constraints (hard rejects) - -The launchers reject these up front to avoid silent miscompares: - -| Constraint | Why | -|---|---| -| `K` must be **even** | The splitk pipeline accumulates a ~3-7% error on odd K (latent K-tail bug); split-barrier independently requires `ceil_div(K, B_K)` even. | -| `A` and `B` dtype = bf16 | a16w16-family kernels lock the input dtype at launch time. | -| `B` is `[N, K]` only when `batch == 1`; otherwise `[batch, N, K]` contiguous | `stride_b_batch == N*K` is hardcoded; broadcast views silently corrupt. | -| pre-shuffled B | not supported; pass plain layout. | -| `bias.dtype == dtype` | match-output; otherwise a host-side `TORCH_CHECK` fires. | -| `bias` shape ∈ {`[M]` (only `batch==1`), `[batch, M]`} | reduce / split-barrier kernels expect this exact layout. | -| GPU arch must be **gfx950 (MI350)** today | opus uses gfx950-only intrinsics (MFMA-32x32x16, ds_read_b64_tr) and the 160 KiB LDS budget. Three-layer enforcement: Python import-time `_detect_arch` swaps `gemm_a16w16_opus` / `opus_gemm_a16w16_tune` for stubs and emits a `RuntimeWarning` on non-gfx950 devices (the import itself succeeds — calling the stubs raises `RuntimeError`); the C++ host dispatcher routes per `gcnArchName` and currently only implements the gfx950 branch (others fail with a clear "pipeline TBD" message); each `__global__` kernel body wraps real code in `#if defined(__gfx950__)` so multi-arch wheels (e.g. `GPU_ARCHS='gfx942;gfx950'`) still compile, but the gfx942 device pass produces an empty kernel stub that is unreachable at runtime. To add support for a new arch, extend `OpusGfxArch` in `csrc/opus_gemm/opus_gemm.cu` and add a per-arch dispatch function. | - -Scale / FP8 paths are handled by other opus submodules (a8w8 / -a8w8_blockscale, landing in follow-up PRs); they share the same -`module_deepgemm_opus` JIT build but expose their own Python entry -under `aiter.ops.opus.*`. Those paths currently reject non-empty -`bias` at the dispatcher. - ---- - -## 2. How Dispatch Works - -When the user calls `gemm_a16w16_opus(A, B)` without an explicit kernel -id, the wrapper does **one** lookup against the global aiter BF16 tuned -CSVs (filtered by `libtype == 'opus'`), then falls through to the C++ -dispatcher: +from aiter.ops.opus import gemm_a16w16_opus, opus_bmm, opus_gemm + +opus_gemm( # XQ [M,K], WQ [N,K], Y [M,N] + XQ, + WQ, + Y, + *, + kid, + layout="plain", + x_scale=None, + w_scale=None, + bias=None, + split_k=0, + workspace=None, +) + +opus_bmm( # XQ [B,M,K], WQ [B,N,K], Y [B,M,N] + XQ, + WQ, + Y, + *, + kid, + layout="plain", + x_scale=None, + w_scale=None, + bias=None, + split_k=0, + workspace=None, +) + +# Retained A16W16 shape-driven API. An explicit kernelId wins; otherwise this +# performs OPUS-only tuned lookup followed by the per-architecture heuristic. +result = gemm_a16w16_opus(A, B, bias=None, dtype=torch.bfloat16) ``` -gemm_a16w16_opus(A, B, bias=...) - ├─ explicit kernelId=N? ───yes──► opus_gemm_a16w16_tune(N, ..., bias) - │ (C++ dispatcher TORCH_CHECKs that kid is bias-aware - │ when bias.has_value()) - │ - ├─ Python-side global-CSV lookup ───hit──► opus_gemm_a16w16_tune(solidx, splitK, bias) - │ (scans aiter/configs/bf16_tuned_gemm.csv + - │ aiter/configs/model_configs/*_bf16_tuned_gemm.csv, - │ filters `libtype=='opus'`, key = - │ (cu_num, M, N, K, bias, dtype, outdtype, scaleAB, - │ bpreshuffle); cached for process lifetime) - │ - └─ miss ──► opus_gemm(..., bias) [C++] - ├─ C++ compile-time (M,N,K) lookup - │ (same global CSV opus rows baked into - │ opus_gemm_lookup.h at JIT-codegen time; - │ key is (M,N,K) only, bias forwarded to - │ the matched launcher) - └─ miss ──► opus_a16w16_heuristic_kid_gfx950 - (M-bucket rule -> integer kid -> tune_lookup - table; always returns a bias-aware kid so - bias is safe to forward unconditionally) -``` - -There is **one** CSV source of truth now: the global aiter BF16 tuned -CSVs. The opus runtime dispatch (`aiter/ops/opus/common.py`) reads opus -rows live every new process; CSV edits take effect immediately on the -Python side. The C++ side bakes the same opus rows into -`opus_gemm_lookup.h` at JIT-codegen time via -`gen_instances.py --tune_files`, and **requires `AITER_REBUILD=1` to -pick up CSV edits**. - -The heuristic-fallback path no longer hardcodes launcher symbol names. -`opus_a16w16_heuristic_kid_gfx950(M, N, K)` returns an integer kid, and -the caller resolves it through `opus_a16w16_tune_dispatch_gfx950<>` (the -same table that powers `opus_gemm_a16w16_tune`). The kids the heuristic -can return are listed in `HEURISTIC_DEFAULT_KIDS` in -`csrc/opus_gemm/opus_gemm_common.py`; `gen_instances.py` asserts they -are all in the subset-compile set `S` before writing -the generated compiled-kid sidecar. - -### Subset compile - -`module_deepgemm_opus` only compiles the kids it actually needs, not -the full `kernels_list`. The compile set `S` is the union of: - -1. Kids referenced by the **global tuned CSVs** with `libtype == 'opus'`. -2. Kids from the last successful **sidecar** at - `{bd_dir}/compiled_kids_opus.json`, outside the per-module build directory. -3. The current target architectures' `HEURISTIC_DEFAULT_KIDS`, so heuristic - fallback always has a compiled kernel. -4. The applicable a8w8 kids, since the opus `.so` also exposes a8w8 dispatch. -5. Additional tuner candidates passed through `--extra_kids`. - -The final set is restricted to valid kids and the target architectures. An -explicit `--extra_kids` request that cannot survive these filters or a -`--kernel_tag` restriction fails codegen; it is not silently omitted. The set -size depends on the target architectures, CSVs and previous tuning requests. - -The tuner synchronously compiles new candidates before starting its workers. -It does not expand the canonical sidecar in advance: JIT first installs the -binary, then publishes the generated sidecar and a receipt binding its contents -to that installed binary. The canonical files survive `clear_build`, while -generated working files live in the module's `blob.staging` directory. Runtime -dispatch uses the CSV/C++ lookup tables, not the sidecar. See -[transactional JIT cache](../../../docs/jit_cache.md) for recovery, permissions -and storage requirements. - -Explicit `kernelId=` bypass exists for tuning, debugging, and future -integrations (e.g. `aiter.tuned_gemm.solMap["opus"]`). The C++ -dispatcher gates `bias` to bias-aware kid ranges (split-barrier 4..9 -or a16w16_flatmm_splitk 200..299) when `bias.has_value()`; passing -bias to a non-bias-aware kid is a hard error. - ---- - -## 3. Tuning Your Shapes - -For production tuning use **gradlib**: it integrates opus alongside -asm / triton / skinny / flydsl / torch / hipblaslt backends and writes -to the global tuned CSV. - -### 3.1 Production: gradlib with `--libtype opus` -```bash -# Tune only opus, single shape (or pass --input_file to sweep a CSV): -python3 gradlib/gemm_tuner.py --libtype opus \ - --input_file aiter/configs/bf16_untuned_gemm.csv - -# Or tune all backends in one pass; gradlib picks the winning libtype -# per shape: -python3 gradlib/gemm_tuner.py --libtype all \ - --input_file aiter/configs/bf16_untuned_gemm.csv - -# Output path follows gradlib's existing --tuned_file / GTUNE_TUNED CLI; -# default is aiter/configs/bf16_tuned_gemm.csv. To write to a sandbox -# during testing: -GTUNE_TUNED=/tmp/test_tuned.csv \ - python3 gradlib/gemm_tuner.py --libtype opus --input_file ... +For `opus_gemm` and `opus_bmm`, `kid` is mandatory. `Y` is caller-owned and is +returned unchanged after the launch. The selected exact function determines +the logical rank, while the resolved family must support that operation; +dtype does not determine it. +Among A8 families, no-scale, blockscale and blockscale-bpreshuffle are +GEMM-only, while MXFP8 is BMM-only. The disjoint architecture id bands and the +merged `kernels_list` form the canonical registry. Both entries call +`kernels_list.get(kid)`; they do not introduce or renumber ids. The returned +instance tag plus the dtype/layout arguments determine the private family +adapter, rather than a second selector or a numeric-range guess. + +## Dispatch model + +```text +caller-resolved final kid + -> opus_gemm (strict 2D) or opus_bmm (strict batch-first 3D) + -> existing kernels_list.get(kid) + -> instance arch/tag metadata + -> family-local dtype/layout/scale checks + -> A16W16 or A8W8 family adapter + -> shared immutable A16 launch plan or A8 family planner + -> family executor + -> unchanged exact-kid C++ family table ``` -gradlib stamps every opus row with `libtype='opus'` so the opus runtime -dispatch picks it up via the libtype filter in -`aiter/ops/opus/common.py`. Other backends' rows (asm / triton / ...) in -the same CSV stay there and are picked up by their respective dispatch -modules. - -**Candidate kid selection**: rather than benchmarking every kid for -every shape, gradlib (via `candidate_kids_for_shape` in -`opus_gemm_common.py`) uses an occupancy heuristic on a `128 × 128` -proxy tile: if `ceil(M/128) * ceil(N/128) < 2 * cu_num`, only splitk -kids are tuned (a non-splitk tile can't fill the device twice); for -larger problems both splitk and non-splitk classes compete. Two -structural fallbacks force splitk-only: K not aligned for non-splitk -launchers (need `K%64==0` and `ceil(K/64)%2==0`), and bias=True when -the candidate set has no bias-aware non-splitk kids. - -**Missing candidates trigger a synchronous rebuild**: the tuner passes new -kids through `--extra_kids` and waits for `module_deepgemm_opus` to finish -building before spawning workers. A matching sidecar alone is insufficient -to skip this step: its receipt must also match the current installed binary. -Missing or stale metadata conservatively triggers a rebuild. An explicit -`AITER_REBUILD` request is honored once in the parent even on a cache hit; -successful preparation sets the environment to `AITER_REBUILD=0` for workers. -On failure the original environment is restored and the canonical sidecar -is not advanced by the failed compile. Build time depends on the requested -set and available compiler resources. - -### 3.2 Debug-only: `opus_gemm_tune.py` (single shape / kid) - -```bash -python3 csrc/opus_gemm/opus_gemm_tune.py \ - -m 128 -n 2880 -k 4096 --dtype bf16 --outdtype bf16 -# default -o is /tmp/opus_debug_tuned.csv (NOT aiter/configs/) +The public operation split does not duplicate kernels, workspace allocation, +or raw bindings. A logical GEMM becomes a batch-one view at the family +boundary. A logical BMM keeps its batch-first public layout; the MXScale +adapter alone converts activation/output tensors to the raw kernel's existing +M-major views with `transpose(0, 1)`, which does not copy storage. The physical +3D raw ABI used by a non-MX A8 GEMM is not exposed as public BMM. + +There is no tuned-CSV lookup, architecture heuristic, redirect, or framework +fallback inside either exact public path. The two shape-driven A16 callers +remain intentionally different: + +```text +aiter.gemm_a16w16 + -> global multi-backend tuned row -> selected backend + -> no valid row -> skinny (eligible gfx90a/gfx942/gfx950) + -> gfx1250 Triton + -> otherwise PyTorch + +gemm_a16w16_opus + -> explicit kernelId -> legacy requested-to-actual resolution + -> otherwise OPUS-only tuned row -> the same compatibility resolution + -> missing/invalid OPUS row -> per-arch OPUS heuristic + -> validate -> local exact A16 GEMM/BMM launcher ``` -This is retained for single-shape smoke / debug runs against a -specific kid. **It never writes to the global aiter/configs/ tree** -unless you explicitly pass `-o aiter/configs/bf16_tuned_gemm.csv`, -which is discouraged -- use gradlib for that. - -Verify winners with the end-to-end test: - -```bash -python3 op_tests/test_opus_a16w16_gemm.py -m 128 -n 256 -k 1024 -b 1 -# expected: allclose passed -``` +The shared OPUS candidate helpers are isolated in `policy.py`. +`tuned_gemm.py` validates tuned candidates and keeps its normal framework +fallback. The OPUS-only compatibility entry warns once and uses its heuristic +for a stale tuned row, but rejects an invalid explicit id. It applies legacy +gfx942 requested-to-actual resolution before calling the local exact launcher; +the strict `opus_gemm`/`opus_bmm` APIs never redirect. ---- +There is currently no high-level A16W16 BF16 BMM wrapper. In particular, +`aiter/ops/batched_gemm_op_bf16.py` contains the existing CK entry points but +does not define `batched_gemm_bf16_OPUS` or a tuned CK/OPUS dispatcher. A16W16 +BMM therefore starts at exact-kid `opus_bmm`: its caller owns `Y`, resolves the +final `kid`/`split_k`, and may provide a Torch workspace. The public router +calls `_launch_a16w16_bmm`, which preserves the batch dimension and forwards +to the same `_execute_a16w16` planner/executor used by A16W16 GEMM. -## 4. API Reference +## Current families -### `gemm_a16w16_opus(A, B, bias=None, dtype=bf16, *, kernelId=None, splitK=None, out=None)` +| Registry family | Current route | Public operation and dtype rules | +|---|---|---| +| `a16w16` | gfx942, gfx950, gfx1250 | GEMM or BMM; BF16 `XQ/WQ`, normally BF16 or FP32 `Y`, plain WQ, optional bias/split-K/Torch workspace; gfx942 BF16-workspace exact kids require BF16 `Y` | +| `a8w8` | gfx950 kid 2 | GEMM only; FP8 `XQ/WQ`, FP32 `Y`, plain WQ, no scales | +| `a8w8_blockscale` | gfx950 kid 1 | GEMM only; FP8 `XQ/WQ`, FP32 `Y`, plain WQ, two FP32 scales | +| `a8w8_blockscale_bpreshuffle` | gfx942 kid 11000 | GEMM only; FP8 `XQ/WQ`, BF16 `Y`, pre-shuffled WQ, two FP32 scales | +| `a8w8_mxscale_bmm` | gfx950 global kids 8000--8653 (45 registered ids) | BMM only; batch-first FP8 inputs, E8M0 scales, BF16 or FP32 output, optional split-K Torch workspace | -Primary user entry. Implemented in -[aiter/ops/opus/gemm_op_a16w16.py](gemm_op_a16w16.py). +Empty family tables on another architecture are valid capability states. A +kid registered for another architecture is rejected before family launch. -| Param | Type | Default | Notes | -|---|---|---|---| -| `A` | Tensor | required | `[M, K]` or `[batch, M, K]` bf16. **K must be even.** | -| `B` | Tensor | required | `[N, K]` (batch=1 only) or contiguous `[batch, N, K]` bf16; broadcast views are rejected (see §1). | -| `bias` | Tensor? | `None` | Optional per-row bias. Shape `[M]` (broadcast across batch; requires `batch == 1`) or `[batch, M]`. dtype must equal `dtype` (match-output). Bias is folded into fp32 acc before cast. Honored on split-barrier (kid 4..9) and splitk (kid 200..299) families; the C++ dispatcher rejects bias on other kids. CSV-miss requests fall through to the heuristic dispatcher, which always returns a bias-aware kid. | -| `dtype` | torch.dtype | `bf16` | Output dtype; both `bf16` and `fp32` are supported on every kid family (the splitk reduce kernel templated on `D_OUT` casts at launch time). | -| `kernelId` | int? | `None` | Override: skip CSV/heuristic and launch this specific instance. With `bias is not None`, must be a kid in `[4, 10) ∪ [200, 300)`. | -| `splitK` | int? | `None` | Only honored with explicit `kernelId`; literal KBatch for splitk kids. | -| `out` | Tensor? | `None` | Reuse a preallocated output buffer. | +## Examples -### `opus_gemm_a16w16_tune(XQ, WQ, Y, bias=None, kernelId=0, splitK=0)` +### A16W16 GEMM and BMM -Low-level id-based dispatcher. Used by the tuner and the high-level -wrapper. Accepts 3D inputs only (`[batch, M, K]`, `[batch, N, K]`, -`[batch, M, N]`) and requires contiguous strides on all three tensors -(`(M*K, K, 1)`, `(N*K, K, 1)`, `(M*N, N, 1)` respectively); a Python -guard raises `NotImplementedError` for broadcast / transpose / slice -views before launching the kernel. `bias` is optional and follows the -same shape / dtype rules documented for `gemm_a16w16_opus`. +```python +import torch +from aiter.ops.opus import opus_bmm, opus_gemm -For backwards compatibility, the legacy 5-arg call form -`opus_gemm_a16w16_tune(XQ, WQ, Y, kernelId, splitK)` (positional `int` -in slot 4) still works: when the 4th positional argument is an `int`, -it is silently reinterpreted as `kernelId` and the rest of the args -shift accordingly. Mixed-style calls -(`..., bias=t, kernelId=k`) keep their kwargs semantics. Prefer -`gemm_a16w16_opus` unless you need explicit control. +XQ = torch.randn((64, 512), device="cuda", dtype=torch.bfloat16) +WQ = torch.randn((64, 512), device="cuda", dtype=torch.bfloat16) +Y = torch.empty((64, 64), device="cuda", dtype=torch.bfloat16) -### Legacy shim +# The caller/tuner has already chosen gfx950 kid 200 and split_k 2. +opus_gemm(XQ, WQ, Y, kid=200, split_k=2) -`aiter.ops.deepgemm.opus_gemm_a16w16_tune` still works and forwards to -`aiter.ops.opus.*` with a `DeprecationWarning`; scheduled for removal -one aiter minor release later. +XQ_b = torch.randn((8, 64, 512), device="cuda", dtype=torch.bfloat16) +WQ_b = torch.randn((8, 64, 512), device="cuda", dtype=torch.bfloat16) +Y_b = torch.empty((8, 64, 64), device="cuda", dtype=torch.bfloat16) +opus_bmm(XQ_b, WQ_b, Y_b, kid=200, split_k=2) +``` -`aiter.ops.deepgemm.deepgemm_opus` (the old aggregate entry that -exposed FP8 grouped + a16w16 no-scale through a single function) has -been **removed** along with any internal opus binding in that module. -Migration: +`opus_gemm` requires 2D tensors; `opus_bmm` requires batch-first 3D tensors. +Inputs are K-contiguous and `Y` is N-contiguous. gfx1250 two-stage workspace +kernels require BMM batch one; pre-built CO kernels support batched inputs. +Exact instances can impose additional tile, output +dtype, bias, or K-loop constraints. The BMM example is a direct exact-API call; +there is no current `batched_gemm_bf16_OPUS` high-level wrapper. -- BF16 no-scale GEMM: use `gemm_a16w16_opus` from this module. -- FP8 grouped GEMM: future `aiter.ops.opus.a8w8*` modules (separate - PR). Until they land, bind `opus_gemm` yourself via `compile_ops` - against `module_deepgemm_opus` / `fc_name="opus_gemm"`. +A16 bias follows the `F.linear` output-feature convention: `[N]` broadcasts +across batch and `[batch,N]` supplies a separate bias for each batch. -`aiter.ops.deepgemm.deepgemm()` is now a thin forwarder to -`deepgemm_ck`; the `AITER_DEEPGEMM_BACKEND=opus` dispatch env is no -longer recognized. +### gfx950 A8W8 without scales ---- +```python +from aiter.ops.opus import opus_gemm -## 5. Testing +Y = torch.empty((M, N), device=XQ.device, dtype=torch.float32) +opus_gemm(XQ, WQ, Y, kid=2) +``` -All tests run inside the project's container -(`docker exec -w /wksp/aiter demon_test bash -lc ...`). +The general `aiter.gemm_a8w8` API remains the scaled CK/Triton operation and +requires both `x_scale` and `w_scale`; omitting scales does not select OPUS. -The single end-to-end test exercises `gemm_a16w16_opus` (shape-driven -API). It supports both single-shape smoke runs and CSV sweeps (e.g. -the gptoss untuned set). +### gfx950 A8W8 blockscale -| Test | Purpose | Pass criterion | -|---|---|---| -| `op_tests/test_opus_a16w16_gemm.py` | End-to-end test of `gemm_a16w16_opus` (shape-driven API); supports single-shape smoke and CSV sweep | `allclose` passes on all shapes | +```python +from aiter.ops.opus import opus_gemm + +opus_gemm( + XQ, + WQ, + Y, + kid=1, + x_scale=x_scale, + w_scale=w_scale, +) +``` -Examples: +The group contract is 1x128x128. GEMM scales are contiguous FP32 +`[M,K/128]` and `[N/128,K/128]` tensors. This family does not accept +`opus_bmm`. The general `aiter.gemm_a8w8_blockscale` dispatcher remains a +BF16/FP16 CK/CKTile/ASM/Triton API; FP32 output is available only through the +explicit OPUS exact-kid call above. -```bash -# single-shape smoke -python3 op_tests/test_opus_a16w16_gemm.py -m 128 -n 256 -k 1024 -b 1 +### gfx950 MXFP8 BMM -# CSV sweep (each row is one (M, N, K, batch) shape) -python3 op_tests/test_opus_a16w16_gemm.py --csv /path/to/shapes.csv +```python +Y = torch.empty((G, M, N), device=XQ.device, dtype=torch.bfloat16) +opus_bmm( + XQ, # [G,M,K], batch-first and K-contiguous + WQ, # [G,N,K], batch-first and K-contiguous + Y, + kid=8311, # exact global id; family-local id 311 + layout="mxscale_bmm", + x_scale=x_scale, # [G,M,K/128], one-byte E8M0 + w_scale=w_scale, # [G,N/128,K/128], one-byte E8M0 + split_k=1, +) ``` ---- +The 45 MXFP8 BMM kernels use global ids `8000 + family_local_kid`; the +family-local ids remain recognizable while sharing the canonical +`kernels_list` without colliding with existing GEMM ids. The public layout +name is strictly `mxscale_bmm`. Internally, the family adapter passes zero-copy +`[M,G,*]` transpose views to the unchanged raw kernel ABI. +The high-level tuned caller remains +`aiter.batched_gemm_a8w8_mxscale`. Its cold-path tuned-row, padded-M, +local-to-global-id and heuristic selection live in `policy.py` beside the A16 +caller policy. +The caller caches the final id/split pair per shape: split-one enters the +checked raw launcher directly, while workspace launches call `opus_bmm`. -## 6. Environment +### gfx942 blockscale bpreshuffle -| Env var | Default | Effect | -|---|---|---| -| `AITER_OPUS_TUNED_CSV_GLOB` | `aiter/configs/bf16_tuned_gemm.csv:aiter/configs/model_configs/*_bf16_tuned_gemm.csv` | Colon-separated glob list of tuned BF16 GEMM CSVs that the opus runtime dispatch (`common.py::lookup_tuned`) and the C++ codegen (`gen_instances.py --tune_files`) read. Each file is filtered by `libtype == 'opus'`. | -| `AITER_OPUS_DEBUG_TUNED_CSV` | `/tmp/opus_debug_tuned.csv` | Default `-o` for the debug-only `opus_gemm_tune.py`. Never set this to a path under `aiter/configs/` -- use gradlib for production tuning. | -| `AITER_REBUILD` | `0` | `1` forces JIT rebuild of `module_deepgemm_opus`. Needed after CSV edits if you want the C++ lookup to pick them up. | -| `GTUNE_TUNED` | `$AITER_CONFIG_GEMM_BF16` (`aiter/configs/bf16_tuned_gemm.csv`) | gradlib output path. Pass `--tuned_file ` to override on the CLI. | -| `FLATMM_HIP_CLANG_PATH` | unset | Optional hipcc override (see `optCompilerConfig.json`). | - -**Removed env vars** (autolog feature deleted in this release): -`AITER_OPUS_A16W16_TUNED_CSV`, `AITER_OPUS_A16W16_UNTUNED_CSV`, -`AITER_OPUS_LOG_UNTUNED`. The autolog code path is gone; collect -untuned shapes via gradlib's standard `--input_file` flow instead. - ---- - -## 7. Under the Hood - -### 7.1 Two-level dispatch (mirrors `csrc/ck_gemm_a8w8/gemm_a8w8.cu`) - -`opus_dispatch_a16w16_gfx950` in -[csrc/opus_gemm/include/gfx950/opus_gemm_arch_gfx950.cuh](../../../csrc/opus_gemm/include/gfx950/opus_gemm_arch_gfx950.cuh) -binary-searches a sorted flat array of `(M, N, K) -> kernel` entries -generated from the global tuned CSV; on miss it routes to the -heuristic-kid path: - -```cpp -template <> -inline OpusA16W16NoscaleKernel -opus_dispatch_a16w16_gfx950(int M, int N, int K, int batch) -{ - static constexpr OpusA16W16RuntimeEntry kLookup[] = { - GENERATE_OPUS_LOOKUP_TABLE_BF16(bf16_t) - }; - OpusA16W16RuntimeEntry needle{{M, N, K}, nullptr}; - auto it = std::lower_bound(kLookup, kLookup + kSize, needle, entry_less); - if (it != kLookup + kSize && entry_eq(*it, needle)) - return it->func; - - // Miss: ask the heuristic for an integer kid, resolve through - // tune_lookup. Splitk kids force (their main kernel only - // has the instantiation; the reduce kernel templated on Y - // dtype handles bf16/fp32 output at launch time). - const int kid = opus_a16w16_heuristic_kid_gfx950(M, N, K); - if (kid_is_splitk(kid)) - return opus_a16w16_tune_dispatch_gfx950(kid); - return opus_a16w16_tune_dispatch_gfx950(kid); -} -``` - -The `` specialization is analogous, using -`GENERATE_OPUS_LOOKUP_TABLE_FP32` and always routing the heuristic kid -through `` (splitk kid is forced; non-splitk kid happens to be -the same in the fp32 lookup table). - -### 7.2 Kernel inventory (see [opus_gemm_common.py](../../../csrc/opus_gemm/opus_gemm_common.py)) - -Two a16w16-class pipelines are compiled today: - -- **Split-barrier a16w16** (kid 4..9): traditional 2-stage double- - buffered pipeline. Both `` and `` instantiations - emitted. Requires even `ceil_div(K, B_K)` (and the cross-family - `K % 2 == 0` rule). Supports per-row bias via two specializations - (`HAS_BIAS = true / false`); the launcher dispatches at runtime on - `bias.has_value()`. -- **Warp-specialized flatmm_splitk** (kid 200..210): 4-wave warp-spec - kernel with runtime splitK (literal KBatch), fp32 workspace, reduce - kernel casts to bf16/fp32 Y. Only `` main-kernel - instantiations emitted; the reduce kernel is templated on `D_OUT` - and dispatches `__bf16` / `float` at launch time, so both bf16 and - fp32 Y are valid. Handles arbitrary even-K / any N via `mask_va_tail` - + reduce-kernel tail path. Bias is folded inside the reduce kernel - via SGPR scalar load (`s_load_dword`), per-row, in fp32 acc before - cast — reduce kernel emits 4 specializations - (`{__bf16, float} × {HAS_BIAS true, false}`) so non-bias callers - pay no bias-add overhead. - -Representative instances (full table lives in -[opus_gemm_common.py](../../../csrc/opus_gemm/opus_gemm_common.py)): - -| kid | Pipeline | Tile (B_M, B_N, B_K) | WG/CU | Notes | -|-----|----------|-----|-------|-------| -| 9 | a16w16 split-barrier | (256, 256, 64) | 2 | Traditional sweet spot for large aligned M/N | -| 200 | flatmm_splitk | (64, 64, 64) | 2 | splitk default for M ≤ 128 | -| 208 | flatmm_splitk | (64, 64, 128) | 1 | Deep K / very skinny M | - -**16 additional a16w16_flatmm kid slots (100..115) are reserved but -currently empty** (`a16w16_flatmm_kernels_list = {}`). Filling them -is orthogonal to this module and does not require changes here. - -### 7.3 Heuristic fallback (bf16 Y path) - -`opus_a16w16_heuristic_kid_gfx950(M, N, K) -> int` in -[csrc/opus_gemm/include/gfx950/opus_gemm_heuristic_dispatch_gfx950.cuh](../../../csrc/opus_gemm/include/gfx950/opus_gemm_heuristic_dispatch_gfx950.cuh) -returns an integer kid based on M-bucket rules; the caller resolves the -kid through the same tune lookup that powers `opus_gemm_a16w16_tune`: - -| M range | kid (oob / nooob) | Pipeline | Rationale | -|---|---|---|---| -| `M ≤ 4` | 208 / 1208 | splitk `(64, 64, 128)` WG=1 | Very skinny M; deep K keeps splitk workspace small | -| `M ≤ 64` | 206 / 1206 | splitk `(64, 32, 128)` WG=2 | cc-recommended mid-M tile | -| `M ≤ 128` | 200 / 1200 | splitk `(64, 64, 64)` WG=2 | splitk sweet spot | -| `M > 128`, N%16 + K%64 + loops even | 300 / 1300 | persistent `(256, 256, 64)` | Persistent + XCD swizzle wins large aligned | -| `M > 128`, misaligned | 200 / 1200 | splitk `(64, 64, 64)` WG=2 | splitk tolerates arbitrary N (per-element tail store) | - -These gfx950 fallback kids are represented in the architecture-specific -heuristic sets in `csrc/opus_gemm/opus_gemm_common.py`. `gen_instances.py` -asserts that all heuristic kids required by the target architectures are in -the subset-compile set `S` before writing the generated sidecar, so fallback -does not select a kid omitted from that build. - -The same heuristic kid function is used for both `` and -`` dispatch specializations; splitk kids force the `` -tune_lookup branch regardless of CDataType (their main kernel only has -``; the reduce kernel handles Y dtype at launch time). -Persistent kid 300/1300 honors CDataType so both bf16 and fp32 output -work. Every kid the heuristic returns supports bias (`HAS_BIAS=true`); -CSV-miss requests with bias are forwarded unchanged. - -### 7.4 The splitk `` trick in the BF16 lookup map - -splitk kid instantiations exist only as `` (their traits -`static_assert(D_C == float)` — the main kernel writes an fp32 -workspace; the reduce kernel then casts to bf16 Y). So the BF16 lookup -map contains mixed template arguments: - -```cpp -// aiter/jit/build/module_deepgemm_opus/blob/opus_gemm_lookup.h (generated) -#define GENERATE_OPUS_LOOKUP_TABLE_BF16(CTYPE) \ - { \ - {{1, 100, 5120}, \ - opus_gemm_flatmm_splitk_256x32x32x64_..._wgpcu2}, \ - {{256, 51200, 5120}, \ - opus_gemm_512x256x256x64_2x4_16x16x32_0x0x0}, \ - ... - } +```python +from aiter.ops.shuffle import shuffle_weight + +WQ_shuffled = shuffle_weight(WQ, layout=(16, 16)) +opus_gemm( + XQ, + WQ_shuffled, + Y, + kid=11000, + layout="bpreshuffle", + x_scale=x_scale, + w_scale=w_scale, +) ``` -`gen_instances.py:gen_lookup_dict` hardcodes `` for splitk kids -regardless of which per-CTYPE map they land in. FP32 map drops splitk -entries entirely (launcher `TORCH_CHECK`s `Y.dtype() == BFloat16`). - -### 7.5 JIT build pipeline - -1. `aiter.ops.opus.gemm_op_a16w16` triggers `compile_ops("module_deepgemm_opus")`. -2. [aiter/jit/optCompilerConfig.json](../../jit/optCompilerConfig.json) - invokes - `csrc/opus_gemm/gen_instances.py --working_path {blob_staging_dir} --tune_files ... --compiled_kids_sidecar={blob_staging_dir}/compiled_kids_opus.json`. - JIT first seeds this staged file from `{bd_dir}/compiled_kids_opus.json`. - The tuner additionally supplies its candidate kids through `--extra_kids`. -3. `gen_instances.py` computes the subset-compile set - `S = (CSV opus rows' solidx) ∪ (sidecar contents) ∪ (extra kids) ∪ HEURISTIC_DEFAULT_KIDS ∪ a8w8_kids`, - applies validity, target-architecture and optional family filters, then - checks that target-specific heuristic kids and all explicit extra kids - remain in `S`. It writes: - - `compiled_kids_opus.json` — the staged sidecar listing every kid in `S` - - `impl/*.cuh` — per-kid kernel launcher templates (one per kid in `S`) - - `instances/all_instances_host.cu` — fused host TU (one per build) - - `instances/{kid_name}_C{bf16_t,fp32_t}.device.cu` — per-(kid, dtype) device TU - - `instances/splitk_reduce.device.cu` — dedicated splitk reduce TU - - `opus_gemm_manifest.h` — forward declarations - - `opus_gemm_a16w16_tune_lookup.h` — int-id → kernel maps for the kids in `S` - - `opus_gemm_lookup.h` — **(M, N, K) → kernel** maps baked from CSV opus rows - (two macros: `_BF16`, `_FP32`) -4. `opus_gemm.cu` is compiled and linked against the generated - instances into the fixed target `module_deepgemm_opus.so`. JIT invokes - Ninja's incremental dependency checks on every build request rather than - skipping through the Python extension versioner. This also permits retries - after a failed compile, header-only changes and missing build outputs. -5. After checking that the codegen generation has not changed, JIT atomically - installs the `.so`, then publishes the canonical sidecar and its `.receipt` - independently of the best-effort source snapshot under `blob/`. A metadata - publication failure does not invalidate a successful compile, but a later - tuner cannot reuse an unmatched receipt. - -Direct generator use without `--compiled_kids_sidecar` retains the default -`{working_path}/compiled_kids.json`; JIT supplies the explicit staged path above. - -### 7.6 Compile-time techniques - -JIT build of `module_deepgemm_opus` is on the user-visible critical -path (first call into any opus entry point on a fresh checkout pays -for it). Five landed rounds of optimization, in order: - -1. **Host/device pass split** -- the codegen-emitted `.cuh` files - guard their `` + launcher body behind - `#if !defined(__HIP_DEVICE_COMPILE__) && !defined(__HIPCC_RTC__)`, - so the device pass parses ~10K lines instead of ~70K. -2. **Fusion** -- 38 per-kid host TUs collapse into one - `all_instances_host.cu`. The heavy `` parse - only happens once per module rebuild instead of 38 times. -3. **Torch removal in launchers** (mirrors PR #2932 for quant) -- - the dispatcher entry points and the codegen-emitted launcher - bodies use `aiter_tensor_t` (POD, defined in - `csrc/include/aiter_tensor.h`) instead of `torch::Tensor`. - `` is replaced with a ~200-line header. The - splitk launcher's fp32 workspace is allocated stream-ordered - via `hipMallocAsync` / `hipFreeAsync` instead of `torch::empty`. -4. **Dispatcher TU torch removal** -- one stale - `#include "py_itfs_common.h"` in `opus_gemm_arch_gfx950.cuh` - was still pulling `` + the full `` stack - into the dispatcher TU even after step 3. Replaced with the - torch-free `opus_gemm_utils.cuh` (same `bf16_t` / `fp32_t` - aliases). Drops dispatcher TU preprocessed input from 401K - lines to 154K. -5. **Lookup map: `unordered_map` + `std::function` -> sorted flat - array + function pointer + `std::lower_bound`** -- the runtime - `(M,N,K) -> kernel` and `kid -> kernel` tables used to be - `std::unordered_map<..., std::function<...>>`, which alone - added ~1s of frontend / template instantiation per dispatcher - TU because of the `std::function` + hashtable templates. The - replacement is a `static constexpr` array of POD entries - `{shape, kernel_function_pointer}` plus - `std::lower_bound`. No template instantiation overhead, no - heap allocation on first call, faster runtime lookup. -6. **`splitk_reduce_kernel` carved out of every splitk kid's - device.cu into one dedicated `splitk_reduce.device.cu`** -- - the 4 reduce specialisations (D_OUT bf16/fp32 x HAS_BIAS - true/false) used to be appended to every splitk kid's - `template __global__` instantiation list, so all 23 splitk - TUs each compiled the 4 reduce kernels redundantly. Linker - deduped the resulting weak symbols, but each TU still paid - the full RA + ISA emit cost on its own compile (~0.3-0.5s - per TU x 23 = ~9s of duplicated CPU work). Now they live in - a single 0.3s TU. -7. **`#pragma unroll` on `tiled_mma_adaptor` MMA-tile loops in - `opus.hpp`** -- the runtime - `for (I = 0; I < EXPAND_K * EXPAND_M * EXPAND_N; I++)` - outer loop and the inner `for (j = 0; j < a_len; j++)` - extract / insert loops in the `vector_t` overload of - `operator()` and `step_k` were relying on clang's default - unroll heuristic to fold trip counts that are all - constexpr. For small tiles (e.g. non-splitk's - 2x2x2 = 8 outer iters) clang did unroll them; but for the - large splitk tiles - (`flatmm_splitk_64x96x64_wgpcu1`'s 4x6x2 = 48 outer iters - x mma_a_len=4 inner, 192 reads in total) the loop was - left as a runtime loop over the `vtype_a` / `vtype_b` / - `vtype_c` register arrays. GFX9 has no "load VGPR by - runtime index" instruction, so LLVM expanded each - `s_a[j] = a[i_a + j]` into an N-way - `s_cmp_eq + s_cselect_b64 + v_cndmask_b32` select tree. - Result on the worst kid: 8931 SGPR spills, 12100 - `s_cselect_b64`, 17956 `v_writelane_b32`, 388 KB ISA, and - 5.2s of LLVM Greedy RA Evict time = 7.7s slowest TU wall - (the build's critical path). Forcing the unroll lets every - index resolve at compile time, eliminates the select trees, - and collapses register pressure. Numbers after the fix on - the same kid: - - | Metric | Before | After | Change | - |---|---:|---:|---:| - | Slowest TU wall | 7.7s | **1.3s** | **-83%** | - | ISA size | 388 KB | 27 KB | -93% | - | `s_cselect_b64` | 12 098 | 2 | -99.98% | - | `v_writelane_b32` | 17 956 | 0 | -100% | - | `sgpr_spill_count` | 8 931 | 0 | -100% | - | `sgpr_count` | 106 | 46 | -57% | - | `vgpr_count` | 423 | 310 | -27% | - | `agpr_count` | 167 | 54 | -68% | - | `private_segment` | 24 704 B | 0 | -100% | - | ASM total lines | 69 426 | 1 958 | -97% | - - This change touches `opus.hpp` (a shared header used by - opus_attn / opus_fmm / opus_gemm). Smaller tile configs - that already unrolled spontaneously are unaffected: they - unroll the same way they did before, just under the - explicit pragma instead of the heuristic. - -**Headline** (128-core demon_test, ROCm 7.2.2, `MAX_JOBS=102`, -3-trial average over `AITER_REBUILD=1` with cleared build dir): - -| Build | wall time | -|---|---:| -| Pre-split baseline (38 .cpp's, each parses torch + has both passes) | **48.4s** (49.2 / 47.2 / 48.9) | -| Host/device split (38 .cpp's, kernel decl in `#ifdef`) | **32.5s** (31.8 / 34.2 / 31.6) | -| Fused host TU + 38 device TUs | **22.3s** (22.5 / 22.0 / 22.4) | -| + torch removal in launchers | **19.4s** (19.5 / 19.5 / 19.2) | -| + dispatcher torch removal + flat-array lookup | **14.0s** (14.1 / 14.0 / 14.0) | -| + dedicated splitk_reduce TU | **14.4s** (14.0 / 14.7 / 14.4) | -| **+ MMA-tile unroll on opus.hpp (current)** | **11.1s** (11.1 / 11.3 / 11.0) | -| **Saving vs. baseline** | **−37.3s (−77%)** | - -Rounds 6 + 7 together flipped the build's critical path. Round -6 (dedicated reduce TU) doesn't move end-to-end wall on its own -because round 7's eventual fix is in the slowest kid's main -kernel, but it cuts ~9s of duplicated reduce codegen across all -splitk TUs (a real win on hardware with smaller MAX_JOBS). -Round 7 cracks the slowest-TU bottleneck the previous five -rounds had left untouched: every splitk TU drops to -~1.2-1.5s wall (vs the worst's 7.7s before), and the new -critical path is the pybind TU (4.7s, mostly pybind11 + libtorch -parse) and the fused host TU (2.4s). Perf on the 24-shape -dsv3+gptoss bf16 benchmark is unchanged across rounds 6+7 -(geomean +0.37% per shape, well within measurement noise; total -+0.27% sum of best-kernel us). - -Functional regression: `op_tests/test_opus_a16w16_gemm.py` end-to-end -shape sweep still passes (`allclose` on every shape). - -**Per-TU breakdown** (single-TU `-ftime-report` wall): - -| TU class | host pass | device pass | TU wall | -|---|---:|---:|---:| -| Pre-split `instance.cpp` | 13.6s | 13.3s | ~26s | -| Host/device-split `instance.cpp` | 11.7s | 1.2s (instantiations only) | ~15s | -| Fused `all_instances_host.cu` (with torch) | ~11.7s (one-time, all 38 launchers) | ~0.4s (device pass empty) | ~12s | -| Fused `all_instances_host.cu` (current, torch-free) | **2.05s** (FE 1.33s 65% / OPT 0.32s 16% / MCG 0.34s 17%) | 0.42s (empty) | **~2.5s** | -| Per-kid `*.device.cu` (typical a16w16) | ~0.10s (RTC) | ~1.5s (single Traits codegen) | ~1.8s | -| Per-kid `*.device.cu` (worst splitk: 64x96x64-wgpcu1, pre-round-7) | ~0.11s | **8.29s (MCG 6.93s 84%, RA 5.74s 77%)** | **~8.5s** | -| Per-kid `*.device.cu` (worst splitk: 96x64x128-wgpcu1, post-round-7) | ~0.10s | ~1.4s (clean codegen, 0 spill) | **~1.5s** | -| Pre-split dispatcher (`opus_gemm.cu`, with torch) | 12.5s | 15.0s | ~27s | -| Split dispatcher (with torch) | 12.5s | 0.4s | ~13s | -| Dispatcher after torch + lookup overhaul (current) | **1.43s** (FE 1.21s 85%) | 0.42s | **~2.0s** | -| Pre-split pybind | ~13s | ~13s | ~26s | -| Split pybind (current) | ~5s | 0.42s | ~5.5s | - -The end-to-end wall is now bounded by **the slowest single -pybind TU (~4.7s, mostly pybind11 + libtorch parse on host -pass)** plus ninja schedule + link + Python startup overhead -(~6s). Critical path breakdown post-round-7: +`layout="bpreshuffle"` is a declaration of WQ content, not something Tensor +shape or strides can prove. Kid 11000 requires batch one, exact 128-wide N/K +tiles, BF16 output, and its registered scale storage contracts. The high-level +`gemm_a8w8_blockscale_bpreshuffle` dispatcher enters this OPUS route only when +the tuned row has `libtype=opus`; CK, CKTile, ASM and Triton rows remain on +their respective backends. gfx950 currently registers zero OPUS kids for this +family, so a gfx950 OPUS validation must report it unavailable rather than run +a non-OPUS fallback as coverage. -``` -opus_gemm_pybind.cu host pass ~4.7s ← pybind11 + libtorch parse -ninja + link + python startup ~6.4s -total wall ~11s -``` +## Tuning compatibility -The slowest device TU now finishes in 1.5s (vs 7.7s before -round 7); every device.cu's GPU codegen is bounded comfortably -under the pybind TU. Round 7's `#pragma unroll` on -`tiled_mma_adaptor` was the breakthrough -- see §7.6.1 below -for the original forensic breakdown of the spill blow-up that -the unroll fixed. +The exact public APIs execute a caller-selected kid. A16 production tuning +continues through `csrc/gemm_a16w16/gemm_a16w16_tune.py`; plain A8W8 and +MXFP8 BMM use `csrc/opus_gemm/opus_gemm_a8w8_tune.py` and +`csrc/opus_gemm/opus_bmm_mxscale_tune.py`, respectively. -#### 7.6.1 Device-pass forensics on the slowest splitk kid (historical) +The CK-owned blockscale tuner remains unchanged. Its legacy +`opus_gemm_a8w8_blockscale_bpreshuffle_tune(...)` import is retained in +`gemm_op_a8w8.py` and calls the bpreshuffle family launcher directly. -> Status: **fixed in round 7**. This section documents the -> diagnostic path for posterity. The numbers below are -> pre-fix and no longer reproducible -- they referred to the -> kernel before the `#pragma unroll` was added to the MMA-tile -> loops in `opus.hpp::tiled_mma_adaptor::operator()` / -> `step_k`. After the fix, this same kid compiles in 1.3s -> with 0 SGPR spills. +### Plain A8W8 GEMM -`flatmm_splitk_64x96x64_wgpcu1` *was* the build's -critical-path TU. `hipcc -ftime-report` on it shows: +`csrc/opus_gemm/opus_gemm_a8w8_tune.py` tunes the gfx950 no-scale and ordinary +blockscale GEMMs. It accepts the existing GEMM tuner options; CSV `scaleAB` +selects the scale mode, so there is no `--family` option. -``` -Pass 1 (DEVICE pass): 8.29s total -├── Front end: 0.66s ( 8%) parse pipeline header + traits -├── Optimizer: 0.61s ( 7%) middle-end opt passes -├── LLVM IR generation: 0.09s ( 1%) -└── Machine code generation: 6.94s (84%) ← AMDGPU backend - ├── Greedy Register Allocator: 5.74s (77% of total) - │ └── Evict sub-pass: 5.22s (99% of RA) - ├── Machine Instruction Sched: 0.56s ( 7%) - └── ~150 other passes: ~0.45s - -Pass 2 (HOST pass): 0.11s ← RTC short-circuits libtorch + libstdc++ +```bash +# Run from the repository root. +export PYTHONPATH="$PWD" +export ROCR_VISIBLE_DEVICES=0 +export HIP_VISIBLE_DEVICES=0 + +cat > /tmp/opus_a8w8_shapes.csv <<'CSV' +M,N,K,dtype,outdtype,bias,scaleAB,bpreshuffle +64,4096,4096,fp8,fp32,False,False,False +128,4096,4096,fp8,fp32,False,False,False +64,4096,4096,fp8,fp32,False,True,False +128,4096,4096,fp8,fp32,False,True,False +CSV + +python3 csrc/opus_gemm/opus_gemm_a8w8_tune.py \ + --input_file /tmp/opus_a8w8_shapes.csv \ + --tuned_file /tmp/opus_a8w8_tuned.csv \ + --libtype opus --mp 1 + +python3 csrc/opus_gemm/opus_gemm_a8w8_tune.py \ + --run_config /tmp/opus_a8w8_tuned.csv --libtype opus --mp 1 ``` -GPU codegen metadata (extracted from the resulting fat binary): - -| Metric | Value | -|---|---:| -| Main kernel ISA size | **388 KB** (vs ~8 KB for non-splitk kids) | -| `vgpr_count` (logical) | 423 | -| `agpr_count` (acc registers) | 167 | -| `sgpr_count` | 106 | -| `sgpr_spill_count` | **8931** | -| `vgpr_spill_count` | 0 | -| `private_segment_fixed_size` (scratch / wave) | 24 704 bytes | -| `group_segment_fixed_size` (LDS / WG) | 144 KB | -| `prefetch_k_iter` | 7 | -| `WG_PER_CU` | 1 (wgpcu1) | - -**Why `Evict` runs for 5.2s** -- ISA size, register pressure -and SGPR spill all stem from one structural choice. `wgpcu1` -sets WG_PER_CU=1, which gives the kernel the entire CU's -registers + LDS. `prefetch_k_iter = LDS_total / per_iter_LDS = 7` -on this shape, so the K-loop carries **7 prefetch buffers' -worth of register tile state simultaneously**. With -`comrep=(2,6)` (2 M x 6 N MFMA tiles per consumer wave) that's -~336 independent vector-register live ranges open at once -across 256 physical VGPRs + 167 AGPRs. LLVM's Greedy RA enters -its `Evict` policy: every conflict triggers an enumeration of -candidate live ranges to spill, with recursive spill-cost -calculation. On this kernel the candidate count and recursion -depth combine into ~5s of pure RA work. - -**Why `sgpr_spill_count = 8931`** -- the K-loop indexes the -prefetch buffers via `slot = issue_k % pfk` (a runtime int). -GFX9 has no "load VGPR by runtime index" instruction -(`v_movrels_b32` exists but is restricted), so for each of the -~192 register tile elements that need slot-indexed access, LLVM -expands the read into a 7-way `s_cmp_eq + s_cselect_b64 + -v_cndmask_b32` select tree: - -```asm -s_cmp_eq_u32 s0, 1 -s_cselect_b64 vcc, -1, 0 -v_cndmask_b32_e32 v130, v6, v183, vcc -s_cmp_eq_u32 s0, 2 -s_cselect_b64 vcc, -1, 0 -v_cndmask_b32_e32 v130, v130, v7, vcc -... ×7 (one per prefetch slot) ×192 tile elements +`-i/--untune_file` and `-o/--tune_file` are equivalent aliases. Missing +`dtype`, `outdtype`, and `scaleAB` columns default to FP8, FP32, and `False`. +`scaleAB=True` uses FP32 scales with the registered 1x128x128 group contract. +Both modes require contiguous inputs/output, no bias or preshuffle, and +`splitK=0`. Candidates come from the canonical registry and are checked against +an independent FP32 dequantize-and-matmul reference. The default `--errRatio 0` +rejects any element outside `rtol=atol=1e-2`. + +The output key includes `gfx,cu_num,M,N,K,dtype,outdtype,scaleAB`, so both +scale modes can coexist for the same shape. `--all` retunes the input rows; +`--profile_file` records all candidates. `--run_config` executes the saved +`kernelId` through `opus_gemm` with preallocated FP32 output; without a path it +reads `--tuned_file`. Callers own this CSV lookup. For M padding, allocate and +zero-pad XQ, pad `x_scale` rows with finite scales (for example 1), call the +saved kid with padded output, then slice back to the original M. + +## A16 Torch workspace + +Workspace ownership is call-scoped and remains in Torch: + +```text +validate exact kid and split_k + -> derive immutable workspace plan from the exact instance + -> reuse caller workspace or torch.empty for this call + -> _launch_a16w16_backend ``` -Every `s_cselect_b64` produces a `vcc` (a 64-bit SGPR pair) -live range. With ~12 000 `s_cselect_b64` and ~12 000 -`s_cmpk_eq_i32 / s_cmp_eq_u32` instructions in the kernel, the -SGPR pressure exceeds the gfx950 physical SGPR cap (~100 -addressable). LLVM's AMDGPU backend handles the overflow by -**spilling SGPRs to VGPR lanes** (`v_writelane_b32 v252-v255, -sN, lane_idx`) instead of going to scratch memory. We see -8931 such spill points statically; the assembly contains 17956 -total `v_writelane` / `v_readlane` instructions because each -spilled SGPR is reloaded once on average. Four whole VGPRs -(v252, v253, v254, v255) are reserved as 64-lane SGPR scratch. +There is no process-global Tensor, pointer registry, HIP allocator, or prewarm +API. The bounded public-contract and A16 launch-plan caches store only registry +metadata, integers, dtypes, option-presence flags and shapes; they never retain +Tensor objects, data pointers, devices, streams or workspaces. -Stack frame metadata confirms it: +Let `padded_M=ceil_div(M,B_M)*B_M` and +`padded_N=ceil_div(N,B_N)*B_N`: -``` -Function: gemm_a16w16_flatmm_splitk_kernel<...wgpcu1...> - private_segment_fixed_size = 24704 bytes - ~96 x 256-byte slots (each = one 64-lane wave-level register - tile spilled to scratch as a fallback when even the 4 VGPR - scratch lanes can't hold a particular live range) +| Architecture/family | Workspace shape | Instance storage | +|---|---|---| +| gfx950 two-stage | `[workspace_capacity_split_k,batch,padded_M,padded_N]` | FP32 | +| gfx942 two-stage | `[workspace_capacity_split_k,batch,padded_M,padded_N]` | exact BF16/FP32 dtype | +| gfx1250 two-stage | `[workspace_capacity_split_k,padded_M,padded_N]` | FP32 | +| gfx1250 pre-built CO direct | none | none | +| gfx1250 fused | not publicly registered | factory/emitter/source retained for repair | + +For gfx942, `abi_split_k` records the converged value passed to the launcher. +`workspace_capacity_split_k` uses the same value, so automatic allocation +reserves one workspace slice per launched split. + +An explicit workspace must be on the XQ device, contiguous, 16-byte aligned, +of the exact instance dtype, and large enough for the final split. Larger +caller-provided workspaces are also accepted. A direct kid requires +`workspace=None`. + +gfx1250 two-stage kids require `M <= 65535`, including `split_k=0/1`, because +their separate reducer places one logical row in each `grid.y` block. The +tuner excludes larger M, policy rejects stale split-K rows, and exact launch +checks the limit before either kernel runs. Larger M can use a compatible +tuned CO kid; the gfx1250 heuristic only selects two-stage kids and requests +tuning when this limit is exceeded. + +gfx1250 clusterlaunch exact kids round the physical launch grid up to complete +clusters; tile-less workgroups exit inside the pipeline. This does not change +the logical workspace shape above. The experimental fused family is disabled, +so no fused kid in `[27000,30000)` can be resolved through the public registry. +The `[21000,27000)` band belongs to pre-built CO kids, which are direct and +therefore never request a workspace. + +gfx942 BF16-workspace kids `10210`, `10213`, and `10216` are exact ids. Their +registered exact-N contract is `{64,128,256,384,512,1024,2048}`. A different N +or an FP32 `Y` is rejected by the strict `opus_gemm`/`opus_bmm` exact APIs; they +never redirect. To preserve the former shape-driven API, `gemm_a16w16_opus` +maps `10210` to `10200` and `10213` to `10203` for a non-exact N. Kid `10216` +has no FP32-workspace sibling and remains rejected. + +## MXFP8 BMM Torch workspace + +MXFP8 BMM follows the same ownership rule: Python either uses the caller's +contiguous FP32 Tensor or creates a call-scoped `torch.empty`; C++ receives a +direct pointer and never retains it. Two-stage split-K uses +`split_k * G * padded_M * padded_N` FP32 elements. The fused family stores its +partials and aligned tile counters in one FP32 Tensor. `split_k == 1` and +families that do not consume workspace reject a supplied Tensor. + +`launch_plan.py` owns the shared immutable workspace specification plus the +family-specific A16W16 and A8W8 plans. Its `A8W8MxscaleBMMPlan` records the +resolved exact kid, the split-K value passed to the ABI and an optional +`WorkspaceSpec`. `gemm_op_a8w8.py` only adapts logical layouts, materializes +that workspace and invokes `_launch_a8w8_backend`. + +For gfx950 kid 8326, only the `split_k > 1`, `D_OUT=void` workspace +specialization sets `PRELOAD_SF_LDS=false` to avoid the ROCm 7.2.4 compiler +failure. Its direct BF16/FP32 `split_k == 1` specializations keep +`PRELOAD_SF_LDS=true`. + +## A8 pybind backend + +All three non-MX A8 GEMM adapters and the MXFP8 BMM executor enter one +low-level facade: + +```text +validated family + resolved kid + physical Tensor views + -> _launch_a8w8_backend + -> no-scale pybind raw launcher + -> plain blockscale pybind raw launcher + -> blockscale-bpreshuffle pybind raw launcher + -> MXFP8 BMM pybind raw launcher ``` -**Bottom line on this TU**: the 8.5s wall is the price of -choosing `prefetch_k_iter = 7` for runtime perf reasons. The -Greedy RA Evict cost and the SGPR spill blow-up are downstream -consequences of that one structural choice. Reducing wall here -requires either: - -1. structurally lowering `pfk` (perf trade-off -- shallower - K-pipeline = less L1/LDS reuse), or -2. teaching the kernel author to write the slot-rotation - without runtime indexing (hand-unroll the slot dispatch), - which removes the `s_cselect` chains entirely and probably - cuts SGPR spill 10x. - -Neither is a build-system change. The B-track flag sweep -(see §7.7) confirms no `-mllvm` knob recovers any meaningful -wall on this kernel. - -**What the codegen actually emits** (post-fusion): - -1. **`csrc/include/opus/hip_minimal.hpp`** — kept torch-free + adds - `__forceinline__` / `__noinline__` keyword fallbacks on top of the - existing `__launch_bounds__` / `__shared__` / `__device__` / - `__global__` / `__host__` set. Pipeline files use the - `opus::thread_id_x()` / `opus::block_id_x()` etc. wrappers from - `opus.hpp` instead of HIP's `threadIdx` / `blockIdx` magic globals, - which lets the device pass skip `` (~100K - preprocessed lines) entirely. - -2. **`csrc/opus_gemm/include/opus_gemm_utils.cuh`** — three include - modes: - * `__HIP_DEVICE_COMPILE__` (any device pass): ``. - * `__HIPCC_RTC__` (RTC mode, set per-source on `*.device.cu`): - `` on both passes. The device TU's host - pass is empty content-wise, so the bare minimal header is - enough; `` would be wasted parse and would - also pull in `` which depends on the - wrapper that `__HIPCC_RTC__` short-circuits. - * Otherwise: the full `` + `` + - ``. Used by `all_instances_host.cu`, - `opus_gemm.cu`, `opus_gemm_pybind.cu`. - -3. **`csrc/opus_gemm/gen_instances.py`** — restructured around three - file shapes: - * `impl/{name}.cuh` (one per kid): Traits aliases + launcher body. - Three guard combinations: skip torch headers when the host pass - is irrelevant (`__HIP_DEVICE_COMPILE__` or `__HIPCC_RTC__` set); - pick `traits header + forward kernel decl` over `pipeline body` - when `OPUS_FUSED_HOST_TU` is set (avoids the ODR clash on - same-named layout helpers between pipeline headers). - * `instances/all_instances_host.cu` (one for the WHOLE module): - defines `OPUS_FUSED_HOST_TU`, includes `aiter_tensor.h` + - `aiter_stream.h` + `` once, includes every kid's - `.cuh`, and emits all `template void xxx(...)` - instantiations. The launcher's `<<<...>>>` calls produce - undefined `__device_stub__` references. Wrapped in - `#ifndef __HIP_DEVICE_COMPILE__` so the device pass sees an - empty TU. - * `instances/{name}_C{dtype}.device.cu` (one per kid, dtype): - includes the kid's `.cuh` with neither `OPUS_FUSED_HOST_TU` nor - `__HIP_DEVICE_COMPILE__` (so the full pipeline header IS - visible), but with `-D__HIPCC_RTC__` from - `flags_extra_hip_per_source` so the host pass takes the lean - branch. Emits `template __global__ void kernel<...>(...)`, - producing the host stub + device GPU IR that the linker pairs - with the fused host TU's undefined references. - -4. **Torch removal across the dispatcher graph** -- mirrors PR #2932 - (`csrc/kernels/quant_kernels.cu`): - * **`csrc/opus_gemm/include/opus_gemm.h`** -- entry-point - signatures take `aiter_tensor_t&` (POD, - `csrc/include/aiter_tensor.h`) instead of `torch::Tensor&`, - return `void`. The header costs ~200 preprocessed lines instead - of ~50K. - * **`csrc/opus_gemm/include/opus_gemm_arch.cuh`** and - **`opus_gemm_arch_gfx950.cuh`** -- `TORCH_CHECK` → - `AITER_CHECK`, `` → `aiter_hip_common.h`. - * **`csrc/opus_gemm/include/gfx950/opus_gemm_heuristic_dispatch_gfx950.cuh`** - -- `OpusA16W16NoscaleKernel` is now - `std::function` so every dispatch - map entry is torch-free. - * **`csrc/opus_gemm/opus_gemm.cu`** -- both entry points - (`opus_gemm`, `opus_gemm_a16w16_tune`) take `aiter_tensor_t`, - use `AiterDtype` enum (`AITER_DTYPE_bf16` / `_fp32` / `_fp8`) - instead of `at::ScalarType::*` / `torch_fp8`, return `void`. - * **`csrc/opus_gemm/gen_instances.py`** -- the codegen-emitted - launcher signatures use `aiter_tensor_t&`, the bias validator - calls `AITER_CHECK` + `bt.is_contiguous() / dtype() / dim() / - size()` (POD accessors that `aiter_tensor_t` provides - PyTorch-compatible by design). The splitk launcher allocates - its fp32 workspace with `hipMallocAsync(stream)` + matching - `hipFreeAsync(stream)` after the reduce kernel, replacing - `torch::empty(... TensorOptions().dtype(kFloat32).device(...))` - while preserving the same stream-ordered lifetime invariant - PyTorch's caching allocator gave us. - * **`aiter/ops/opus/gemm_op_a16w16.py`** -- - `@compile_ops("module_deepgemm_opus", develop=True)` on both - `_opus_gemm_a16w16_tune_raw` and `_opus_gemm_bf16_dispatch`. - `develop=True` makes the JIT wrapper (a) inject the current - torch CUDA stream into the C++ `aiter::getCurrentHIPStream` - thread-local via `module._set_current_hip_stream` before the - call and (b) auto-convert any `torch.Tensor` arg to - `aiter_tensor_t` via `torch_to_aiter_pybind`. Because the C++ - side now returns `void`, `opus_gemm_a16w16_tune` keeps its - `return Y` contract by returning the in-place tensor directly. - -5. **`csrc/opus_gemm/opus_gemm.cu` and `csrc/pybind/opus_gemm_pybind.cu`** - — entire-file `#ifndef __HIP_DEVICE_COMPILE__` skip. Pure host - code with no `__global__` / `<<<>>>`; their device pass is dead - weight (12.5–15s → 0.4s). - `opus_gemm_pybind.cu` additionally registers - `AITER_SET_STREAM_PYBIND` so Python can call - `module._set_current_hip_stream(...)`. - -6. **Per-file flag plumbing in `aiter/jit/`** — `core.py` reads the - new `flags_extra_hip_per_source` dict from - `optCompilerConfig.json` and forwards it through `_jit_compile` - into `_write_ninja_file`, which emits a per-build - `cuda_post_cflags = $cuda_post_cflags ` override on - matching ninja rules. Used by opus to apply `-D__HIPCC_RTC__` to - `*.device.cu` only — the dispatcher / pybind TUs would break with - it because they transitively pull in ck_tile / pybind11, both of - which depend on the wrapper that RTC short-circuits. - -**Why fusion was the right move on this hardware**: in the -host/device-split layout, the critical path was ~15s (each -`instance.cpp` had to parse `` AND run device -codegen, in series inside one hipcc invocation). The fused host TU -detaches those: launcher instantiations all live in one .cu that -parses headers ONCE but skips device codegen entirely; per-kid -device codegen lives in 38 tiny self-contained .cu's that finish in -~2s each in parallel. After torch removal the fused TU's parse -drops from ~12s to ~7s, and that's now the end-to-end critical -path. - -**What did NOT help** on this hardware: - -- **MAX_JOBS tweaks**: aiter already auto-sets it to - `min(80% × cpu_count, free_mem / 0.5 GB)` = 102 on the test host; - CPU saturation isn't the bottleneck — the host TU's serial parse - of `` is. - -### 7.7 Future compile-time work - -The current 11.1s floor is now set by **the pybind TU at ~4.7s** -(mostly pybind11 + libtorch parse, neither of which the -dispatcher / launcher refactors can touch since the pybind -layer is the boundary with Python). The previous champion -- -the slowest device TU's GPU codegen on -`flatmm_splitk_64x96x64_wgpcu1` -- went from 7.7s down to 1.3s -in round 7 once we found the SGPR spill root cause and added -`#pragma unroll` (see §7.6 round 7 + §7.6.1). Remaining attack -surface, in rough order of expected payoff vs. invasiveness: - -1. **(MEASURED, NEGATIVE)** AMDGPU-side `-mllvm` flag sweep -- - the build's device passes inherit five aiter-global -mllvm - flags (`--amdgpu-kernarg-preload-count=16`, - `--lsr-drop-solution=1`, `-amdgpu-early-inline-all=true`, - `-amdgpu-function-calls=false`, `-enable-post-misched=0`) - plus opus-private `--amdgpu-mfma-vgpr-form`. Earlier - speculation was that `-amdgpu-early-inline-all=true` + - `-amdgpu-function-calls=false` (which together force every - call to inline into one mega-function for the RA) might be - responsible for the long Greedy RA Evict pass. We measured - five build configurations: - - | Config | Slowest TU wall | Build wall (3-trial avg) | Perf vs baseline (geomean of 24 dsv3+gptoss bf16 shapes) | - |---|---:|---:|---:| - | baseline (all 6 flags on) | 7.66s | 18.1s | reference | - | `-amdgpu-early-inline-all=false` override | 7.68s | 17.9s | +0.24% (slower by 0.24%) | - | `-amdgpu-function-calls=true` override | 7.74s | 17.8s | -0.10% (faster by 0.10%) | - | `--amdgpu-mfma-vgpr-form=false` override | 7.75s | 18.3s | +0.04% | - | all three off | 7.71s | 17.9s | +0.06% | - - Both build wall and perf differences are within measurement - noise (~1%). We also tried `-O1`, `--amdgpu-igrouplp-exact-solver=0`, - and `-greedy-regalloc-eviction-max-iterations=2` (last one - doesn't exist in ROCm 7.2.2 LLVM and was rejected). None - reduced the slowest TU's wall by more than measurement - noise. Conclusion: **the Greedy RA Evict cost on this kid - is fundamental to the IR's register pressure**, not gated - by any user-tunable -mllvm flag in this LLVM revision. See - §7.6.1 for the underlying SGPR-spill / pfk=7 analysis. The - five flags should stay in the global config because they - improve perf on smaller kids elsewhere in aiter. - -2. **Trim or split heavy splitk kid instantiations** (untried, - high potential) -- the slowest 4-5 splitk kids (`*_64x96x*` - and `*_96x64x*` family with `wgpcu1`) eat the entire ninja - schedule's tail. Empirical sweep of tuned-CSV winners would - reveal which of these are actually selected for production - shapes; un-selected kids can be dropped at codegen time, - removing them from the build entirely. Saving: depends on - CSV coverage; potentially -3 to -5s end-to-end if the - slowest 1-2 kids turn out unused. - -3. **(LANDED, round 7)** Force unroll on tiled_mma_adaptor's - MMA-tile loops -- root cause of the §7.6.1 SGPR spill - blow-up turned out to be `opus.hpp`'s - `for (I = 0; I < EXPAND_K * EXPAND_M * EXPAND_N; I++)` - relying on clang's heuristic to unroll. For the worst - splitk tile (4x6x2 = 48 outer iters x mma_a_len=4 inner) - the heuristic gave up, the loop ran at runtime, and the - `a[i_a + j]` reads compiled to N-way s_cselect select - trees. Adding `#pragma unroll` to the five overloads in - `tiled_mma_adaptor` (plus the inner extract / insert loops - in the vector_t path) eliminated the spill problem - entirely: slowest TU 7.7s -> 1.3s, 8931 spills -> 0, - ASM 69k lines -> 2k. End-to-end wall 14.4s -> 11.1s. - -4. **`ccache` / `sccache` integration** -- reuse `.cuda.o` - across rebuilds when `gen_instances.py` produces a - byte-identical TU. Pure infra change, complementary to all - other items. The first build still pays full freight; - subsequent rebuilds (e.g. `AITER_REBUILD=1` after a CSV-only - edit) drop to seconds. This was deferred earlier because - parsing was the bottleneck and parses don't compose well - across rebuilds; with parse gone, MCG dominates and MCG - output is much more cacheable. - -5. **Header structure cleanup** (low priority) -- `opus.hpp` is - 3055 lines, parsed by every device TU on its device pass and - by the fused host TU on its host pass. The fused TU now - parses it in ~1.3s and that's already off the critical path; - the per-device TU host pass is even shorter (~0.1s with RTC). - Only worth doing for cleanliness, not for time. - -Practical ceiling on this hardware (post-round-7): - -- The new bottleneck is the **pybind TU at 4.7s** (mostly - pybind11 + libtorch parse, neither easy to remove without - abandoning the python binding). -- Fused host TU is 2.4s, dispatcher TU 1.8s, slowest device - TU 1.5s -- the gap between pybind and the next slowest TU - is ~2.3s, leaving room for ninja schedule / link / Python - startup (~6s observed). -- If item 2 (CSV-driven kid trimming) lands and removes - ~5 unused splitk variants: probably **~10s** end-to-end - (slow TUs are already fast, savings are linear in TU count - and amortized across MAX_JOBS=102 parallelism). -- Removing the pybind TU entirely (e.g. via a C-API shim - similar to the dispatcher's torch-free refactor) would - bring this to ~7-8s but requires wider-scope changes to - the Python wrapper layer. - -Items 1 (-mllvm flag sweep) and 3 (MMA-tile unroll) are -closed -- 1 measured negative, 3 landed in round 7. Items -4 + 5 are infrastructure / cleanliness, not on the wall -critical path. - ---- - -## 8. Troubleshooting - -### CSV edits don't seem to take effect - -Python-side lookup in `common.py` is read lazily per process -(`functools.lru_cache(maxsize=1)`). Restart the process — that is -enough for Python-layer routing. - -The **C++ compile-time lookup** (`opus_gemm_lookup.h`) only picks up -CSV changes on JIT rebuild: +## Graphs and streams + +Automatic `torch.empty` during graph capture uses the graph-private pool. + +A CO image is opened and registered on the first call to its launcher. That +first load must happen before graph capture; warm-up followed by capture/replay +is supported, while first-ever loading inside capture is not. `import aiter` +sets `OPUS_GEN_CO_DIR` to the packaged `csrc/opus_gemm/gen_co` directory, and an +explicit environment value overrides it for testing locally rebuilt images. + +## Build-time subset compile + +Tuned CSVs, the last successful compiled-kids sidecar, and additional tuner +candidates passed through `--extra_kids` are build inputs only. Their valid +non-BMM OPUS ids are unioned with: + +- `DEFAULT_COMPILED_KIDS_BY_ARCH`, the exact-id compile floor containing every + A16 caller-side heuristic result; +- mandatory A8 ids (`gfx950: {1,2}`, `gfx942: {11000}`). + +This controls which launchers enter a subset `.so`. The high-level A16 caller +may read a tuned row at runtime, but the public/C++ path receives only its +resolved id. Calling a known non-BMM registry kid that was omitted from a +subset build produces an uncompiled-id error. A gfx950 build emits all 45 +MXFP8 BMM routes as one deduplicated family so every registered BMM id remains +exact-routable. A gfx1250 build keeps all 219 available CO host launchers in its +default compile floor; their device code remains in the packaged `.co` files. +The sidecar records all emitted ids, including the deduplicated BMM family. +An explicit `--extra_kids` request that is unknown, outside the target +architectures, or excluded by `--kernel_tag` fails codegen before the sidecar +is updated. + +The canonical sidecar is `{bd_dir}/compiled_kids_opus.json`, outside the +per-module build directory so it survives `clear_build`. Tuners synchronously +build candidates before spawning workers. They pass requests through +`--extra_kids` without expanding the canonical sidecar in advance. JIT uses +`blob.staging` for generated working files, installs the binary, then publishes +the generated sidecar and a receipt binding its contents to that binary. +Runtime exact dispatch does not read this sidecar. + +A tuner skips rebuilding only when the sidecar and its receipt match the +required kids and installed binary. Missing or stale metadata triggers a +rebuild. An explicit `AITER_REBUILD` request runs once in the parent even on +a cache hit; successful preparation sets `AITER_REBUILD=0` for workers. A +failed compile restores the original environment and preserves the previous +successful metadata. See [transactional JIT cache](../../../docs/jit_cache.md) +for recovery and storage requirements. + +## Migration + +For new exact-id integrations, allocate `Y`, resolve the final id in the +caller, use `opus_gemm` for logical 2D calls, and use `opus_bmm` for +batch-first 3D calls. The retained `gemm_a16w16_opus` entry preserves the +former A16W16 shape-driven behavior: explicit id, then OPUS-only tuned lookup, +then the migrated per-architecture heuristic. Do not infer the operation from +dtype or expose the physical 3D raw ABI of a GEMM-only A8 family as public +BMM. The A8 family module exports only its legacy tuner compatibility name. + +## Validation + +Run the retained OPUS numerical tests on matching target GPUs rather than +treating architecture skips as coverage: ```bash -AITER_REBUILD=1 python3 -c "from aiter.ops.opus import gemm_a16w16_opus" +pytest -q op_tests/test_opus_a16w16_gemm.py +PYTHONPATH=. python3 op_tests/test_opus_a8w8_bmm.py \ + -g 2 -s 16,1024,4096 -d bf16 ``` -The whole rebuild takes ~11s on dev hardware (128-core, -ROCm 7.2.2; see [§7.6 Compile-time techniques](#76-compile-time-techniques) -for the seven-stage optimization stack that drops it from ~48s). - -### `RuntimeError: K=... must be even` - -The a16w16-family launchers reject odd `K` because the splitk pipeline -silently accumulates a ~3-7% maxdelta on odd K (latent K-tail bug; -e.g. `K=257` / `513`). Even K is unaffected. Pad / round your `K` to -an even number (typically 4-aligned for VEC_A=8 layout) or wait for -the K-tail handling fix. - -### `RuntimeError: bias is currently only supported on a16w16 split-barrier kids [4, 10) or a16w16_flatmm_splitk kids [200, 300)` - -Triggered when an explicit `kernelId` outside the bias-aware ranges is -passed together with a non-empty `bias`. Pick a kid in `[4, 10) ∪ -[200, 300)`, or drop the explicit override and let the dispatcher pick -a bias-aware kid. - -### `Kernel id N not found in a16w16 ... tune lookup table` - -The CSV references a kid that the current JIT build didn't compile -(usually a flatmm kid 100..115 from an older tuning run, since -`a16w16_flatmm_kernels_list` is currently empty). Re-tune the affected -shapes against the current build, or remove those rows from the CSV. - -### Why the cross-family `K % 2 == 0` rule exists - -Two independent K-tail problems on the a16w16 family motivate the -launcher-side `TORCH_CHECK(K % 2 == 0, ...)`: - -- Split-barrier (kid 4..9): the prefetched double-buffer reads one - tile past the valid K range and corrupts the accumulator on - `ceil_div(K, B_K)` odd. The launcher additionally enforces - `loops_ % 2 == 0`, which already covers most cases (B_K is 32 or 64, - so K must be a multiple of B_K; K must therefore be 64 / 128 aligned - in practice). -- Splitk (kid 200..299): on odd K (e.g. 257 / 513) the - `mask_va_tail` + reduce-tail interplay yields a 3-7% maxdelta vs. - reference, while even K stays at the bf16 noise floor. The exact - root cause is still under investigation. - -The launchers reject odd K uniformly to give callers a clear error -instead of silent miscompares; relax once the underlying handling is -fixed. - -### HIP graph compatibility - -splitk kernels allocate a fresh fp32 workspace via `torch::empty` per -call (same pattern as triton `gemm_a16w16` uses for `y_pp`). This works -under `torch.cuda.graph` capture + replay. See splitk plan §5 for the -design notes. - ---- +In particular, gfx942 and gfx1250 validation must run on matching hardware; a +skip on another architecture is not a pass for that target. -## 9. File Map +## Source map | Path | Role | |---|---| -| [aiter/ops/opus/gemm_op_a16w16.py](gemm_op_a16w16.py) | `gemm_a16w16_opus` wrapper + low-level `opus_gemm_a16w16_tune` pybind + private `_opus_gemm_bf16_dispatch` fallback binding | -| [aiter/ops/opus/common.py](common.py) | Python tuned-CSV lookup against `aiter/configs/bf16_tuned_gemm.csv` (+ `model_configs/*_bf16_tuned_gemm.csv`), filtered by `libtype=='opus'` | -| [aiter/ops/opus/__init__.py](__init__.py) | Public symbol aggregator | -| [aiter/configs/bf16_tuned_gemm.csv](../../configs/bf16_tuned_gemm.csv) | Global tuned BF16 GEMM CSV. Opus rows live here (`libtype=='opus'`) alongside asm / triton / skinny / flydsl / torch / hipblaslt rows. | -| [aiter/configs/model_configs/](../../configs/model_configs/) | Per-model tuned BF16 GEMM CSVs (gptoss / dsv4 / glm5 / kimik2 / qwen / ...). Same schema; same `libtype` filter. | -| [aiter/ops/deepgemm.py](../deepgemm.py) | CK backend (`deepgemm_ck` + `deepgemm()` forwarder). Also hosts the `opus_gemm_a16w16_tune` deprecation shim. | -| [csrc/opus_gemm/opus_gemm_common.py](../../../csrc/opus_gemm/opus_gemm_common.py) | Kernel instance metadata, architecture-specific heuristic sets and `_opus_sidecar_path()` | -| [csrc/opus_gemm/opus_gemm_tune.py](../../../csrc/opus_gemm/opus_gemm_tune.py) | Candidate/shape helpers and synchronous `_ensure_kids_compiled()` shared with gradlib; also the debug tuner entry point | -| [gradlib/gradlib/GemmTuner.py](../../../gradlib/gradlib/GemmTuner.py) | Production tuner; `--libtype opus` adds opus to the candidate sweep alongside other backends. | -| [csrc/opus_gemm/gen_instances.py](../../../csrc/opus_gemm/gen_instances.py) | Subset codegen: `--tune_files` drives the (M,N,K) lookup, while CSV/sidecar/heuristic/extra kids determine `S`; writes the staged sidecar for JIT publication | -| [csrc/opus_gemm/opus_gemm.cu](../../../csrc/opus_gemm/opus_gemm.cu) | Pybind entries (`opus_gemm`, `opus_gemm_a16w16_tune`) + per-arch router | -| [csrc/opus_gemm/include/gfx950/opus_gemm_arch_gfx950.cuh](../../../csrc/opus_gemm/include/gfx950/opus_gemm_arch_gfx950.cuh) | gfx950 dispatch: (M,N,K) lookup + heuristic-kid fallback | -| [csrc/opus_gemm/include/gfx950/opus_gemm_heuristic_dispatch_gfx950.cuh](../../../csrc/opus_gemm/include/gfx950/opus_gemm_heuristic_dispatch_gfx950.cuh) | `opus_a16w16_heuristic_kid_gfx950(M,N,K) -> int` (single source: integer kid only, no launcher symbol names) | -| [csrc/opus_gemm/include/gfx950/](../../../csrc/opus_gemm/include/gfx950/) | Kernel source (a16w16, flatmm, flatmm_splitk, persistent) for gfx950 | -| [op_tests/test_opus_a16w16_gemm.py](../../../op_tests/test_opus_a16w16_gemm.py) | End-to-end `gemm_a16w16_opus` (single-shape + CSV sweep) | - ---- - -## 10. Related Plans - -- [splitk_flatmm_aiter](/.cursor/plans/splitk_flatmm_aiter_446c6aa0.plan.md) — splitk kernel integration (kid 200..210, 17-column CSV schema, validation matrix). -- [opus_a16w16_refactor](/.cursor/plans/opus_a16w16_refactor_71298e24.plan.md) — this module's refactor (PR1: layout + shim; PR2: two-level dispatch + Python wrapper). - -Future work (separate plans / PRs): - -- Fill the a8w8 / a8w8_blockscale Python interfaces under - `aiter/ops/opus/`, mirroring this module's shape; extend bias support - through them. -- Fix the splitk K-tail accumulation on odd K and lift the `K % 2 == 0` - launcher assert. -- Optionally repopulate `a16w16_flatmm_kernels_list` (kid 100..115) - with a bias-aware warp-spec epilogue (currently empty; the splitk - pipeline with `splitK=0` covers the same shapes bit-identically). +| `__init__.py` | thin public `opus_gemm`/`opus_bmm` delegates and lazy `gemm_a16w16_opus` compatibility entry | +| `dispatch.py` | public contract validation and strict exact-kid family routing | +| `_arch.py` | per-explicit-device architecture/CU scalar cache | +| `policy.py` | A16 tuned/heuristic candidate selection plus MXFP8 tuned CSV discovery, padded-M lookup, local-to-global kid normalization and heuristic fallback | +| `launch_plan.py` | shared `WorkspaceSpec`, A16 exact-kid/split-K planning, and A8 family contract/MXFP8 BMM planning | +| `gemm_op_a8w8.py` | three non-MX A8 GEMM adapters, the legacy bpreshuffle tuner compatibility entry, MXFP8 BMM workspace materialization, and the unified `_launch_a8w8_backend` over four pybind raw bindings | +| `csrc/opus_gemm/opus_gemm_a8w8_tune.py` | plain A8W8 no-scale/blockscale tuner and saved-kid CSV replay | +| `moe_stage1_a8w4.py` | A8W4 MoE stage-1 runtime binding and launcher | +| `moe_stage2_a8w4.py` | A8W4 MoE stage-2 runtime bindings and launchers | +| `../gemm_op_a8w8.py` | general scaled CK/CKTile/ASM/Triton A8 dispatchers plus the tuned-row OPUS bpreshuffle route | +| `../batched_gemm_op_bf16.py` | existing high-level CK BF16 BMM path; it is not an OPUS A16W16 BMM wrapper | +| `../batched_gemm_op_a8w8.py` | MXFP8 high-level caller, scalar launch cache, output allocation and split-one/workspace execution choice | +| `../../../csrc/opus_gemm/` | canonical registry, C++ family launchers, codegen, traits and pipelines | diff --git a/aiter/ops/opus/__init__.py b/aiter/ops/opus/__init__.py index ac5a4a38af..ee13863f1d 100644 --- a/aiter/ops/opus/__init__.py +++ b/aiter/ops/opus/__init__.py @@ -1,82 +1,101 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2025-2026, Advanced Micro Devices, Inc. All rights reserved. -"""opus kernel Python user-facing API. +"""Public OPUS GEMM/BMM interfaces backed by shared exact-kid launchers.""" -Public API: `gemm_a16w16_opus` (CSV lookup + C++ heuristic) and -`opus_gemm_a16w16_tune` (id-based binding). The gfx942 A8W8 blockscale -bpreshuffle entry is an explicit tune API. -""" +from __future__ import annotations -from ._arch import _detect_arch +import torch +from torch import Tensor -_SUPPORTED = {"gfx950", "gfx942", "gfx1250"} -_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 " - "GPU_ARCHS to one of these (or run on a matching device) to use this " - "module." -) +from .dispatch import _opus_dispatch -_arch_ok, _detected_arch = _detect_arch(_SUPPORTED) +def opus_gemm( + XQ: Tensor, + WQ: Tensor, + Y: Tensor, + *, + kid: int, + layout: str = "plain", + x_scale: Tensor | None = None, + w_scale: Tensor | None = None, + bias: Tensor | None = None, + split_k: int = 0, + workspace: Tensor | None = None, +) -> Tensor: + """Launch logical 2D ``[M,K] x [N,K] -> [M,N]`` by exact ``kid``. -def _make_unsupported_arch_stub(name: str): - """Build a callable that always raises with the detected-arch context.""" - - def _stub(*_args, **_kwargs): - raise RuntimeError( - f"{name} requires GPU arch in {sorted(_SUPPORTED)}; " - f"detected {_detected_arch!r}. {_HINT}" - ) - - _stub.__name__ = name - _stub.__qualname__ = name - _stub.__doc__ = f"Stub: {_FEATURE} unavailable on {_detected_arch!r}." - return _stub + ``Y`` is caller-owned and returned. ``layout='bpreshuffle'`` declares a + transformed WQ content layout that Tensor metadata cannot prove. + """ + return _opus_dispatch( + "opus_gemm", + 2, + XQ, + WQ, + Y, + kid=kid, + layout=layout, + x_scale=x_scale, + w_scale=w_scale, + bias=bias, + split_k=split_k, + workspace=workspace, + ) -if _arch_ok: - from .bmm_op import bmm_a8w8_mxscale_opus - from .gemm_op_a16w16 import ( - gemm_a16w16_opus, - opus_gemm_a16w16_tune, - opus_gemm_workspace_init, - opus_gemm_workspace_release, - opus_gemm_workspace_release_all, +def opus_bmm( + XQ: Tensor, + WQ: Tensor, + Y: Tensor, + *, + kid: int, + layout: str = "plain", + x_scale: Tensor | None = None, + w_scale: Tensor | None = None, + bias: Tensor | None = None, + split_k: int = 0, + workspace: Tensor | None = None, +) -> Tensor: + """Launch batch-first ``[B,M,K] x [B,N,K] -> [B,M,N]`` by exact kid.""" + return _opus_dispatch( + "opus_bmm", + 3, + XQ, + WQ, + Y, + kid=kid, + layout=layout, + x_scale=x_scale, + w_scale=w_scale, + bias=bias, + split_k=split_k, + workspace=workspace, ) - 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) +def gemm_a16w16_opus( + A: Tensor, + B: Tensor, + bias: Tensor | None = None, + dtype: torch.dtype = torch.bfloat16, + *, + kernelId: int | None = None, + splitK: int | None = None, + out: Tensor | None = None, +) -> Tensor: + """Run the legacy shape-driven A16W16 OPUS selection path.""" + from .gemm_op_a16w16 import gemm_a16w16_opus as _impl -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") - bmm_a8w8_mxscale_opus = _make_unsupported_arch_stub("bmm_a8w8_mxscale_opus") - 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") - opus_gemm_workspace_release = _make_unsupported_arch_stub( - "opus_gemm_workspace_release" - ) - opus_gemm_workspace_release_all = _make_unsupported_arch_stub( - "opus_gemm_workspace_release_all" + return _impl( + A, + B, + bias, + dtype, + kernelId=kernelId, + splitK=splitK, + out=out, ) -__all__ = [ - "bmm_a8w8_mxscale_opus", - "gemm_a16w16_opus", - "opus_gemm_a8w8_blockscale_bpreshuffle_tune", - "opus_gemm_a16w16_tune", - "opus_gemm_workspace_init", - "opus_gemm_workspace_release", - "opus_gemm_workspace_release_all", -] +__all__ = ["gemm_a16w16_opus", "opus_bmm", "opus_gemm"] diff --git a/aiter/ops/opus/_arch.py b/aiter/ops/opus/_arch.py index 35640a5f41..676cda3274 100644 --- a/aiter/ops/opus/_arch.py +++ b/aiter/ops/opus/_arch.py @@ -37,8 +37,17 @@ import os from collections.abc import Iterable +import torch + logger = logging.getLogger("aiter.ops.opus._arch") +GFX942 = "gfx942" +GFX950 = "gfx950" +GFX1250 = "gfx1250" +SUPPORTED_OPUS_ARCHES = frozenset((GFX942, GFX950, GFX1250)) + +_DEVICE_INFO_CACHE: dict[torch.device, tuple[str, int]] = {} + def _detect_arch( supported: Iterable[str], @@ -134,3 +143,50 @@ def _check_arch( if hint: msg = f"{msg} {hint}" raise ImportError(msg) + + +def _normalize_device(device: torch.device | str | int) -> torch.device: + """Return an explicit device so cache entries never follow current_device.""" + if isinstance(device, int): + device = torch.device("cuda", device) + elif not isinstance(device, torch.device): + device = torch.device(device) + if device.type == "cuda" and device.index is None: + device = torch.device("cuda", torch.cuda.current_device()) + return device + + +def _read_device_arch_and_cu(device: torch.device) -> tuple[str, int]: + """Read immutable architecture properties for one explicit GPU.""" + if device.type != "cuda": + raise RuntimeError(f"OPUS GEMM requires a GPU tensor; got device {device}") + props = torch.cuda.get_device_properties(device) + raw_arch = str(props.gcnArchName).strip() + arch = raw_arch.split(":", 1)[0].lower() + if not arch.startswith("gfx"): + try: + from ...jit.utils.chip_info import get_gfx_runtime + + arch = get_gfx_runtime().lower() + except Exception as exc: + raise RuntimeError( + f"cannot determine the AMD gfx architecture for device {device}" + ) from exc + return arch, int(props.multi_processor_count) + + +def _device_arch_and_cu( + device: torch.device | str | int, +) -> tuple[str, int]: + """Return cached arch/CU metadata scoped to an explicit device.""" + explicit = _normalize_device(device) + info = _DEVICE_INFO_CACHE.get(explicit) + if info is None: + info = _read_device_arch_and_cu(explicit) + _DEVICE_INFO_CACHE[explicit] = info + return info + + +def _device_arch(device: torch.device | str | int) -> str: + """Return the runtime gfx architecture for one tensor device.""" + return _device_arch_and_cu(device)[0] diff --git a/aiter/ops/opus/bmm_op.py b/aiter/ops/opus/bmm_op.py deleted file mode 100644 index 7f79c5d765..0000000000 --- a/aiter/ops/opus/bmm_op.py +++ /dev/null @@ -1,175 +0,0 @@ -# SPDX-License-Identifier: MIT -# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. -"""Opus batched-BMM Python bindings. - -This module is intentionally separate from `gemm_op_a16w16.py`: BMM callers use -batch-in-the-middle or grouped layouts (for example DSV4 `wo_a`) while the -underlying kernels still live in the shared opus GEMM backend. -""" - -import functools - -import torch - -from ...jit.core import compile_ops - - -def _gen_bmm_a8w8_scale_fake_tensors( - x: torch.Tensor, - wo_a: torch.Tensor, - Y: torch.Tensor, - x_scale: torch.Tensor, - w_scale: torch.Tensor, - splitK: int = 2, - kernelId: int = 0, -) -> None: - # In-place mutation of ``Y``; fake must mirror the void C++ op (full arg - # list + None return) so torch.compile registers a mutating op, not a - # tensor-producing one. - return None - - -# mmajor fp8 e8m0 mxscale BMM raw binding: x/Y are [M, batch, *], wo_a + w_scale -# batch-major (zero-copy DSV4 wo_a). kid-dispatched; driven by -# bmm_a8w8_mxscale_opus below. -@compile_ops( - "module_deepgemm_opus", - fc_name="opus_bmm_a8w8_mxscale", - gen_fake=_gen_bmm_a8w8_scale_fake_tensors, - develop=True, -) -def _opus_bmm_a8w8_mxscale_raw( - x: torch.Tensor, - wo_a: torch.Tensor, - Y: torch.Tensor, - x_scale: torch.Tensor, - w_scale: torch.Tensor, - splitK: int = 2, - kernelId: int = 0, -) -> None: - # In-place: result written into ``Y``, void return (``-> None`` keeps it - # torch.compile-safe as a mutating op). Callers read ``Y``. - ... - - -# ---- Shape-driven mxscale flatmm BMM (tuned row + heuristic fallback) ------ -# The raw binding has no tuning of its own (kernelId=0 -> slow k32 fused). This -# wrapper adds selection: explicit kernelId -> verbatim; else the tuned row the -# family entry looked up; else M-split for large unaligned M; else a coarse M/G -# heuristic. - - -@functools.cache -def _mxscale_kid_m_align() -> dict[int, int]: - """kid -> M multiple its launcher requires (1 == it masks a partial M tile). - - Comes from the codegen instance table, which is also what the tuner filters - candidates on. This used to be a hand-kept kid allowlist here and a second - hand-kept m_align column in the tuner, and the two disagreed: kid326 was - dispatched at unaligned M by this file while the tuner never tuned it there, - which cost ~9% at the wo_a decode shapes. - """ - from csrc.opus_gemm.opus_gemm_common import a8w8_mxscale_bmm_kernel_lists - - return { - int(kid): int(inst.m_align) - for fam in a8w8_mxscale_bmm_kernel_lists - for kid, inst in fam.items() - } - - -def _kid_runs_m(kid: int, m: int) -> bool: - """True iff kid's launcher accepts this M (unknown kid -> assume it does not). - - Only a tuned row found at a padded M can name a kernel that rejects the - real, smaller M, so this is what an incoming id is checked against below. - No tuned winner needs alignment today (all 11 mask their partial M tile), - but 10 of the 45 codegen instances require M % 128 or % 256, so a re-tune - can put one in the CSV. - """ - align = _mxscale_kid_m_align().get(int(kid)) - return align is not None and m % align == 0 - - -def _heuristic_mxscale_kid(g: int, m: int, n: int, k: int) -> int: - """Coarse M/G kid picker for shapes not in the tuned CSV. - - kid 158 (512x256 preload pipeline) for large-M/high-G, falling back to kid 150 - (256x256 plain) for K>8192 where 158 early-returns; kid 320/640 for small-M; - kid 653 the general strong mid/small-M pick; kid 0 (k32 fused) for shapes that - are not tile-aligned in N or K. - """ - - def div(a: int, b: int) -> bool: - return a % b == 0 - - if div(n, 256) and div(k, 128) and (m >= 2048 or (m >= 1024 and g >= 8)): - # Large M: the preload pipeline (kid158) is the tuned winner across this - # whole region (CSV picks 158 for every aligned m>=2048). No M alignment - # needed -- the pipeline family masks its partial trailing tile via buffer - # OOB. kid158 stages the SFA/SFB scales into LDS and early-returns for - # K>8192 (SFA_K_MAX), so gate the preload pick at K<=8192 and fall back to - # the plain 256x256 (kid150) for K>8192. Measured on g=2,n=1024,k=4096: - # kid150 was 34-51% slower than 158 at the untuned m=2560/3072/3584 - # buckets, and on unaligned M a single kid158 launch beats the sub-tile - # kid653 by 13-34% (g2/m2624, g8/m1000, g16/m600). - return 158 if 4096 <= k <= 8192 else 150 - # Sub-tile M: B_M=32/64 tiles mask partial M via buffer OOB, so run any M - # (no m-alignment needed -- verified 653/321/... run arbitrary unaligned M). - if m < 64: - return 640 if (div(n, 64) and div(k, 256)) else 653 - if m <= 256 and k <= 1024 and div(n, 32) and div(k, 256): - return 320 - if div(n, 64) and div(k, 128): - return 653 - return 0 # nothing tile-aligned: k32 fused runs arbitrary shapes - - -def bmm_a8w8_mxscale_opus( - x: torch.Tensor, - wo_a: torch.Tensor, - x_scale: torch.Tensor, - w_scale: torch.Tensor, - out: torch.Tensor | None = None, - dtype: torch.dtype = torch.bfloat16, - kernelId: int | None = None, - splitK: int | None = None, -) -> torch.Tensor: - """Opus fp8 e8m0 mxscale (block-scale) BMM by kernel id. - - mmajor DSV4 wo_a layout: ``x`` [M, G, K] fp8, ``wo_a`` [G, N, K] fp8, - ``x_scale`` [M, G, K/128], ``w_scale`` [G, N/128, K/128], ``out`` optional - [M, G, N]. Returns the [M, G, N] output. - - ``kernelId`` None falls back to the shape heuristic: the tuned CSV is read - one layer up, in batched_gemm_a8w8_mxscale, which hands the tuned id down. - An id this backend cannot run at this M gets the heuristic too, so the - caller never has to know the alignment rules; _opus_bmm_a8w8_mxscale_raw is - the entry that launches an id verbatim. ``splitK`` defaults to 1. - """ - m, g, k = int(x.shape[0]), int(x.shape[1]), int(x.shape[2]) - n = int(wo_a.shape[1]) - - if out is not None: - Y = out - else: - Y = torch.empty((m, g, n), dtype=dtype, device=x.device) - - # A tuned row found at a padded M can name a kernel whose launcher rejects - # the real, smaller M; drop its splitK along with it and let the heuristic - # pick instead of letting the launcher throw. - if kernelId is not None and not _kid_runs_m(int(kernelId), m): - kernelId = splitK = None - if kernelId is None: - kernelId = _heuristic_mxscale_kid(g, m, n, k) - if splitK is None: - splitK = 1 - - _opus_bmm_a8w8_mxscale_raw(x, wo_a, Y, x_scale, w_scale, int(splitK), int(kernelId)) - return Y - - -__all__ = [ - "_opus_bmm_a8w8_mxscale_raw", - "bmm_a8w8_mxscale_opus", -] diff --git a/aiter/ops/opus/common.py b/aiter/ops/opus/common.py deleted file mode 100644 index 747af9b1e7..0000000000 --- a/aiter/ops/opus/common.py +++ /dev/null @@ -1,240 +0,0 @@ -# SPDX-License-Identifier: MIT -# Copyright (C) 2025-2026, Advanced Micro Devices, Inc. All rights reserved. -""" -Opus a16w16 tuned-CSV lookup against the **global** aiter BF16 GEMM CSVs. - -Source of truth: - - aiter/configs/bf16_tuned_gemm.csv - aiter/configs/model_configs/*_bf16_tuned_gemm.csv - -These are the same CSVs that the aiter-global `gemm_a16w16` dispatcher -reads; opus rows are stamped with `libtype == 'opus'` and coexist with -`asm` / `triton` / `skinny` / `flydsl` / `torch` / `hipblaslt` rows for -the same (cu_num, M, N, K, ...) keys. We filter by `libtype == 'opus'` -here, so the opus runtime dispatch only returns a tuned winner when one -of those CSVs has an opus row matching the shape. - -Schema (matches gradlib/GemmTuner.py output): - - gfx, cu_num, M, N, K, bias, dtype, outdtype, scaleAB, bpreshuffle, - libtype, solidx, splitK, us, kernelName, err_ratio, tflops, bw - -`gfx` is optional (the legacy opus-private CSV did not have it); when -absent we tolerate it. Rows missing any of the 9 key columns -(cu_num/M/N/K/bias/dtype/outdtype/scaleAB/bpreshuffle) are skipped. - -Configuration: - - AITER_OPUS_TUNED_CSV_GLOB - Colon-separated list of glob patterns for tuned CSVs. Default - includes both the global BF16 GEMM CSV and the per-model CSVs - under aiter/configs/model_configs/. - - (Removed in this rewrite: AITER_OPUS_A16W16_TUNED_CSV, - AITER_OPUS_A16W16_UNTUNED_CSV, AITER_OPUS_LOG_UNTUNED, and the - autolog feature. Untuned-shape collection is no longer supported; - use gradlib/gemm_tuner.py --libtype opus to tune shapes offline.) -""" - -from __future__ import annotations - -import functools -import glob -import os - -import pandas as pd -import torch - -from aiter.jit.core import AITER_ROOT_DIR - -# ---- Env / default paths -------------------------------------------------- - -# Colon-separated list of glob patterns; each pattern is expanded with -# glob.glob() and the results concatenated. Order does not matter -- on -# duplicate keys we keep the row with the smallest `us` (best winner) -# across all files. -_DEFAULT_TUNED_CSV_GLOB = ( - f"{AITER_ROOT_DIR}/aiter/configs/bf16_tuned_gemm.csv" - f":{AITER_ROOT_DIR}/aiter/configs/model_configs/*_bf16_tuned_gemm.csv" -) - -OPUS_TUNED_CSV_GLOB = os.getenv("AITER_OPUS_TUNED_CSV_GLOB", _DEFAULT_TUNED_CSV_GLOB) - - -# ---- Tuned CSV lookup ----------------------------------------------------- - -_KEY_COLUMNS = ( - "cu_num", - "M", - "N", - "K", - "bias", - "dtype", - "outdtype", - "scaleAB", - "bpreshuffle", -) - - -def _resolve_csv_paths() -> list[str]: - """Expand OPUS_TUNED_CSV_GLOB into a deduplicated list of file paths.""" - paths: list[str] = [] - seen: set[str] = set() - for pattern in OPUS_TUNED_CSV_GLOB.split(os.pathsep): - pattern = pattern.strip() - if not pattern: - continue - for path in sorted(glob.glob(pattern)): - if path in seen: - continue - seen.add(path) - paths.append(path) - return paths - - -@functools.lru_cache(maxsize=1) -def _load_tuned_dict() -> dict: - """Load opus-flagged rows from all configured tuned CSVs into a dict. - - Returns a mapping `key -> {'solidx', 'splitK', 'kernelName'}` where key - is the 9-tuple from `_KEY_COLUMNS`. When multiple CSVs report a winner - for the same key, the one with the smallest `us` is retained (best - timing wins). - - Cached for the process lifetime. Call `_load_tuned_dict.cache_clear()` - if a fresh CSV is dropped in between invocations (rare in production). - """ - paths = _resolve_csv_paths() - if not paths: - return {} - - frames: list[pd.DataFrame] = [] - for path in paths: - try: - df = pd.read_csv(path) - except (pd.errors.EmptyDataError, FileNotFoundError): - continue - if "libtype" not in df.columns: - # CSVs without a `libtype` column predate the multi-backend - # schema; they cannot contain opus rows by definition. Skip - # rather than misclassify their rows as opus. - continue - df = df[df["libtype"] == "opus"] - if df.empty: - continue - missing = [c for c in _KEY_COLUMNS if c not in df.columns] - if missing: - # Malformed / partial-schema CSV. Skip rather than crash. - continue - frames.append(df) - - if not frames: - return {} - combined = pd.concat(frames, ignore_index=True).drop_duplicates() - - # Conflict resolution: same 9-tuple key from multiple files -> keep - # the row with the smallest `us` (best timing). If `us` is missing, - # fall back to first-write-wins. - has_us = "us" in combined.columns - if has_us: - combined = combined.sort_values("us", ascending=True, kind="mergesort") - out: dict = {} - for _, row in combined.iterrows(): - key = tuple(row[c] for c in _KEY_COLUMNS) - if key in out: - continue # already kept the better one - try: - out[key] = { - "solidx": int(row["solidx"]), - "splitK": int(row["splitK"]), - "kernelName": str(row.get("kernelName", "")), - } - except (KeyError, ValueError, TypeError): - # Missing solidx / splitK or non-int values; skip this row. - continue - return out - - -def _key_from_runtime( - M: int, - N: int, - K: int, - bias: bool, - dtype: torch.dtype, - outdtype: torch.dtype, - scaleAB: bool = False, - bpreshuffle: bool = False, -) -> tuple: - """Build the 9-tuple lookup key using the current device's cu_num.""" - cu_num = torch.cuda.get_device_properties( - torch.cuda.current_device() - ).multi_processor_count - return ( - int(cu_num), - int(M), - int(N), - int(K), - bool(bias), - str(dtype), - str(outdtype), - bool(scaleAB), - bool(bpreshuffle), - ) - - -# Mono-tile kid → (B_M, B_N, B_K). Must stay in lock-step with -# csrc/opus_gemm/opus_gemm_common.py:_MONO_TILE_TILES; the runtime guard -# below uses it to validate (N, K) alignment for CSV-picked mono kids, -# since tuned_gemm.get_padded_m pads the lookup key by M only and can -# surface a kid whose B_N / B_K does not divide the actual N / K. -_MONO_TILE_KID_TILES = { - 1400: (192, 256, 64), - 1401: (128, 256, 64), - 1402: (192, 128, 64), - 1403: (128, 128, 64), - 1404: (64, 128, 64), -} - - -def mono_kid_shape_ok(kid: int, N: int, K: int) -> bool: - """Return True iff `kid` is a mono-tile kid whose B_N / B_K divides N / K. - - Returns True for non-mono kids (out of range) so callers can use this - as an unconditional gate without having to special-case the kid range. - B_M is intentionally NOT checked: the mono-tile launcher now handles - M-tail via the bounded gmem descriptor (commit 41e2d482a), so M may - be non-tile-aligned. N and K must still be tile-aligned -- the kernel - has no N-tail mask (column writes would spill into the next row) and - no K-tail mask. - """ - bm_bn_bk = _MONO_TILE_KID_TILES.get(int(kid)) - if bm_bn_bk is None: - return True - _, B_N, B_K = bm_bn_bk - return (int(N) % B_N == 0) and (int(K) % B_K == 0) - - -def lookup_tuned( - M: int, - N: int, - K: int, - bias: bool, - dtype: torch.dtype, - outdtype: torch.dtype, - scaleAB: bool = False, - bpreshuffle: bool = False, -) -> dict | None: - """Look up a tuned winner for this shape; returns dict or None. - - Dict contains 'solidx' (kernelId), 'splitK', 'kernelName'. - """ - key = _key_from_runtime(M, N, K, bias, dtype, outdtype, scaleAB, bpreshuffle) - return _load_tuned_dict().get(key) - - -__all__ = [ - "OPUS_TUNED_CSV_GLOB", - "lookup_tuned", - "mono_kid_shape_ok", -] diff --git a/aiter/ops/opus/dispatch.py b/aiter/ops/opus/dispatch.py new file mode 100644 index 0000000000..2d3168f392 --- /dev/null +++ b/aiter/ops/opus/dispatch.py @@ -0,0 +1,255 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2025-2026, Advanced Micro Devices, Inc. All rights reserved. +"""Contract validation and exact-kid dispatch for public OPUS operations.""" + +from __future__ import annotations + +from functools import lru_cache + +import torch +from torch import Tensor + +from csrc.opus_gemm.opus_gemm_common import ( + OpusGemmInstance, + get_kernel_instance, + kernels_list, +) + +from . import gemm_op_a8w8 as _a8w8_family +from . import gemm_op_a16w16 as _a16w16_family +from . import launch_plan +from ._arch import GFX950 + + +def _validate_a16w16_public_contract( + *, + kid: int, + instance: OpusGemmInstance, + input_dtype: torch.dtype, + weight_dtype: torch.dtype, + output_dtype: torch.dtype, + layout: str, + has_x_scale: bool, + has_w_scale: bool, +) -> None: + """Validate A16W16-only options shared by both public routers.""" + if input_dtype != weight_dtype: + raise ValueError( + f"OPUS requires matching XQ/WQ dtypes; got " f"{input_dtype}/{weight_dtype}" + ) + if input_dtype != torch.bfloat16: + raise ValueError(f"OPUS kid {kid} requires bf16 XQ/WQ; got {input_dtype}") + if layout != "plain": + raise ValueError( + f"OPUS kid {kid} belongs to family a16w16 and requires " + f"layout='plain'; got {layout!r}" + ) + if has_x_scale or has_w_scale: + raise ValueError("OPUS a16w16 does not accept x_scale/w_scale") + arch = (instance.arch_prefix or GFX950).lower() + if get_kernel_instance(arch, "a16w16", kid, output_dtype) is None: + raise ValueError(f"OPUS kid {kid} does not support Y.dtype={output_dtype}") + + +@lru_cache(maxsize=4096) +def _resolve_contract( + kid: int, + input_dtype: torch.dtype, + weight_dtype: torch.dtype, + output_dtype: torch.dtype, + layout: str, + has_x_scale: bool, + has_w_scale: bool, + has_bias: bool, + has_workspace: bool, + split_k: int, +) -> tuple[str, object, object]: + """Validate/cache the public contract and its resolved family module.""" + instance = kernels_list.get(kid) + if instance is None: + raise ValueError(f"unknown OPUS kid {kid}") + + if instance.kernel_tag.startswith("a16w16"): + _validate_a16w16_public_contract( + kid=kid, + instance=instance, + input_dtype=input_dtype, + weight_dtype=weight_dtype, + output_dtype=output_dtype, + layout=layout, + has_x_scale=has_x_scale, + has_w_scale=has_w_scale, + ) + return "a16w16", instance, _a16w16_family + + family = launch_plan._validate_a8w8_public_contract( + kernel_tag=instance.kernel_tag, + kid=kid, + input_dtype=input_dtype, + weight_dtype=weight_dtype, + output_dtype=output_dtype, + layout=layout, + has_x_scale=has_x_scale, + has_w_scale=has_w_scale, + has_bias=has_bias, + has_workspace=has_workspace, + split_k=split_k, + ) + return family, instance, _a8w8_family + + +def _opus_dispatch( + operation: str, + rank: int, + XQ: Tensor, + WQ: Tensor, + Y: Tensor, + *, + kid: int, + layout: str = "plain", + x_scale: Tensor | None = None, + w_scale: Tensor | None = None, + bias: Tensor | None = None, + split_k: int = 0, + workspace: Tensor | None = None, +) -> Tensor: + bad_name = None + bad_value = None + if not isinstance(XQ, Tensor): + bad_name, bad_value = "XQ", XQ + elif not isinstance(WQ, Tensor): + bad_name, bad_value = "WQ", WQ + elif not isinstance(Y, Tensor): + bad_name, bad_value = "Y", Y + elif x_scale is not None and not isinstance(x_scale, Tensor): + bad_name, bad_value = "x_scale", x_scale + elif w_scale is not None and not isinstance(w_scale, Tensor): + bad_name, bad_value = "w_scale", w_scale + if bad_name is not None: + raise TypeError( + f"{operation}: {bad_name} must be a Tensor, got {type(bad_value)!r}" + ) + + bad_rank = None + if XQ.dim() != rank: + bad_rank = "XQ" + elif WQ.dim() != rank: + bad_rank = "WQ" + elif Y.dim() != rank: + bad_rank = "Y" + elif x_scale is not None and x_scale.dim() != rank: + bad_rank = "x_scale" + elif w_scale is not None and w_scale.dim() != rank: + bad_rank = "w_scale" + if bad_rank is not None: + expected = "logical 2D" if rank == 2 else "batch-first 3D" + raise ValueError( + f"{operation} expects {expected} {bad_rank}; the selected kid " + "family must also support that operation" + ) + + if type(kid) is not int: + raise ValueError(f"OPUS kid must be an integer id, got {kid!r}") + if type(split_k) is not int: + raise ValueError(f"OPUS split_k must be an integer, got {split_k!r}") + if split_k < 0: + raise ValueError(f"OPUS split_k must be non-negative, got {split_k}") + if layout not in ( + "plain", + "bpreshuffle", + "mxscale_bmm", + ): + raise ValueError( + f"unsupported OPUS weight layout {layout!r}; expected " + "'plain', 'bpreshuffle' or 'mxscale_bmm'" + ) + if operation == "opus_gemm" and layout == "mxscale_bmm": + raise ValueError("layout='mxscale_bmm' is only supported by opus_bmm") + + has_x_scale = x_scale is not None + has_w_scale = w_scale is not None + family, instance, family_module = _resolve_contract( + kid, + XQ.dtype, + WQ.dtype, + Y.dtype, + layout, + has_x_scale, + has_w_scale, + bias is not None, + workspace is not None, + split_k, + ) + route_arch = (instance.arch_prefix or GFX950).lower() + + if family == "a16w16": + launch = ( + family_module._launch_a16w16_gemm + if operation == "opus_gemm" + else family_module._launch_a16w16_bmm + ) + return launch( + XQ, + WQ, + Y, + bias, + kid=kid, + split_k=split_k, + workspace=workspace, + route_arch=route_arch, + instance=instance, + ) + + if family == "a8w8_mxscale_bmm": + if operation != "opus_bmm": + raise ValueError("OPUS a8w8_mxscale_bmm supports opus_bmm only") + assert x_scale is not None and w_scale is not None + return family_module._launch_a8w8_mxscale_bmm( + XQ, + WQ, + Y, + x_scale, + w_scale, + kid=kid, + split_k=split_k, + workspace=workspace, + route_arch=route_arch, + instance=instance, + ) + + if operation != "opus_gemm": + raise ValueError(f"OPUS family {family} is GEMM-only; use opus_gemm") + + if family == "a8w8": + return family_module._launch_a8w8_gemm( + XQ, + WQ, + Y, + kid=kid, + route_arch=route_arch, + instance=instance, + ) + assert x_scale is not None and w_scale is not None + if family == "a8w8_blockscale": + return family_module._launch_a8w8_blockscale_gemm( + XQ, + WQ, + Y, + x_scale, + w_scale, + kid=kid, + route_arch=route_arch, + instance=instance, + ) + if family == "a8w8_blockscale_bpreshuffle": + return family_module._launch_a8w8_blockscale_bpreshuffle_gemm( + XQ, + WQ, + x_scale, + w_scale, + Y, + kid=kid, + route_arch=route_arch, + instance=instance, + ) + raise RuntimeError(f"unsupported canonical OPUS family {family!r}") diff --git a/aiter/ops/opus/gemm_op_a16w16.py b/aiter/ops/opus/gemm_op_a16w16.py index 3a36b8fd16..5c3c636df0 100644 --- a/aiter/ops/opus/gemm_op_a16w16.py +++ b/aiter/ops/opus/gemm_op_a16w16.py @@ -1,712 +1,482 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2025-2026, Advanced Micro Devices, Inc. All rights reserved. -""" -Opus a16w16 Python user-facing API. - -Public entry points: - -* `gemm_a16w16_opus(A, B, bias=None, dtype=bf16, *, kernelId=None, splitK=None, out=None)` - Shape-driven wrapper. The typical user writes `gemm_a16w16_opus(A, B)` - and never sees a kid number. Internal path: - - 1. Reshape A/B to 3D, allocate Y, validate (bias allowed across the - split-barrier / splitk kid families; bpreshuffle and non-bf16 A/B - unsupported). - 2. If `kernelId` is given explicitly -> opus_gemm_a16w16_tune (bias - is forwarded; the C++ dispatcher rejects non-bias-aware kids). - 3. Otherwise query the global aiter BF16 tuned CSVs via - aiter.ops.opus.common (filtered by `libtype == 'opus'`, key - includes bias=True/False); on hit -> opus_gemm_a16w16_tune - with the tuned (solidx, splitK). - 4. On miss -> fall through to the private bf16 no-scale binding - `_opus_gemm_bf16_dispatch`, which forwards to the C++ entry - `opus_gemm` whose bf16 branch does its own lookup + heuristic - dispatch (see csrc/opus_gemm/opus_gemm.cu). bias is forwarded - through this path: the C++ entry skips its bias-agnostic lookup - map when bias is present and goes straight to the heuristic - dispatcher (which always returns a bias-aware kid). - -* `opus_gemm_a16w16_tune(XQ, WQ, Y, bias, kernelId, splitK)` - Low-level pybind binding to the id-based tune dispatcher. Exposes a - specific kernel instance by `kernelId` plus optional literal KBatch - via `splitK` and an optional bias tensor (D_OUT-typed, [N] or - [batch, N]; F.linear convention). Intended for the tuner, for debugging a specific kid, - and for aiter-global integrations (e.g. future tuned_gemm.solMap). - -All entry points share the JIT module `module_deepgemm_opus`, which -still hosts bindings for other opus kernel families (a8w8 etc.). The -Python surface is deliberately per-dtype: a16w16 here, a8w8 in its own -module when that lands. -""" - -import functools -import logging -import os -import sys +"""A16W16 exact launch, Torch workspace, and shape-driven compatibility.""" + +from functools import lru_cache import torch -from torch import Tensor -from ...jit.core import AITER_ROOT_DIR, compile_ops -from . import common as _opus_common +from aiter import logger +from csrc.opus_gemm.opus_gemm_common import OpusGemmInstance -logger = logging.getLogger("aiter") +from ...jit.core import compile_ops +from ._arch import _device_arch_and_cu +from .launch_plan import _get_cached_a16w16_launch_plan -# ---- Low-level pybind bindings -------------------------------------------- +# ---- Low-level A16W16 backend -------------------------------------------- -def _gen_opus_gemm_a16w16_tune_fake_tensors( +def _gen_opus_gemm_a16w16_launch_fake_tensors( XQ: torch.Tensor, WQ: torch.Tensor, Y: torch.Tensor, - bias: torch.Tensor | None = None, - workspace: torch.Tensor | None = None, - kernelId: int = 0, - splitK: int = 0, + bias: torch.Tensor | None, + workspace: torch.Tensor | None, + kid: int, + split_k: int, ) -> torch.Tensor: return Y -# Raw pybind binding to the C++ id-based dispatcher. We wrap it in a Python -# function below to add a stride-layout guard before the C++ call -- the -# launcher hardcodes stride_b_batch == N*K and reads gpu memory directly, -# so a broadcast / non-contiguous WQ silently corrupts results or faults -# the GPU. Keep `gen_fake` and `fc_name` on the raw binding so dynamo and -# torch.library see the underlying op. @compile_ops( "module_deepgemm_opus", - fc_name="opus_gemm_a16w16_tune", - gen_fake=_gen_opus_gemm_a16w16_tune_fake_tensors, + fc_name="opus_gemm_a16w16_launch", + gen_fake=_gen_opus_gemm_a16w16_launch_fake_tensors, develop=True, ) -def _opus_gemm_a16w16_tune_raw( +def _opus_gemm_a16w16_launch_raw( XQ: torch.Tensor, WQ: torch.Tensor, Y: torch.Tensor, - bias: torch.Tensor | None = None, - workspace: torch.Tensor | None = None, - kernelId: int = 0, - splitK: int = 0, + bias: torch.Tensor | None, + workspace: torch.Tensor | None, + kid: int, + split_k: int, ) -> torch.Tensor: ... -def _check_a16w16_tune_layout(XQ: torch.Tensor, WQ: torch.Tensor, Y: torch.Tensor): - """Reject layouts that the opus launcher's hardcoded strides cannot serve. +def _launch_a16w16_backend( + XQ: torch.Tensor, + WQ: torch.Tensor, + Y: torch.Tensor, + bias: torch.Tensor | None, + workspace: torch.Tensor | None, + kid: int, + split_k: int, +) -> None: + if ( + torch.compiler.is_compiling() + or XQ.is_meta + or getattr(XQ, "fake_mode", None) is not None + ): + _opus_gemm_a16w16_launch_raw( + XQ, + WQ, + Y, + bias, + workspace, + kid, + split_k, + ) + return + with torch.cuda.device(XQ.device): + _opus_gemm_a16w16_launch_raw( + XQ, + WQ, + Y, + bias, + workspace, + kid, + split_k, + ) - Mirrors the kargs setup in csrc/opus_gemm/gen_instances.py - (_gen_flatmm_splitk_instance et al.): - kargs.stride_a = K - kargs.stride_b = K - kargs.stride_c = N - kargs.stride_a_batch = M * K - kargs.stride_b_batch = N * K - kargs.stride_c_batch = M * N - The kernel reads memory at `ptr + batch_id * stride_*_batch + ...` - directly. Any broadcast view (batch stride == 0), transpose, or - sliced layout will hit garbage / unmapped memory. - Cheap to run (a handful of integer comparisons); only raised on real - misuse so the hot path pays nothing. - """ - for name, t in (("XQ", XQ), ("WQ", WQ), ("Y", Y)): - if t.dim() != 3: +def _check_a16w16_launch_layout( + XQ: torch.Tensor, + WQ: torch.Tensor, + Y: torch.Tensor, +) -> None: + """Validate launcher-required 3D shapes and physical strides.""" + for name, tensor in (("XQ", XQ), ("WQ", WQ), ("Y", Y)): + if tensor.dim() != 3: raise ValueError( - f"opus_gemm_a16w16_tune: {name} must be 3D (got " - f"{name}.shape={tuple(t.shape)}). The C++ launcher reads " - f"`{name}.size(0)` as batch and indexes with hardcoded " - f"stride_*_batch == size(1)*size(2)." + f"opus_gemm_a16w16_launch: {name} must be 3D " + f"(got {name}.shape={tuple(tensor.shape)}). " + "The C++ launcher reads size(0) as batch and uses " + "hardcoded dense batch strides." ) batch, M, K = XQ.shape b_w, N, K_w = WQ.shape - b_y, M_y, N_y = Y.shape + expected_wq = (batch, N, K) + expected_y = (batch, M, N) if (b_w, K_w) != (batch, K): raise ValueError( - f"opus_gemm_a16w16_tune: WQ shape mismatch (got " - f"WQ.shape={tuple(WQ.shape)}, expected " - f"({batch}, N, {K})); XQ.shape={tuple(XQ.shape)}" + "opus_gemm_a16w16_launch: WQ shape mismatch " + f"(got {tuple(WQ.shape)}, expected {expected_wq}); " + f"XQ.shape={tuple(XQ.shape)}" ) - if (b_y, M_y, N_y) != (batch, M, N): + if tuple(Y.shape) != expected_y: raise ValueError( - f"opus_gemm_a16w16_tune: Y shape mismatch (got " - f"Y.shape={tuple(Y.shape)}, expected ({batch}, {M}, {N}))" + "opus_gemm_a16w16_launch: Y shape mismatch " + f"(got {tuple(Y.shape)}, expected {expected_y})" ) - # XQ / WQ: the K (innermost / contraction) dimension may be padded -- the - # launcher passes the tensor's real leading stride as kargs.stride_a/stride_b - # and the kernels use it as the lda for BOTH addressing and the gmem buffer - # bound, so a row pitch > K (e.g. a 2880-wide tensor stored at lda 3072) is - # served correctly. We only require: - # * innermost stride == 1 (the kernel layout hardcodes the K stride to 1) - # * row pitch (stride[1]) >= K - # * batch stride == rows * row pitch (or batch == 1) -- rejects broadcast - # (stride 0) and transposed / overlapping views. - for name, t, rows in (("XQ", XQ, M), ("WQ", WQ, N)): - s0, s1, s2 = t.stride() - k_inner = t.shape[2] - ok = s2 == 1 and s1 >= k_inner and (batch == 1 or s0 == rows * s1) - if not ok: + # XQ/WQ allow padded rows but require contiguous K and dense batches. + for name, tensor, rows in (("XQ", XQ, M), ("WQ", WQ, N)): + stride_batch, stride_row, stride_k = tensor.stride() + if ( + stride_k != 1 + or stride_row < K + or (batch != 1 and stride_batch != rows * stride_row) + ): raise NotImplementedError( - f"opus_gemm_a16w16_tune: {name} must be K-contiguous with an " - f"optional padded leading dim -- need stride[2]==1, " - f"stride[1]>={k_inner}, and stride[0]==size(1)*stride[1] (or " - f"batch==1). Got {name}.stride()={tuple(t.stride())}, " - f"{name}.shape={tuple(t.shape)}. Broadcast / transpose / " - f"non-K-contiguous slices are not supported; materialize with " - f"`{name} = {name}.contiguous()` before calling." + f"opus_gemm_a16w16_launch: {name} must be K-contiguous " + "with an optional padded leading dimension; need " + "stride[2]==1, stride[1]>=K, and " + "stride[0]==size(1)*stride[1] when batch>1. " + f"Got {name}.stride()={tuple(tensor.stride())}, " + f"{name}.shape={tuple(tensor.shape)}. " + f"Materialize with `{name} = {name}.contiguous()`." ) - # Y is the output: the launcher hardcodes stride_c == N and - # stride_c_batch == M*N, so it must be fully contiguous. - y_want = (M * N, N, 1) - if tuple(Y.stride()) != y_want: + + # Y must match the launcher's contiguous output strides. + expected_y_stride = (M * N, N, 1) + if Y.stride() != expected_y_stride: raise NotImplementedError( - f"opus_gemm_a16w16_tune: Y must have contiguous strides {y_want} " - f"(got Y.stride()={tuple(Y.stride())}, Y.shape={tuple(Y.shape)}). " - f"The launcher hardcodes stride_c == N and stride_c_batch == M*N; " - f"materialize with `Y = Y.contiguous()` before calling." + "opus_gemm_a16w16_launch: Y must have contiguous strides " + f"{expected_y_stride} (got {tuple(Y.stride())}, " + f"Y.shape={tuple(Y.shape)}). " + "Materialize with `Y = Y.contiguous()`." ) -# gfx1250 split-K workspace element dtype (kid property) -> torch dtype. The -# main kernel WRITES partials in this dtype and the reduce kernel READS them -# back, so the buffer MUST be allocated with the matching element type: a -# bf16-sized buffer handed to an fp32-workspace kid is half the bytes the kernel -# writes -> global-memory overrun -> machine hang. Sizing is therefore done by -# ELEMENT COUNT in this dtype, never by a raw byte size. -_OPUS_WS_TORCH_DTYPE = { - "bf16_t": torch.bfloat16, - "fp32_t": torch.float32, -} - - -@functools.lru_cache(maxsize=1) -def _gfx1250_kids() -> dict: - """Lazily load the opus kid table (csrc/opus_gemm/opus_gemm_common.py). - - Lets the split-K workspace be sized from each kid's ACTUAL kernel - definition (tile B_M/B_N, split_k, workspace dtype) instead of a byte - guess. The module is pure-Python (no torch/JIT deps); returns ``{}`` if it - can't be located so the caller can fall back to a safe over-estimate. - """ - csrc = os.path.join(AITER_ROOT_DIR, "csrc", "opus_gemm") - if csrc not in sys.path: - sys.path.insert(0, csrc) - try: - from opus_gemm_common import kernels_list # type: ignore[import-not-found] - - return kernels_list - except Exception: # noqa: BLE001 - return {} - - -def _get_opus_workspace( - device: torch.device, ws_shape: tuple, dtype: torch.dtype -) -> torch.Tensor: - """Split-K workspace, allocated per call. - - Allocated with its natural ``[batch, split_k, padded_M, padded_N]`` element - shape (never a raw byte count) so the tensor is self-describing and matches - the kernel's ``ptr_ws`` layout; the launcher only reads ``data_ptr()``, so - the extra dims cost nothing. - - A single torch.empty path serves eager AND capture: torch's caching - allocator is HIP graph-capture aware, so a torch.empty issued while a - capture is active is drawn from the graph's mempool and gets an address that - stays valid on replay -- exactly how a captured graph allocates all of its - other intermediates. - - This used to be @functools.cache'd, to pin one buffer per - (device, shape, dtype) for the process lifetime. That is bounded for - inference, which sees a handful of shapes, but the tuner sweeps every - (tile padding x split_k) pair: 1065 distinct buffers totalling 450 GiB on a - 432 GiB part, which is what exhausted VRAM mid-tune. Nothing needs the - buffers to outlive the call -- only one is live at a time -- so they are not - kept. - """ - return torch.empty(ws_shape, dtype=dtype, device=device) - - -def _alloc_splitk_workspace( - kernelId: int, - batch: int, - M: int, - N: int, - splitK: int, - device: torch.device, +def _execute_a16w16( + XQ: torch.Tensor, + WQ: torch.Tensor, + Y: torch.Tensor, + bias: torch.Tensor | None = None, + *, + kid: int, + split_k: int = 0, + workspace: torch.Tensor | None = None, + route_arch: str | None = None, + instance: OpusGemmInstance | None = None, ) -> torch.Tensor: - """Allocate the gfx1250 split-K partial workspace by ELEMENT COUNT in - ``[batch, split_k, padded_M, padded_N]`` -- never by a raw byte size. - - Extents and element dtype come from the selected kid's own kernel - definition, so the buffer is exactly what the kernel writes and reads back: - - * fuse kids (``a16w16_clusterlaunch_tdm_splitk_fuse``): ``split_k`` and the - workspace dtype are COMPILE-TIME per kid and the runtime ``splitK`` arg - is IGNORED, so the buffer is sized from ``fuse_split_k`` / - ``fuse_ws_dtype`` (a bf16- or fp32-workspace kid). - * ws-variant kids (``*_tdm_splitk_ws``): fp32 workspace with a runtime - ``split_k`` the launcher clamps DOWN from ``splitK`` (so ``splitK`` is a - safe upper bound). - - Falls back to a safe over-estimate (fp32 element, split_k=16, 128x512 tile - padding -- the widest gfx1250 split-K tile is B_M<=128, B_N<=256) when the - kid table can't be loaded. - """ - inst = _gfx1250_kids().get(int(kernelId)) - if inst is not None: - b_m, b_n = int(inst.B_M), int(inst.B_N) - if inst.kernel_tag == "a16w16_clusterlaunch_tdm_splitk_fuse": - split_k = max(int(inst.fuse_split_k), 1) - ws_dtype = _OPUS_WS_TORCH_DTYPE.get(inst.fuse_ws_dtype, torch.float32) - else: - # ws-variant: the partial type is the kid's own (its traits D_C, which - # the main kernel stores and the reduce reads), so read it off the - # instance rather than assuming fp32 -- a bf16-partial kid handed an - # fp32 buffer would be walked at the wrong stride by both ends. - # split_k is runtime, and the launcher clamps it DOWN from splitK, so - # splitK is a safe upper bound for sizing. - split_k = max(1, int(splitK)) - ws_dtype = _OPUS_WS_TORCH_DTYPE.get( - getattr(inst, "splitk_workspace_dtype", "fp32_t"), torch.float32 - ) + """Validate, plan, and launch one exact 3D A16W16 operation.""" + _check_a16w16_launch_layout(XQ, WQ, Y) + batch, M, K = XQ.shape + N = Y.shape[2] + + use_gfx950_caller_workspace_fast_path = ( + route_arch == "gfx950" + and workspace is not None + and split_k > 0 + and instance is not None + and instance.splitk_workspace_dtype is not None + ) + if use_gfx950_caller_workspace_fast_path: + # A caller-owned gfx950 workspace and the public registry route avoid + # re-reading device metadata. Explicit gfx950 plans do not consult the + # CU count, so one is a safe cache-key placeholder here. + arch, cu_num = route_arch, 1 else: - # Kid table unavailable: widest-element (fp32) upper bound. split_k must - # cover a fuse kid's max baked split_k (15) and any runtime splitK. - b_m, b_n = 128, 512 - split_k = max(int(splitK), 16) - ws_dtype = torch.float32 + arch, cu_num = _device_arch_and_cu(XQ.device) + + plan = _get_cached_a16w16_launch_plan( + arch, + M, + N, + K, + batch, + cu_num, + bias is not None, + XQ.dtype, + Y.dtype, + int(kid), + int(split_k), + ) + workspace_spec = plan.workspace_spec + if use_gfx950_caller_workspace_fast_path and workspace_spec is None: + raise RuntimeError( + f"OPUS gfx950 kid {plan.resolved_kid} unexpectedly has no " + "caller-workspace plan" + ) + if workspace_spec is None: + if workspace is not None: + raise ValueError( + "opus_gemm_a16w16_launch: " + f"kid {plan.resolved_kid} does not use an external workspace" + ) + elif workspace is None: + workspace = torch.empty( + workspace_spec.shape, + dtype=workspace_spec.dtype, + device=XQ.device, + ) - padded_M = ((int(M) + b_m - 1) // b_m) * b_m - padded_N = ((int(N) + b_n - 1) // b_n) * b_n - ws_shape = (int(batch), split_k, padded_M, padded_N) - return _get_opus_workspace(device, ws_shape, ws_dtype) + _launch_a16w16_backend( + XQ, + WQ, + Y, + bias, + workspace, + plan.resolved_kid, + plan.abi_split_k, + ) + return Y -def opus_gemm_a16w16_tune( +def _launch_a16w16_gemm( XQ: torch.Tensor, WQ: torch.Tensor, Y: torch.Tensor, - bias=None, - kernelId: int = 0, - splitK: int = 0, + bias: torch.Tensor | None = None, + *, + kid: int, + split_k: int = 0, + workspace: torch.Tensor | None = None, + route_arch: str | None = None, + instance: OpusGemmInstance | None = None, ) -> torch.Tensor: - """Low-level id-based dispatcher (Python guard + C++ launch). - - See module docstring. This Python wrapper checks XQ/WQ/Y layout up - front (rejecting broadcast / transpose / slice views that the C++ - kernel would happily run with garbage data); on success it forwards - to the underlying pybind binding. - - Parameters - ---------- - bias : optional D_OUT-typed bias tensor, accepted shapes: - [M] (broadcast across batch; requires batch==1) or [batch, M]. - Only honored on bias-aware kid ranges (split-barrier kid 4..9 - and a16w16_flatmm_splitk kid 200..299); the C++ dispatcher - rejects bias on other kids. - - Backwards-compatibility note - ---------------------------- - Older callers used ``opus_gemm_a16w16_tune(XQ, WQ, Y, kernelId, splitK)`` - with positional args (no bias slot). When the 4th positional argument - is an int, we silently treat it as kernelId and shift remaining args - accordingly so existing tuner / test scripts keep working without an - edit. Mixed-style calls (``..., bias=t, kernelId=k``) keep their kwargs - semantics. - """ - # Positional-int back-compat: opus_gemm_a16w16_tune(XQ, WQ, Y, kid, splitK). - # When `bias` arrives as an int (which torch_library would otherwise - # reject as not Optional[Tensor]), reinterpret as kernelId. - if isinstance(bias, int) and not isinstance(bias, bool): - # Positional int means "this was meant to be kernelId"; treat the - # next positional (kernelId) as splitK and the original splitK - # (default 0) as truly unset. - if splitK != 0 and kernelId == 0: - # Shouldn't happen in old call sites, but be defensive. - new_splitK = splitK - else: - new_splitK = kernelId - kernelId = bias - splitK = new_splitK - bias = None - _check_a16w16_tune_layout(XQ, WQ, Y) - # split-K kids need a workspace tensor allocated externally (torch.empty) - # and passed to the C++ launcher. Asked through is_splitk_kid() rather than - # a literal range: the gfx1250 band is not contiguous (the pre-compiled .co - # kids at [21000, 23000) sit inside it and take no workspace at all), and a - # second spelling of the range is a second thing to keep in sync. - workspace = None - if is_splitk_kid(kernelId): - batch, M, N = Y.shape - workspace = _alloc_splitk_workspace(kernelId, batch, M, N, splitK, XQ.device) - # Mono-tile kid guard: the launcher requires N / K to be tile-aligned - # (the kernel has no N-tail mask and no K-tail mask; M-tail IS handled - # via the bounded gmem desc). A CSV winner picked through - # tuned_gemm.get_padded_m can surface a mono kid whose B_N / B_K does - # not divide the actual N / K -- the launcher would AITER_CHECK abort - # the process. Reroute to opus's own bf16 heuristic dispatch instead; - # it never returns a mono kid, so it always picks something that can - # run the shape. - _, _, N = Y.shape - _, _, K = XQ.shape - if not _opus_common.mono_kid_shape_ok(kernelId, N, K): - logger.warning( - "opus_gemm_a16w16_tune: mono-tile kid %d requires N/K aligned " - "to its tile; got N=%d K=%d -- rerouting to opus bf16 heuristic.", - kernelId, - N, - K, + """Launch logical 2D ``[M,K] x [N,K] -> [M,N]`` A16W16 GEMM.""" + if instance is None and (XQ.dim() != 2 or WQ.dim() != 2 or Y.dim() != 2): + raise ValueError( + "opus_gemm A16W16 expects 2D XQ/WQ/Y; use opus_bmm for " + "batch-first 3D tensors" ) - _opus_gemm_bf16_dispatch(XQ, WQ, Y, None, None, None, bias) - return Y - # C++ launcher is in-place on Y (returns void after PR #2932-style - # refactor to aiter_tensor_t). Keep the wrapper's `return Y` - # contract so callers that did `Y = opus_gemm_a16w16_tune(...)` - # still see the populated Y. - _opus_gemm_a16w16_tune_raw(XQ, WQ, Y, bias, workspace, kernelId, splitK) + _execute_a16w16( + XQ.unsqueeze(0), + WQ.unsqueeze(0), + Y.unsqueeze(0), + bias, + kid=kid, + split_k=split_k, + workspace=workspace, + route_arch=route_arch, + instance=instance, + ) return Y -# Private bf16 no-scale dispatch binding, used only by gemm_a16w16_opus -# as the CSV-miss fallback path. Wraps the same C++ function (opus_gemm) -# that used to be exposed via the legacy aiter.ops.deepgemm.deepgemm_opus -# entry, but deliberately hides its scale / group_layout arguments so -# callers of the a16w16 module do not see FP8-grouped concepts. The C++ -# side's bf16 branch handles lookup + heuristic dispatch internally. -# -# Parameter annotations match the C++ signature exactly; torch_library's -# infer_schema requires every parameter be typed even though we always -# pass None for the last three. -def _gen_opus_gemm_bf16_dispatch_fake_tensors( +def _launch_a16w16_bmm( XQ: torch.Tensor, WQ: torch.Tensor, Y: torch.Tensor, - group_layout: torch.Tensor | None = None, - x_scale: torch.Tensor | None = None, - w_scale: torch.Tensor | None = None, bias: torch.Tensor | None = None, + *, + kid: int, + split_k: int = 0, + workspace: torch.Tensor | None = None, + route_arch: str | None = None, + instance: OpusGemmInstance | None = None, ) -> torch.Tensor: - return Y + """Launch batch-first ``[B,M,K] x [B,N,K] -> [B,M,N]`` A16W16 BMM.""" + if instance is None and (XQ.dim() != 3 or WQ.dim() != 3 or Y.dim() != 3): + raise ValueError( + "opus_bmm A16W16 expects batch-first 3D XQ/WQ/Y; use " + "opus_gemm for logical 2D tensors" + ) + return _execute_a16w16( + XQ, + WQ, + Y, + bias, + kid=kid, + split_k=split_k, + workspace=workspace, + route_arch=route_arch, + instance=instance, + ) -@compile_ops( - "module_deepgemm_opus", - fc_name="opus_gemm", - gen_fake=_gen_opus_gemm_bf16_dispatch_fake_tensors, - develop=True, -) -def _opus_gemm_bf16_dispatch( +def opus_gemm_a16w16_tune( XQ: torch.Tensor, WQ: torch.Tensor, Y: torch.Tensor, - group_layout: torch.Tensor | None = None, - x_scale: torch.Tensor | None = None, - w_scale: torch.Tensor | None = None, - bias: torch.Tensor | None = None, -) -> torch.Tensor: ... - - -# ---- High-level shape-driven API ----------------------------------------- - -# splitk kids main kernel only has the instantiation (traits -# static_assert D_C==float, fp32 workspace). The reduce kernel -# (splitk_reduce_kernel) is templated on D_OUT and dispatches to either -# __bf16 or float at launch time based on Y.dtype(), so both bf16 and fp32 -# outputs are valid. The dispatch code below no longer needs to special-case -# Y.dtype against splitk kids. -# -# splitk kid ranges, one half-open [lo, hi) interval per device family. Kept -# in exact lockstep with the C++ authority `opus_kid_is_splitk` in -# csrc/opus_gemm/opus_gemm.cu -- adding a new device's splitk band means -# appending one row HERE and there. Consumed by is_splitk_kid() below, which -# gates the split-K workspace prewarm in aiter/tuned_gemm.py (non-splitk kids -# never touch the workspace, so warming it for them is pure waste). -_SPLITK_KID_RANGES = ( - (200, 300), # gfx950 base - (1200, 1300), # gfx950 non-OOB mirror (+1000) - (10200, 10300), # gfx942 (+10000) - (20000, 21000), # gfx1250 cluster/TDM split-K - # [21000, 27000) is the pre-compiled .co family -- NOT split-K, no workspace. - (27000, 30000), # gfx1250 fused single-kernel split-K (currently unregistered) -) - - -def is_splitk_kid(kid: int) -> bool: - """True iff `kid` selects a split-K opus a16w16 kernel (fp32 workspace + - reduce). Mirrors C++ `opus_kid_is_splitk`; keep the two in sync.""" - kid = int(kid) - return any(lo <= kid < hi for lo, hi in _SPLITK_KID_RANGES) - - -# Back-compat: the old gfx950-only single-band constants some callers imported. -_SPLITK_KID_MIN = 200 -_SPLITK_KID_MAX = 299 + kernelId: int = 0, + splitK: int = 0, +) -> torch.Tensor: + """Launch the legacy A16W16 GEMM interface by exact kernel id.""" + return _execute_a16w16( + XQ, + WQ, + Y, + kid=int(kernelId), + split_k=int(splitK), + ) -def _validate_and_reshape(A: Tensor, B: Tensor, bias, dtype, out): +def _prepare_shape_driven_a16w16( + A: torch.Tensor, + B: torch.Tensor, + bias: torch.Tensor | None, + output_dtype: torch.dtype, + out: torch.Tensor | None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, bool]: + """Normalize the legacy 2D/3D caller contract to batch-first tensors.""" + if not isinstance(A, torch.Tensor) or not isinstance(B, torch.Tensor): + raise TypeError("gemm_a16w16_opus requires Tensor A and B") if A.dtype != torch.bfloat16 or B.dtype != torch.bfloat16: raise NotImplementedError( - f"gemm_a16w16_opus only supports bf16 A/B " - f"(got A.dtype={A.dtype}, B.dtype={B.dtype})." + "gemm_a16w16_opus only supports bf16 A/B " + f"(got A.dtype={A.dtype}, B.dtype={B.dtype})" ) - if dtype not in (torch.bfloat16, torch.float32): + if output_dtype not in (torch.bfloat16, torch.float32): raise NotImplementedError( - f"gemm_a16w16_opus only supports bf16/fp32 output dtype, got {dtype}" + "gemm_a16w16_opus only supports bf16/fp32 output dtype, " + f"got {output_dtype}" + ) + if A.device != B.device: + raise ValueError( + f"gemm_a16w16_opus requires A/B on one device; got {A.device}/{B.device}" ) - # Resolve A first so we know `batch`. - if A.dim() == 2: - M, K = A.shape + is_gemm = A.dim() == 2 + if is_gemm: + M, K = map(int, A.shape) batch = 1 - XQ = A.unsqueeze(0) - reshape_out_to_2d = True elif A.dim() == 3: - batch, M, K = A.shape - XQ = A - reshape_out_to_2d = False + batch, M, K = map(int, A.shape) else: raise ValueError(f"A must be 2D or 3D, got shape {tuple(A.shape)}") - # B accepted shapes: - # * [N, K] - allowed only when batch == 1 - # * [batch, N, K] real-strided - allowed for any batch - # - # The opus a16w16-family launchers hardcode `kargs.stride_b_batch = N * K` - # (csrc/opus_gemm/gen_instances.py around lines 531/634/735/865) and the - # device kernel computes `ptr_b + batch_id * stride_b_batch` directly, - # ignoring the tensor's reported stride. A `B.unsqueeze(0).expand(batch, - # -1, -1)` view has batch_stride == 0, so the kernel reads garbage past - # B's real allocation -- this manifests as NaN, large numerical errors, - # or HIP "Memory access fault by GPU node-1" depending on what the - # caching allocator parked next to B. Reject the broken case at the - # Python boundary rather than letting it through. if B.dim() == 2: - N, K_b = B.shape - if K_b != K: - raise ValueError(f"K dimension mismatch: A has K={K}, B has K={K_b}") - if batch > 1: + N, K_b = map(int, B.shape) + if batch != 1: raise NotImplementedError( - f"gemm_a16w16_opus: B must be 3D [batch, N, K] when A is " - f"batched (got A.shape={tuple(A.shape)}, " - f"B.shape={tuple(B.shape)}). The opus a16w16 launchers " - f"assume stride_b_batch == N*K (see " - f"csrc/opus_gemm/gen_instances.py), which is incompatible " - f"with the batch_stride=0 view a B.unsqueeze(0)." - f"expand(batch, -1, -1) would produce. Two valid fixes:\n" - f" 1. Broadcast explicitly: B = B.expand({batch}, -1, " - f"-1).contiguous()\n" - f" 2. Pass a real 3D weight: B with shape ({batch}, N, K)" + "gemm_a16w16_opus requires a real 3D [batch,N,K] weight " + f"when A is batched; got A.shape={tuple(A.shape)}, " + f"B.shape={tuple(B.shape)}" ) - WQ = B.unsqueeze(0) # batch == 1 here; kernel never reads stride_b_batch. elif B.dim() == 3: - b_b, N, K_b = B.shape - if K_b != K: - raise ValueError(f"K dimension mismatch: A has K={K}, B has K={K_b}") + b_b, N, K_b = map(int, B.shape) if b_b != batch: - raise ValueError( - f"B batch mismatch: A has batch={batch}, B has batch={b_b}" - ) - # Reject expand-style broadcast views (batch_stride=0) up front. Any - # other layout (contiguous, transposed N/K, etc.) is still rejected - # below by the elements-per-row check; the launcher requires - # B[b].stride(0) == N*K and B[b].stride(1) == K. - bs0, bs1, bs2 = B.stride() - if bs0 != N * K or bs1 != K or bs2 != 1: - raise NotImplementedError( - f"gemm_a16w16_opus: B must be a contiguous 3D tensor with " - f"strides (N*K, K, 1) (got B.shape={tuple(B.shape)}, " - f"B.stride()={tuple(B.stride())}). The opus launchers " - f"hardcode stride_b_batch == N*K and stride_b == K; any " - f"non-standard layout (broadcast view, transpose, slice) " - f"will produce wrong results or a memory access fault. " - f"Materialize via B = B.contiguous() first." - ) - WQ = B + raise ValueError(f"B batch mismatch: expected {batch}, got {b_b}") else: raise ValueError( - f"B must be 2D [N, K] or 3D [batch, N, K] (got shape {tuple(B.shape)})" + f"B must be 2D [N,K] or 3D [batch,N,K], got shape {tuple(B.shape)}" + ) + if K_b != K: + raise ValueError(f"K dimension mismatch: A has K={K}, B has K={K_b}") + + Y = out + if Y is None: + Y = torch.empty((batch, M, N), dtype=output_dtype, device=A.device) + elif not isinstance(Y, torch.Tensor): + raise TypeError(f"gemm_a16w16_opus out must be a Tensor, got {type(Y)!r}") + elif Y.device != A.device or Y.dtype != output_dtype: + raise ValueError( + "gemm_a16w16_opus out must match A.device and dtype; " + f"got {Y.device}/{Y.dtype}, expected {A.device}/{output_dtype}" ) - if out is not None: - Y = out - else: - Y = torch.empty(batch, M, N, dtype=dtype, device=A.device) - - # Bias validation. Bias may be fp32 OR match the output dtype: the gfx1250 - # splitk main kernel always writes an fp32 workspace and the reduce kernel - # folds bias in fp32 before the final cast to Y, so an fp32 bias is exact - # and free regardless of Y dtype (the common accuracy-friendly case for a - # bf16 output). Bias is per-output-feature [N] (F.linear convention): - # * [N] -> stride_bias_batch = 0 (broadcast across batch) - # * [batch, N] -> stride_bias_batch = N - # Matches the C++-side gfx1250 bias validation in gen_instances_gfx1250.py. + XQ = A.unsqueeze(0) if is_gemm else A + WQ = B.unsqueeze(0) if B.dim() == 2 else B + _check_a16w16_launch_layout(XQ, WQ, Y) if bias is not None: - if bias.dtype not in (dtype, torch.float32): + if not isinstance(bias, torch.Tensor): + raise TypeError("gemm_a16w16_opus bias must be a Tensor") + if bias.device != A.device or bias.dtype not in (output_dtype, torch.float32): raise ValueError( - f"gemm_a16w16_opus: bias dtype must be fp32 or match output " - f"dtype (got bias.dtype={bias.dtype}, dtype={dtype})" + "gemm_a16w16_opus bias must be on A.device and use fp32 or " + f"the output dtype; got {bias.device}/{bias.dtype}" ) if not bias.is_contiguous(): + raise ValueError("gemm_a16w16_opus bias must be contiguous") + if tuple(bias.shape) not in ((N,), (batch, N)): raise ValueError( - f"gemm_a16w16_opus: bias must be contiguous (got " - f"bias.stride()={tuple(bias.stride())})" - ) - if bias.dim() == 1: - if bias.shape[0] != N: - raise ValueError( - f"gemm_a16w16_opus: 1D bias length must equal N (got " - f"bias.shape={tuple(bias.shape)}, N={N})" - ) - elif bias.dim() == 2: - if tuple(bias.shape) != (batch, N): - raise ValueError( - f"gemm_a16w16_opus: 2D bias must be [batch, N] (got " - f"bias.shape={tuple(bias.shape)}, batch={batch}, N={N})" - ) - else: - raise ValueError( - f"gemm_a16w16_opus: bias must be 1D [N] or 2D [batch, N] " - f"(got bias.shape={tuple(bias.shape)})" + f"gemm_a16w16_opus bias must have shape [{N}] or [{batch},{N}], " + f"got shape {tuple(bias.shape)}" ) - return XQ, WQ, Y, M, N, K, batch, reshape_out_to_2d + return XQ, WQ, Y, is_gemm -def _finalize_output(Y: Tensor, reshape_out_to_2d: bool) -> Tensor: - return Y.squeeze(0) if reshape_out_to_2d else Y +@lru_cache(maxsize=256) +def _warn_invalid_a16w16_tuned_row(message: str) -> None: + logger.warning(message) def gemm_a16w16_opus( - A: Tensor, - B: Tensor, - bias: Tensor | None = None, + A: torch.Tensor, + B: torch.Tensor, + bias: torch.Tensor | None = None, dtype: torch.dtype = torch.bfloat16, *, kernelId: int | None = None, splitK: int | None = None, - out: Tensor | None = None, -) -> Tensor: - """Shape-driven opus a16w16 GEMM. - - Parameters - ---------- - A : [M, K] or [batch, M, K], bf16 - B : bf16 weight, plain layout (not pre-shuffled). Two accepted shapes: - * [N, K] -- requires batch == 1 (i.e. A is 2D, or A is - 3D with leading dim 1). - * [batch, N, K] -- contiguous strides (N*K, K, 1) only. - Broadcast views (e.g. ``B.unsqueeze(0). - expand(batch, -1, -1)``) are rejected - because the opus launcher assumes - ``stride_b_batch == N*K``; pass - ``.contiguous()`` if you need to broadcast - a single-batch weight across A. - bias : optional per-output-feature bias (F.linear convention), dtype - must equal `dtype` (match_d_out). Accepted shapes: - * [N] -- broadcast across batch. - * [batch, N] -- per-batch bias vector. - bias is consumed by the a16w16 split-barrier (kid 4..9) and the - a16w16_flatmm_splitk (kid 200..299) families. CSV-miss requests - with bias fall back to the C++ heuristic dispatcher (which only - returns bias-aware kids), so any (M, N, K) is supported even - without a tuned bias-aware winner -- accuracy is preserved at - whatever the heuristic kid achieves; performance may not be - optimal until the shape is re-tuned with `--bias`. - dtype : output dtype, bf16 or fp32 (any kernel family supports either) - kernelId : optional explicit override. When given, bypass CSV / C++ - dispatch and launch this specific tuned instance via - opus_gemm_a16w16_tune. - splitK : optional literal KBatch; only honored when kernelId is set. - out : optional preallocated [batch, M, N] output; reused instead of - allocating a fresh tensor. - - Returns - ------- - Tensor with shape [M, N] when A was 2D, [batch, M, N] when A was 3D. - """ - XQ, WQ, Y, M, N, K, _batch, reshape_out_to_2d = _validate_and_reshape( - A, B, bias, dtype, out - ) + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Shape-driven A16W16 compatibility API over the exact-kid launchers. - # 1) Explicit-kid override path. The C++ dispatcher gates non-bias-aware - # kids when bias is present, so we just forward. - if kernelId is not None: - opus_gemm_a16w16_tune(XQ, WQ, Y, bias, int(kernelId), int(splitK or 0)) - return _finalize_output(Y, reshape_out_to_2d) - - # 2) Default path: opus-private tuned CSV lookup. lookup_tuned() keys - # on bias=True/False as part of its 9-column tuple, so bias=True - # only matches rows that were tuned with the bias path. CSV miss on - # bias=True falls through to the explicit error below; we never - # silently route bias to the no-bias fallback. - cfg = _opus_common.lookup_tuned( - M=M, - N=N, - K=K, - bias=(bias is not None), - dtype=A.dtype, - outdtype=dtype, - scaleAB=False, - bpreshuffle=False, + An explicit ``kernelId`` wins over tuned lookup. Explicit and tuned ids + still pass through legacy requested-to-actual compatibility resolution; + a missing or invalid tuned row uses the migrated architecture heuristic. + """ + XQ, WQ, Y, is_gemm = _prepare_shape_driven_a16w16(A, B, bias, dtype, out) + from .policy import ( + lookup_a16w16_opus_config, + resolve_a16w16_heuristic_candidate, + resolve_a16w16_tuned_candidate, ) - if cfg is not None: - kid = cfg["solidx"] - # Both bf16 and fp32 Y are now valid for splitk kids (the reduce - # kernel handles the cast / passthrough), so no Y.dtype gating is - # needed here -- always honor the tuned winner. - opus_gemm_a16w16_tune(XQ, WQ, Y, bias, kid, int(cfg["splitK"])) - return _finalize_output(Y, reshape_out_to_2d) - - # 3) CSV miss: fall through to the C++ heuristic dispatcher via - # opus_gemm. Bias is forwarded through; the C++ entry skips its - # bias-agnostic lookup map when bias is present and routes - # directly to the heuristic (which only ever returns bias-aware - # split-barrier / splitk kids). - # - # (Note: this used to call `_opus_common.maybe_log_untuned_shape` - # to autolog the missed shape to a private CSV for offline tuning. - # The autolog feature has been removed -- collect untuned shapes - # via gradlib's standard --input_file flow instead.) - _opus_gemm_bf16_dispatch(XQ, WQ, Y, None, None, None, bias) - return _finalize_output(Y, reshape_out_to_2d) - - -# Per-stream splitk workspace init. Call once inside `with torch.cuda.stream(s):` -# (eagerly, before HIP graph capture) to register a workspace handle for that -# stream. Needed under vLLM/sglang-style TBO where two CPU threads drive two -# streams concurrently -- each captured graph must bake in its own buffer -# pointer; the prior thread_local cache would fail capture on the second -# stream. After init, run the largest expected gemm eagerly on the same -# stream to grow the buffer, then capture. -@compile_ops("module_deepgemm_opus", fc_name="opus_gemm_workspace_init", develop=True) -def opus_gemm_workspace_init() -> None: ... - - -# Free the per-stream splitk workspace registered by opus_gemm_workspace_init -# (and grown by the splitk launchers). Call inside `with torch.cuda.stream(s):` -# in eager mode (not during HIP graph capture) to reclaim the GPU buffer + -# handles for that stream; no-op if the stream was never registered. Use this -# for explicit teardown of streams the framework will not reuse. -@compile_ops( - "module_deepgemm_opus", fc_name="opus_gemm_workspace_release", develop=True -) -def opus_gemm_workspace_release() -> None: ... + arch, cu_num = _device_arch_and_cu(A.device) + batch, M, K = map(int, XQ.shape) + N = int(WQ.shape[1]) + lookup_args = { + "arch": arch, + "cu_num": cu_num, + "M": M, + "N": N, + "K": K, + "has_bias": bias is not None, + "input_dtype": A.dtype, + "output_dtype": Y.dtype, + } + if kernelId is None: + config = lookup_a16w16_opus_config(**lookup_args) + plan = None + if config is not None: + plan = resolve_a16w16_tuned_candidate( + batch=batch, + requested_kid=config.get("solidx"), + requested_split_k=config.get("splitK"), + **lookup_args, + ) + if plan is None: + _warn_invalid_a16w16_tuned_row( + "Ignoring invalid OPUS A16W16 tuned row for " + f"gfx={arch}, cu_num={cu_num}, " + f"shape=({batch},{M},{N},{K}), " + f"kid={config.get('solidx')!r}, " + f"splitK={config.get('splitK')!r}; " + "using OPUS heuristic fallback" + ) + if plan is None: + plan = resolve_a16w16_heuristic_candidate(batch=batch, **lookup_args) + if plan is None: + raise RuntimeError( + "gemm_a16w16_opus found no valid OPUS kernel for " + f"arch={arch}, shape=({batch},{M},{N},{K})" + ) + kid, split_k = plan.resolved_kid, 0 + else: + kid = plan.resolved_kid + split_k = int(config["splitK"]) + else: + kid = int(kernelId) + split_k = int(splitK or 0) + plan = resolve_a16w16_tuned_candidate( + batch=batch, + requested_kid=kid, + requested_split_k=split_k, + **lookup_args, + ) + if plan is not None: + kid = plan.resolved_kid + + if is_gemm: + return _launch_a16w16_gemm( + XQ.squeeze(0), + WQ.squeeze(0), + Y.squeeze(0), + kid=kid, + bias=bias, + split_k=split_k, + ) -# Free the splitk workspace for all registered streams and clear the registry. -# Eager mode only. Use for a full teardown before the framework reclaims its -# stream pool / at process shutdown. -@compile_ops( - "module_deepgemm_opus", fc_name="opus_gemm_workspace_release_all", develop=True -) -def opus_gemm_workspace_release_all() -> None: ... + return _launch_a16w16_bmm(XQ, WQ, Y, kid=kid, bias=bias, split_k=split_k) -__all__ = [ - "gemm_a16w16_opus", - "is_splitk_kid", - "opus_gemm_a16w16_tune", - "opus_gemm_workspace_init", - "opus_gemm_workspace_release", - "opus_gemm_workspace_release_all", -] +__all__ = ["gemm_a16w16_opus", "opus_gemm_a16w16_tune"] diff --git a/aiter/ops/opus/gemm_op_a8w8.py b/aiter/ops/opus/gemm_op_a8w8.py index e4cf523b3b..c500633e95 100644 --- a/aiter/ops/opus/gemm_op_a8w8.py +++ b/aiter/ops/opus/gemm_op_a8w8.py @@ -1,37 +1,115 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2025-2026, Advanced Micro Devices, Inc. All rights reserved. -"""Low-level Opus gfx942 A8W8 blockscale bpreshuffle entry points.""" +"""Private OPUS A8W8 exact-kid launch APIs.""" + +# Keep annotations eager in this binding module. ``torch_compile_guard`` uses +# the first parameter's concrete ``torch.Tensor`` identity to decide whether a +# custom op already has a Tensor dispatch key. Postponed string annotations +# make it add a dummy CUDA Tensor to every raw launch, which is observable host +# overhead on short BMM kernels. import torch from torch import Tensor from ...jit.core import compile_ops +from ._arch import _device_arch +from .launch_plan import ( + _A8W8_BLOCKSCALE_FAMILY, + _A8W8_BPRESHUFFLE_FAMILY, + _A8W8_FAMILY, + _A8W8_MXSCALE_BMM_FAMILY, + _FP8_DTYPES, + _get_cached_a8w8_mxscale_bmm_plan, + _require_registered_kid, +) + +_E8M0_DTYPES = frozenset( + dtype + for dtype in ( + torch.uint8, + getattr(torch, "float8_e8m0fnu", None), + ) + if dtype is not None +) + + +# ---- Low-level A8W8 backend ---------------------------------------------- + + +def _gen_opus_gemm_a8w8_launch_fake_tensors( + XQ: Tensor, + WQ: Tensor, + Y: Tensor, + kid: int, +) -> Tensor: + return Y + + +@compile_ops( + "module_deepgemm_opus", + fc_name="opus_gemm_a8w8_launch", + gen_fake=_gen_opus_gemm_a8w8_launch_fake_tensors, + develop=True, +) +def _opus_gemm_a8w8_launch_raw( + XQ: Tensor, + WQ: Tensor, + Y: Tensor, + kid: int, +) -> Tensor: ... + + +def _gen_opus_gemm_a8w8_blockscale_launch_fake_tensors( + XQ: Tensor, + WQ: Tensor, + Y: Tensor, + x_scale: Tensor, + w_scale: Tensor, + kid: int, +) -> Tensor: + return Y + + +@compile_ops( + "module_deepgemm_opus", + fc_name="opus_gemm_a8w8_blockscale_launch", + gen_fake=_gen_opus_gemm_a8w8_blockscale_launch_fake_tensors, + develop=True, +) +def _opus_gemm_a8w8_blockscale_launch_raw( + XQ: Tensor, + WQ: Tensor, + Y: Tensor, + x_scale: Tensor, + w_scale: Tensor, + kid: int, +) -> Tensor: ... -def _gen_opus_a8w8_blockscale_bpreshuffle_fake_tensors( +def _gen_opus_gemm_a8w8_blockscale_bpreshuffle_launch_fake_tensors( XQ: Tensor, WQ: Tensor, x_scale: Tensor, w_scale: Tensor, Y: Tensor, - kernelId: int, + kid: 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, + fc_name="opus_gemm_a8w8_blockscale_bpreshuffle_launch", + gen_fake=_gen_opus_gemm_a8w8_blockscale_bpreshuffle_launch_fake_tensors, develop=True, ) -def _opus_gemm_a8w8_blockscale_bpreshuffle_tune_raw( +def _opus_gemm_a8w8_blockscale_bpreshuffle_launch_raw( XQ: Tensor, WQ: Tensor, x_scale: Tensor, w_scale: Tensor, Y: Tensor, - kernelId: int, + kid: int, ) -> Tensor: ... @@ -43,13 +121,426 @@ def opus_gemm_a8w8_blockscale_bpreshuffle_tune( Y: Tensor | None = None, kernelId: int = 11000, ) -> Tensor: - """Run one gfx942 Opus A8W8 blockscale bpreshuffle kernel by explicit id.""" + """Compatibility entry for the existing A8W8 blockscale tuner.""" if Y is None: Y = torch.empty( - (XQ.shape[-2], WQ.shape[-2]), device=XQ.device, dtype=torch.bfloat16 + (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 _launch_a8w8_blockscale_bpreshuffle_gemm( + XQ, + WQ, + x_scale, + w_scale, + Y, + kid=int(kernelId), + ) + + +def _gen_opus_gemm_a8w8_mxscale_bmm_launch_fake_tensors( + XQ: Tensor, + WQ: Tensor, + Y: Tensor, + x_scale: Tensor, + w_scale: Tensor, + workspace: Tensor | None, + kid: int, + split_k: int, +) -> Tensor: + return Y + + +@compile_ops( + "module_deepgemm_opus", + fc_name="opus_gemm_a8w8_mxscale_bmm_launch", + gen_fake=_gen_opus_gemm_a8w8_mxscale_bmm_launch_fake_tensors, + develop=True, +) +def _opus_gemm_a8w8_mxscale_bmm_launch_raw( + XQ: Tensor, + WQ: Tensor, + Y: Tensor, + x_scale: Tensor, + w_scale: Tensor, + workspace: Tensor | None, + kid: int, + split_k: int, +) -> Tensor: ... + + +def _launch_a8w8_backend( + XQ: Tensor, + WQ: Tensor, + Y: Tensor, + x_scale: Tensor | None, + w_scale: Tensor | None, + workspace: Tensor | None, + family: str, + kid: int, + split_k: int, +) -> None: + if family == _A8W8_FAMILY: + if x_scale is not None or w_scale is not None: + raise RuntimeError("A8W8 no-scale backend received scale tensors") + if workspace is not None or split_k != 0: + raise RuntimeError("A8W8 no-scale backend received split-K state") + _opus_gemm_a8w8_launch_raw(XQ, WQ, Y, kid) + return + + if x_scale is None or w_scale is None: + raise RuntimeError(f"A8W8 backend family {family!r} requires both scales") + + if family == _A8W8_BLOCKSCALE_FAMILY: + if workspace is not None or split_k != 0: + raise RuntimeError("A8W8 blockscale backend received split-K state") + _opus_gemm_a8w8_blockscale_launch_raw( + XQ, + WQ, + Y, + x_scale, + w_scale, + kid, + ) + return + + if family == _A8W8_BPRESHUFFLE_FAMILY: + if workspace is not None or split_k != 0: + raise RuntimeError( + "A8W8 blockscale-bpreshuffle backend received split-K state" + ) + _opus_gemm_a8w8_blockscale_bpreshuffle_launch_raw( + XQ, + WQ, + x_scale, + w_scale, + Y, + kid, + ) + return + + if family == _A8W8_MXSCALE_BMM_FAMILY: + _opus_gemm_a8w8_mxscale_bmm_launch_raw( + XQ, + WQ, + Y, + x_scale, + w_scale, + workspace, + kid, + split_k, + ) + return + + raise RuntimeError(f"unsupported A8W8 backend family {family!r}") + + +# ---- A8W8 execution and logical adapters --------------------------------- + + +def _launch_a8w8_gemm( + XQ: Tensor, + WQ: Tensor, + Y: Tensor, + *, + kid: int, + route_arch: str | None = None, + instance: object | None = None, +) -> Tensor: + """Launch logical 2D no-scale FP8 ``[M,K] x [N,K] -> FP32 [M,N]``.""" + resolved_kid = kid + if instance is None: + if XQ.dim() != 2 or WQ.dim() != 2 or Y.dim() != 2: + raise ValueError( + "opus_gemm A8W8 expects logical 2D XQ/WQ/Y; " "this family is GEMM-only" + ) + arch = route_arch or _device_arch(XQ.device) + resolved_kid = _require_registered_kid( + arch=arch, + family=_A8W8_FAMILY, + kid=kid, + output_dtype=Y.dtype, + ) + _launch_a8w8_backend( + XQ.unsqueeze(0), + WQ.unsqueeze(0), + Y.unsqueeze(0), + None, + None, + None, + _A8W8_FAMILY, + resolved_kid, + 0, + ) + return Y + + +def _launch_a8w8_blockscale_gemm( + XQ: Tensor, + WQ: Tensor, + Y: Tensor, + x_scale: Tensor, + w_scale: Tensor, + *, + kid: int, + route_arch: str | None = None, + instance: object | None = None, +) -> Tensor: + """Launch logical 2D blockscale A8W8 GEMM with 2D scales.""" + resolved_kid = kid + if instance is None: + if any(tensor.dim() != 2 for tensor in (XQ, WQ, Y, x_scale, w_scale)): + raise ValueError( + "opus_gemm A8W8 blockscale expects logical 2D " + "XQ/WQ/Y/x_scale/w_scale; this family is GEMM-only" + ) + arch = route_arch or _device_arch(XQ.device) + resolved_kid = _require_registered_kid( + arch=arch, + family=_A8W8_BLOCKSCALE_FAMILY, + kid=kid, + output_dtype=Y.dtype, + ) + _launch_a8w8_backend( + XQ.unsqueeze(0), + WQ.unsqueeze(0), + Y.unsqueeze(0), + x_scale, + w_scale, + None, + _A8W8_BLOCKSCALE_FAMILY, + resolved_kid, + 0, + ) + return Y + + +def _launch_a8w8_blockscale_bpreshuffle_gemm( + XQ: Tensor, + WQ: Tensor, + x_scale: Tensor, + w_scale: Tensor, + Y: Tensor, + *, + kid: int, + route_arch: str | None = None, + instance: object | None = None, +) -> Tensor: + """Launch logical 2D bpreshuffled blockscale A8W8 GEMM. + + ``WQ`` pre-shuffle is a content/layout semantic. It + cannot be proven from Tensor shape or strides. Build it with + ``shuffle_weight(WQ, layout=(16, 16))``. The generated launcher checks + output dtype, scale layout, batch and tile alignment. + """ + resolved_kid = kid + if instance is None: + if any(tensor.dim() != 2 for tensor in (XQ, WQ, Y, x_scale, w_scale)): + raise ValueError( + "opus_gemm A8W8 blockscale bpreshuffle expects logical 2D " + "XQ/WQ/Y/x_scale/w_scale; this family is GEMM-only" + ) + arch = route_arch or _device_arch(XQ.device) + resolved_kid = _require_registered_kid( + arch=arch, + family=_A8W8_BPRESHUFFLE_FAMILY, + kid=kid, + output_dtype=Y.dtype, + ) + _launch_a8w8_backend( + XQ.unsqueeze(0), + WQ.unsqueeze(0), + Y.unsqueeze(0), + x_scale, + w_scale, + None, + _A8W8_BPRESHUFFLE_FAMILY, + resolved_kid, + 0, + ) + return Y + + +def _validate_a8w8_mxscale_bmm_tensors( + XQ: Tensor, + WQ: Tensor, + Y: Tensor, + x_scale: Tensor, + w_scale: Tensor, +) -> None: + entry = "opus_gemm_a8w8_mxscale_bmm_launch" + tensors = (XQ, WQ, Y, x_scale, w_scale) + if any(tensor.dim() != 3 for tensor in tensors): + raise ValueError(f"{entry}: all inputs and Y must be 3D") + # This path validates inputs before Python allocates split-K workspace; + # ordinary launches leave the same-device contract to the checked C++ ABI. + device = XQ.device + if any(tensor.device != device for tensor in tensors[1:]): + devices = {tensor.device for tensor in tensors} + raise ValueError( + f"{entry}: all tensors must be on one device; got " + f"{sorted(map(str, devices))}" + ) + if XQ.dtype not in _FP8_DTYPES or WQ.dtype != XQ.dtype: + raise ValueError(f"{entry}: XQ and WQ must have the same FP8 dtype") + if Y.dtype not in (torch.bfloat16, torch.float32): + raise ValueError(f"{entry}: Y must be BF16 or FP32") + if x_scale.dtype not in _E8M0_DTYPES or w_scale.dtype not in _E8M0_DTYPES: + raise ValueError( + f"{entry}: x_scale and w_scale must contain one-byte E8M0 values" + ) + + M, batch, K = map(int, XQ.shape) + w_batch, N, w_K = map(int, WQ.shape) + if min(M, batch, N, K) <= 0: + raise ValueError(f"{entry}: M, batch, N and K must be positive") + if N % 128 or K % 128: + raise ValueError(f"{entry}: N and K must be multiples of 128; got N={N}, K={K}") + if (w_batch, w_K) != (batch, K): + raise ValueError( + f"{entry}: WQ must have shape [{batch},N,{K}], got {tuple(WQ.shape)}" + ) + if tuple(Y.shape) != (M, batch, N): + raise ValueError( + f"{entry}: Y must have shape {(M, batch, N)}, got {tuple(Y.shape)}" + ) + expected_x_scale = (M, batch, K // 128) + expected_w_scale = (batch, N // 128, K // 128) + if tuple(x_scale.shape) != expected_x_scale: + raise ValueError( + f"{entry}: x_scale must have shape {expected_x_scale}, " + f"got {tuple(x_scale.shape)}" + ) + if tuple(w_scale.shape) != expected_w_scale: + raise ValueError( + f"{entry}: w_scale must have shape {expected_w_scale}, " + f"got {tuple(w_scale.shape)}" + ) + if any(tensor.stride(-1) != 1 for tensor in (XQ, WQ, Y, x_scale, w_scale)): + raise ValueError(f"{entry}: every tensor must be contiguous in its last axis") + + +def _launch_a8w8_mxscale_bmm( + XQ: Tensor, + WQ: Tensor, + Y: Tensor, + x_scale: Tensor, + w_scale: Tensor, + *, + kid: int, + split_k: int, + workspace: Tensor | None, + route_arch: str | None = None, + instance: object | None = None, +) -> Tensor: + """Launch batch-first MXFP8 BMM through the shared physical launcher. + + Public tensors use ``[B,M,K]``, ``[B,N,K]`` and ``[B,M,N]``. The raw + kernels retain their established M-major ``[M,B,*]`` activation/output + ABI; transpose views bridge the two contracts without copying storage. + """ + if instance is None and any( + tensor.dim() != 3 for tensor in (XQ, WQ, Y, x_scale, w_scale) + ): + raise ValueError( + "opus_bmm A8W8 mxscale expects batch-first 3D " "XQ/WQ/Y/x_scale/w_scale" + ) + launch_x = XQ.transpose(0, 1) + launch_y = Y.transpose(0, 1) + launch_x_scale = x_scale.transpose(0, 1) + + if instance is not None and workspace is None and split_k <= 1: + # The checked C++ entry owns the dynamic tensor contract. Public + # routing already validated the immutable family/kid contract, so the + # common split-one path need not repeat the Python registry/planner. + _launch_a8w8_backend( + launch_x, + WQ, + launch_y, + launch_x_scale, + w_scale, + None, + _A8W8_MXSCALE_BMM_FAMILY, + kid, + max(1, split_k), + ) + return Y + + M, batch, K = map(int, launch_x.shape) + N = int(WQ.shape[1]) + arch = route_arch or _device_arch(XQ.device) + plan = _get_cached_a8w8_mxscale_bmm_plan( + arch, + int(kid), + Y.dtype, + M, + batch, + N, + K, + int(split_k), + ) + + launch_workspace = workspace + workspace_spec = plan.workspace_spec + if workspace_spec is not None: + # Workspace sizing reads Tensor dimensions before entering C++. Keep + # the full Python contract here so malformed metadata cannot drive an + # allocation. The split-one hot path above remains C++-checked only. + _validate_a8w8_mxscale_bmm_tensors( + launch_x, + WQ, + launch_y, + launch_x_scale, + w_scale, + ) + required_numel = workspace_spec.shape[0] + if launch_workspace is None: + launch_workspace = torch.empty( + workspace_spec.shape, + dtype=workspace_spec.dtype, + device=XQ.device, + ) + else: + if not isinstance(launch_workspace, Tensor): + raise TypeError( + "opus_gemm_a8w8_mxscale_bmm_launch: workspace must be a Tensor" + ) + if launch_workspace.device != XQ.device: + raise ValueError( + "opus_gemm_a8w8_mxscale_bmm_launch: workspace must be on " + f"{XQ.device}, got {launch_workspace.device}" + ) + if launch_workspace.dtype != workspace_spec.dtype: + raise ValueError( + "opus_gemm_a8w8_mxscale_bmm_launch: workspace must be FP32" + ) + if not launch_workspace.is_contiguous(): + raise ValueError( + "opus_gemm_a8w8_mxscale_bmm_launch: workspace must be contiguous" + ) + if launch_workspace.numel() < required_numel: + raise ValueError( + "opus_gemm_a8w8_mxscale_bmm_launch: workspace capacity is " + f"{launch_workspace.numel()}, but {required_numel} elements " + "are required" + ) + elif launch_workspace is not None: + raise ValueError( + f"OPUS BMM kid {plan.resolved_kid} with split_k={split_k} " + "does not use workspace" + ) + + _launch_a8w8_backend( + launch_x, + WQ, + launch_y, + launch_x_scale, + w_scale, + launch_workspace, + _A8W8_MXSCALE_BMM_FAMILY, + plan.resolved_kid, + plan.abi_split_k, ) return Y diff --git a/aiter/ops/opus/launch_plan.py b/aiter/ops/opus/launch_plan.py new file mode 100644 index 0000000000..61801eb2e1 --- /dev/null +++ b/aiter/ops/opus/launch_plan.py @@ -0,0 +1,767 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2025-2026, Advanced Micro Devices, Inc. All rights reserved. +"""Immutable exact-kid contracts and workspace planning for OPUS.""" + +from __future__ import annotations + +from dataclasses import dataclass +from functools import cache, lru_cache + +import torch + +from csrc.opus_gemm.opus_gemm_common import ( + BIAS_AWARE_KIDS, + GFX942_BF16WS_EXACT_N, + GFX942_EVEN_LOOP_SPLITK_TAGS, + GFX942_MAX_AUTO_SPLIT_K, + GFX942_MIN_ITERS_PER_SPLIT, + SPLITK_KIDS, + OpusGemmInstance, + a8w8_mxscale_flatmm_prefetch_k_iter, + a16w16_flatmm_prefetch_k_iter, + get_kernel_instance, +) + +from ._arch import GFX942, GFX950, GFX1250 + +_WORKSPACE_DTYPES = { + "bf16_t": torch.bfloat16, + "fp32_t": torch.float32, +} +_GFX1250_FUSED_SPLITK_TAG = "a16w16_clusterlaunch_tdm_splitk_fuse" +_GFX1250_CO_TAGS = frozenset({"a16w16_4wave_co", "a16w16_4wave_wl_co"}) + + +@dataclass(frozen=True) +class WorkspaceSpec: + """Immutable metadata required to materialize one launch workspace.""" + + shape: tuple[int, ...] + dtype: torch.dtype + + +@dataclass(frozen=True) +class A16W16LaunchPlan: + """One validated exact A16W16 launch with immutable workspace metadata.""" + + registry_arch: str + resolved_kid: int + workspace_capacity_split_k: int + abi_split_k: int + workspace_spec: WorkspaceSpec | None + + +# ---- A16W16 exact launch planning --------------------------------------- + + +def _supports_a16w16_shape( + instance: OpusGemmInstance, + *, + M: int, + N: int, + K: int, +) -> bool: + if instance.kernel_tag == _GFX1250_FUSED_SPLITK_TAG: + split_k = int(instance.fuse_split_k) + n_cluster = int(instance.fuse_m_cluster) + if K % 2 != 0 or N % instance.B_N != 0: + return False + if split_k < 2 or split_k * n_cluster > 16: + return False + num_tiles_n = N // instance.B_N + if num_tiles_n % n_cluster != 0: + return False + return split_k <= (K + instance.B_K - 1) // instance.B_K + + if instance.kernel_tag == "a16w16_mono_tile": + return N % instance.B_N == 0 and K % instance.B_K == 0 + + if not instance.has_oob: + return M % instance.B_M == 0 and N % instance.B_N == 0 + return True + + +def _plan_gfx942_split_k( + instance: OpusGemmInstance, + *, + M: int, + N: int, + K: int, + batch: int, + cu_num: int, + requested: int, +) -> int: + """Return the converged ABI split-K matching the gfx942 launcher.""" + if requested > 0: + abi_split_k = requested + else: + tiles_mn = ( + (M + instance.B_M - 1) + // instance.B_M + * ((N + instance.B_N - 1) // instance.B_N) + * batch + ) + tiles_mn = max(1, tiles_mn) + target_wg = (2 * cu_num) if instance.kernel_tag.endswith("_p1") else cu_num + abi_split_k = (target_wg + tiles_mn - 1) // tiles_mn + abi_split_k = min(GFX942_MAX_AUTO_SPLIT_K, max(1, abi_split_k)) + + total_iters = (K + instance.B_K - 1) // instance.B_K + if total_iters < GFX942_MIN_ITERS_PER_SPLIT: + raise ValueError( + f"K={K} is too small for gfx942 kid B_K={instance.B_K}; " + f"need at least {instance.B_K * GFX942_MIN_ITERS_PER_SPLIT}" + ) + + require_even = instance.kernel_tag in GFX942_EVEN_LOOP_SPLITK_TAGS + while abi_split_k > 1: + iters_full = (total_iters + abi_split_k - 1) // abi_split_k + last_loops = total_iters - (abi_split_k - 1) * iters_full + parity_ok = not require_even or (iters_full % 2 == 0 and last_loops % 2 == 0) + if ( + iters_full >= GFX942_MIN_ITERS_PER_SPLIT + and last_loops >= GFX942_MIN_ITERS_PER_SPLIT + and parity_ok + ): + break + abi_split_k -= 1 + + if require_even: + iters_full = (total_iters + abi_split_k - 1) // abi_split_k + last_loops = total_iters - (abi_split_k - 1) * iters_full + if iters_full % 2 != 0 or last_loops % 2 != 0: + raise ValueError( + f"gfx942 kid {instance.name} needs even loops per split; " + f"K={K}, split_k={abi_split_k}, " + f"loops=({iters_full},{last_loops})" + ) + return abi_split_k + + +def _plan_gfx950_split_k( + instance: OpusGemmInstance, + *, + K: int, + requested: int, +) -> int: + """Return the converged ABI split-K matching the gfx950 launcher.""" + total_iters = (K + instance.B_K - 1) // instance.B_K + prefetch_k_iter = a16w16_flatmm_prefetch_k_iter(instance) + if total_iters < prefetch_k_iter: + raise ValueError( + f"K={K} is too small for gfx950 kid B_K={instance.B_K}; " + f"need at least {instance.B_K * prefetch_k_iter}" + ) + abi_split_k = min(max(1, requested), total_iters // prefetch_k_iter) + while abi_split_k > 1: + iters_full = (total_iters + abi_split_k - 1) // abi_split_k + last_loops = total_iters - (abi_split_k - 1) * iters_full + if iters_full >= prefetch_k_iter and last_loops >= prefetch_k_iter: + break + abi_split_k -= 1 + return abi_split_k + + +def _plan_gfx1250_split_k( + instance: OpusGemmInstance, + *, + K: int, + requested: int, +) -> int: + """Return the converged ABI split-K matching the gfx1250 launcher.""" + total_iters = (K + instance.B_K - 1) // instance.B_K + abi_split_k = min(max(1, requested), total_iters) + while abi_split_k > 1: + iters_full = (total_iters + abi_split_k - 1) // abi_split_k + if (abi_split_k - 1) * iters_full < total_iters: + break + abi_split_k -= 1 + return abi_split_k + + +def _build_a16w16_workspace_spec( + instance: OpusGemmInstance, + *, + registry_arch: str, + resolved_kid: int, + workspace_capacity_split_k: int, + batch: int, + M: int, + N: int, +) -> WorkspaceSpec | None: + """Build workspace metadata from the already-resolved registry instance.""" + needs_workspace = resolved_kid in SPLITK_KIDS + declares_workspace = instance.splitk_workspace_dtype is not None + if needs_workspace != declares_workspace: + raise RuntimeError( + "inconsistent OPUS a16w16 workspace registry for " + f"{registry_arch} kid {resolved_kid}" + ) + if not needs_workspace: + return None + + is_fused = instance.kernel_tag == _GFX1250_FUSED_SPLITK_TAG + split_k = ( + int(instance.fuse_split_k) if is_fused else int(workspace_capacity_split_k) + ) + if split_k <= 0: + raise ValueError( + "opus_gemm_a16w16_launch: workspace capacity split_k must be " + f"positive, got {split_k}" + ) + + block_m = int(instance.B_M) + block_n = int(instance.B_N) + + if registry_arch == GFX1250: + if batch != 1: + raise ValueError( + "opus_gemm_a16w16_launch: gfx1250 workspace kids require " + f"batch=1; got batch={batch}" + ) + num_tiles_m = (M + block_m - 1) // block_m + num_tiles_n = (N + block_n - 1) // block_n + if is_fused: + if split_k < 2: + raise ValueError( + "opus_gemm_a16w16_launch: " + f"gfx1250 fused kid {resolved_kid} must declare " + f"compile-time SplitK >= 2, got {split_k}" + ) + shape = ( + num_tiles_m, + num_tiles_n, + split_k - 1, + block_m, + block_n, + ) + else: + padded_m = num_tiles_m * block_m + padded_n = num_tiles_n * block_n + shape = (split_k, padded_m, padded_n) + else: + padded_m = ((M + block_m - 1) // block_m) * block_m + padded_n = ((N + block_n - 1) // block_n) * block_n + shape = (split_k, batch, padded_m, padded_n) + + dtype_token = instance.splitk_workspace_dtype + try: + dtype = _WORKSPACE_DTYPES[dtype_token] + except KeyError as exc: + raise ValueError( + "opus_gemm_a16w16_launch: " + f"workspace kid {resolved_kid} must declare bf16_t or fp32_t " + f"storage, got {dtype_token!r}" + ) from exc + + required_numel = 1 + max_numel = (2**63 - 1) // int(dtype.itemsize) + for extent in shape: + if extent <= 0 or required_numel > max_numel // extent: + raise OverflowError( + "opus_gemm_a16w16_launch: " + f"workspace shape {shape} exceeds the supported tensor size " + f"for dtype {dtype}" + ) + required_numel *= extent + + return WorkspaceSpec(shape=shape, dtype=dtype) + + +def _build_a16w16_launch_plan( + *, + arch: str, + M: int, + N: int, + K: int, + batch: int, + cu_num: int, + has_bias: bool, + input_dtype: object, + output_dtype: object, + kid: int, + split_k: int, +) -> A16W16LaunchPlan: + """Validate one exact kid and build all immutable launch metadata.""" + registry_arch = str(arch).lower().split(":", 1)[0] + M, N, K, batch, cu_num = map(int, (M, N, K, batch, cu_num)) + if min(M, N, K, batch, cu_num) <= 0: + raise ValueError("M, N, K, batch, and cu_num must all be positive") + try: + resolved_kid = int(kid) + requested_split_k = int(split_k) + except (TypeError, ValueError) as exc: + raise ValueError( + f"OPUS a16w16 kid/split_k must be integers, got {kid!r}/{split_k!r}" + ) from exc + if requested_split_k < 0: + raise ValueError( + "OPUS a16w16 split_k must be non-negative, " f"got {requested_split_k}" + ) + if input_dtype != torch.bfloat16: + raise ValueError( + f"OPUS a16w16 requires bf16 XQ/WQ, got input dtype {input_dtype}" + ) + + instance = get_kernel_instance(registry_arch, "a16w16", resolved_kid, output_dtype) + if instance is None: + if get_kernel_instance(registry_arch, "a16w16", resolved_kid) is None: + raise ValueError( + f"OPUS kid {resolved_kid} is not an a16w16 kernel for " + f"runtime arch {registry_arch}" + ) + raise ValueError( + f"OPUS kid {resolved_kid} does not support output dtype " f"{output_dtype}" + ) + + if instance.max_m is not None and M > instance.max_m: + raise ValueError( + f"OPUS kid {resolved_kid} requires M <= {instance.max_m}; got M={M}" + ) + + needs_workspace = resolved_kid in SPLITK_KIDS + if instance.kernel_tag in _GFX1250_CO_TAGS and requested_split_k > 1: + raise ValueError( + f"gfx1250 CO kid {resolved_kid} does not support split-K; " + f"got split_k={requested_split_k}" + ) + if ( + registry_arch == GFX942 + and needs_workspace + and instance.splitk_workspace_dtype == "bf16_t" + and N not in GFX942_BF16WS_EXACT_N + ): + raise ValueError( + f"gfx942 exact kid {resolved_kid} requires N in " + f"{sorted(GFX942_BF16WS_EXACT_N)}; got N={N}" + ) + if registry_arch == GFX1250 and needs_workspace and batch != 1: + raise ValueError( + "opus_gemm_a16w16_launch: gfx1250 workspace kids require " + f"batch=1; got batch={batch}" + ) + if not _supports_a16w16_shape(instance, M=M, N=N, K=K): + raise ValueError( + f"OPUS kid {resolved_kid} is incompatible with " + f"shape (batch={batch}, M={M}, N={N}, K={K})" + ) + if has_bias and resolved_kid not in BIAS_AWARE_KIDS: + raise ValueError(f"OPUS kid {resolved_kid} does not support bias") + if has_bias and instance.kernel_tag == _GFX1250_FUSED_SPLITK_TAG: + raise ValueError( + "gfx1250 splitk_fuse has a narrower bf16 [N] bias contract than " + "the public OPUS interfaces can represent" + ) + if has_bias and registry_arch == GFX942 and needs_workspace: + raise ValueError( + "the current gfx942 a16w16 launch rejects bias on split-K kernels" + ) + + workspace_capacity_split_k = 1 + abi_split_k = requested_split_k + if needs_workspace: + if instance.kernel_tag == _GFX1250_FUSED_SPLITK_TAG: + abi_split_k = int(instance.fuse_split_k) + elif registry_arch == GFX942: + abi_split_k = _plan_gfx942_split_k( + instance, + M=M, + N=N, + K=K, + batch=batch, + cu_num=cu_num, + requested=requested_split_k, + ) + elif registry_arch == GFX950: + abi_split_k = _plan_gfx950_split_k( + instance, + K=K, + requested=requested_split_k, + ) + elif registry_arch == GFX1250: + abi_split_k = _plan_gfx1250_split_k( + instance, + K=K, + requested=requested_split_k, + ) + workspace_capacity_split_k = max(1, abi_split_k) + + # Validate the launch split-K independently of workspace sizing. + launch_split_k = max(1, abi_split_k) + block_k = int(instance.B_K) + max_useful_split_k = (K + block_k - 1) // block_k + if launch_split_k > max_useful_split_k: + raise ValueError( + "opus_gemm_a16w16_launch: " + f"launch split_k={launch_split_k} exceeds the per-kid " + f"K-tile limit {max_useful_split_k} for K={K}, B_K={block_k}" + ) + + workspace_spec = _build_a16w16_workspace_spec( + instance, + registry_arch=registry_arch, + resolved_kid=resolved_kid, + workspace_capacity_split_k=workspace_capacity_split_k, + batch=batch, + M=M, + N=N, + ) + return A16W16LaunchPlan( + registry_arch=registry_arch, + resolved_kid=resolved_kid, + workspace_capacity_split_k=workspace_capacity_split_k, + abi_split_k=abi_split_k, + workspace_spec=workspace_spec, + ) + + +@lru_cache(maxsize=256) +def _get_cached_a16w16_launch_plan( + arch: str, + M: int, + N: int, + K: int, + batch: int, + cu_num: int, + has_bias: bool, + input_dtype: object, + output_dtype: object, + kid: int, + split_k: int, +) -> A16W16LaunchPlan: + """Return a scalar-only cached exact-kid launch plan.""" + return _build_a16w16_launch_plan( + arch=arch, + M=M, + N=N, + K=K, + batch=batch, + cu_num=cu_num, + has_bias=has_bias, + input_dtype=input_dtype, + output_dtype=output_dtype, + kid=kid, + split_k=split_k, + ) + + +# ---- A8W8 contracts and MXFP8 BMM planning ------------------------------ + +_A8W8_FAMILY = "a8w8" +_A8W8_BLOCKSCALE_FAMILY = "a8w8_blockscale" +_A8W8_BPRESHUFFLE_FAMILY = "a8w8_blockscale_bpreshuffle" +_A8W8_MXSCALE_BMM_FAMILY = "a8w8_mxscale_bmm" + +_A8W8_MXSCALE_BMM_TAGS = frozenset( + { + "a8w8_mxscale_bmm_flatmm_splitk", + "a8w8_mxscale_bmm_fused", + "a8w8_mxscale_bmm_minterleave", + "a8w8_mxscale_bmm_mouter", + "a8w8_mxscale_bmm_mouter_tunable", + "a8w8_mxscale_bmm_pipeline", + "a8w8_mxscale_bmm_wave8n2", + "a8w8_mxscale_bmm_wave4m2_selfload", + } +) +_A8W8_MXSCALE_BMM_WORKSPACE_TAGS = frozenset( + { + "a8w8_mxscale_bmm_flatmm_splitk", + "a8w8_mxscale_bmm_fused", + } +) +_A8W8_MXSCALE_BMM_PREFETCH_TAGS = _A8W8_MXSCALE_BMM_WORKSPACE_TAGS | frozenset( + { + "a8w8_mxscale_bmm_minterleave", + "a8w8_mxscale_bmm_mouter", + "a8w8_mxscale_bmm_mouter_tunable", + } +) +_A8W8_FAMILY_BY_TAG = { + "a8w8": _A8W8_FAMILY, + "a8w8_scale": _A8W8_BLOCKSCALE_FAMILY, + "a8w8_blockscale_bpreshuffle_singlebuf": _A8W8_BPRESHUFFLE_FAMILY, + **{tag: _A8W8_MXSCALE_BMM_FAMILY for tag in _A8W8_MXSCALE_BMM_TAGS}, +} +_A8W8_FAMILY_LAYOUT = { + _A8W8_FAMILY: "plain", + _A8W8_BLOCKSCALE_FAMILY: "plain", + _A8W8_BPRESHUFFLE_FAMILY: "bpreshuffle", + _A8W8_MXSCALE_BMM_FAMILY: "mxscale_bmm", +} +_FP8_DTYPES = frozenset( + dtype + for dtype in ( + getattr(torch, "float8_e4m3fnuz", None), + getattr(torch, "float8_e4m3fn", None), + ) + if dtype is not None +) + + +@dataclass(frozen=True) +class A8W8MxscaleBMMPlan: + """One resolved MXFP8 BMM launch and its optional workspace.""" + + registry_arch: str + resolved_kid: int + abi_split_k: int + workspace_spec: WorkspaceSpec | None + + +def _validate_a8w8_public_contract( + *, + kernel_tag: str, + kid: int, + input_dtype: torch.dtype, + weight_dtype: torch.dtype, + output_dtype: torch.dtype, + layout: str, + has_x_scale: bool, + has_w_scale: bool, + has_bias: bool, + has_workspace: bool, + split_k: int, +) -> str: + """Validate immutable options for the public A8W8 operation routers.""" + try: + family = _A8W8_FAMILY_BY_TAG[kernel_tag] + except KeyError as exc: + raise ValueError( + f"OPUS kid {kid} has unsupported registry tag {kernel_tag!r}" + ) from exc + + if input_dtype != weight_dtype: + raise ValueError( + f"OPUS requires matching XQ/WQ dtypes; got " f"{input_dtype}/{weight_dtype}" + ) + if input_dtype not in _FP8_DTYPES: + raise ValueError(f"OPUS kid {kid} requires FP8 XQ/WQ; got {input_dtype}") + + if family == _A8W8_MXSCALE_BMM_FAMILY: + if output_dtype not in (torch.bfloat16, torch.float32): + raise ValueError( + f"OPUS kid {kid} requires BF16 or FP32 Y; got {output_dtype}" + ) + else: + expected_output_dtype = ( + torch.bfloat16 if family == _A8W8_BPRESHUFFLE_FAMILY else torch.float32 + ) + if output_dtype != expected_output_dtype: + raise ValueError( + f"OPUS kid {kid} does not support Y.dtype={output_dtype}; " + f"expected {expected_output_dtype}" + ) + + expected_layout = _A8W8_FAMILY_LAYOUT[family] + if layout != expected_layout: + raise ValueError( + f"OPUS kid {kid} belongs to family {family} and requires " + f"layout={expected_layout!r}; got {layout!r}" + ) + if has_x_scale != has_w_scale: + raise ValueError("OPUS requires x_scale and w_scale together") + if has_bias: + raise ValueError(f"OPUS family {family} does not support bias") + + if family == _A8W8_MXSCALE_BMM_FAMILY: + if not has_x_scale: + raise ValueError("OPUS a8w8_mxscale_bmm requires x_scale and w_scale") + if split_k == 0 and has_workspace: + raise ValueError("OPUS a8w8_mxscale_bmm split_k=0 does not use workspace") + return family + + if has_workspace: + raise ValueError(f"OPUS family {family} does not use workspace") + if split_k != 0: + raise ValueError(f"OPUS family {family} does not accept split_k") + if family == _A8W8_FAMILY: + if has_x_scale: + raise ValueError("OPUS a8w8 kid does not accept scales") + elif not has_x_scale: + raise ValueError(f"OPUS family {family} requires x_scale and w_scale") + return family + + +@cache +def _require_registered_kid_cached( + arch: str, + family: str, + resolved_kid: int, + output_dtype: torch.dtype, +) -> int: + """Validate one A8W8 registry entry and cache successful lookups.""" + if get_kernel_instance(arch, family, resolved_kid, output_dtype) is None: + raise ValueError( + "no registered OPUS kernel for " + f"(arch={arch!r}, family={family!r}, kid={resolved_kid}, " + f"Y.dtype={output_dtype})" + ) + return resolved_kid + + +def _require_registered_kid( + *, + arch: str, + family: str, + kid: object, + output_dtype: torch.dtype, +) -> int: + """Normalize a kid and require the matching A8W8 registry entry.""" + try: + resolved_kid = int(kid) + except (TypeError, ValueError) as exc: + raise ValueError(f"OPUS {family} kid must be an integer, got {kid!r}") from exc + return _require_registered_kid_cached(arch, family, resolved_kid, output_dtype) + + +def _build_a8w8_mxscale_bmm_plan( + *, + arch: str, + kid: int, + output_dtype: torch.dtype, + M: int, + batch: int, + N: int, + K: int, + split_k: int, +) -> A8W8MxscaleBMMPlan: + """Resolve one MXFP8 BMM kid, ABI split-K and FP32 workspace.""" + registry_arch = str(arch).lower().split(":", 1)[0] + resolved_kid = int(kid) + M, batch, N, K = map(int, (M, batch, N, K)) + requested_split_k = int(split_k) + + instance = get_kernel_instance( + registry_arch, + _A8W8_MXSCALE_BMM_FAMILY, + resolved_kid, + output_dtype, + ) + if instance is None: + raise ValueError( + "no registered OPUS kernel for " + f"(arch={registry_arch!r}, " + f"family={_A8W8_MXSCALE_BMM_FAMILY!r}, " + f"kid={resolved_kid}, Y.dtype={output_dtype})" + ) + + tag = instance.kernel_tag + if tag not in _A8W8_MXSCALE_BMM_TAGS: + raise ValueError(f"OPUS kid {resolved_kid} is not an MXFP8 BMM kernel") + + if requested_split_k < 0 or requested_split_k > (1 << 31) - 1: + raise ValueError( + f"OPUS BMM kid {resolved_kid} requires 0 <= split_k <= 2147483647; " + f"got {requested_split_k}" + ) + abi_split_k = max(1, requested_split_k) + if min(M, batch, N, K) <= 0: + raise ValueError( + "OPUS BMM requires positive M, batch, N and K; " + f"got M={M}, batch={batch}, N={N}, K={K}" + ) + m_align = max(1, int(instance.m_align)) + if M % m_align: + raise ValueError( + f"OPUS BMM kid {resolved_kid} requires M % {m_align} == 0; got M={M}" + ) + n_align = int(instance.B_N) * (2 if tag == "a8w8_mxscale_bmm_wave8n2" else 1) + if N % n_align: + raise ValueError( + f"OPUS BMM kid {resolved_kid} requires N % {n_align} == 0; got N={N}" + ) + k_align = int(instance.B_K) + if K % k_align: + raise ValueError( + f"OPUS BMM kid {resolved_kid} requires K % {k_align} == 0; got K={K}" + ) + if (instance.k1024_only or instance.k1024_lb1) and K != 1024: + raise ValueError(f"OPUS BMM kid {resolved_kid} requires K == 1024; got K={K}") + if instance.direct_only and abi_split_k != 1: + raise ValueError(f"OPUS BMM kid {resolved_kid} requires split_k <= 1") + + if tag in _A8W8_MXSCALE_BMM_PREFETCH_TAGS: + total_iters = K // k_align + prefetch_k_iter = a8w8_mxscale_flatmm_prefetch_k_iter(instance) + if tag in _A8W8_MXSCALE_BMM_WORKSPACE_TAGS: + if abi_split_k > total_iters: + raise ValueError( + f"OPUS BMM kid {resolved_kid} split_k={abi_split_k} exceeds " + f"the K-tile count {total_iters} for K={K}" + ) + iters_full = (total_iters + abi_split_k - 1) // abi_split_k + last_loops = total_iters - (abi_split_k - 1) * iters_full + if last_loops < prefetch_k_iter: + raise ValueError( + f"OPUS BMM kid {resolved_kid} requires every split to have " + f"at least {prefetch_k_iter} K-tiles; K={K}, " + f"split_k={abi_split_k}, last split has {last_loops}" + ) + elif total_iters < prefetch_k_iter: + raise ValueError( + f"OPUS BMM kid {resolved_kid} requires at least " + f"{prefetch_k_iter} K-tiles; K={K} gives {total_iters}" + ) + + workspace_numel = 0 + if tag in _A8W8_MXSCALE_BMM_WORKSPACE_TAGS: + if abi_split_k > 1: + tiles_m = (M + instance.B_M - 1) // instance.B_M + tiles_n = (N + instance.B_N - 1) // instance.B_N + padded_m = tiles_m * instance.B_M + padded_n = tiles_n * instance.B_N + partial_numel = abi_split_k * batch * padded_m * padded_n + if tag == "a8w8_mxscale_bmm_fused": + counter_offset = (partial_numel * 4 + 255) & ~255 + counter_bytes = batch * tiles_m * tiles_n * 4 + workspace_numel = (counter_offset + counter_bytes + 3) // 4 + else: + workspace_numel = partial_numel + elif tag == "a8w8_mxscale_bmm_mouter_tunable": + pass + elif abi_split_k != 1: + raise ValueError(f"OPUS BMM kid {resolved_kid} requires split_k <= 1") + + workspace_spec = ( + WorkspaceSpec(shape=(workspace_numel,), dtype=torch.float32) + if workspace_numel + else None + ) + return A8W8MxscaleBMMPlan( + registry_arch=registry_arch, + resolved_kid=resolved_kid, + abi_split_k=abi_split_k, + workspace_spec=workspace_spec, + ) + + +@lru_cache(maxsize=4096) +def _get_cached_a8w8_mxscale_bmm_plan( + arch: str, + kid: int, + output_dtype: torch.dtype, + M: int, + batch: int, + N: int, + K: int, + split_k: int, +) -> A8W8MxscaleBMMPlan: + """Return a scalar-only cached MXFP8 BMM launch plan.""" + return _build_a8w8_mxscale_bmm_plan( + arch=arch, + kid=kid, + output_dtype=output_dtype, + M=M, + batch=batch, + N=N, + K=K, + split_k=split_k, + ) + + +__all__ = [ + "A8W8MxscaleBMMPlan", + "A16W16LaunchPlan", + "WorkspaceSpec", +] diff --git a/aiter/ops/opus/policy.py b/aiter/ops/opus/policy.py new file mode 100644 index 0000000000..c460de76d0 --- /dev/null +++ b/aiter/ops/opus/policy.py @@ -0,0 +1,799 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2025-2026, Advanced Micro Devices, Inc. All rights reserved. +"""Caller-side selection policy for OPUS A16W16 and MXFP8 BMM. + +This module runs before the exact-kid :func:`opus_gemm` or :func:`opus_bmm` +entry. A16W16 policy validates tuned candidates and supplies explicit +per-architecture heuristic candidates. The gfx950 MXFP8 BMM policy owns tuned +CSV discovery, legacy-id normalization and its fallback kid/split selection. +Exact launch, workspace materialization and C++ dispatch remain outside this +module; no policy path replaces a caller-supplied exact kid inside the public +launcher. +""" + +from __future__ import annotations + +import math +from functools import cache, lru_cache + +import pandas as pd + +from aiter import logger +from csrc.opus_gemm.opus_gemm_common import ( + DEFAULT_COMPILED_KIDS_BY_ARCH, + GFX942_BF16WS_EXACT_N, + canonical_output_dtype, + get_kernel_instance, +) + +from ...jit.core import AITER_CONFIGS, AITER_LOG_TUNED_CONFIG +from ...jit.utils.chip_info import get_gfx_runtime as get_gfx +from ..gemm_op_common import get_padded_m +from ._arch import GFX942, GFX950, GFX1250 +from .launch_plan import ( + A16W16LaunchPlan, + _get_cached_a8w8_mxscale_bmm_plan, + _get_cached_a16w16_launch_plan, +) + +# ---- A16W16 tuned-candidate and heuristic policy ------------------------- + +_A16W16_TUNED_KEY_COLUMNS = ( + "gfx", + "cu_num", + "M", + "N", + "K", + "bias", + "dtype", + "outdtype", + "scaleAB", + "bpreshuffle", +) + + +@cache +def _load_a16w16_opus_tuned() -> dict: + """Load only OPUS rows from the merged global A16W16 tuned config.""" + path = AITER_CONFIGS.AITER_CONFIG_GEMM_BF16_FILE + try: + df = pd.read_csv(path).drop_duplicates() + except (FileNotFoundError, pd.errors.EmptyDataError): + return {} + except (OSError, UnicodeDecodeError, pd.errors.ParserError) as exc: + logger.warning( + "Ignoring unreadable A16W16 tuned CSV %r; OPUS will use its " + "heuristic fallback: %s", + path, + exc, + ) + return {} + + required = set(_A16W16_TUNED_KEY_COLUMNS) | {"libtype", "solidx", "splitK"} + missing = required.difference(df.columns) + if missing: + logger.warning( + "Ignoring A16W16 tuned CSV %r; missing columns %s", + path, + sorted(missing), + ) + return {} + + df = df[df["libtype"].eq("opus")] + if df.empty: + return {} + + integer_columns = ["solidx", "splitK"] + numeric = df[integer_columns].apply(pd.to_numeric, errors="coerce") + valid = ( + numeric.notna().all(axis=1) + & numeric.ge(0).all(axis=1) + & numeric.lt(float("inf")).all(axis=1) + & numeric.eq(numeric.round()).all(axis=1) + ) + invalid_rows = int((~valid).sum()) + if invalid_rows: + logger.warning( + "Skipping %d malformed OPUS row(s) in A16W16 tuned CSV %r: " + "solidx and splitK must be non-negative integers", + invalid_rows, + path, + ) + df = df.loc[valid].copy() + if df.empty: + return {} + df[integer_columns] = numeric.loc[valid].astype("int64") + + if "us" in df.columns: + df["_opus_sort_us"] = pd.to_numeric(df["us"], errors="coerce") + df = df.sort_values("_opus_sort_us", kind="stable", na_position="last").drop( + columns="_opus_sort_us" + ) + keys = list(_A16W16_TUNED_KEY_COLUMNS) + return df.drop_duplicates(keys).set_index(keys).to_dict("index") + + +@lru_cache(maxsize=4096) +def lookup_a16w16_opus_config( + *, + arch: str, + cu_num: int, + M: int, + N: int, + K: int, + has_bias: bool, + input_dtype: object, + output_dtype: object, +) -> dict | None: + """Return the exact-shape OPUS tuned row used by the legacy OPUS caller.""" + key = ( + str(arch).lower().split(":", 1)[0], + int(cu_num), + int(M), + int(N), + int(K), + bool(has_bias), + str(input_dtype), + str(output_dtype), + False, + False, + ) + config = _load_a16w16_opus_tuned().get(key) + return None if config is None else dict(config) + + +def _heuristic_a16w16_kid_gfx950( + M: int, + N: int, + K: int, + batch: int = 1, + has_bias: bool = False, + output_dtype: object = "bf16", +) -> int: + """Return the original gfx950 no-tuned-row fallback kid.""" + del batch, output_dtype + M, N, K = map(int, (M, N, K)) + split_barrier_ok = N % 16 == 0 and K % 64 == 0 and (K // 64) % 2 == 0 + + if M <= 4: + if M % 64 == 0 and N % 64 == 0 and K % 128 == 0: + return 1208 + return 208 + if M <= 64: + if M % 64 == 0 and N % 32 == 0 and K % 128 == 0: + return 1206 + return 206 + if M <= 128: + if M % 64 == 0 and N % 64 == 0 and K % 64 == 0: + return 1200 + return 200 + if split_barrier_ok and not has_bias: + if M % 256 == 0 and N % 256 == 0 and K % 64 == 0: + return 1300 + return 300 + if M % 64 == 0 and N % 64 == 0 and K % 64 == 0: + return 1200 + return 200 + + +def _heuristic_a16w16_kid_gfx1250( + M: int, + N: int, + K: int, + batch: int = 1, + has_bias: bool = False, + output_dtype: object = "bf16", +) -> int: + """Return the original gfx1250 no-tuned-row fallback kid.""" + del K, batch, has_bias, output_dtype + M, N = map(int, (M, N)) + if M % 32 == 0: + if N % 128 == 0: + return 20007 + if N % 64 == 0: + return 20006 + if N % 32 == 0: + return 20005 + if N % 128 == 0: + return 20004 + if N % 64 == 0: + return 20003 + return 20000 + + +@lru_cache(maxsize=1) +def _gfx942_heuristic_symbol_to_kid() -> dict[str, int]: + """Build the canonical gfx942 launcher-symbol mapping from the registry.""" + result: dict[str, int] = {} + for kid in DEFAULT_COMPILED_KIDS_BY_ARCH[GFX942]: + instance = get_kernel_instance(GFX942, "a16w16", kid) + if instance is None: + raise RuntimeError(f"gfx942 heuristic kid {kid} has no a16w16 instance") + previous = result.setdefault(instance.name, int(kid)) + if previous != kid: + raise RuntimeError( + f"duplicate gfx942 launcher symbol {instance.name!r}: " + f"kids {previous} and {kid}" + ) + return result + + +def _gfx942_heuristic_kid_for_symbol(symbol: str) -> int: + try: + return _gfx942_heuristic_symbol_to_kid()[symbol] + except KeyError as exc: + raise RuntimeError( + f"gfx942 heuristic returned unknown launcher symbol {symbol!r}" + ) from exc + + +def _gfx942_heuristic_split_barrier_ok(N: int, K: int) -> bool: + loops = (K + 63) // 64 + return N % 16 == 0 and K % 64 == 0 and loops >= 2 and loops % 2 == 0 + + +def _gfx942_heuristic_bf16ws_band(M: int, N: int, K: int) -> bool: + return ( + K >= 4096 and K % 64 == 0 and 104 <= M <= 608 and (N == 256 or 512 <= N <= 2048) + ) + + +def _gfx942_heuristic_bf16_symbol(M: int, N: int, K: int) -> str: + """Port of the original gfx942 BF16 launcher-choice ordering.""" + k64_ok = K % 64 == 0 + k32_ok = K % 32 == 0 + wkc_bk64_ok = K >= 4096 and K % 512 == 0 + p1_ok = K % 128 == 0 + sb_ok = _gfx942_heuristic_split_barrier_ok(N, K) + + if K == 4096: + if p1_ok and (M in (48, 64) and N == 1024): + return "opus_gemm_gfx942_splitk_p1_bk128_bf16ws_256x64x64x128_2x2_16x16x16_0x0x0" + if p1_ok and ((M == 128 and N == 512) or (M == 256 and N == 256)): + return "opus_gemm_gfx942_splitk_p1_bk128_bf16ws_256x64x64x128_2x2_16x16x16_0x0x0" + if p1_ok and M == 512 and N == 256: + return "opus_gemm_gfx942_splitk_p1_bk128_256x64x64x128_2x2_16x16x16_0x0x0" + if M in (48, 64) and 1536 <= N <= 2048: + return "opus_gemm_gfx942_splitk_legacy_512x64x128x64_2x4_16x16x16_0x0x0" + if (M == 128 and N == 1024) or (M == 256 and N == 512): + return "opus_gemm_gfx942_splitk_legacy_512x64x128x64_2x4_16x16x16_0x0x0" + if ( + (M == 128 and 1536 <= N <= 2048) + or (M == 256 and N == 1024) + or (M == 512 and N == 512) + ): + return "opus_gemm_gfx942_splitk_legacy_512x128x128x64_2x4_16x16x16_0x0x0" + + if K >= 1024 and k32_ok and N >= 1536 and M <= 32: + if M <= 4 and N >= 4096: + return "opus_gemm_gfx942_wkc_512x16x16x64_1x1_16x16x16_0x0x0" + if M <= 16: + if wkc_bk64_ok: + return "opus_gemm_gfx942_wkc_512x16x32x64_1x1_16x16x16_0x0x0" + return "opus_gemm_gfx942_wkc_512x16x32x32_1x1_16x16x16_0x0x0" + if M == 32 and K == 4096 and wkc_bk64_ok: + return "opus_gemm_gfx942_wkc_512x16x32x64_1x1_16x16x16_0x0x0" + return "opus_gemm_gfx942_wkc_256x32x32x64_1x1_16x16x16_0x0x0" + + if ( + K >= 512 + and k64_ok + and (N <= 64 or (M <= 128 and N <= 1024) or (M <= 8 and N <= 1536)) + ): + if N <= 64 and M > 128: + return "opus_gemm_gfx942_wkc_512x32x16x64_1x1_16x16x16_0x0x0" + if N <= 256 or M <= 8 or (M <= 16 and N <= 800): + return "opus_gemm_gfx942_wkc_512x16x16x64_1x1_16x16x16_0x0x0" + return "opus_gemm_gfx942_wkc_512x32x16x64_1x1_16x16x16_0x0x0" + + if _gfx942_heuristic_bf16ws_band(M, N, K): + return "opus_gemm_gfx942_splitk_legacy_bf16ws_512x128x128x64_2x4_16x16x16_0x0x0" + + if N == 384 and K >= 4096: + if M <= 128: + return "opus_gemm_gfx942_wkc_512x32x16x64_1x1_16x16x16_0x0x0" + if M <= 224: + return "opus_gemm_gfx942_splitk_p1_256x64x64x64_2x2_16x16x16_0x0x0" + if 392 <= M <= 512: + return "opus_gemm_gfx942_splitk_em3en4_lds1_pgr2_256x128x96x128_2x2_16x16x16_0x0x0" + return "opus_gemm_gfx942_splitk_legacy_512x128x128x64_2x4_16x16x16_0x0x0" + + if k64_ok and N >= 4096 and K <= 3200: + if K <= 640 and M <= 128: + return "opus_gemm_gfx942_p1_256x64x64x64_2x2_16x16x16_0x0x0" + return "opus_gemm_gfx942_512x128x128x64_2x4_16x16x16_0x0x0" + + if sb_ok and M >= 128: + return "opus_gemm_gfx942_512x128x128x64_2x4_16x16x16_0x0x0" + if N <= 256 and p1_ok: + return "opus_gemm_gfx942_splitk_p1_256x64x64x64_2x2_16x16x16_0x0x0" + return "opus_gemm_gfx942_splitk_legacy_512x128x128x64_2x4_16x16x16_0x0x0" + + +def _heuristic_a16w16_kid_gfx942( + M: int, + N: int, + K: int, + batch: int = 1, + has_bias: bool = False, + output_dtype: object = "bf16", +) -> int: + """Return the original gfx942 no-tuned-row fallback kid.""" + del batch + M, N, K = map(int, (M, N, K)) + if canonical_output_dtype(output_dtype) == "bf16_t" and not has_bias: + symbol = _gfx942_heuristic_bf16_symbol(M, N, K) + elif N <= 256 and K % 128 == 0: + symbol = "opus_gemm_gfx942_splitk_p1_256x64x64x64_2x2_16x16x16_0x0x0" + else: + symbol = "opus_gemm_gfx942_splitk_legacy_512x128x128x64_2x4_16x16x16_0x0x0" + return _gfx942_heuristic_kid_for_symbol(symbol) + + +_A16W16_HEURISTICS = { + GFX942: _heuristic_a16w16_kid_gfx942, + GFX950: _heuristic_a16w16_kid_gfx950, + GFX1250: _heuristic_a16w16_kid_gfx1250, +} + +_UINT32_MAX_BYTES = (1 << 32) - 1 + + +def _check_a16w16_heuristic_4g( + *, + arch: str, + M: int, + N: int, + K: int, + output_dtype: object, +) -> None: + """Mirror the legacy C++ heuristic guard for 32-bit buffer descriptors.""" + if arch not in (GFX950, GFX1250): + return + + output_itemsize = {"bf16_t": 2, "fp32_t": 4}.get( + canonical_output_dtype(output_dtype) + ) + if output_itemsize is None: + return + + M, N, K = map(int, (M, N, K)) + if max(M * K * 2, N * K * 2, M * N * output_itemsize) <= _UINT32_MAX_BYTES: + return + + reason = ( + "legacy kids require a tuned 4g_safe kid" + if arch == GFX950 + else "launcher gmem descriptors are 32-bit" + ) + raise RuntimeError( + f"opus {arch} a16w16 heuristic refuses >4 GiB shape " + f"(M={M} N={N} K={K}): {reason}" + ) + + +def select_a16w16_heuristic_kid( + *, + arch: str, + M: int, + N: int, + K: int, + batch: int, + has_bias: bool, + output_dtype: object, +) -> int: + """Select one baseline-parity A16 kid before the exact public call.""" + arch = str(arch).lower().split(":", 1)[0] + heuristic = _A16W16_HEURISTICS.get(arch) + if heuristic is None: + raise ValueError(f"no OPUS a16w16 heuristic for runtime arch {arch}") + kid = int(heuristic(M, N, K, batch, has_bias, output_dtype)) + if kid not in DEFAULT_COMPILED_KIDS_BY_ARCH.get(arch, frozenset()): + raise RuntimeError( + f"{arch} a16w16 heuristic returned kid {kid}, which is not in " + "DEFAULT_COMPILED_KIDS_BY_ARCH" + ) + instance = get_kernel_instance(arch, "a16w16", kid) + if instance is not None and instance.max_m is not None and int(M) > instance.max_m: + raise RuntimeError( + f"opus {arch} a16w16 heuristic refuses M={M}: " + f"kid {kid} requires M <= {instance.max_m}. " + "Tune this shape for a non-split-K kernel." + ) + return kid + + +def _resolve_a16w16_candidate( + *, + arch: str, + M: int, + N: int, + K: int, + batch: int, + cu_num: int, + has_bias: bool, + input_dtype: object, + output_dtype: object, + kid: int, + split_k: int, +) -> A16W16LaunchPlan | None: + """Apply caller-only redirects, then validate one exact candidate.""" + if arch == GFX942 and N not in GFX942_BF16WS_EXACT_N: + if kid == 10210: + kid = 10200 + elif kid == 10213: + kid = 10203 + elif kid == 10216: + return None + + try: + return _get_cached_a16w16_launch_plan( + arch, + M, + N, + K, + batch, + cu_num, + has_bias, + input_dtype, + output_dtype, + kid, + split_k, + ) + except (ValueError, OverflowError): + return None + + +def resolve_a16w16_tuned_candidate( + *, + arch: str, + M: int, + N: int, + K: int, + batch: int, + cu_num: int, + has_bias: bool, + input_dtype: object, + output_dtype: object, + requested_kid: object, + requested_split_k: object = 0, +) -> A16W16LaunchPlan | None: + """Validate one tuned candidate without invoking a heuristic.""" + arch = str(arch).lower().split(":", 1)[0] + try: + split_k = int(requested_split_k) + kid = int(requested_kid) + except (TypeError, ValueError): + return None + return _resolve_a16w16_candidate( + arch=arch, + M=M, + N=N, + K=K, + batch=batch, + cu_num=cu_num, + has_bias=has_bias, + input_dtype=input_dtype, + output_dtype=output_dtype, + kid=kid, + split_k=split_k, + ) + + +def resolve_a16w16_heuristic_candidate( + *, + arch: str, + M: int, + N: int, + K: int, + batch: int, + cu_num: int, + has_bias: bool, + input_dtype: object, + output_dtype: object, + requested_split_k: object = 0, +) -> A16W16LaunchPlan | None: + """Select and validate one per-architecture heuristic candidate.""" + arch = str(arch).lower().split(":", 1)[0] + try: + split_k = int(requested_split_k) + _check_a16w16_heuristic_4g( + arch=arch, + M=M, + N=N, + K=K, + output_dtype=output_dtype, + ) + preferred = select_a16w16_heuristic_kid( + arch=arch, + M=M, + N=N, + K=K, + batch=batch, + has_bias=has_bias, + output_dtype=output_dtype, + ) + except (TypeError, ValueError): + return None + return _resolve_a16w16_candidate( + arch=arch, + M=M, + N=N, + K=K, + batch=batch, + cu_num=cu_num, + has_bias=has_bias, + input_dtype=input_dtype, + output_dtype=output_dtype, + kid=preferred, + split_k=split_k, + ) + + +def resolve_a16w16_caller_candidate( + *, + arch: str, + M: int, + N: int, + K: int, + batch: int, + cu_num: int, + has_bias: bool, + input_dtype: object, + output_dtype: object, + requested_kid: object | None, + requested_split_k: object = 0, +) -> A16W16LaunchPlan | None: + """Resolve a tuned kid when supplied, otherwise use the arch heuristic.""" + if requested_kid is None: + return resolve_a16w16_heuristic_candidate( + arch=arch, + M=M, + N=N, + K=K, + batch=batch, + cu_num=cu_num, + has_bias=has_bias, + input_dtype=input_dtype, + output_dtype=output_dtype, + requested_split_k=requested_split_k, + ) + return resolve_a16w16_tuned_candidate( + arch=arch, + M=M, + N=N, + K=K, + batch=batch, + cu_num=cu_num, + has_bias=has_bias, + input_dtype=input_dtype, + output_dtype=output_dtype, + requested_kid=requested_kid, + requested_split_k=requested_split_k, + ) + + +# ---- gfx950 MXFP8 BMM tuned-row and heuristic policy --------------------- + +_MXSCALE_BMM_KID_OFFSET = 8000 +_MXSCALE_BMM_LOCAL_KID_MAX = 653 +_TUNED_PERF_COLUMNS = ("us", "tflops", "bw", "errRatio") +_C_INT_MAX = (1 << 31) - 1 + + +def _parse_mxscale_bmm_tuned_split_k(value: object) -> int: + """Require one saved split-K value that can cross the C++ int ABI.""" + if isinstance(value, bool): + raise TypeError(f"splitK must be a positive integer, got {value!r}") + try: + numeric = float(value) + except (TypeError, ValueError, OverflowError) as exc: + raise ValueError(f"splitK must be a positive integer, got {value!r}") from exc + if ( + not math.isfinite(numeric) + or not numeric.is_integer() + or numeric < 1 + or numeric > _C_INT_MAX + ): + raise ValueError(f"splitK must be a positive integer, got {value!r}") + return int(numeric) + + +@cache +def _load_mxscale_bmm_tuned(libtype: str | None = None) -> dict: + path = AITER_CONFIGS.AITER_CONFIG_BATCHED_GEMM_A8W8_BLOCKSCALE_MXSCALE_FILE + try: + df = pd.read_csv(path).drop_duplicates() + except FileNotFoundError: + logger.warning("MXFP8 BMM tuned CSV was not found at %s", path) + return {} + + required = {"gfx", "b", "m", "n", "k", "kernelId", "splitK"} + missing = required.difference(df.columns) + if missing: + raise ValueError(f"MXFP8 BMM tuned CSV is missing columns {sorted(missing)}") + + if libtype is not None and "libtype" in df.columns: + df = df[df["libtype"] == libtype].copy() + + opus_rows = ( + df["libtype"].eq("opus") + if "libtype" in df.columns + else pd.Series(True, index=df.index, dtype=bool) + ) + + # Normalize legacy local ids, then validate each complete saved launch. + opus_kids = pd.to_numeric(df.loc[opus_rows, "kernelId"], errors="coerce") + legacy_rows = opus_kids.between(0, _MXSCALE_BMM_LOCAL_KID_MAX) + opus_kids.loc[legacy_rows] += _MXSCALE_BMM_KID_OFFSET + integer_kids = ( + opus_kids.notna() & opus_kids.lt(float("inf")) & opus_kids.eq(opus_kids.round()) + ) + valid_opus_rows = pd.Series(False, index=opus_kids.index, dtype=bool) + for index in opus_kids.index[integer_kids]: + kid = int(opus_kids.at[index]) + arch = str(df.at[index, "gfx"]).lower().split(":", 1)[0] + try: + split_k = _parse_mxscale_bmm_tuned_split_k(df.at[index, "splitK"]) + batch, m, n, k = ( + int(df.at[index, column]) for column in ("b", "m", "n", "k") + ) + _get_cached_a8w8_mxscale_bmm_plan( + arch, + kid, + "bf16", + m, + batch, + n, + k, + split_k, + ) + except (TypeError, ValueError, OverflowError): + continue + else: + df.at[index, "kernelId"] = kid + df.at[index, "splitK"] = split_k + valid_opus_rows.at[index] = True + + invalid_opus_rows = opus_rows.copy() + invalid_opus_rows.loc[opus_rows] = ~valid_opus_rows + if invalid_opus_rows.any(): + logger.warning( + "Skipping %d invalid OPUS row(s) in MXFP8 BMM tuned CSV %r: " + "kernelId, splitK and shape must form a compatible registered " + "MXFP8 BMM launch", + int(invalid_opus_rows.sum()), + path, + ) + df = df.loc[~invalid_opus_rows].copy() + + shape_keys = ["gfx", "b", "m", "n", "k"] + duplicate_shapes = df.duplicated(subset=shape_keys, keep=False) + if duplicate_shapes.any(): + rows = df.loc[duplicate_shapes, shape_keys].drop_duplicates().to_dict("records") + raise RuntimeError(f"duplicate shapes across MXFP8 BMM tuned CSV files: {rows}") + return df.set_index(shape_keys).to_dict("index") + + +@lru_cache(maxsize=1024) +def lookup_mxscale_bmm_config( + b: int, + m: int, + n: int, + k: int, + *, + libtype: str | None = None, +): + """Return the exact or existing padded-M tuned row for one shape.""" + gfx = get_gfx() + 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( + "shape B:%s M:%s N:%s K:%s has no MXFP8 BMM tuned row", + b, + m, + n, + k, + ) + return None + if AITER_LOG_TUNED_CONFIG: + cfg = { + key: value for key, value in row.items() if key not in _TUNED_PERF_COLUMNS + } + logger.info( + "shape B:%s M:%s N:%s K:%s uses padded_M:%s MXFP8 config %s", + b, + m, + n, + k, + padded_m, + cfg, + ) + return row + + +def _heuristic_mxscale_bmm_kid(g: int, m: int, n: int, k: int) -> int: + """Choose a final global kid only when the tuned table has no usable row.""" + + def divisible(value: int, divisor: int) -> bool: + return value % divisor == 0 + + if ( + divisible(n, 256) + and divisible(k, 128) + and (m >= 2048 or (m >= 1024 and g >= 8)) + ): + return 8158 if 4096 <= k <= 8192 else 8150 + if m < 64: + return 8640 if divisible(n, 64) and divisible(k, 256) else 8653 + if m <= 256 and k <= 1024 and divisible(n, 32) and divisible(k, 256): + return 8320 + if divisible(n, 64) and divisible(k, 128): + return 8653 + return 8000 + + +def resolve_a8w8_mxscale_bmm_plan( + g: int, + m: int, + n: int, + k: int, +) -> tuple[int, int]: + """Resolve one final global kid/split pair for the high-level caller.""" + config = lookup_mxscale_bmm_config(g, m, n, k) + libtype = config.get("libtype", "opus") if config is not None else "opus" + if libtype != "opus": + raise NotImplementedError( + f"MXFP8 BMM tuned row requests unsupported backend {libtype!r}" + ) + + if config is not None: + try: + kid = int(config["kernelId"]) + split_k = _parse_mxscale_bmm_tuned_split_k(config["splitK"]) + plan = _get_cached_a8w8_mxscale_bmm_plan( + get_gfx(), + kid, + "bf16", + m, + g, + n, + k, + split_k, + ) + except (TypeError, ValueError, OverflowError) as exc: + logger.warning( + "Ignoring invalid OPUS MXFP8 BMM tuned row for " + "B:%s M:%s N:%s K:%s (kernelId=%r, splitK=%r): %s; " + "using heuristic fallback", + g, + m, + n, + k, + config.get("kernelId"), + config.get("splitK"), + exc, + ) + else: + return plan.resolved_kid, plan.abi_split_k + + kid = _heuristic_mxscale_bmm_kid(g, m, n, k) + return kid, 1 + + +__all__ = [ + "lookup_a16w16_opus_config", + "lookup_mxscale_bmm_config", + "resolve_a8w8_mxscale_bmm_plan", + "resolve_a16w16_caller_candidate", + "resolve_a16w16_heuristic_candidate", + "resolve_a16w16_tuned_candidate", + "select_a16w16_heuristic_kid", +] diff --git a/aiter/tuned_gemm.py b/aiter/tuned_gemm.py index 9f047bd876..61cb5b0e0b 100644 --- a/aiter/tuned_gemm.py +++ b/aiter/tuned_gemm.py @@ -30,9 +30,9 @@ from aiter.ops.gemm_op_common import get_padded_m try: - from aiter.ops.opus.gemm_op_a16w16 import opus_gemm_a16w16_tune as _opus_tune + from aiter.ops.opus import opus_gemm as _opus_launch except Exception: # noqa: BLE001 blanket catch is intentional here - _opus_tune = None + _opus_launch = None @functools.lru_cache(maxsize=1) @@ -42,23 +42,8 @@ def _get_flydsl_gemm_kernels(): return gemm_kernels -# NOTE: gfx1250 split-K kids allocate their partial-sum workspace as a plain -# torch.empty tensor (see aiter.ops.opus.gemm_op_a16w16._get_opus_workspace) -# passed explicitly to the launcher. torch's caching allocator is HIP graph- -# capture aware, so that single torch.empty path serves both eager and capture -# (a buffer first touched inside capture comes from the graph mempool with a -# replay-stable address) and no eager pre-warm of the shape is required. (The -# old per-stream hipMalloc registry -- opus_gemm_workspace_init / -# opus_splitk_ws_get -- used by the gfx942/gfx950 a16w16 split-K path still needs -# an eager warm before capture; if that path is ever exercised under cudagraphs, -# warm it via aiter.opus_gemm_workspace_init() on the capture stream. It fails -# loudly ("splitk workspace not initialized") rather than silently corrupting, -# so its absence here is safe to detect.) - - this_dir = os.path.dirname(os.path.abspath(__file__)) - extensions_created = False untune_path = f"{this_dir}/configs/bf16_untuned_gemm.csv" tune_path = AITER_CONFIGS.AITER_CONFIG_GEMM_BF16_FILE @@ -139,6 +124,7 @@ def get_GEMM_A16W16_config( padded_M = M config = None gfx = get_gfx() + warned_invalid_opus = set() for gl in [None, 0, 1]: padded_M = M if gl is None else get_padded_m(M, N, K, gl) config = cfg.get( @@ -173,6 +159,52 @@ def get_GEMM_A16W16_config( config = None if config is None: continue + if config["libtype"] == "opus": + if _opus_launch is None: + resolved = None + else: + from aiter.ops.opus.policy import ( + resolve_a16w16_tuned_candidate, + ) + + resolved = resolve_a16w16_tuned_candidate( + arch=gfx, + M=M, + N=N, + K=K, + batch=1, + cu_num=cu_num, + has_bias=bias, + input_dtype=eval(dtype), + output_dtype=eval(otype), + requested_kid=config.get("solidx"), + requested_split_k=config.get("splitK"), + ) + if resolved is None: + # Discard the whole stale (kid, split-K) pair before + # trying another padded row or the default fallback. + invalid_row = ( + padded_M, + config.get("solidx"), + config.get("splitK"), + ) + if invalid_row not in warned_invalid_opus: + logger.warning( + "Ignoring invalid OPUS tuned row for gfx=%s, " + "shape=(%d,%d,%d), kid=%r, splitK=%r; trying " + "the next padded row or default backend", + gfx, + padded_M, + N, + K, + config.get("solidx"), + config.get("splitK"), + ) + warned_invalid_opus.add(invalid_row) + config = None + continue + config = dict(config) + config["solidx"] = int(resolved.resolved_kid) if AITER_LOG_TUNED_CONFIG: kernelName = ( config["kernelName"] if config["libtype"] != "hipblaslt" else "" @@ -207,6 +239,7 @@ def get_GEMM_A16W16_config( elif gfx in ("gfx90a", "gfx942", "gfx950") and is_skinny_default_shape( M, N, K, dtype, cu_num ): + # soltype, solution_idx = 3, 2 default_config["libtype"] = "skinny" default_config["solidx"] = 2 default_config["kernelName"] = "" @@ -531,7 +564,8 @@ def opus_gemm( bpreshuffle: bool | None = False, config: dict | None = None, ): - if _opus_tune is None: + """Run one tuned OPUS A16W16 row through the exact-kid interface.""" + if _opus_launch is None: logger.warning( "opus tuned config found but opus is not available; falling back to torch" ) @@ -554,22 +588,16 @@ def opus_gemm( splitK = int(config.get("splitK", 0)) if config is not None else 0 m, _k = inp.shape n = weights.shape[0] - # The split-K workspace (if any) is allocated capture-safely inside - # opus_gemm_a16w16_tune -> _get_opus_workspace; no eager pre-warm needed. Y = torch.empty(m, n, dtype=otype or inp.dtype, device=inp.device) - _opus_tune( - inp.unsqueeze(0), - weights.unsqueeze(0), - Y.unsqueeze(0), + _opus_launch( + inp, + weights, + Y, + kid=int(solidx), bias=bias, - kernelId=int(solidx), - splitK=splitK, + split_k=splitK, ) - # NOTE: do NOT add bias again here -- the opus splitk reduce kernel already - # folds `bias` into the fp32 accumulator before the bf16/fp32 cast (HAS_BIAS - # path). The previous `Y = Y + bias` double-counted bias (output = A@B^T + - # 2*bias), causing ~54% miscompare (maxabs ~= bias range) for every bias!=None - # opus shape under tgemm (e.g. ATOM's bf16 linear). + # The OPUS launcher already applies bias, including split-K reduction. return Y diff --git a/csrc/gemm_a16w16/gemm_a16w16_tune.py b/csrc/gemm_a16w16/gemm_a16w16_tune.py index 5c7a41f7ff..0b44f319ae 100644 --- a/csrc/gemm_a16w16/gemm_a16w16_tune.py +++ b/csrc/gemm_a16w16/gemm_a16w16_tune.py @@ -67,13 +67,11 @@ kid_rejects_shape as _opus_kid_rejects_shape, ) - from aiter.ops.opus.gemm_op_a16w16 import ( - opus_gemm_a16w16_tune as _opus_gemm_a16w16_tune, - ) + from aiter.ops.opus import opus_gemm as _opus_gemm _opus_all_kernels = dict(_opus_kernels_list) except Exception as _opus_exc: # noqa: BLE001 - _opus_gemm_a16w16_tune = None + _opus_gemm = None _opus_all_kernels = None _opus_splitk_kids = frozenset() _opus_candidate_kids_for_shape = None @@ -203,16 +201,14 @@ def run_triton_gemm_bf16(input, weight, bias=None, otype=dtypes.bf16): def run_opus_gemm_bf16(inp, weight, out, bias=None, kid=0, splitK=0): - inp3 = inp.unsqueeze(0) - weight3 = weight.unsqueeze(0) - out3 = out.unsqueeze(0) - _opus_gemm_a16w16_tune( - inp3, - weight3, - out3, + """Launch one OPUS kid and run the existing maximum-error check.""" + _opus_gemm( + inp, + weight, + out, bias=bias, - kernelId=kid, - splitK=splitK, + kid=kid, + split_k=splitK, ) if torch.cuda.is_current_stream_capturing(): return out @@ -227,13 +223,13 @@ def run_opus_gemm_bf16(inp, weight, out, bias=None, kid=0, splitK=0): ) if cache_key in _opus_max_delta_checked: return out - ref_fp32 = torch.bmm(inp3.float(), weight3.float().transpose(-1, -2)) + # ``opus_gemm`` is the strict logical-2D public entry point. Keep this + # catastrophic-error guard in the same shape domain instead of referring + # to the pre-split adapter's removed 3D views. + ref_fp32 = F.linear(inp.float(), weight.float()) if bias is not None: - if bias.dim() == 1: - ref_fp32 = ref_fp32 + bias.float().view(1, 1, -1) - else: - ref_fp32 = ref_fp32 + bias.float().unsqueeze(1) - max_delta = (out3.float() - ref_fp32).abs().max().item() + ref_fp32 = ref_fp32 + bias.float() + max_delta = (out.float() - ref_fp32).abs().max().item() max_ref = ref_fp32.abs().max().item() bound = max(max_ref * 0.1, 1.0) if max_delta > bound: @@ -683,7 +679,7 @@ def _get_asm_tasks( def _get_opus_tasks( self, info_keys, has_bias, indtype, outdtype, scaleAB, is_shuffle, run_kwargs ): - if _opus_gemm_a16w16_tune is None: + if _opus_gemm is None: logger.warning(f"opus not available, skip. reason: {OPUS_TUNE_ERROR}") return [] if scaleAB or indtype != dtypes.bf16: diff --git a/csrc/include/rocm_ops.hpp b/csrc/include/rocm_ops.hpp index 6290c96889..9b24af3cc9 100644 --- a/csrc/include/rocm_ops.hpp +++ b/csrc/include/rocm_ops.hpp @@ -290,72 +290,62 @@ namespace py = pybind11; py::arg("x_scale") = std::nullopt, \ py::arg("w_scale") = std::nullopt); -#define OPUS_GEMM_PYBIND \ - m.def("opus_gemm", \ - &opus_gemm, \ - "opus_gemm", \ - py::arg("XQ"), \ - py::arg("WQ"), \ - py::arg("Y"), \ - py::arg("group_layout") = std::nullopt, \ - py::arg("x_scale") = std::nullopt, \ - py::arg("w_scale") = std::nullopt, \ - py::arg("bias") = std::nullopt); - -#define OPUS_GEMM_A16W16_TUNE_PYBIND \ - m.def("opus_gemm_a16w16_tune", \ - &opus_gemm_a16w16_tune, \ - "opus_gemm_a16w16_tune", \ +// OPUS exact-kid bindings; blockscale scale tensors are required. +#define OPUS_GEMM_A16W16_LAUNCH_PYBIND \ + m.def("opus_gemm_a16w16_launch", \ + &opus_gemm_a16w16_launch, \ + "opus_gemm_a16w16_launch", \ py::arg("XQ"), \ py::arg("WQ"), \ py::arg("Y"), \ - py::arg("bias") = std::nullopt, \ - py::arg("workspace") = std::nullopt, \ - py::arg("kernelId") = 0, \ - py::arg("splitK") = 0); - -#define OPUS_BMM_A8W8_MXSCALE_PYBIND \ - m.def("opus_bmm_a8w8_mxscale", \ - &opus_bmm_a8w8_mxscale, \ - "mmajor fp8 e8m0 mxscale (block-scale) BMM with native " \ - "scaled MFMA; kid-dispatched flatmm split-K backend", \ - py::arg("O"), \ - py::arg("wo_a"), \ - py::arg("Y"), \ - py::arg("x_scale"), \ - py::arg("w_scale"), \ - py::arg("splitK") = 2, \ - py::arg("kernelId") = 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("bias"), \ + py::arg("workspace"), \ + py::arg("kid"), \ + py::arg("split_k")); + +#define OPUS_GEMM_A8W8_LAUNCH_PYBIND \ + m.def("opus_gemm_a8w8_launch", \ + &opus_gemm_a8w8_launch, \ + "opus_gemm_a8w8_launch", \ + py::arg("XQ"), \ + py::arg("WQ"), \ + py::arg("Y"), \ + py::arg("kid")); + +#define OPUS_GEMM_A8W8_BLOCKSCALE_LAUNCH_PYBIND \ + m.def("opus_gemm_a8w8_blockscale_launch", \ + &opus_gemm_a8w8_blockscale_launch, \ + "opus_gemm_a8w8_blockscale_launch", \ + py::arg("XQ"), \ + py::arg("WQ"), \ + py::arg("Y"), \ + py::arg("x_scale"), \ + py::arg("w_scale"), \ + py::arg("kid")); + +#define OPUS_GEMM_A8W8_BLOCKSCALE_BPRESHUFFLE_LAUNCH_PYBIND \ + m.def("opus_gemm_a8w8_blockscale_bpreshuffle_launch", \ + &opus_gemm_a8w8_blockscale_bpreshuffle_launch, \ + "opus_gemm_a8w8_blockscale_bpreshuffle_launch", \ + py::arg("XQ"), \ + py::arg("WQ"), \ + py::arg("x_scale"), \ + py::arg("w_scale"), \ + py::arg("Y"), \ + py::arg("kid")); + +#define OPUS_GEMM_A8W8_MXSCALE_BMM_LAUNCH_PYBIND \ + m.def("opus_gemm_a8w8_mxscale_bmm_launch", \ + &opus_gemm_a8w8_mxscale_bmm_launch, \ + "opus_gemm_a8w8_mxscale_bmm_launch", \ py::arg("XQ"), \ py::arg("WQ"), \ + py::arg("Y"), \ 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, \ - "Register a splitk fp32 workspace handle for the current " \ - "CUDA stream. Call once per stream eagerly (outside HIP " \ - "graph capture) before capturing graphs that include " \ - "opus_gemm splitk kernels under TBO."); - -#define OPUS_GEMM_WORKSPACE_RELEASE_PYBIND \ - m.def("opus_gemm_workspace_release", \ - &opus_gemm_workspace_release, \ - "Free the splitk workspace (buffer + handles + registry " \ - "entry) for the current CUDA stream. Eager mode only; " \ - "no-op if the stream was never registered."); \ - m.def("opus_gemm_workspace_release_all", \ - &opus_gemm_workspace_release_all, \ - "Free the splitk workspace for all registered streams and " \ - "clear the registry. Eager mode only. Use for explicit " \ - "teardown before a framework reclaims its stream pool."); + py::arg("workspace"), \ + py::arg("kid"), \ + py::arg("split_k")); #define OPUS_MOE_PYBIND \ m.def("opus_moe_stage2_a8w4_decode_fwd", \ diff --git a/csrc/opus_gemm/README.md b/csrc/opus_gemm/README.md index f58612efb0..6f756ea8e6 100644 --- a/csrc/opus_gemm/README.md +++ b/csrc/opus_gemm/README.md @@ -1,332 +1,214 @@ -# Opus GEMM (C++ side) - -The user-facing documentation for the opus a16w16 GEMM lives at -[**aiter/ops/opus/README.md**](../../aiter/ops/opus/README.md). It -covers Quick Start, dispatch architecture, tuning workflow, env vars, -testing, internals, and troubleshooting. - -This directory holds the C++ / JIT build inputs only. - -## Layout - -| File | Role | -|---|---| -| `opus_gemm.cu` | Top-level entry points (`opus_gemm()` / `opus_gemm_a16w16_tune()`) and arch routers that switch on `opus_get_gfx_arch()` | -| `opus_gemm_common.py` | Kernel instance metadata — all kids (a16w16 split-barrier, flatmm, flatmm_splitk) live here | -| `gen_instances.py` | JIT codegen driver; `--tune_file` bakes the tuned CSV into `opus_gemm_lookup.h` | -| `opus_gemm_tune.py` | Offline tuner CLI (see `aiter/ops/opus/README.md` §3 for usage) | -| `opus_bmm_mxscale_tune.py` | Offline tuner CLI for the a8w8 mxscale BMM (DSV4 wo_a), writing `dsv4_batched_gemm_a8w8_blockscale_mxscale_tuned.csv`. Its candidate pool is `_TUNE_POLICY` (kid -> split-K factors); tile shape, kernelName and M alignment come from `opus_gemm_common.py`, so a kid is never tuned on a shape its launcher rejects | -| `include/opus_gemm.h`, `include/opus_gemm_arch.cuh` | Cross-arch declarations + `OpusGfxArch` enum + `opus_get_arch_info()` probe | -| `include/opus_gemm_common.cuh`, `include/opus_gemm_utils.cuh` | Cross-arch traits umbrella + opus.hpp shim | -| `include/gfx950/*.cuh` | gfx950-specific pipelines (a16w16 split-barrier / flatmm / flatmm_splitk, a8w8 noscale / scale), traits, splitk reduce, heuristic dispatch (`opus_a16w16_heuristic_dispatch_gfx950`), and the dispatch glue (`opus_gemm_arch_gfx950.cuh`). | - -The dispatch flow (gfx950 today): - -``` -opus_gemm() / opus_gemm_a16w16_tune() [opus_gemm.cu] - │ - ├─ opus_get_gfx_arch() ──────────► OpusGfxArch::{Gfx950, ...} [opus_gemm_arch.cuh] - │ - └─ switch (arch) { - case Gfx950: opus_dispatch_a16w16_gfx950(...) [gfx950/opus_gemm_arch_gfx950.cuh] - │ - ├─ tuned (M,N,K) lookup map (baked from CSV) - └─ opus_a16w16_heuristic_dispatch_gfx950(...) [gfx950/opus_gemm_heuristic_dispatch_gfx950.cuh] - } -``` - -Per-arch headers carry an `_` suffix on every file, every kargs -struct, every traits class, and the heuristic dispatch function so two -arches' headers can be visible in the same TU without ODR collisions. -The launcher symbol names (e.g. `opus_gemm_512x256x256x64_2x4_16x16x32_0x0x0`) -are not suffixed because each arch picks a different valid tile set, so -collisions cannot happen unless a future arch picks an identical tile; -if that happens, prefix `gen_instances.py`'s emitted launcher names with -the arch tag at that point. - -## Adding a new arch - -The codebase is staged so a new arch (e.g. `gfx942`) can be brought up -without touching gfx950 code. The Python and JIT layers are -arch-aware; only the kernel pipelines / traits are gfx950-specific -today. - -### 1. Arch enum and runtime probe - -Edit [`include/opus_gemm_arch.cuh`](include/opus_gemm_arch.cuh): - -```cpp -enum class OpusGfxArch -{ - Unknown = 0, - Gfx950, - Gfx942, // (1) add the enum value -}; - -// inside opus_get_arch_info(): -if (name.rfind("gfx950", 0) == 0) a = OpusGfxArch::Gfx950; -else if (name.rfind("gfx942", 0) == 0) a = OpusGfxArch::Gfx942; // (2) prefix-match -``` - -### 2. Per-arch headers - -Create `include/gfx942/` and mirror the gfx950 layout: - -``` -include/gfx942/ -├── opus_gemm_arch_gfx942.cuh # dispatch glue (lookup + heuristic wrapper) -├── opus_gemm_heuristic_dispatch_gfx942.cuh # M-bucket → launcher symbol heuristic -├── opus_gemm_traits_a16w16_gfx942.cuh # traits + 5 kargs structs -├── opus_gemm_traits_a8w8_noscale_gfx942.cuh -├── opus_gemm_traits_a8w8_scale_gfx942.cuh -├── opus_gemm_pipeline_a16w16_gfx942.cuh # __global__ kernel bodies -├── opus_gemm_pipeline_a16w16_flatmm_gfx942.cuh -├── opus_gemm_pipeline_a16w16_flatmm_splitk_gfx942.cuh -├── opus_gemm_pipeline_a8w8_noscale_gfx942.cuh -├── opus_gemm_pipeline_a8w8_scale_gfx942.cuh -└── splitk_reduce_gfx942.cuh -``` - -Naming rules (mirror gfx950, replace the suffix): - -- File names: `opus_gemm_*_gfx942.cuh` (one suffix per file). -- Traits / kargs structs: `opus_gemm_a16w16_traits_gfx942`, - `opus_gemm_noscale_kargs_gfx942`, `opus_gemm_flatmm_kargs_gfx942`, - `opus_gemm_flatmm_splitk_kargs_gfx942`, - `opus_gemm_a16w16_flatmm_traits_gfx942`, - `opus_flatmm_splitk_traits_gfx942`, - `opus_gemm_a8w8_noscale_traits_gfx942`, - `opus_gemm_a8w8_scale_traits_gfx942`, - `opus_gemm_scale_kargs_gfx942`. -- Shared-ABI kargs guard macro: `OPUS_GEMM_NOSCALE_KARGS_GFX942_DEFINED`. -- Heuristic dispatch function: `opus_a16w16_heuristic_dispatch_gfx942`. -- Per-arch dispatch glue: - - `opus_dispatch_a16w16_gfx942(int M, int N, int K, int batch)` - - `opus_a16w16_tune_dispatch_gfx942(int id)` - -If gfx942 reuses the same launcher tile sizes, also rename the emitted -launcher symbols (see step 5) to avoid ODR collisions in the manifest. - -### 3. Cross-arch umbrella - -Edit [`include/opus_gemm_common.cuh`](include/opus_gemm_common.cuh) to -include the new arch's traits headers alongside gfx950's: - -```cpp -#include "gfx950/opus_gemm_traits_a8w8_scale_gfx950.cuh" -#include "gfx950/opus_gemm_traits_a8w8_noscale_gfx950.cuh" -#include "gfx950/opus_gemm_traits_a16w16_gfx950.cuh" -#include "gfx942/opus_gemm_traits_a8w8_scale_gfx942.cuh" // new -#include "gfx942/opus_gemm_traits_a8w8_noscale_gfx942.cuh" // new -#include "gfx942/opus_gemm_traits_a16w16_gfx942.cuh" // new -``` - -The `_gfx942` suffix on every struct keeps the two arches' definitions -from clashing in the same TU. - -### 4. Arch routers in `opus_gemm.cu` - -Edit [`opus_gemm.cu`](opus_gemm.cu) — add the include and one `case` -per router: - -```cpp -#include "gfx950/opus_gemm_arch_gfx950.cuh" -#include "gfx942/opus_gemm_arch_gfx942.cuh" // new - -template -OpusA16W16NoscaleKernel opus_dispatch_a16w16(int M, int N, int K, int batch) -{ - switch (opus_get_gfx_arch()) { - case OpusGfxArch::Gfx950: - return opus_dispatch_a16w16_gfx950(M, N, K, batch); - case OpusGfxArch::Gfx942: // new - return opus_dispatch_a16w16_gfx942(M, N, K, batch); // new - default: { /* TORCH_CHECK with arch_info */ } - } -} +# OPUS GEMM C++ and code generation + +The public Python contract is documented in +[`aiter/ops/opus/README.md`](../../aiter/ops/opus/README.md). C++ keeps five +family launch ABIs. They are shared private implementation boundaries for the +Python `opus_gemm(..., kid=...)` and `opus_bmm(..., kid=...)` entries; the +public operation split does not duplicate C++ launchers or kernels. + +## Exact-id architecture + +Kernel identity is `(arch, logical family, kid, Y dtype)`. Python resolves a +bare final id through the merged `kernels_list`; C++ receives an already +resolved family call and performs strict lookup in the current architecture's +typed table. + +```text +caller final kid + -> strict 2D opus_gemm or batch-first 3D opus_bmm + -> Python canonical registry route and family adapter + -> family C++ entry + -> runtime architecture + output-dtype table + -> exact kid lookup + -> generated launcher checks ``` -Same edit for `opus_a16w16_tune_dispatch` (id-based router) and the -`a8w8` block (it currently `TORCH_CHECK`s on `arch == Gfx950`; widen -the check or move a8w8 to its own arch router). - -### 5. Codegen tables - -`gen_instances.py` keeps four arch-tagged tables that drive launcher -emission. Today they hard-code gfx950; the cleanest extension is to -make them dispatch on a per-`OpusGemmInstance` `arch` field: - -| Table | What it controls | gfx950 entry today | -|---|---|---| -| `PIPELINE_HEADER_MAP` | `#include "{pipeline_header}"` in each launcher TU | `gfx950/opus_gemm_pipeline_*_gfx950.cuh` | -| `TRAITS_NAME_MAP` | `using Traits = {traits_name}<...>` | `opus_gemm_*_traits_gfx950` | -| `KARGS_NAME_MAP` | `{kargs_name} kargs{};` | `opus_gemm_*_kargs_gfx950` | -| `KERNEL_FUNC_MAP` | `__global__` template name (unchanged across archs) | `gemm_*_kernel` | - -Two implementation options: +C++ does not choose a default kid, read a CSV, run a shape heuristic, redirect +an id, allocate a workspace, or fall back to another backend. -- **Per-arch dicts** (smallest surface change): add - `PIPELINE_HEADER_MAP_GFX942`, etc., and pick the right one inside the - `opus_gemm_codegen` methods based on the instance's arch. -- **Tagged keys**: keep one dict but key by `(arch, kid_tag)`. Cleaner - long-term; needs more adapter code. +The Python layer retains two distinct A16 shape-driven flows. The generic +`aiter.gemm_a16w16` dispatcher uses the global multi-backend tuned result and, +on a miss or invalid OPUS row, keeps its original skinny, gfx1250 Triton, or +PyTorch fallback. It does not run an OPUS heuristic. The OPUS-only +`gemm_a16w16_opus` compatibility entry instead uses an explicit id when +provided, otherwise tries a tuned row and falls back to its heuristic for a +missing or invalid row. All selections pass through legacy compatibility +resolution before the local exact launcher; the strict `opus_gemm`/`opus_bmm` +APIs never redirect. Reusable policy helpers live in +`aiter/ops/opus/policy.py`. -Also extend `OpusGemmInstance` (in `opus_gemm_common.py`) with an -`arch: str` field, and make the launcher symbol name include the arch -suffix when a launcher with the same tile already exists for another -arch — otherwise the manifest will see two prototypes with the same -function name. - -### 6. Python import-time guard - -Edit [`aiter/ops/opus/__init__.py`](../../aiter/ops/opus/__init__.py) -to widen the supported arch set: - -```python -_SUPPORTED = {"gfx950", "gfx942"} # new -``` - -The probe helper at -[`aiter/ops/opus/_arch.py`](../../aiter/ops/opus/_arch.py) is -non-raising: ``_detect_arch(supported)`` returns ``(ok, detected)`` so -the package can install stubs and emit a ``RuntimeWarning`` instead of -breaking ``from aiter.ops.opus import *`` (which sits inside the -swallow-ImportError ``try`` block in ``aiter/__init__.py``). Calling a -stub raises ``RuntimeError`` with the detected arch and the supported -set. ``_check_arch`` (raising variant) is still available for callers -that prefer hard failure. - -### 7. Tuning data - -If gfx942 needs its own tuned CSV (different tile choices), either: - -- Co-locate per-arch CSV files (e.g. - `aiter/ops/opus/configs/opus_gemm_a16w16_tuned_gfx942.csv`) and have - `gen_instances.py --tune_file` consume the right one for the active - arch (or both, baked into separate macros consumed by the per-arch - glue header). -- Keep one CSV with an `arch` column and filter at codegen time. - -### 8. Multi-arch wheel build - -The Layer-3 device-pass guard -(`#if defined(__gfx950__)` wrapping the kernel body) is per-arch and -already in place for gfx950. Add the matching guard at the top of each -new gfx942 kernel body so multi-arch wheels (e.g. -`GPU_ARCHS=gfx950;gfx942`) compile cleanly: +## Family entries ```cpp -__global__ void gemm_a16w16_kernel_gfx942(...) { -#ifdef __HIP_DEVICE_COMPILE__ -#if defined(__gfx942__) - /* real body */ -#else - /* empty stub: unreachable at runtime; here so other arches' device pass - does not try to instantiate gfx942-only intrinsics */ -#endif -#endif -} -``` - -### 9. Validation - -Run the standard regression on a gfx950 box: - -```bash -GPU_ARCHS=gfx950 python op_tests/test_opus_a16w16_gemm.py -m 128 -n 256 -k 1024 -b 1 -# CSV sweep (optional, if you have a shapes file): -GPU_ARCHS=gfx950 python op_tests/test_opus_a16w16_gemm.py --csv /path/to/shapes.csv +void opus_gemm_a16w16_launch( + aiter_tensor_t& XQ, + aiter_tensor_t& WQ, + aiter_tensor_t& Y, + std::optional bias, + std::optional workspace, + int kid, + int split_k); + +void opus_gemm_a8w8_launch( + aiter_tensor_t& XQ, + aiter_tensor_t& WQ, + aiter_tensor_t& Y, + int kid); + +void opus_gemm_a8w8_blockscale_launch( + aiter_tensor_t& XQ, + aiter_tensor_t& WQ, + aiter_tensor_t& Y, + aiter_tensor_t& x_scale, + aiter_tensor_t& w_scale, + int kid); + +void opus_gemm_a8w8_blockscale_bpreshuffle_launch( + aiter_tensor_t& XQ, + aiter_tensor_t& WQ, + aiter_tensor_t& x_scale, + aiter_tensor_t& w_scale, + aiter_tensor_t& Y, + int kid); + +void opus_gemm_a8w8_mxscale_bmm_launch( + aiter_tensor_t& XQ, + aiter_tensor_t& WQ, + aiter_tensor_t& Y, + aiter_tensor_t& x_scale, + aiter_tensor_t& w_scale, + std::optional workspace, + int kid, + int split_k); ``` -Then, on the new arch hardware, repeat with `GPU_ARCHS=gfx942` and -provide a tuned CSV (if any) to populate the lookup map. - -For a multi-arch wheel sanity check, do a full rebuild: - -```bash -rm -f aiter/jit/module_deepgemm_opus.so -AITER_REBUILD=1 GPU_ARCHS="gfx942;gfx950" python -c \ - "from aiter.ops.opus import gemm_a16w16_opus; print('ok')" +## Registry and capability + +| Family | gfx942 | gfx950 | gfx1250 | +|---|---|---|---| +| `a16w16` | direct + two-stage | direct + two-stage | two-stage + pre-built BF16 direct; fused source retained but unregistered | +| `a8w8` | empty | kid 2, FP32 Y | empty | +| `a8w8_blockscale` | empty | kid 1, FP32 Y | empty | +| `a8w8_blockscale_bpreshuffle` | kid 11000, BF16 Y | empty | empty | +| `a8w8_mxscale_bmm` | empty | 45 exact ids in 8000--8653, BF16/FP32 Y | empty | + +Empty tables are explicit capability states. The merged registry currently +contains 925 final ids, including 219 pre-built gfx1250 A16W16 CO ids. Those CO +ids currently occupy 21016--21315 inside the reserved `[21000,27000)` band. +The MXFP8 BMM ids are +`8000 + family_local_kid`, which places them in an unused global band while +preserving family-local tuning/debug correlation. Historical child-dictionary +collisions are resolved by the final merge; runtime routing always follows the +resulting `kernels_list[kid]` instance and never a numeric interval. + +The gfx942 BF16-workspace A16 exact kids (`10210`, `10213`, `10216`) are the +one workspace-output exception: their current exact-N reducer requires BF16 +`Y`. The canonical Python registry rejects FP32 `Y` before launch, matching the +generated host guard. + +## Generated tables + +Generated roots are: + +```text +opus_gemm_a16w16_kid_dispatch.h +opus_gemm_a8w8_kid_dispatch.h +opus_bmm_mxscale_kid_dispatch.h +opus_gemm_manifest.h +opus_build_archs.h ``` -The build must finish without errors; both `--offload-arch=gfx950` and -`--offload-arch=gfx942` must appear in the hipcc invocation; and the -runtime call on whichever device the host machine has must produce -correct results (no clean way to cross-test on a single-arch host). - -## Splitk workspace and TBO (two-batch overlap) - -The splitk kid family (kid 200..299, a16w16_flatmm_splitk) writes per-split -partials into an fp32 scratch buffer that the reduce kernel folds into the -final output. That buffer is owned host-side as a stable -`opus_splitk_ws_handle*` slot: the kernel reads `slot->ptr` at launch, so -captured HIP graphs bake in the slot address, not the raw buffer. The -launcher grows `slot->ptr` lazily (4 MiB-aligned), draining outstanding work -with `hipDeviceSynchronize` before freeing the old buffer. - -### Ownership: per-stream, not per-thread - -The slot is registered in a **process-global, mutex-protected map keyed by -`hipStream_t`** (see `opus_splitk_ws_get` in `opus_gemm.cu`). vLLM/sglang- -style TBO drives two CUDA streams from two CPU threads; each captured graph -must bake in its own buffer pointer so concurrent replays write disjoint -scratch. A per-thread cache would either share one buffer between the two -streams (corrupting concurrent replays) or fail the in-capture grow guard on -the second thread's first call. - -### Framework usage - -```python -import aiter - -compute = torch.cuda.Stream() -comm = torch.cuda.Stream() - -with torch.cuda.stream(compute): - aiter.opus_gemm_workspace_init() # register handle for this stream - _ = gemm_a16w16_opus(A_max, B_max) # warm: grow buffer to max size - # ... capture compute graph ... - -with torch.cuda.stream(comm): - aiter.opus_gemm_workspace_init() - _ = gemm_a16w16_opus(A_max, B_max) - # ... capture comm graph ... - -# Replay both graphs concurrently from two threads -- each sees its own -# workspace; no aborts, no cross-stream interference. -``` - -Rules: - -* Call `opus_gemm_workspace_init()` once per TBO stream, **eagerly** (outside - HIP graph capture). Calling it during capture raises; capture cannot run - `hipHostMalloc`. -* Before capture, run the largest expected splitk shape on that stream once, - so the buffer grows to its final size. Grow inside capture is illegal - (`hipMalloc` / `hipFree` are stream-capture-illegal) and aborts with a - message pointing back here. -* Single-stream, single-thread (non-TBO) callers do **not** need to call - `opus_gemm_workspace_init` -- the registry lazy-creates a handle on the - first eager call on each new stream. Init only becomes mandatory when the - first call on a stream is inside HIP graph capture. - -### When the buffer is freed - -Not automatically during steady-state. Each registered stream holds one -`opus_splitk_ws_handle` plus its current `hipMalloc` buffer; on grow the old -buffer is `hipFree`d (after a device sync) before the larger one is allocated, -so grow never leaks. Streams the torch CUDAStream pool reuses map to the same -entry; streams the pool never reclaims would otherwise retain their handle + -buffer until process exit. - -For explicit teardown, call one of: - -* `opus_gemm_workspace_release()` -- frees the buffer, host/device handles and - registry entry for the **current** stream (call inside - `with torch.cuda.stream(s):`). No-op if the stream was never registered. -* `opus_gemm_workspace_release_all()` -- frees every registered stream's - workspace and clears the registry. - -Both must run in eager mode (frees are stream-capture-illegal); `_release` -synchronizes the target stream and `_release_all` does a device sync before -freeing so no in-flight kernel references a buffer being freed. +A16 tables separate direct BF16/FP32 launchers from workspace launchers. A8 +tables are family and output-dtype scoped. Every macro has a `_SIZE`; an empty +capability produces `std::array` without referencing a missing +launcher. + +Full canonical A16 counts are: + +| Architecture | Direct BF16 | Direct FP32 | Workspace | +|---|---:|---:|---:| +| gfx942 | 14 | 1 | 8 | +| gfx950 | 92 | 92 | 48 | +| gfx1250 | 219 | 0 | 496 | + +`gen_instances.py` treats tuned CSV ids, the sidecar, the per-architecture +default compile floor, and mandatory A8 ids as build availability. It emits no +runtime shape table. All 45 gfx950 MXFP8 BMM ids are emitted as one family and +deduplicated by generated symbol name rather than entering the ordinary +per-kid subset. All available gfx1250 CO ids are in the gfx1250 compile floor; +codegen emits their five-argument host launchers but no device translation +units. The device bodies come from `gen_co/gfx1250/.co`. + +## A16 workspace checks + +Torch owns every workspace Tensor. Generated launchers validate the final +launch inputs after architecture-specific split resolution: + +- XQ/WQ/Y shape, dtype, stride and batch rules; +- exact instance workspace dtype; +- same device, contiguous storage and 16-byte alignment; +- overflow-checked extent and byte-span arithmetic; +- sufficient capacity for the final effective split; +- exact-kid bias support. + +Two-stage layouts are split-major. gfx1250 exact kids currently use BF16 +workspace storage; the generated launcher/reducer ABI remains typed for either +BF16 or FP32. C++ never owns or retains a Tensor or pointer. + +The gfx1250 TDM pipelines use the policy-tag, element-unit API. Clusterlaunch +rounds only the physical grid to `(cluster_wg_m, cluster_wg_n)` multiples; +surplus workgroups arrive at the required cluster barrier and leave through the +uniform `tile_oob` path. Logical tile counts and workspace strides remain +unrounded. The separate reducer dispatches runtime split-K to compile-time +specializations `SPLIT_K_=1..16`, with `SPLIT_K_=0` as the runtime fallback, +using the VEC=8/BLOCK=128 geometry. + +The fused gfx1250 factory, emitter and device pipeline remain in-tree for repair, +but `GFX1250_SPLITK_FUSE_ENABLED` is `False`. No fused kid is registered, the +unified capability tables cannot return one, and its `[27000,30000)` band is +unclaimed. The preceding `[21000,27000)` band is reserved for CO ids. + +gfx942 continues to wave-uniformize both halves of the direct 64-bit workspace +pointer with `__builtin_amdgcn_readfirstlane` in main and reduce kernels. + +## A8 input checks + +The family router owns common device/dtype checks. Generated exact-instance +launchers own tile and storage details: + +- gfx950 no-scale kid 2: matching 3D FP8 inputs, contiguous FP32 output and + valid K-loop depth/parity; +- gfx950 blockscale kid 1: the same tensors plus contiguous FP32 1x128x128 + scales and exact scale shapes; +- gfx942 bpreshuffle kid 11000: batch one, BF16 output, exact 128-wide N/K + tiles, registered scale layouts and truly pre-shuffled WQ content. + +MXFP8 BMM is gfx950-only. `opus_bmm.cu` first applies the shared FP8/E8M0 +shape, stride, device and output checks, then performs an exact lookup in +`opus_bmm_mxscale_kid_dispatch.h`. Unknown ids fail immediately. Generated +launchers enforce their own M/tile/K restrictions; they never redirect to kid +8000 or another family. + +For two-stage BMM split-K, the caller supplies a direct FP32 partial-buffer +pointer. Fused split-K stores partials and aligned tile counters in the same +caller Tensor. The reduce kernel also receives the direct pointer. No BMM +launcher allocates, frees, registers or retains workspace memory. + +For global kid 8326 (family-local kid 326), codegen sets +`PRELOAD_SF_LDS=false` only on the `split_k > 1`, `D_OUT=void` workspace +specialization that writes partial sums. Its direct BF16/FP32 +`split_k == 1` specializations keep `PRELOAD_SF_LDS=true`. + +## Source layout + +| Path | Role | +|---|---| +| `opus_bmm.cu` / `include/opus_bmm.h` | MXFP8 BMM exact-kid family entry and Torch-workspace forwarding | +| `opus_gemm_common.py` | canonical registry, unique route map and compile-floor constants | +| `gen_instances.py` | subset selection, manifests and typed dispatch generation | +| `codegen/gen_instances_gfx*.py` | exact-instance host launchers and generated input checks | +| `gen_co/` | offline CO manifest/builder, build metadata and packaged gfx1250 ELF images | +| `include/gfx950/opus_bmm_*` | MXFP8 BMM traits, launchers and pipelines | +| `include/gfx1250/opus_co_launch_gfx1250.cuh` | first-use CO loader and cluster launcher | +| `include/gfx*/opus_gemm_arch_*.cuh` | sorted exact-kid tables | +| `include/gfx*/**/opus_gemm_traits*.cuh` | kernel arguments and traits | diff --git a/csrc/opus_gemm/codegen/common.py b/csrc/opus_gemm/codegen/common.py index 47b684e194..9d7abc053e 100644 --- a/csrc/opus_gemm/codegen/common.py +++ b/csrc/opus_gemm/codegen/common.py @@ -30,27 +30,22 @@ ) + _NOSPLIT ) -# Pre-compiled (.co) a16w16 families: device code comes from an offline build -# (csrc/opus_gemm/gen_co), not from the JIT. They are a16w16 tags like any other -# -- same input dtypes, same tune-lookup machinery -- but they get their own -# launcher signature and their own dispatch table, because none of them needs -# the workspace argument the gfx1250 split-K families carry. A future split-K -# .co family would drop out of this tuple and back into the workspace ABI. _A16W16_CO_TAGS = ("a16w16_4wave_co", "a16w16_4wave_wl_co") - _A16W16_TAGS = ( "a16w16", "a16w16_flatmm", "a16w16_flatmm_splitk", "a16w16_persistent", "a16w16_mono_tile", - # gfx1250 cluster/TDM split-K (fp32 workspace + reduce kernel). + # gfx1250 cluster/TDM split-K (exact-kid typed workspace + reduce kernel). "a16w16_cluster_tdm_splitk_ws", - # gfx1250 CLUSTER-LAUNCH (multicast) TDM split-K (fp32 workspace + reduce). + # gfx1250 CLUSTER-LAUNCH (multicast) TDM split-K (typed workspace + reduce). "a16w16_clusterlaunch_tdm_splitk_ws", - # gfx1250 FUSED single-kernel in-cluster split-K reduce (no reduce kernel); - # B multicast + GL2-resident partial workspace + cluster-barrier sync. + # gfx1250 fused in-cluster reduction source (currently unregistered). When + # enabled it consumes caller-owned typed workspace without a second reduce. "a16w16_clusterlaunch_tdm_splitk_fuse", + # Pre-built gfx1250 device images. Their generated host launchers reuse the + # ordinary five-argument non-workspace exact-kid ABI. *_A16W16_CO_TAGS, ) + _GFX942_A16W16_TAGS @@ -61,6 +56,23 @@ # default maps. ARCH_MAP_REGISTRY = {} +_SPLITK_WORKSPACE_TYPES = { + "bf16_t": ("bf16_t", "__bf16", "AITER_DTYPE_bf16"), + "fp32_t": ("fp32_t", "float", "AITER_DTYPE_fp32"), +} + + +def splitk_workspace_type(k): + """Return C++ storage, pointer, and Aiter dtype tokens declared by a kid.""" + dtype = k.splitk_workspace_dtype + try: + return _SPLITK_WORKSPACE_TYPES[dtype] + except KeyError as exc: + raise ValueError( + f"workspace instance {getattr(k, 'name', '')} must declare " + f"splitk_workspace_dtype as bf16_t or fp32_t, got {dtype!r}" + ) from exc + def register_arch_map(arch, map_name, mapping): key = (arch, map_name) @@ -76,7 +88,7 @@ def get_arch_map(arch, map_name): def kid_arch(k): """Resolve a kid's target arch_prefix (defaults to gfx950 for legacy kids).""" - return (getattr(k, "arch_prefix", "") or "gfx950").lower() + return (k.arch_prefix or "gfx950").lower() def register_emit(arch, kernel_tag, fn): diff --git a/csrc/opus_gemm/codegen/gen_instances_gfx1250.py b/csrc/opus_gemm/codegen/gen_instances_gfx1250.py index bdd58a3b9c..353969a2f0 100644 --- a/csrc/opus_gemm/codegen/gen_instances_gfx1250.py +++ b/csrc/opus_gemm/codegen/gen_instances_gfx1250.py @@ -1,21 +1,11 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2025-2026, Advanced Micro Devices, Inc. All rights reserved. -"""gfx1250 codegen -- emit launchers for gfx1250-targeted kid families. - -Wires the a16w16 cluster/TDM split-K pipeline that reduces via an fp32 -WORKSPACE + a separate REDUCE kernel (no atomic_add), mirroring the gfx950 -flatmm-splitk launcher (workspace + main kernel + reduce -kernel). The main kernel is always instantiated (it writes the fp32 -workspace); the reduce kernel casts the fp32 partials to the runtime Y dtype -(bf16 / fp32) and folds bias once. - -Self-registers each emit into codegen.common.EMIT_REGISTRY at import time. -""" +"""Generate gfx1250 OPUS A16W16 launchers.""" import os from pathlib import Path -from codegen.common import register_arch_map, register_emit +from codegen.common import register_arch_map, register_emit, splitk_workspace_type # ---------------- gfx1250 arch-override maps ---------------- @@ -29,11 +19,9 @@ "a16w16_clusterlaunch_tdm_splitk_fuse": ( "gfx1250/opus_gemm_pipeline_a16w16_clusterlaunch_tdm_splitk_fuse_gfx1250.cuh" ), - # 4wave_co has a real pipeline header, but NOTHING in the JIT build may - # include it: it uses pin builtins a release toolchain does not have, and - # its device code comes from a pre-built .co instead. Point the entry at the - # traits header so any generic consumer of this map stays on safe ground -- - # the co emit below never uses it at all. + # CO device bodies are built offline. Point generic codegen lookup at the + # release-toolchain-safe traits header; the host-only emitter never includes + # either pipeline header. "a16w16_4wave_co": "gfx1250/opus_gemm_traits_a16w16_gfx1250.cuh", "a16w16_4wave_wl_co": "gfx1250/opus_gemm_traits_a16w16_gfx1250.cuh", } @@ -50,9 +38,6 @@ "a16w16_cluster_tdm_splitk_ws": "gemm_a16w16_cluster_tdm_splitk_ws_kernel_gfx1250", "a16w16_clusterlaunch_tdm_splitk_ws": "gemm_a16w16_clusterlaunch_tdm_splitk_ws_kernel_gfx1250", "a16w16_clusterlaunch_tdm_splitk_fuse": "gemm_a16w16_splitk_fuse_kernel_gfx1250", - # Device-side body name. The JIT never instantiates it (the .co does); it is - # here because gen_instance() resolves this map for every kid, and because - # build_co.py reads it to spell the stub TU's call. "a16w16_4wave_co": "gemm_a16w16_4wave_compute_body_gfx1250", "a16w16_4wave_wl_co": "gemm_a16w16_4wave_wl_body_gfx1250", } @@ -74,22 +59,17 @@ } -# 4wave_co traits argument list, shared by the JIT launcher emit below and the -# offline stub TU (gen_co/build_co.py imports this). One spelling, so the .co's -# baked traits and the host launcher's tile constants cannot disagree. def co_traits_args(k): + """Return the traits arguments shared by the offline and host generators.""" d_a, d_b, d_c, d_acc = k.co_dtypes - return ( + args = ( f"{k.BLOCK_SIZE}, {k.B_M}, {k.B_N}, {k.B_K}, {k.num_slots}, " f"{d_a}, {d_b}, {d_c}, {d_acc}, " f"{k.cluster_wg_m}, {k.cluster_wg_n}" - # The wave-layout family takes two more: how the 4 waves tile the block. - + ( - f", {k.co_wave_layout[0]}, {k.co_wave_layout[1]}" - if k.kernel_tag == "a16w16_4wave_wl_co" - else "" - ) ) + if k.kernel_tag == "a16w16_4wave_wl_co": + args += f", {k.co_wave_layout[0]}, {k.co_wave_layout[1]}" + return args # fuse workspace storage dtype -> (C type, byte size) for the fuse kernel instantiation. @@ -97,21 +77,20 @@ def co_traits_args(k): def splitk_reduce_extra_device_instantiations(): - # gfx1250 only: fp32 bias with a bf16 output (D_OUT=__bf16, D_BIAS=float). - # The main kernel writes the partial workspace, so an fp32 bias folds in fp32 in - # the reduce before the cast to bf16. The baseline instantiations cover the - # matched-dtype cases; this adds the bf16-out + fp32-bias mix. Emitted for - # every compile-time split_k (0=runtime fallback, 1..16=unrolled) and - # HAS_OOB, and for BOTH partial types -- which one a kid uses is its - # splitk_workspace_dtype, so both have to exist. Same kernel NAME/ABI. - out = "// fp32-bias + bf16-out (gfx1250 f32 bias support), per split_k + D_WS\n" - for d_ws in ("__bf16", "float"): - for has_oob in ("true", "false"): - for sk in range(17): + # The shared generator emits matched output/bias combinations for both + # physical workspace types and every split_k specialization. gfx1250 also + # accepts fp32 bias with bf16 output, so add that mixed combination here. + out = "// gfx1250 fp32-bias + bf16-output reduce variants\n" + for has_oob in ("true", "false"): + for split_k in range(17): + for workspace_type in ("__bf16", "float"): out += ( - f"template __global__ void splitk_reduce_kernel_gfx1250<8, 128, __bf16, true, float, {has_oob}, {sk}, {d_ws}>(\n" + "template __global__ void " + "splitk_reduce_kernel_gfx1250<" + f"8, 128, __bf16, true, float, {has_oob}, " + f"{split_k}, {workspace_type}>(\n" " const void*, __bf16*, int, int, int, int, int, int,\n" - " const float*, int);\n" + " const float*, int);\n" ) return out @@ -147,14 +126,13 @@ def gen_cluster_tdm_splitk_ws_instance( BIAS_HOST_VALIDATE="", **_unused, ): - """gfx1250 a16w16 TDM split-K (workspace + reduce) launcher emit. - - NO-CLUSTER grid: grid = (M/B_M, N/B_N, split_k); each WG owns one - B_M x B_N tile (so M %% B_M == 0, N %% B_N == 0). The main kernel writes - its split's fp32 partial into ws[split, padded_M, padded_N]; the reduce - kernel sums split_k slices, folds bias, casts to Y dtype. batch handled by - a per-batch host launch (sequential on stream -> workspace reuse is safe). - """ + """Emit a checked gfx1250 two-stage split-K launcher.""" + workspace_dtype, workspace_ptr_type, workspace_aiter_dtype = splitk_workspace_type( + k + ) + # The final #4246 reducer uses the same coalesced VEC=8/BLOCK=128 geometry + # for either physical workspace type. + reduce_vec, reduce_bs = 8, 128 layout_int = _LAYOUT_INT[getattr(k, "ctdm_layout", "tileN")] has_oob_str = "true" if k.has_oob else "false" enable_bias_str = "true" if getattr(k, "enable_bias", False) else "false" @@ -187,31 +165,24 @@ def gen_cluster_tdm_splitk_ws_instance( if is_clusterlaunch else "" ) - # Cluster round-up emitted before the grid launch: the runtime rejects a grid - # that is not a whole number of clusters. The surplus workgroups own no tile and - # return right after their one cluster-barrier arrival (tile_oob in the pipeline), - # so any (M, N) is launchable with any cluster dims -- no exact-fill assert. + # A cluster launch grid must contain whole clusters. Round only the physical + # launch grid up; logical tile counts and workspace strides remain unrounded. cluster_grid_roundup = "" grid_m_expr, grid_n_expr = "num_tiles_m", "num_tiles_n" if is_clusterlaunch: cluster_grid_roundup = ( - f" // CLUSTER-LAUNCH: the grid must be a whole number of " - f"{cwm}x{cwn} clusters, so\n" - f" // round the tile counts up. The surplus workgroups have no tile and " - f"leave at\n" - f" // the pipeline's tile_oob exit; the workspace strides below stay on " - f"the\n" - f" // UNROUNDED counts, so the reduce kernel is unaffected by the " - f"padding.\n" + f" // Round the physical grid to complete {cwm}x{cwn} clusters.\n" + f" // Surplus WGs take the pipeline tile_oob exit; workspace layout\n" + f" // continues to use the unrounded logical tile counts.\n" f" int grid_tiles_m = (num_tiles_m + {cwm} - 1) / {cwm} * {cwm};\n" f" int grid_tiles_n = (num_tiles_n + {cwn} - 1) / {cwn} * {cwn};\n" ) grid_m_expr, grid_n_expr = "grid_tiles_m", "grid_tiles_n" # gfx1250-specific bias validation (does NOT use the shared BIAS_HOST_VALIDATE, - # which forces bias.dtype == Y.dtype). The main kernel always writes an fp32 - # workspace and the reduce kernel folds bias in fp32 before the final cast to - # Y, so an fp32 bias is exact for ANY Y dtype (bf16 or fp32). We therefore + # which forces bias.dtype == Y.dtype). The reduce kernel folds bias into its + # fp32 accumulator before the final cast to Y, regardless of workspace + # storage, so an fp32 bias is exact for ANY Y dtype (bf16 or fp32). We therefore # accept bias.dtype in {{fp32, Y.dtype}} and record bias_is_fp32_ so the reduce # launch below can pick the matching D_BIAS template. (Double C++ braces are # intentional -- this string is inserted verbatim into the f-string template.) @@ -253,7 +224,7 @@ def gen_cluster_tdm_splitk_ws_instance( using {k.name}_Traits = {traits_name}<{k.BLOCK_SIZE}, {k.B_M}, {k.B_N}, {k.B_K}, {layout_int}, - {da}, {db}, D_C, fp32_t, + {da}, {db}, {workspace_dtype}, fp32_t, {enable_bias_str}, {num_slots}, {wg_per_cu}{cluster_traits_args}>; """ @@ -264,6 +235,7 @@ def gen_cluster_tdm_splitk_ws_instance( #if !defined(__HIP_DEVICE_COMPILE__) && !defined(__HIPCC_RTC__) #include "aiter_tensor.h" #include "aiter_stream.h" +#include "opus_gemm_common.cuh" #include #endif #ifdef OPUS_FUSED_HOST_TU @@ -279,10 +251,8 @@ def gen_cluster_tdm_splitk_ws_instance( #endif {traits_aliases} #if !defined(__HIP_DEVICE_COMPILE__) && !defined(__HIPCC_RTC__) -// Reduce kernel forward-decl + split_k -> compile-time-instance launch -// dispatcher (opus_splitk_reduce_launch_gfx1250). The reduce kernel definition -// lives in gfx1250/splitk_reduce_gfx1250.cuh; explicit instantiations (per -// SPLIT_K + D_WS) live in the dedicated splitk_reduce_gfx1250.device.cu TU. +// Host launch helper dispatches runtime split_k to the matching compile-time +// reducer specialization. The device definitions remain in the per-arch TU. #include "gfx1250/splitk_reduce_launch_gfx1250.cuh" template @@ -295,21 +265,8 @@ def gen_cluster_tdm_splitk_ws_instance( std::optional bias, int splitK) {{{{ - // D_C is the split-K PARTIAL type (this kernel stores it, the reduce below - // reads it and casts to Y), picked per kid by splitk_workspace_dtype. Y is - // independent and may be bf16 or fp32 either way. - static_assert(std::is_same::value || std::is_same::value, - "cluster_tdm_splitk_ws split-K partial must be fp32_t or bf16_t"); - - // The host sizes this buffer from the kid table's splitk_workspace_dtype - // while D_C was baked in at build time, so a .so that is stale with respect - // to the table makes the two disagree. Catch it here: unchecked, a narrower - // buffer than D_C is a GPU page fault with no hint of where it came from. - AITER_CHECK(workspace.element_size() == sizeof(D_C), - "split-K workspace is ", workspace.element_size(), - "-byte but this kernel stores ", sizeof(D_C), - "-byte partials -- rebuild module_deepgemm_opus after changing " - "splitk_workspace_dtype"); + static_assert(std::is_same::value, + "cluster_tdm_splitk_ws uses the fp32 launch-dispatch specialization"); int batch = XQ.size(0); int M = XQ.size(1); @@ -321,7 +278,7 @@ def gen_cluster_tdm_splitk_ws_instance( // M / N need NOT be multiples of B_M / B_N: the grid is padded to // ceil(M/B_M) x ceil(N/B_N) tiles, the main kernel TDM-clamps OOB global // reads to the real (M, N) tensor extents (tensor_dim1 = m - tile_row / - // n - tile_col), partials for padded rows/cols land in the padded fp32 + // n - tile_col), partials for padded rows/cols land in the padded typed // workspace, and the reduce kernel only iterates m in [0, M) and writes // n in [0, N) (HAS_OOB tail). So M=49 transparently runs as a padded // M=64 tile, etc. @@ -329,6 +286,9 @@ def gen_cluster_tdm_splitk_ws_instance( "K=", K, " must be even (a16w16 family rejects odd K)"); AITER_CHECK(M >= 1 && N >= 1 && K >= 1 && batch >= 1, "M, N, K, batch must be >= 1"); + AITER_CHECK(batch == 1, + "gfx1250 cluster_tdm_splitk_ws supports batch == 1 only; got batch=", + batch); {gfx1250_bias_validate} using Traits = {k.name}_Traits; @@ -342,32 +302,41 @@ def gen_cluster_tdm_splitk_ws_instance( split_k--; }}}} - int num_tiles_m = (M + {k.B_M} - 1) / {k.B_M}; - int num_tiles_n = (N + {k.B_N} - 1) / {k.B_N}; - int padded_M = num_tiles_m * {k.B_M}; - int padded_N = num_tiles_n * {k.B_N}; - + // The reducer uses one logical row per grid.y block. Reject stale tuned + // rows and explicit kids before either kernel can launch an invalid grid. + AITER_CHECK(M <= {k.max_m}, + "{k.name}: split-K reduce requires M <= {k.max_m}; got M=", M); + + int num_tiles_m = 1 + (M - 1) / {k.B_M}; + int num_tiles_n = 1 + (N - 1) / {k.B_N}; + const size_t padded_M_size = opus_checked_extent_product( + {{static_cast(num_tiles_m), static_cast({k.B_M})}}, + "{k.name}"); + const size_t padded_N_size = opus_checked_extent_product( + {{static_cast(num_tiles_n), static_cast({k.B_N})}}, + "{k.name}"); + const size_t workspace_slice_numel = opus_checked_extent_product( + {{padded_M_size, padded_N_size}}, "{k.name}"); + AITER_CHECK(padded_M_size <= static_cast(std::numeric_limits::max()) + && padded_N_size <= static_cast(std::numeric_limits::max()) + && workspace_slice_numel <= static_cast(std::numeric_limits::max()), + "{k.name}: padded workspace extents exceed 32-bit kernel stride limits"); + int padded_M = static_cast(padded_M_size); + int padded_N = static_cast(padded_N_size); + + // One-batch layout: [split_k, padded_M, padded_N]. + const size_t required_numel = opus_checked_extent_product( + {{static_cast(split_k), workspace_slice_numel}}, + "{k.name}"); + void* workspace_ptr_ = opus_validate_workspace( + workspace, XQ, {workspace_aiter_dtype}, required_numel, 16, "{k.name}"); auto stream = aiter::getCurrentHIPStream(); - void* ws_ptr_ = workspace.data_ptr(); {cluster_grid_roundup} dim3 grid_main({grid_m_expr}, {grid_n_expr}, split_k); dim3 block_main({k.BLOCK_SIZE}); - // VEC=8 -> each lane owns one dwordx4 of bf16 so the wave stores 512B fully - // contiguous with no cross-lane shuffle (100% write-transaction efficiency), - // and the fp32 workspace load drops from a 64B to a 32B lane stride. BLOCK=128 - // (4 waves) is the tuned reduce block; grid.x = ceil(N, VEC*BLOCK) is unchanged - // vs the old VEC=16/BS=64 (both 1024 N per block). - constexpr int REDUCE_VEC = 8; - constexpr int REDUCE_BS = 128; - // The reduce carries one row per grid.y block, and grid.y stops at 65535. - // Past that it does not fail -- it writes garbage (NaN at M=65536, measured) - // until M is large enough that the launch itself is rejected. The tuner will - // not pick a split-K kid here, so reaching this means a stale tuned CSV or an - // explicit kernelId; say so rather than returning wrong numbers. - AITER_CHECK(M <= 65535, - "split-K reduce puts one row per grid.y block and grid.y is capped at " - "65535; M=", M, " needs a non-split-K kernel"); + constexpr int REDUCE_VEC = {reduce_vec}; + constexpr int REDUCE_BS = {reduce_bs}; dim3 grid_reduce((N + REDUCE_VEC * REDUCE_BS - 1) / (REDUCE_VEC * REDUCE_BS), M, 1); dim3 block_reduce(REDUCE_BS); @@ -379,7 +348,7 @@ def gen_cluster_tdm_splitk_ws_instance( {kargs_name} kargs{{{{}}}}; kargs.ptr_a = XQ.data_ptr(); kargs.ptr_b = WQ.data_ptr(); - kargs.ptr_ws = workspace.data_ptr(); + kargs.ptr_ws = workspace_ptr_; kargs.ptr_c = Y.data_ptr(); kargs.ptr_bias = ptr_bias_; kargs.m = M; kargs.n = N; kargs.k = K; kargs.batch = 1; kargs.split_k = split_k; @@ -389,15 +358,13 @@ def gen_cluster_tdm_splitk_ws_instance( kargs.stride_c = N; kargs.stride_a_batch = XQ.stride(0); kargs.stride_b_batch = WQ.stride(0); - kargs.stride_ws_batch = padded_M * padded_N; + kargs.stride_ws_batch = static_cast(workspace_slice_numel); kargs.stride_c_batch = M * N; kargs.stride_bias_batch = stride_bias_batch_; {kernel_func}<<>>(kargs); - // Reduce reads the split-K workspace the main kernel wrote. D_C is that - // partial type for BOTH ends -- the main kernel stores Traits::DataC and - // this reads the same D_C -- so a per-kid choice cannot desynchronise them, + // Reduce reads the exact kid's typed partials (D_WS={workspace_ptr_type}), // re-accumulates in fp32, folds bias, casts to Y dtype. split_k is dispatched // to a compile-time (unrolled) reduce instance by the launch helper. if (Y.dtype() == AITER_DTYPE_bf16) {{{{ @@ -405,31 +372,31 @@ def gen_cluster_tdm_splitk_ws_instance( if (ptr_bias_ && bias_is_fp32_) {{{{ // fp32 bias + bf16 output: fold the exact fp32 bias in the // reduce (D_BIAS=float), then cast the fp32 sum to bf16. - opus_splitk_reduce_launch_gfx1250( + opus_splitk_reduce_launch_gfx1250( grid_reduce, block_reduce, stream, - ws_ptr_, y_ptr, split_k, M, N, 1, padded_M, padded_N, + workspace_ptr_, y_ptr, split_k, M, N, 1, padded_M, padded_N, reinterpret_cast(ptr_bias_), stride_bias_batch_); }}}} else if (ptr_bias_) {{{{ - opus_splitk_reduce_launch_gfx1250( + opus_splitk_reduce_launch_gfx1250( grid_reduce, block_reduce, stream, - ws_ptr_, y_ptr, split_k, M, N, 1, padded_M, padded_N, + workspace_ptr_, y_ptr, split_k, M, N, 1, padded_M, padded_N, reinterpret_cast(ptr_bias_), stride_bias_batch_); }}}} else {{{{ - opus_splitk_reduce_launch_gfx1250( + opus_splitk_reduce_launch_gfx1250( grid_reduce, block_reduce, stream, - ws_ptr_, y_ptr, split_k, M, N, 1, padded_M, padded_N, nullptr, 0); + workspace_ptr_, y_ptr, split_k, M, N, 1, padded_M, padded_N, nullptr, 0); }}}} }}}} else {{{{ float* y_ptr = reinterpret_cast(Y.data_ptr()); if (ptr_bias_) {{{{ - opus_splitk_reduce_launch_gfx1250( + opus_splitk_reduce_launch_gfx1250( grid_reduce, block_reduce, stream, - ws_ptr_, y_ptr, split_k, M, N, 1, padded_M, padded_N, + workspace_ptr_, y_ptr, split_k, M, N, 1, padded_M, padded_N, reinterpret_cast(ptr_bias_), stride_bias_batch_); }}}} else {{{{ - opus_splitk_reduce_launch_gfx1250( + opus_splitk_reduce_launch_gfx1250( grid_reduce, block_reduce, stream, - ws_ptr_, y_ptr, split_k, M, N, 1, padded_M, padded_N, nullptr, 0); + workspace_ptr_, y_ptr, split_k, M, N, 1, padded_M, padded_N, nullptr, 0); }}}} }}}} }}}} @@ -437,20 +404,9 @@ def gen_cluster_tdm_splitk_ws_instance( """ Path(os.path.join(cg.impl_path, f"{k.name}.cuh")).write_text(INSTANCE_IMPL) - # The kid's template slot is the split-K PARTIAL type (traits D_C), not its - # output dtype -- the reduce decides the output. Instantiate the partial type - # the kid declares, since the dispatch table now references it by that type - # and the host sizes the workspace from the same field. The fuse family keeps - # output_dtypes: it reduces in-kernel and never lands a partial for us. - _ws_partial_tags = ( - "a16w16_cluster_tdm_splitk_ws", - "a16w16_clusterlaunch_tdm_splitk_ws", - ) - if k.kernel_tag in _ws_partial_tags: - inst_dtypes = (getattr(k, "splitk_workspace_dtype", "fp32_t"),) - else: - inst_dtypes = tuple(k.output_dtypes) - for CDtype in inst_dtypes: + # The token is the host launch-dispatch specialization. The physical + # workspace type is independently embedded in the Traits alias above. + for CDtype in k.output_dtypes: host_decl = ( f"template void\n" f"{k.name}<{CDtype}>(\n" @@ -486,44 +442,32 @@ def gen_splitk_fuse_instance( BIAS_HOST_VALIDATE="", **_unused, ): - """gfx1250 FUSED single-kernel in-cluster split-K reduce launcher emit. - - No reduce kernel: the last split WG folds bias + reduces the partials in-kernel - (cluster-barrier sync) and writes C directly. The kernel is templated on - ; SplitK / MClusterWg are compile- - time (cluster dims), so each kid bakes one (tile, split_k, m_cluster, ws_dtype) - combo. The launcher (instantiated for the split-K lookup ABI) picks - D_OUT = __bf16 / float at runtime from Y.dtype. Requires M %% B_M == 0, - N %% B_N == 0 (no OOB C-store mask), ceil(M/B_M) %% MClusterWg == 0, and a - compile-time SplitK with no empty trailing K-slice for the runtime K. - """ + """Emit the fused gfx1250 split-K launcher and workspace checks.""" + del BIAS_HOST_VALIDATE + workspace_dtype, workspace_ptr_type, workspace_aiter_dtype = splitk_workspace_type( + k + ) layout_int = _LAYOUT_INT[getattr(k, "ctdm_layout", "tileN")] enable_bias_str = "true" if getattr(k, "enable_bias", False) else "false" num_slots = getattr(k, "num_slots", 3) wg_per_cu = getattr(k, "wg_per_cu", 2) - split_k = getattr(k, "fuse_split_k", 2) - # fuse_m_cluster field holds the cluster's 2nd-dim WG count; for this pipeline - # it groups N-tile peers (cluster.y, A-multicast), so expose it as n_cluster. - n_cluster = getattr(k, "fuse_m_cluster", 1) - ws_dtype = getattr(k, "fuse_ws_dtype", "bf16_t") - ws_ctype, _ws_bytes_elem = _FUSE_WS_CTYPE[ws_dtype] - - # Traits: 11-arg form (default cluster dims; the fuse kernel drives its own - # __cluster_dims__(SplitK, MClusterWg, 1) and only uses the traits for tile - # geometry / WindowA/B, not the traits cluster args). + split_k = int(k.fuse_split_k) + # Historical field name retained for compatibility; physically cluster.y + # groups N-tile peers sharing A. + n_cluster = int(k.fuse_m_cluster) + if split_k < 2: + raise ValueError(f"fused instance {k.name} must declare fuse_split_k >= 2") + traits_aliases = f""" template using {k.name}_Traits = {traits_name}<{k.BLOCK_SIZE}, {k.B_M}, {k.B_N}, {k.B_K}, {layout_int}, - {da}, {db}, D_C, fp32_t, + {da}, {db}, {workspace_dtype}, fp32_t, {enable_bias_str}, {num_slots}, {wg_per_cu}>; """ - # Host expansion of __cluster_dims__ (the fused HOST TU includes hip_runtime.h, - # not hip_minimal, so the attribute macro is otherwise not in scope -> the - # launch would not form the cluster -> multicast + cluster barrier stall). cluster_dims_host_def = ( "#ifndef __cluster_dims__\n" "#define __cluster_dims__(...) __attribute__((cluster_dims(__VA_ARGS__)))\n" @@ -536,16 +480,13 @@ def gen_splitk_fuse_instance( #if !defined(__HIP_DEVICE_COMPILE__) && !defined(__HIPCC_RTC__) #include "aiter_tensor.h" #include "aiter_stream.h" +#include "opus_gemm_common.cuh" #include #endif #ifdef OPUS_FUSED_HOST_TU #include "{traits_header}" -{cluster_dims_host_def}// Forward decl for the host <<<>>> launch stub. The __cluster_dims__ attribute -// uses this kid's CONCRETE (split_k, m_cluster) -- NOT the template params -- so -// the host launch site actually sets the cluster geometry (a template-parameter -// attribute does not propagate to the launch config; mirrors the ws clusterlaunch -// stub which also bakes concrete cluster dims). -template +{cluster_dims_host_def}// Concrete cluster geometry must be visible on the host launch stub. +template __global__ __launch_bounds__(128, 1) __cluster_dims__({split_k}, {n_cluster}, 1) void {kernel_func}({kargs_name} kargs); @@ -565,98 +506,98 @@ def gen_splitk_fuse_instance( int splitK) {{{{ static_assert(std::is_same::value, - "splitk_fuse launcher uses the split-K lookup ABI (D_C=fp32 traits;" - " Y dtype is chosen at runtime as D_OUT)"); - (void)splitK; // SplitK is compile-time ({split_k}); runtime splitK ignored. + "splitk_fuse uses the fp32 workspace-dispatch specialization"); + (void)splitK; // SplitK is compile-time ({split_k}) for this exact kid. int batch = XQ.size(0); int M = XQ.size(1); int N = WQ.size(1); int K = XQ.size(2); - AITER_CHECK(batch == 1, "splitk_fuse is batch==1 only (got batch=", batch, ")"); + AITER_CHECK(batch == 1, + "gfx1250 splitk_fuse supports batch == 1 only; got batch=", batch); + AITER_CHECK(M >= 1 && N >= 1 && K >= 1, + "splitk_fuse requires positive M, N, and K"); AITER_CHECK(Y.dtype() == AITER_DTYPE_bf16 || Y.dtype() == AITER_DTYPE_fp32, "splitk_fuse requires Y dtype bf16 or fp32"); AITER_CHECK(K % 2 == 0, "K=", K, " must be even"); AITER_CHECK(N % {k.B_N} == 0, - "splitk_fuse writes full-N C tiles (no N OOB mask): N must be a " - "multiple of B_N={k.B_N} (got N=", N, "). Ragged M is OK: the last " - "M-tile's OOB rows fall past the C buffer num_records and are dropped."); + "splitk_fuse writes full-N C tiles: N must be a multiple of B_N={k.B_N}; got N=", + N, ". Ragged M remains supported by the bounded C descriptor."); - int num_tiles_m = (M + {k.B_M} - 1) / {k.B_M}; // ceil: last M-tile may be partial (OOB rows dropped by C buffer num_records) + int num_tiles_m = 1 + (M - 1) / {k.B_M}; int num_tiles_n = N / {k.B_N}; - // N-direction cluster (cluster.y groups {n_cluster} N-tile peers, A-multicast): - // ceil(N/B_N) must be a multiple of the cluster N-peer count (exact fill; an - // OOB tail WG would still be named in the multicast mask and stall the barrier). AITER_CHECK(num_tiles_n % {n_cluster} == 0, - "splitk_fuse kid n_cluster={n_cluster}: ceil(N/B_N)=", num_tiles_n, - " must be a multiple of n_cluster (cluster.y N-peer fill)"); + "splitk_fuse kid n_cluster={n_cluster}: N/B_N=", num_tiles_n, + " must exactly fill cluster.y"); int k_steps_tot = (K + {k.B_K} - 1) / {k.B_K}; - // Balanced K-tile split (see pipeline): every split WG gets >=1 tile as long as - // split_k <= k_steps_tot, so no WG is empty (the K tail is TDM-clamped, not - // handled by emptying a WG). Only reject when there are fewer whole B_K tiles - // than splits. AITER_CHECK({split_k} <= k_steps_tot, - "splitk_fuse kid split_k={split_k} exceeds k_steps_tot=", k_steps_tot, - " for K=", K, " (more splits than whole B_K tiles -> some WG would be empty);" - " pick a kid with a smaller split_k for this K"); + "splitk_fuse kid split_k={split_k} exceeds K-tile count ", k_steps_tot, + " for K=", K, " and B_K={k.B_K}"); - // Bias: read as bf16 in-kernel; require bf16 (or absent) for round-1. + // #4246 round-1 bias contract: contiguous bf16 [N]. const void* ptr_bias_ = nullptr; int stride_bias_batch_ = 0; if (bias.has_value()) {{{{ const auto& bt = bias.value(); AITER_CHECK(bt.is_contiguous(), "splitk_fuse bias must be contiguous"); AITER_CHECK(bt.dtype() == AITER_DTYPE_bf16, - "splitk_fuse bias must be bf16 (got ", AiterDtype_to_str(bt.dtype()), ")"); - if (bt.dim() == 1) {{{{ - AITER_CHECK(bt.size(0) == N, "splitk_fuse 1D bias length must equal N"); - stride_bias_batch_ = 0; - }}}} else {{{{ - AITER_CHECK(false, "splitk_fuse round-1 supports only 1D [N] bias"); - }}}} + "splitk_fuse bias must be bf16; got ", AiterDtype_to_str(bt.dtype())); + AITER_CHECK(bt.dim() == 1 && bt.size(0) == N, + "splitk_fuse bias must have shape [N]; got dim=", bt.dim()); ptr_bias_ = bt.data_ptr(); }}}} - using Traits = {k.name}_Traits; + // Physical layout: [num_tiles_m, num_tiles_n, SplitK-1, B_M, B_N]. + const size_t tile_numel = opus_checked_extent_product( + {{static_cast({k.B_M}), static_cast({k.B_N})}}, + "{k.name}"); + const size_t required_numel = opus_checked_extent_product( + {{static_cast(num_tiles_m), + static_cast(num_tiles_n), + static_cast({split_k - 1}), + tile_numel}}, + "{k.name}"); + void* workspace_ptr_ = opus_validate_workspace( + workspace, XQ, {workspace_aiter_dtype}, required_numel, 16, "{k.name}"); + using Traits = {k.name}_Traits; auto stream = aiter::getCurrentHIPStream(); {kargs_name} kargs{{{{}}}}; - kargs.ptr_a = XQ.data_ptr(); - kargs.ptr_b = WQ.data_ptr(); - kargs.ptr_ws = workspace.data_ptr(); - kargs.ptr_c = Y.data_ptr(); - kargs.ptr_bias = ptr_bias_; - kargs.m = M; kargs.n = N; kargs.k = K; kargs.batch = 1; kargs.split_k = {split_k}; - kargs.stride_a = XQ.stride(1); - kargs.stride_b = WQ.stride(1); - kargs.stride_c = N; - kargs.stride_a_batch = XQ.stride(0); - kargs.stride_b_batch = WQ.stride(0); - kargs.stride_c_batch = M * N; + kargs.ptr_a = XQ.data_ptr(); + kargs.ptr_b = WQ.data_ptr(); + kargs.ptr_ws = workspace_ptr_; + kargs.ptr_c = Y.data_ptr(); + kargs.ptr_bias = ptr_bias_; + kargs.m = M; kargs.n = N; kargs.k = K; kargs.batch = 1; + kargs.split_k = {split_k}; + kargs.stride_a = XQ.stride(1); + kargs.stride_b = WQ.stride(1); + kargs.stride_c = N; + kargs.stride_a_batch = XQ.stride(0); + kargs.stride_b_batch = WQ.stride(0); + kargs.stride_c_batch = M * N; kargs.stride_bias_batch = stride_bias_batch_; kargs.num_tiles_m = num_tiles_m; kargs.num_tiles_n = num_tiles_n; - // N-direction cluster: N-tiles on grid.y so cluster.y groups the {n_cluster} - // N-peers (A-multicast); M-tiles on grid.z. cluster = ({split_k}, {n_cluster}, 1). + // cluster = (SplitK, N peers, 1); M tiles occupy grid.z. dim3 grid_main({split_k}, num_tiles_n, num_tiles_m); dim3 block_main({k.BLOCK_SIZE}); if (Y.dtype() == AITER_DTYPE_bf16) {{{{ - {kernel_func} + {kernel_func} <<>>(kargs); }}}} else {{{{ - {kernel_func} + {kernel_func} <<>>(kargs); }}}} }}}} -#endif // launcher only on regular host pass +#endif """ Path(os.path.join(cg.impl_path, f"{k.name}.cuh")).write_text(INSTANCE_IMPL) - # Host launcher: only (split-K lookup ABI). Device kernel: both D_OUT. host_decl = ( f"template void\n" f"{k.name}(\n" @@ -673,8 +614,8 @@ def gen_splitk_fuse_instance( for d_out in ("__bf16", "float"): device_decl = ( f"template __global__ void {kernel_func}<\n" - f" {k.name}_Traits, {split_k}, {ws_ctype}, {n_cluster}, {d_out}>" - f"({kargs_name});\n" + f" {k.name}_Traits, {split_k}, {workspace_ptr_type}, " + f"{n_cluster}, {d_out}>({kargs_name});\n" ) cg._device_instantiations.append( {"kid_name": k.name, "dtype": d_out, "device_decl": device_decl} @@ -685,62 +626,32 @@ def gen_4wave_co_instance( cg, k, traits_header, - da, - db, traits_name, kargs_name, **_unused, ): - """gfx1250 symmetric 4-wave compute launcher emit -- PRE-COMPILED .co kid. - - This is the one emit that does not produce a device kernel. It differs from - an ordinary kid in three places: - - * the launcher takes the ordinary 5-arg a16w16 signature, NOT the 6-arg - workspace-carrying one the gfx1250 split-K kids use: this pipeline has - no split-K, no partial buffer and no reduce kernel, so there is nothing - to put in a workspace. It therefore dispatches through its own pair of - tables (opus_a16w16_co_tune_dispatch_gfx1250 by kid, - opus_a16w16_co_dispatch_gfx1250 by shape) rather than the arch's shared - ones. A split-K .co variant would move back to the workspace ABI. - - * it includes the TRAITS header only, never the pipeline header. The - launcher needs nothing but Traits::kBlockM-style constants, and pulling - the pipeline in would drag TDM builtins and pin builtins into a release - compile unit that cannot handle them. - * it appends NOTHING to cg._device_instantiations, so gen_instances.py's - `{name}_C{dtype}.device.cu` loop skips this kid entirely. That single - omission IS the compile bypass; everything downstream (manifest, tune - lookup, dispatch, tuned CSV) is unchanged. - - The device side is loaded at runtime from gen_co/gfx1250/{k.name}.co, whose - filename and extern "C" symbol both equal the launcher name, so no sidecar - is needed to connect them. The whole family -- tile, cluster, VGPR budget, - launch bounds, device flags -- comes from gen_co/co_kernels.json. + """Emit a host launcher for one pre-built gfx1250 CO exact kid. + + The launcher uses the repository's existing five-argument non-workspace + A16W16 ABI. No device instantiation is emitted: the matching device image is + loaded from ``gen_co/gfx1250/.co`` on first use. """ traits_alias = f""" -// Baked as a plain (non-template) alias rather than being parameterised on the -// launcher's D_C, so it can never be instantiated with a configuration the .co -// was not built for. +// The pre-built image fixes this complete traits configuration. using {k.name}_Traits = {traits_name}<{co_traits_args(k)}>; """ - INSTANCE_IMPL = f"""// SPDX-License-Identifier: MIT + instance_impl = f"""// SPDX-License-Identifier: MIT // Copyright (C) 2025-2026, Advanced Micro Devices, Inc. All rights reserved. // // Auto-generated. Do not edit. See codegen/gen_instances_gfx1250.py. -// -// Pre-compiled (.co) kid: no device TU is emitted for this launcher. The kernel -// image is built offline by csrc/opus_gemm/gen_co/build_co.py. +// Pre-compiled CO kid: host launcher only; no device TU is emitted. #pragma once #if !defined(__HIP_DEVICE_COMPILE__) && !defined(__HIPCC_RTC__) #include "aiter_tensor.h" #include "aiter_stream.h" #include #endif -// Traits only -- see the emit docstring. There is deliberately no -// OPUS_FUSED_HOST_TU branch and no pipeline include: this header is identical -// on every pass. #include "{traits_header}" {traits_alias} #if !defined(__HIP_DEVICE_COMPILE__) && !defined(__HIPCC_RTC__) @@ -756,68 +667,44 @@ def gen_4wave_co_instance( int splitK) {{ static_assert(std::is_same::value, - "the 4wave_co kernel writes bf16 C directly -- there is no reduce " - "kernel and no fp32 workspace slot to dispatch through"); + "gfx1250 4wave CO kernels write bf16 output directly"); int batch = XQ.size(0); int M = XQ.size(1); int N = WQ.size(1); int K = XQ.size(2); - using Traits = {k.name}_Traits; AITER_CHECK(Y.dtype() == AITER_DTYPE_bf16, - "4wave_co writes bf16 C directly (no reduce kernel to cast): Y must be " - "bf16, got ", AiterDtype_to_str(Y.dtype())); - AITER_CHECK(!bias.has_value(), "4wave_co does not support bias"); - AITER_CHECK(splitK <= 1, - "4wave_co has no split-K (got splitK=", splitK, ")"); + "4wave CO output must be bf16, got ", AiterDtype_to_str(Y.dtype())); + AITER_CHECK(!bias.has_value(), "4wave CO kernels do not support bias"); + AITER_CHECK(splitK == 0 || splitK == 1, + "4wave CO kernels require splitK in {{0,1}} (got splitK=", splitK, ")"); AITER_CHECK(M >= 1 && N >= 1 && K >= 1 && batch >= 1, - "M, N, K, batch must be >= 1"); - // No M/N/K alignment guard, and deliberately so. Every tail is handled by - // the TDM descriptor: tdm::make_descriptor() clamps EVERY dimension with - // saturating_sub(extent, origin), and the C store window's fastest axis IS - // N, so the hardware writes min(B_N, n - col) columns of each row. Ragged - // M and ragged K ride the same mechanism. - // - // (An earlier revision asserted N % B_N == 0 on the theory that the - // epilogue bounded the store by bytes-remaining-in-the-matrix and an N tail - // would spill into the next row. That is make_gmem's num_records semantics, - // not the TDM's -- this pipeline builds no gmem descriptor at all.) + "M, N, K and batch must be >= 1"); {kargs_name} kargs{{}}; kargs.ptr_a = XQ.data_ptr(); kargs.ptr_b = WQ.data_ptr(); kargs.ptr_c = Y.data_ptr(); - kargs.m = M; kargs.n = N; kargs.k = K; - kargs.stride_a = XQ.stride(1); - kargs.stride_b = WQ.stride(1); - kargs.stride_c = N; - // 64-bit: a batch stride is an ELEMENT count over a whole matrix, so it - // passes 2^31 at 4 GiB of bf16 and the kernel offsets by it per batch. kargs.stride_a_batch = XQ.stride(0); kargs.stride_b_batch = WQ.stride(0); - // kargs carries no batch count (grid.z is it) and no C batch stride (the - // kernel derives m * stride_c), which is what keeps the struct at 64 B. - - // Round the tile counts up to whole clusters: the runtime rejects a cluster - // launch whose grid is not a multiple of the cluster dims, so a half-full - // cluster does not exist and the surplus workgroups WILL be dispatched. The - // kernel absorbs them -- they fail the tile bound check and leave right - // after the one cluster-scope barrier they owe their peers, having touched - // neither LDS nor the TDM. So any (M, N) is launchable with any cluster - // dims, and there is deliberately no divisibility AITER_CHECK here. + kargs.m = M; + kargs.n = N; + kargs.k = K; + kargs.stride_a = XQ.stride(1); + kargs.stride_b = WQ.stride(1); + kargs.stride_c = Y.stride(1); + int grid_m = (M + Traits::kBlockM - 1) / Traits::kBlockM; int grid_n = (N + Traits::kBlockN - 1) / Traits::kBlockN; - grid_m = (grid_m + Traits::kClusterWgM - 1) / Traits::kClusterWgM * Traits::kClusterWgM; - grid_n = (grid_n + Traits::kClusterWgN - 1) / Traits::kClusterWgN * Traits::kClusterWgN; - - // Resolved (file read + module registration) on first call only: the symbol - // is this launcher's own name, so one static per launcher is exact. - static AiterAsmKernelFast& kernel = opus_gfx1250_co::co_kernel("{k.name}"); + grid_m = (grid_m + Traits::kClusterWgM - 1) / + Traits::kClusterWgM * Traits::kClusterWgM; + grid_n = (grid_n + Traits::kClusterWgN - 1) / + Traits::kClusterWgN * Traits::kClusterWgN; - // One launch covers the whole batch (grid.z), unlike the ws pipeline's host - // batch loop: the kernel offsets A/B/C by workgroup_id_z * stride_*_batch. + static AiterAsmKernelFast& kernel = + opus_gfx1250_co::co_kernel("{k.name}"); opus_co_launch_gfx1250( kernel, dim3(grid_m, grid_n, batch), @@ -825,24 +712,24 @@ def gen_4wave_co_instance( kargs, aiter::getCurrentHIPStream()); }} -#endif // launcher only on regular host pass +#endif """ - Path(os.path.join(cg.impl_path, f"{k.name}.cuh")).write_text(INSTANCE_IMPL) + Path(os.path.join(cg.impl_path, f"{k.name}.cuh")).write_text(instance_impl) - # Host launcher only. NOTHING is appended to cg._device_instantiations -- - # that omission is the whole bypass (see the docstring). - for CDtype in k.output_dtypes: + # Intentionally append host instantiations only. This is the compile bypass + # that keeps release hipcc away from the pin-VGPR device pipeline. + for c_dtype in k.output_dtypes: host_decl = ( - f"template void\n" - f"{k.name}<{CDtype}>(\n" - f" aiter_tensor_t &XQ,\n" - f" aiter_tensor_t &WQ,\n" - f" aiter_tensor_t &Y,\n" - f" std::optional,\n" - f" int);\n" + "template void\n" + f"{k.name}<{c_dtype}>(\n" + " aiter_tensor_t &XQ,\n" + " aiter_tensor_t &WQ,\n" + " aiter_tensor_t &Y,\n" + " std::optional,\n" + " int);\n" ) cg._host_instantiations.append( - {"kid_name": k.name, "dtype": CDtype, "host_decl": host_decl} + {"kid_name": k.name, "dtype": c_dtype, "host_decl": host_decl} ) @@ -856,7 +743,7 @@ def gen_4wave_co_instance( "gfx1250", "a16w16_clusterlaunch_tdm_splitk_fuse", gen_splitk_fuse_instance ) # CLUSTER-LAUNCH variant shares the same emit (it branches on k.kernel_tag to add -# __cluster_dims__, the cluster-fill check, and the CLUSTER_WG_M/N traits args). +# __cluster_dims__, physical-grid round-up, and the CLUSTER_WG_M/N traits args). register_emit( "gfx1250", "a16w16_clusterlaunch_tdm_splitk_ws", gen_cluster_tdm_splitk_ws_instance ) diff --git a/csrc/opus_gemm/codegen/gen_instances_gfx942.py b/csrc/opus_gemm/codegen/gen_instances_gfx942.py index 5d8f053ffb..76e379d728 100644 --- a/csrc/opus_gemm/codegen/gen_instances_gfx942.py +++ b/csrc/opus_gemm/codegen/gen_instances_gfx942.py @@ -1,11 +1,18 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2025-2026, Advanced Micro Devices, Inc. All rights reserved. -"""gfx942 codegen -- emit launchers for gfx942-targeted kid families.""" +"""Generate gfx942 OPUS launchers.""" import os from pathlib import Path -from opus_gemm_common import OpusGemmInstance +from opus_gemm_common import ( + GFX942_BF16WS_EXACT_N, + GFX942_EVEN_LOOP_SPLITK_TAGS, + GFX942_MAX_AUTO_SPLIT_K, + GFX942_MIN_ITERS_PER_SPLIT, + GFX942_QUAD_MFMA32_SPLITK_TAG, + OpusGemmInstance, +) from codegen.common import ( _GFX942_A16W16_TAGS, @@ -15,6 +22,7 @@ WARP_SIZE, register_arch_map, register_emit, + splitk_workspace_type, ) @@ -41,13 +49,7 @@ def _gfx942_pipeline(tag): }, } -GFX942_QUAD_MFMA32_SPLITK_TAG = "a16w16_quad_mfma32_kbuf1_sk" GFX942_SPLITK_TAGS = _SPLITK + ("a16w16_em3en4_lds1_pgr2_sk",) -GFX942_EVEN_LOOP_SPLITK_TAGS = ( - "a16w16_kbuf2v_sk", - "a16w16_kbuf2v_bk128_sk", - GFX942_QUAD_MFMA32_SPLITK_TAG, -) _GFX942_A8W8_TAGS = ("a8w8_blockscale_bpreshuffle_singlebuf",) @@ -59,17 +61,6 @@ def _splitk_traits_geometry(k): return trait_bm, trait_bn, override["lds_depth"] -def _splitk_workspace_types(k): - dtype = getattr(k, "splitk_workspace_dtype", "fp32_t") - if dtype == "bf16_t": - return "bf16_t", "__bf16" - return "fp32_t", "float" - - -def _uses_bf16_workspace(k): - return getattr(k, "splitk_workspace_dtype", "fp32_t") == "bf16_t" - - PIPELINE_HEADER_MAP = { "a8w8_blockscale_bpreshuffle_singlebuf": "gfx942/a8w8/opus_gemm_pipeline_a8w8_blockscale_bpreshuffle.cuh", "a16w16_em3en4_lds1_pgr2_sk": _gfx942_pipeline("a16w16_em3en4_lds1_pgr2_sk"), @@ -136,20 +127,25 @@ def _uses_bf16_workspace(k): (8, 256, 1), # N=2048, 1 row/wg ) +assert ( + frozenset(vec * nvec for vec, nvec, _ in EXACT_N_ROWBLOCK_REDUCE_CONFIGS) + == GFX942_BF16WS_EXACT_N +) + def splitk_reduce_extra_forward_decls(): return ( "template\n" "__global__ void splitk_reduce_kernel_bf16ws_fallback(\n" - " const opus_splitk_ws_handle* ws_handle, D_OUT* c_out,\n" + " const void* ws_ptr, D_OUT* c_out,\n" " int split_k, int M, int N, int batch,\n" " int padded_M, int padded_N,\n" " const D_BIAS_* bias, int stride_bias_batch);\n" "template\n" "__global__ void splitk_reduce_kernel_exact_n_rowblock(\n" - " const opus_splitk_ws_handle* ws_handle, D_OUT* c_out,\n" + " const void* ws_ptr, D_OUT* c_out,\n" " int M, int N, int batch,\n" " int padded_M, int padded_N,\n" " const D_BIAS_* bias, int stride_bias_batch);\n" @@ -161,10 +157,10 @@ def splitk_reduce_extra_device_instantiations(): for out_type in ("__bf16", "float"): contents += ( f"template __global__ void splitk_reduce_kernel_bf16ws_fallback<16, 64, {out_type}, true, {out_type}, true>(\n" - f" const opus_splitk_ws_handle*, {out_type}*, int, int, int, int, int, int,\n" + f" const void*, {out_type}*, int, int, int, int, int, int,\n" f" const {out_type}*, int);\n" f"template __global__ void splitk_reduce_kernel_bf16ws_fallback<16, 64, {out_type}, false, {out_type}, true>(\n" - f" const opus_splitk_ws_handle*, {out_type}*, int, int, int, int, int, int,\n" + f" const void*, {out_type}*, int, int, int, int, int, int,\n" f" const {out_type}*, int);\n" ) for vec, nvec, rows in EXACT_N_ROWBLOCK_REDUCE_CONFIGS: @@ -172,7 +168,7 @@ def splitk_reduce_extra_device_instantiations(): for ws_type in ("float", "__bf16"): contents += ( f"template __global__ void splitk_reduce_kernel_exact_n_rowblock<{sk}, {nvec}, {rows}, {vec}, {ws_type}, __bf16, false, __bf16>(\n" - " const opus_splitk_ws_handle*, __bf16*, int, int, int, int, int,\n" + " const void*, __bf16*, int, int, int, int, int,\n" " const __bf16*, int);\n" ) return contents @@ -203,7 +199,7 @@ def gen_splitk_gfx942_instance( kargs_name, kargs_template_vars, BIAS_HOST_VALIDATE, - A16W16_TUNE_HOST_EXTRA, + A16W16_WORKSPACE_LAUNCH_HOST_EXTRA, make_host_decl, make_device_decl, record_one_instantiation, @@ -218,8 +214,10 @@ def gen_splitk_gfx942_instance( kargs_explicit_param = f", {k.GROUP_M}, opus_gemm_splitk_kargs" fwd_decl_kargs_tpl = ", int COL_MAJOR_GROUP_M, typename Kargs" fwd_decl_kargs_fnarg = "Kargs" - bf16ws = _uses_bf16_workspace(k) - workspace_dtype, workspace_ptr_type = _splitk_workspace_types(k) + workspace_dtype, workspace_ptr_type, workspace_aiter_dtype = splitk_workspace_type( + k + ) + bf16ws = workspace_dtype == "bf16_t" # gfx942 a16w16_traits: 7 params . trait_bm, trait_bn, lds_depth = _splitk_traits_geometry(k) traits_aliases = f""" @@ -281,7 +279,7 @@ def reduce_rowblock_branch(hasbias): dim3 block_rowblock({block_size}); splitk_reduce_kernel_exact_n_rowblock<{sk}, {nvec}, {rows}, {vec}, {workspace_ptr_type}, __bf16, {{hb}}, __bf16> <<>>( - ws_handle_device_, + workspace_ptr_, reinterpret_cast<__bf16*>(Y.data_ptr()), M, N, batch, padded_M, padded_N, {{bias_arg}}); @@ -307,7 +305,7 @@ def _baseline_call(dtype, hasbias, indent): return ( f"{indent}{reduce_kernel}\n" f"{indent} <<>>(\n" - f"{indent} ws_handle_device_,\n" + f"{indent} workspace_ptr_,\n" f"{indent} reinterpret_cast<{dtype}*>(Y.data_ptr()),\n" f"{indent} split_k, M, N, batch, padded_M, padded_N," f"{bias_args}" @@ -318,43 +316,17 @@ def _baseline_call(dtype, hasbias, indent): fp32_t = _baseline_call("float", True, " ") fp32_f = _baseline_call("float", False, " ") bf16_y_check = "" - bf16ws_fallback_decl = "" - bf16ws_host_redirect = "" + bf16ws_host_guard = "" if bf16ws: - fp32ws_name = k.name.replace("_bf16ws", "") exact_reduce_shape_conditions = " ||\n ".join( - f"(N == {n_exact})" - for n_exact in sorted( - {vec * nvec for vec, nvec, _ in EXACT_N_ROWBLOCK_REDUCE_CONFIGS} - ) + f"(N == {n_exact})" for n_exact in sorted(GFX942_BF16WS_EXACT_N) ) - if is_quad_mfma32_splitk: - bf16ws_host_redirect = f""" + bf16ws_host_guard = f""" const bool bf16ws_exact_reduce_shape = {exact_reduce_shape_conditions}; AITER_CHECK(bf16ws_exact_reduce_shape, "{err_label} bf16 workspace currently supports only exact-N rowblock " "reduce shapes"); -""" - else: - bf16ws_fallback_decl = f""" -#if !defined(__HIP_DEVICE_COMPILE__) && !defined(__HIPCC_RTC__) -template -void {fp32ws_name}( - aiter_tensor_t &XQ, - aiter_tensor_t &WQ, - aiter_tensor_t &Y, - std::optional bias, - int splitK); -#endif -""" - bf16ws_host_redirect = f""" - const bool bf16ws_exact_reduce_shape = - {exact_reduce_shape_conditions}; - if (!bf16ws_exact_reduce_shape) {{ - {fp32ws_name}(XQ, WQ, Y, bias, splitK); - return; - }} """ bf16_y_check = ( " AITER_CHECK(Y.dtype() == AITER_DTYPE_bf16,\n" @@ -392,6 +364,7 @@ def _baseline_call(dtype, hasbias, indent): #if !defined(__HIP_DEVICE_COMPILE__) && !defined(__HIPCC_RTC__) #include "aiter_tensor.h" #include "aiter_stream.h" +#include "opus_gemm_common.cuh" #include #include #endif @@ -401,7 +374,6 @@ def _baseline_call(dtype, hasbias, indent): #else #include "{pipeline_header}" #endif -{bf16ws_fallback_decl} {traits_aliases} #if !defined(__HIP_DEVICE_COMPILE__) && !defined(__HIPCC_RTC__) template @@ -410,18 +382,19 @@ def _baseline_call(dtype, hasbias, indent): aiter_tensor_t &XQ, aiter_tensor_t &WQ, aiter_tensor_t &Y, + aiter_tensor_t &workspace, std::optional bias, int splitK) {{{{ static_assert(std::is_same::value, - "{err_label} splitK launcher uses the fp32 tune-dispatch table"); + "{err_label} split_k launcher uses the fp32 launch-dispatch table"); int batch = XQ.size(0); int M = XQ.size(1); int N = WQ.size(1); int K = XQ.size(2); -{bf16ws_host_redirect} +{bf16ws_host_guard} {bf16_y_check} AITER_CHECK(Y.dtype() == AITER_DTYPE_bf16 || Y.dtype() == AITER_DTYPE_fp32, @@ -467,14 +440,18 @@ def _baseline_call(dtype, hasbias, indent): }}}} if (cu_cached <= 0) cu_cached = 64; // safe gfx942 lower bound }}}} - int tiles_mn = ((M + {k.B_M} - 1) / {k.B_M}) - * ((N + {k.B_N} - 1) / {k.B_N}) * batch; - if (tiles_mn <= 0) tiles_mn = 1; + const size_t tiles_mn = opus_checked_extent_product( + {{static_cast(1 + (M - 1) / {k.B_M}), + static_cast(1 + (N - 1) / {k.B_N}), + static_cast(batch)}}, + "{k.name} launch grid"); // P1 variant wants 2 wg/CU co-residency for TLP -> aim for 2x cu_num grid. int target_wg_dbuf2 = {"2 * cu_cached" if k.kernel_tag.endswith("_p1") else "cu_cached"}; - split_k = (target_wg_dbuf2 + tiles_mn - 1) / tiles_mn; + split_k = static_cast( + 1 + (static_cast(target_wg_dbuf2) - 1) / tiles_mn); if (split_k < 1) split_k = 1; - if (split_k > 16) split_k = 16; // matches tuner enumeration ceiling + if (split_k > {GFX942_MAX_AUTO_SPLIT_K}) + split_k = {GFX942_MAX_AUTO_SPLIT_K}; // tuner enumeration ceiling }}}} // Host-side auto-clamp: split-barrier pipeline requires at least 2 @@ -482,7 +459,7 @@ def _baseline_call(dtype, hasbias, indent): // both caller-pinned and auto-picked split_k. P1 (depth=2 K-dbuf) additionally // requires loops even per split. int total_iters = (K + {k.B_K} - 1) / {k.B_K}; - constexpr int min_iters_per_split = 2; + constexpr int min_iters_per_split = {GFX942_MIN_ITERS_PER_SPLIT}; constexpr bool require_even_loops_dbuf2 = {"true" if k.kernel_tag in GFX942_EVEN_LOOP_SPLITK_TAGS else "false"}; while (split_k > 1) {{{{ int iters_full = (total_iters + split_k - 1) / split_k; @@ -503,49 +480,35 @@ def _baseline_call(dtype, hasbias, indent): " split_k=", split_k, " gives loops=(", iters_full, ",", last_loops, ")"); }}}} - int num_tiles_m = (M + {k.B_M} - 1) / {k.B_M}; - int num_tiles_n = (N + {k.B_N} - 1) / {k.B_N}; - int padded_M = num_tiles_m * {k.B_M}; - int padded_N = num_tiles_n * {k.B_N}; - + int num_tiles_m = 1 + (M - 1) / {k.B_M}; + int num_tiles_n = 1 + (N - 1) / {k.B_N}; + const size_t padded_M_size = opus_checked_extent_product( + {{static_cast(num_tiles_m), static_cast({k.B_M})}}, + "{k.name}"); + const size_t padded_N_size = opus_checked_extent_product( + {{static_cast(num_tiles_n), static_cast({k.B_N})}}, + "{k.name}"); + const size_t workspace_slice_numel = opus_checked_extent_product( + {{padded_M_size, padded_N_size}}, "{k.name}"); + AITER_CHECK(padded_M_size <= static_cast(std::numeric_limits::max()) + && padded_N_size <= static_cast(std::numeric_limits::max()) + && workspace_slice_numel <= static_cast(std::numeric_limits::max()), + "{k.name}: padded workspace extents exceed 32-bit kernel stride limits"); + int padded_M = static_cast(padded_M_size); + int padded_N = static_cast(padded_N_size); + + const size_t required_numel = opus_checked_extent_product( + {{static_cast(split_k), static_cast(batch), + workspace_slice_numel}}, + "{k.name}"); + void* workspace_ptr_ = opus_validate_workspace( + workspace, XQ, {workspace_aiter_dtype}, required_numel, 16, "{k.name}"); auto stream = aiter::getCurrentHIPStream(); - hipStreamCaptureStatus capture_status = hipStreamCaptureStatusNone; - HIP_CALL(hipStreamIsCapturing(stream, &capture_status)); - const bool capturing = (capture_status != hipStreamCaptureStatusNone); - extern opus_splitk_ws_handle* opus_splitk_ws_get(hipStream_t, bool); - extern const opus_splitk_ws_handle* opus_splitk_ws_device_handle(hipStream_t, bool); - extern void opus_splitk_ws_sync_to_device(hipStream_t); - auto* ws_handle_ = opus_splitk_ws_get(stream, /*allow_create=*/!capturing); - - size_t ws_bytes = (size_t)split_k * (size_t)batch - * (size_t)padded_M * (size_t)padded_N * sizeof(typename Traits::D_C); - if (ws_handle_->ptr == nullptr || ws_bytes > ws_handle_->bytes) - {{ - AITER_CHECK(!capturing, - "{err_label} workspace grow inside HIP graph capture is not " - "supported. Call aiter.opus_gemm_workspace_init() on the capture " - "stream and warm with the largest expected GEMM before capturing."); - - if (ws_handle_->ptr != nullptr) - {{ - HIP_CALL(hipDeviceSynchronize()); - HIP_CALL(hipFree(ws_handle_->ptr)); - }} - const size_t kGrowAlign = (size_t)4 * 1024 * 1024; - size_t grow_bytes = ((ws_bytes + kGrowAlign - 1) / kGrowAlign) * kGrowAlign; - void* new_ptr = nullptr; - HIP_CALL(hipMalloc(&new_ptr, grow_bytes)); - ws_handle_->ptr = new_ptr; - ws_handle_->bytes = grow_bytes; - opus_splitk_ws_sync_to_device(stream); - }} - const auto* ws_handle_device_ = - opus_splitk_ws_device_handle(stream, /*allow_create=*/!capturing); {kargs_name} kargs{{{{}}}}; kargs.ptr_a = XQ.data_ptr(); kargs.ptr_b = WQ.data_ptr(); - kargs.ws_handle = ws_handle_device_; + kargs.ptr_ws = workspace_ptr_; kargs.ptr_c = Y.data_ptr(); kargs.ptr_bias = ptr_bias_; kargs.m = M; kargs.n = N; kargs.k = K; kargs.batch = batch; @@ -556,7 +519,7 @@ def _baseline_call(dtype, hasbias, indent): kargs.stride_c = N; kargs.stride_a_batch = XQ.stride(0); kargs.stride_b_batch = WQ.stride(0); - kargs.stride_ws_batch = padded_M * padded_N; + kargs.stride_ws_batch = static_cast(workspace_slice_numel); kargs.stride_c_batch = M * N; kargs.stride_bias_batch = stride_bias_batch_; dim3 grid_main(num_tiles_m * num_tiles_n * split_k, 1, batch); @@ -573,7 +536,7 @@ def _baseline_call(dtype, hasbias, indent): k, kernel_func, kargs_name, - A16W16_TUNE_HOST_EXTRA, + A16W16_WORKSPACE_LAUNCH_HOST_EXTRA, kargs_explicit_param, ) @@ -590,7 +553,7 @@ def _emit_a16w16_nosplit_launcher( kargs_name, instance_impl_preamble, instance_impl_host_tu_split, - A16W16_TUNE_TAGS, + A16W16_KID_DISPATCH_TAGS, fwd_decl_kargs_tpl, fwd_decl_kargs_fnarg, traits_extra, @@ -600,15 +563,15 @@ def _emit_a16w16_nosplit_launcher( device_decl_for_dtype, ): extra_param = ( - ",\n std::optional bias," "\n int /*splitK*/" - if k.kernel_tag in A16W16_TUNE_TAGS + ",\n std::optional bias," "\n int /*split_k*/" + if k.kernel_tag in A16W16_KID_DISPATCH_TAGS else "" ) bias_kargs_block = ( " AITER_CHECK(!bias.has_value(),\n" ' "bias not supported on this a16w16 kid");\n' - if k.kernel_tag in A16W16_TUNE_TAGS + if k.kernel_tag in A16W16_KID_DISPATCH_TAGS else "" ) @@ -672,7 +635,7 @@ def _emit_a16w16_nosplit_launcher( inst_extra_param = ( ",\n std::optional,\n int" - if k.kernel_tag in A16W16_TUNE_TAGS + if k.kernel_tag in A16W16_KID_DISPATCH_TAGS else "" ) for CDtype in k.output_dtypes: @@ -709,8 +672,8 @@ def gen_a16w16_quad_mfma32_gfx942_instance( instance_impl_preamble, instance_impl_host_tu_split, record_one_instantiation, - A16W16_TUNE_HOST_EXTRA, - A16W16_TUNE_TAGS, + A16W16_LAUNCH_HOST_EXTRA, + A16W16_KID_DISPATCH_TAGS, **_unused, ): """gfx942 quad MFMA32 launcher emit.""" @@ -766,7 +729,7 @@ def device_decl_for_dtype(CDtype): kargs_name, instance_impl_preamble, instance_impl_host_tu_split, - A16W16_TUNE_TAGS, + A16W16_KID_DISPATCH_TAGS, fwd_decl_kargs_tpl, fwd_decl_kargs_fnarg, traits_extra, @@ -791,8 +754,8 @@ def gen_a16w16_nosplit_gfx942_instance( instance_impl_preamble, instance_impl_host_tu_split, record_one_instantiation, - A16W16_TUNE_HOST_EXTRA, - A16W16_TUNE_TAGS, + A16W16_LAUNCH_HOST_EXTRA, + A16W16_KID_DISPATCH_TAGS, **_unused, ): """gfx942 a16w16 non-splitK launcher emit (kbuf2v / kbuf2v_bk128 / @@ -831,7 +794,17 @@ def gen_a16w16_nosplit_gfx942_instance( AITER_CHECK(M >= 1 && N >= 1, "M and N must be >= 1"); """ - launch_block = f""" + if is_wkc_accum: + launch_block = f""" + auto stream = aiter::getCurrentHIPStream(); + auto memset_status = hipMemsetAsync( + Y.data_ptr(), 0, static_cast(batch) * M * N * sizeof(D_C), stream); + AITER_CHECK(memset_status == hipSuccess, + "hipMemsetAsync failed before gfx942 wave-K accumulate launch: ", + hipGetErrorString(memset_status)); + {kernel_func}<{k.name}_Traits><<>>(kargs);""" + else: + launch_block = f""" auto stream = aiter::getCurrentHIPStream(); {kernel_func}<{k.name}_Traits><<>>(kargs);""" if is_wkc_accum: @@ -859,7 +832,7 @@ def device_decl_for_dtype(CDtype): kargs_name, instance_impl_preamble, instance_impl_host_tu_split, - A16W16_TUNE_TAGS, + A16W16_KID_DISPATCH_TAGS, fwd_decl_kargs_tpl, fwd_decl_kargs_fnarg, traits_extra, @@ -884,15 +857,11 @@ def gen_a8w8_blockscale_bpreshuffle_gfx942_instance( instance_impl_preamble, instance_impl_host_tu_split, record_one_instantiation, - A8W8_SCALE_HOST_EXTRA, + A8W8_BLOCKSCALE_HOST_EXTRA, + make_a8w8_bpreshuffle_host_decl, **_unused, ): - """gfx942 A8W8 blockscale bpreshuffle launcher emit. - - This is an explicit tune path. The public C++ wrapper dispatches by - integer kid through opus_gemm_a8w8_tune_lookup.h; production opus_gemm() - fp8 dispatch remains gfx950-only. - """ + """Emit the checked gfx942 A8W8 bpreshuffle launcher.""" info = _validate_a8w8_blockscale_bpreshuffle_gfx942(k) print( f" {k.name}: E=({info['E_M']},{info['E_N']},{info['E_K']})" @@ -937,32 +906,48 @@ def gen_a8w8_blockscale_bpreshuffle_gfx942_instance( {k.name}( aiter_tensor_t &XQ, aiter_tensor_t &WQ, - aiter_tensor_t &Y, - std::optional x_scale, - std::optional w_scale) + aiter_tensor_t &x_scale, + aiter_tensor_t &w_scale, + aiter_tensor_t &Y) {{{{ AITER_CHECK((XQ.dim() == 2 || XQ.dim() == 3), - "gfx942 a8w8 expects XQ [M,K] or [B,M,K]"); + "opus_gemm_a8w8_blockscale_bpreshuffle_launch: XQ must be " + "[M,K] or [B,M,K]"); AITER_CHECK((WQ.dim() == 2 || WQ.dim() == 3), - "gfx942 a8w8 expects WQ [N,K] or [B,N,K]"); + "opus_gemm_a8w8_blockscale_bpreshuffle_launch: WQ must be " + "[N,K] or [B,N,K]"); AITER_CHECK((Y.dim() == 2 || Y.dim() == 3), - "gfx942 a8w8 expects Y [M,N] or [B,M,N]"); - AITER_CHECK(x_scale.has_value() && w_scale.has_value(), - "gfx942 a8w8 blockscale requires x_scale and w_scale"); + "opus_gemm_a8w8_blockscale_bpreshuffle_launch: Y must be " + "[M,N] or [B,M,N]"); + AITER_CHECK(XQ.dtype() == AITER_DTYPE_fp8 && WQ.dtype() == AITER_DTYPE_fp8, + "opus_gemm_a8w8_blockscale_bpreshuffle_launch: expected fp8 XQ/WQ"); + AITER_CHECK(Y.dtype() == AITER_DTYPE_bf16, + "opus_gemm_a8w8_blockscale_bpreshuffle_launch: expected bf16 Y"); + AITER_CHECK(XQ.is_contiguous() && WQ.is_contiguous() && Y.is_contiguous(), + "opus_gemm_a8w8_blockscale_bpreshuffle_launch: expects contiguous " + "XQ/WQ/Y"); int batch = XQ.dim() == 3 ? XQ.size(0) : 1; int M = XQ.dim() == 3 ? XQ.size(1) : XQ.size(0); int K = XQ.dim() == 3 ? XQ.size(2) : XQ.size(1); int N = WQ.dim() == 3 ? WQ.size(1) : WQ.size(0); AITER_CHECK(batch == 1, - "gfx942 a8w8 tune path currently supports batch=1 only"); + "opus_gemm_a8w8_blockscale_bpreshuffle_launch: gfx942 currently " + "supports batch=1 only"); + AITER_CHECK((WQ.dim() == 2 || WQ.size(0) == batch) && + (Y.dim() == 2 || Y.size(0) == batch), + "opus_gemm_a8w8_blockscale_bpreshuffle_launch: batch dimensions " + "must match"); AITER_CHECK(WQ.size(WQ.dim() - 1) == K, - "WQ K dim must match XQ K dim"); + "opus_gemm_a8w8_blockscale_bpreshuffle_launch: WQ K dim must " + "match XQ K dim"); AITER_CHECK((Y.dim() == 3 ? Y.size(1) : Y.size(0)) == M && (Y.dim() == 3 ? Y.size(2) : Y.size(1)) == N, - "Y shape must be [M,N] or [B,M,N]"); + "opus_gemm_a8w8_blockscale_bpreshuffle_launch: Y shape must be " + "[M,N] or [B,M,N]"); AITER_CHECK(N % {k.B_N} == 0 && K % {k.B_K} == 0, - "gfx942 a8w8 tune path requires exact N/K tiles: N%", + "opus_gemm_a8w8_blockscale_bpreshuffle_launch: gfx942 requires " + "exact N/K tiles: N%", {k.B_N}, "=0 K%", {k.B_K}, "=0"); int GROUP_N = {k.GROUP_N}; @@ -970,14 +955,22 @@ def gen_a8w8_blockscale_bpreshuffle_gfx942_instance( int num_groups_n = N / GROUP_N; int num_groups_k = K / GROUP_K; - const auto& xs = x_scale.value(); - const auto& ws = w_scale.value(); + const auto& xs = x_scale; + const auto& ws = w_scale; AITER_CHECK(xs.dtype() == AITER_DTYPE_fp32 && ws.dtype() == AITER_DTYPE_fp32, - "gfx942 a8w8 blockscale expects fp32 scales"); + "opus_gemm_a8w8_blockscale_bpreshuffle_launch: expects fp32 scales"); + AITER_CHECK(xs.is_contiguous() && ws.is_contiguous(), + "opus_gemm_a8w8_blockscale_bpreshuffle_launch: expects contiguous " + "scales"); + AITER_CHECK(xs.device_id == XQ.device_id && ws.device_id == XQ.device_id, + "opus_gemm_a8w8_blockscale_bpreshuffle_launch: scales must be on " + "the XQ device"); AITER_CHECK(xs.dim() == 2 && xs.size(0) == M && xs.size(1) == num_groups_k, - "x_scale must use CK bpreshuffle layout [K/128,M] flattened as [M,K/128]"); + "opus_gemm_a8w8_blockscale_bpreshuffle_launch: x_scale must use " + "the transposed storage contract with shape [M,K/128]"); AITER_CHECK(ws.dim() == 2 && ws.size(0) == num_groups_n && ws.size(1) == num_groups_k, - "w_scale must be row-major [N/128,K/128]"); + "opus_gemm_a8w8_blockscale_bpreshuffle_launch: w_scale must be " + "row-major [N/128,K/128]"); int num_tiles_m = (M + {k.B_M} - 1) / {k.B_M}; int num_tiles_n = N / {k.B_N}; @@ -1008,8 +1001,9 @@ def gen_a8w8_blockscale_bpreshuffle_gfx942_instance( k, kernel_func, kargs_name, - A8W8_SCALE_HOST_EXTRA, + A8W8_BLOCKSCALE_HOST_EXTRA, kargs_explicit_param, + make_a8w8_bpreshuffle_host_decl, ) @@ -1047,7 +1041,7 @@ def gen_a8w8_blockscale_bpreshuffle_gfx942_instance( def _validate_a8w8_blockscale_bpreshuffle_gfx942(k: OpusGemmInstance): - """Validate gfx942 A8W8 blockscale bpreshuffle tune instances.""" + """Validate one gfx942 A8W8 bpreshuffle kernel entry.""" errors = [] sizeof_da = 1 # fp8 if k.BLOCK_SIZE != k.T_M * k.T_N * WARP_SIZE: diff --git a/csrc/opus_gemm/codegen/gen_instances_gfx950.py b/csrc/opus_gemm/codegen/gen_instances_gfx950.py index c161795be7..e3ae162535 100644 --- a/csrc/opus_gemm/codegen/gen_instances_gfx950.py +++ b/csrc/opus_gemm/codegen/gen_instances_gfx950.py @@ -1,10 +1,6 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2025-2026, Advanced Micro Devices, Inc. All rights reserved. -"""gfx950 codegen -- emit launchers for gfx950-targeted kid families. - -Free functions taking the parent opus_gemm_codegen instance as first arg. -Self-registers each emit into codegen.common.EMIT_REGISTRY at import time. -""" +"""Generate gfx950 OPUS launchers.""" import os from pathlib import Path @@ -15,6 +11,7 @@ WARP_SIZE, register_arch_map, register_emit, + splitk_workspace_type, ) # ---------------- gfx950 arch-override maps ---------------- @@ -142,10 +139,10 @@ def splitk_reduce_extra_device_instantiations(): return ( "// mmajor BMM reduce (a8w8_mxscale split-K launchers)\n" "template __global__ void opus_bmm_splitk_reduce_kernel<__bf16, 8, 128>(\n" - " const opus_splitk_ws_handle*, __bf16*,\n" + " const void*, __bf16*,\n" " int, int, int, int, int, int, int, int);\n" "template __global__ void opus_bmm_splitk_reduce_kernel(\n" - " const opus_splitk_ws_handle*, float*,\n" + " const void*, float*,\n" " int, int, int, int, int, int, int, int);\n" ) @@ -633,7 +630,7 @@ def gen_persistent_instance( instance_impl_preamble, instance_impl_host_tu_split, record_one_instantiation, - A16W16_TUNE_HOST_EXTRA, + A16W16_LAUNCH_HOST_EXTRA, **_unused, ): """gfx950 a16w16_persistent launcher emit. See gen_instances.opus_gemm_codegen._gen_persistent_instance.""" @@ -720,7 +717,7 @@ def gen_persistent_instance( aiter_tensor_t &WQ, aiter_tensor_t &Y, std::optional bias, - int /*splitK*/) // persistent ignores splitK; shares tune-lookup slot signature + int /*splitK*/) // persistent ignores splitK; shares launch-table signature {{{{ int batch = XQ.size(0); int M = XQ.size(1); @@ -753,7 +750,7 @@ def gen_persistent_instance( #endif // launcher only on regular host pass """ Path(os.path.join(cg.impl_path, f"{k.name}.cuh")).write_text(INSTANCE_IMPL) - record_one_instantiation(cg, k, kernel_func, kargs_name, A16W16_TUNE_HOST_EXTRA) + record_one_instantiation(cg, k, kernel_func, kargs_name, A16W16_LAUNCH_HOST_EXTRA) def gen_scale_instance( @@ -770,10 +767,10 @@ def gen_scale_instance( instance_impl_preamble, instance_impl_host_tu_split, record_one_instantiation, - A8W8_SCALE_HOST_EXTRA, + A8W8_BLOCKSCALE_HOST_EXTRA, **_unused, ): - """gfx950 a8w8_scale launcher emit.""" + """Emit the checked gfx950 A8W8 blockscale launcher.""" _kargs_explicit_param, fwd_decl_kargs_tpl, fwd_decl_kargs_fnarg = ( kargs_template_vars(k.kernel_tag, kargs_name) ) @@ -804,13 +801,48 @@ def gen_scale_instance( aiter_tensor_t &XQ, aiter_tensor_t &WQ, aiter_tensor_t &Y, - std::optional x_scale, - std::optional w_scale) + aiter_tensor_t &x_scale, + aiter_tensor_t &w_scale) {{{{ + AITER_CHECK(XQ.dim() == 3 && WQ.dim() == 3 && Y.dim() == 3, + "opus_gemm_a8w8_blockscale_launch: XQ/WQ/Y must be 3D"); + AITER_CHECK(XQ.dtype() == AITER_DTYPE_fp8 && WQ.dtype() == AITER_DTYPE_fp8, + "opus_gemm_a8w8_blockscale_launch: expected fp8 XQ/WQ"); + AITER_CHECK(Y.dtype() == AITER_DTYPE_fp32, + "opus_gemm_a8w8_blockscale_launch: expected fp32 Y"); + AITER_CHECK(XQ.is_contiguous() && WQ.is_contiguous() && Y.is_contiguous(), + "opus_gemm_a8w8_blockscale_launch: XQ/WQ must be K-contiguous " + "and Y must be contiguous"); int batch = XQ.size(0); int M = XQ.size(1); int N = WQ.size(1); int K = XQ.size(2); + AITER_CHECK(batch >= 1 && WQ.size(0) == batch && Y.size(0) == batch, + "opus_gemm_a8w8_blockscale_launch: batch dimensions must match"); + AITER_CHECK(WQ.size(2) == K && Y.size(1) == M && Y.size(2) == N, + "opus_gemm_a8w8_blockscale_launch: tensor shapes must be " + "[B,M,K], [B,N,K], [B,M,N]"); + AITER_CHECK(M >= 1 && N >= 1 && K >= {k.B_K}, + "opus_gemm_a8w8_blockscale_launch: requires positive M/N and " + "K >= {k.B_K}"); + AITER_CHECK(M % {k.GROUP_M} == 0 && N % {k.GROUP_N} == 0 && + K % {k.GROUP_K} == 0, + "opus_gemm_a8w8_blockscale_launch: M/N/K must be divisible by " + "scale groups ", + {k.GROUP_M}, "/", {k.GROUP_N}, "/", {k.GROUP_K}); + + // The pipeline primes two K tiles, then advances in pairs. Rejecting a + // one-tile or odd-tile launch here prevents a negative final tile and an + // out-of-range prefetch in device code. + int loops_ = (K + {k.B_K} - 1) / {k.B_K}; + AITER_CHECK(loops_ >= 2, + "opus_gemm_a8w8_blockscale_launch: ceil_div(K, B_K)=", loops_, + " must be >= 2 (K=", K, ", B_K=", {k.B_K}, ")"); + AITER_CHECK(loops_ % 2 == 0, + "opus_gemm_a8w8_blockscale_launch: ceil_div(K, B_K)=", loops_, + " must be even (prefetch constraint)"); + AITER_CHECK(K % 2 == 0, + "opus_gemm_a8w8_blockscale_launch: K must be even; got K=", K); using Traits = {k.name}_Traits; @@ -821,6 +853,31 @@ def gen_scale_instance( int num_groups_n = N / GROUP_N; int num_groups_k = K / GROUP_K; + AITER_CHECK(x_scale.dtype() == AITER_DTYPE_fp32 && + w_scale.dtype() == AITER_DTYPE_fp32, + "opus_gemm_a8w8_blockscale_launch: expects fp32 scales"); + AITER_CHECK(x_scale.is_contiguous() && w_scale.is_contiguous(), + "opus_gemm_a8w8_blockscale_launch: expects contiguous scales"); + AITER_CHECK(x_scale.device_id == XQ.device_id && + w_scale.device_id == XQ.device_id, + "opus_gemm_a8w8_blockscale_launch: scales must be on the XQ device"); + const bool x_scale_2d = x_scale.dim() == 2 && batch == 1 && + x_scale.size(0) == M && x_scale.size(1) == num_groups_k; + const bool x_scale_3d = x_scale.dim() == 3 && + x_scale.size(0) == batch && x_scale.size(1) == M && + x_scale.size(2) == num_groups_k; + const bool w_scale_2d = w_scale.dim() == 2 && batch == 1 && + w_scale.size(0) == num_groups_n && w_scale.size(1) == num_groups_k; + const bool w_scale_3d = w_scale.dim() == 3 && + w_scale.size(0) == batch && w_scale.size(1) == num_groups_n && + w_scale.size(2) == num_groups_k; + AITER_CHECK(x_scale_2d || x_scale_3d, + "opus_gemm_a8w8_blockscale_launch: x_scale must be " + "[B,M,K/128], or [M,K/128] when B=1"); + AITER_CHECK(w_scale_2d || w_scale_3d, + "opus_gemm_a8w8_blockscale_launch: w_scale must be " + "[B,N/128,K/128], or [N/128,K/128] when B=1"); + {kargs_name} kargs{{}}; kargs.ptr_a = XQ.data_ptr(); kargs.ptr_b = WQ.data_ptr(); @@ -836,8 +893,8 @@ def gen_scale_instance( kargs.stride_b_batch = N * K; kargs.stride_c_batch = M * N; - kargs.ptr_sfa = x_scale.value().data_ptr(); - kargs.ptr_sfb = w_scale.value().data_ptr(); + kargs.ptr_sfa = x_scale.data_ptr(); + kargs.ptr_sfb = w_scale.data_ptr(); kargs.stride_sfa = num_groups_k; kargs.stride_sfb = num_groups_k; kargs.stride_sfa_batch = num_groups_m * num_groups_k; @@ -855,84 +912,7 @@ def gen_scale_instance( #endif // launcher only on regular host pass """ Path(os.path.join(cg.impl_path, f"{k.name}.cuh")).write_text(INSTANCE_IMPL) - record_one_instantiation(cg, k, kernel_func, kargs_name, A8W8_SCALE_HOST_EXTRA) - - # "_mmajor" sibling: A(XQ)/Y are [M, batch, *] (dim0=M, dim1=batch) and - # x_scale is [M, batch, K/GROUP_K] (per-token M) so the DSV4 wo_a activation - # o=[num_tokens, n_groups, K] feeds in with NO caller-side transpose. Weight - # (WQ) and its scale (w_scale) stay batch-major [batch, N, K] / - # [batch, N/GROUP_N, K/GROUP_K]. Same kernel/traits; the launcher just reads - # A/Y/sfa strides from the tensors instead of hardcoding batch-major. - INSTANCE_IMPL_MMAJOR = f""" -#if !defined(__HIP_DEVICE_COMPILE__) && !defined(__HIPCC_RTC__) -template -void -{k.name}_mmajor( - aiter_tensor_t &XQ, - aiter_tensor_t &WQ, - aiter_tensor_t &Y, - std::optional x_scale, - std::optional w_scale) -{{{{ - int M = XQ.size(0); - int batch = XQ.size(1); - int N = WQ.size(1); - int K = XQ.size(2); - - int GROUP_N = {k.GROUP_N}; - int GROUP_K = {k.GROUP_K}; - int num_groups_n = N / GROUP_N; - int num_groups_k = K / GROUP_K; - - {kargs_name} kargs{{}}; - kargs.ptr_a = XQ.data_ptr(); - kargs.ptr_b = WQ.data_ptr(); - kargs.ptr_c = Y.data_ptr(); - kargs.m = M; - kargs.n = N; - kargs.k = K; - kargs.batch = batch; - // mmajor A/Y (dim0=M, dim1=batch); weight WQ stays batch-major. - kargs.stride_a = (int)XQ.stride(0); - kargs.stride_b = (int)WQ.stride(1); - kargs.stride_c = (int)Y.stride(0); - kargs.stride_a_batch = (int)XQ.stride(1); - kargs.stride_b_batch = (int)WQ.stride(0); - kargs.stride_c_batch = (int)Y.stride(1); - - kargs.ptr_sfa = x_scale.value().data_ptr(); - kargs.ptr_sfb = w_scale.value().data_ptr(); - // x_scale mmajor [M, batch, num_groups_k]; w_scale batch-major. - kargs.stride_sfa = (int)x_scale.value().stride(0); - kargs.stride_sfa_batch = (int)x_scale.value().stride(1); - kargs.stride_sfb = num_groups_k; - kargs.stride_sfb_batch = num_groups_n * num_groups_k; - - int num_tiles_m = (M + {k.B_M} - 1) / {k.B_M}; - int num_tiles_n = (N + {k.B_N} - 1) / {k.B_N}; - dim3 grid(num_tiles_m * num_tiles_n, 1, batch); - dim3 block({k.BLOCK_SIZE}); - - auto stream = aiter::getCurrentHIPStream(); - {kernel_func}<{k.name}_Traits><<>>(kargs); - -}}}} -#endif // launcher only on regular host pass -""" - with open(os.path.join(cg.impl_path, f"{k.name}.cuh"), "a") as _f: - _f.write(INSTANCE_IMPL_MMAJOR) - - for CDtype in k.output_dtypes: - host_decl_mmajor = ( - f"template void\n" - f"{k.name}_mmajor<{CDtype}>(\n" - f" aiter_tensor_t &XQ,\n" - f" aiter_tensor_t &WQ,\n" - f" aiter_tensor_t &Y{A8W8_SCALE_HOST_EXTRA});\n" - ) - cg._host_instantiations.append( - {"kid_name": k.name, "dtype": CDtype, "host_decl": host_decl_mmajor} - ) + record_one_instantiation(cg, k, kernel_func, kargs_name, A8W8_BLOCKSCALE_HOST_EXTRA) def gen_noscale_instance_gfx950( @@ -949,11 +929,10 @@ def gen_noscale_instance_gfx950( instance_impl_preamble, instance_impl_host_tu_split, BIAS_HOST_VALIDATE, - A16W16_TUNE_TAGS, + A16W16_KID_DISPATCH_TAGS, **_unused, ): - """gfx950 noscale launcher emit: a16w16 split-barrier (bias-aware double-traits) - and a8w8 noscale (single traits). a8w8 falls through the else branch.""" + """Emit a gfx950 A16W16 or A8W8 no-scale launcher.""" kargs_explicit_param, fwd_decl_kargs_tpl, fwd_decl_kargs_fnarg = ( kargs_template_vars(k.kernel_tag, kargs_name) ) @@ -969,7 +948,9 @@ def gen_noscale_instance_gfx950( ) min_k = 2 * k.B_K - k_check = f""" + shape_preamble = "" + if is_a16w16_split_barrier: + k_check = f""" int loops_ = (K + {k.B_K} - 1) / {k.B_K}; AITER_CHECK(loops_ >= 2, "K=", K, " too small for B_K={k.B_K}, need K >= {min_k}"); @@ -980,10 +961,37 @@ def gen_noscale_instance_gfx950( "latent K-tail accumulation bug; pass an even K)"); AITER_CHECK(M >= 1 && N >= 1, "M and N must be >= 1"); """ + else: + shape_preamble = """ + AITER_CHECK(XQ.dim() == 3 && WQ.dim() == 3 && Y.dim() == 3, + "opus_gemm_a8w8_launch: XQ/WQ/Y must be 3D"); + AITER_CHECK(XQ.dtype() == AITER_DTYPE_fp8 && WQ.dtype() == AITER_DTYPE_fp8, + "opus_gemm_a8w8_launch: expected fp8 XQ/WQ"); + AITER_CHECK(Y.dtype() == AITER_DTYPE_fp32, + "opus_gemm_a8w8_launch: expected fp32 Y"); + AITER_CHECK(XQ.is_contiguous() && WQ.is_contiguous() && Y.is_contiguous(), + "opus_gemm_a8w8_launch: XQ/WQ must be K-contiguous and Y must be " + "contiguous"); +""" + k_check = f""" + AITER_CHECK(batch >= 1 && WQ.size(0) == batch && Y.size(0) == batch, + "opus_gemm_a8w8_launch: batch dimensions must match"); + AITER_CHECK(WQ.size(2) == K && Y.size(1) == M && Y.size(2) == N, + "opus_gemm_a8w8_launch: tensor shapes must be " + "[B,M,K], [B,N,K], [B,M,N]"); + int loops_ = (K + {k.B_K} - 1) / {k.B_K}; + AITER_CHECK(loops_ >= 2, + "K=", K, " too small for B_K={k.B_K}, need K >= {min_k}"); + AITER_CHECK(loops_ % 2 == 0, + "ceil_div(K, {k.B_K})=", loops_, " must be even (prefetch constraint)"); + AITER_CHECK(K % 2 == 0, + "opus_gemm_a8w8_launch: K must be even; got K=", K); + AITER_CHECK(M >= 1 && N >= 1, "M and N must be >= 1"); +""" - if k.kernel_tag in A16W16_TUNE_TAGS: + if k.kernel_tag in A16W16_KID_DISPATCH_TAGS: extra_param = ( - ",\n std::optional bias," "\n int /*splitK*/" + ",\n std::optional bias," "\n int /*split_k*/" ) else: extra_param = "" @@ -996,7 +1004,7 @@ def gen_noscale_instance_gfx950( + " kargs.ptr_bias = ptr_bias_;\n" + " kargs.stride_bias_batch = stride_bias_batch_;\n" ) - elif k.kernel_tag in A16W16_TUNE_TAGS: + elif k.kernel_tag in A16W16_KID_DISPATCH_TAGS: bias_kargs_block = ( " AITER_CHECK(!bias.has_value(),\n" ' "bias not supported on this a16w16 kid");\n' @@ -1068,6 +1076,7 @@ def gen_noscale_instance_gfx950( aiter_tensor_t &WQ, aiter_tensor_t &Y{extra_param}) {{{{ +{shape_preamble} int batch = XQ.size(0); int M = XQ.size(1); int N = WQ.size(1); @@ -1099,7 +1108,7 @@ def gen_noscale_instance_gfx950( """ Path(os.path.join(cg.impl_path, f"{k.name}.cuh")).write_text(INSTANCE_IMPL) - if k.kernel_tag in A16W16_TUNE_TAGS: + if k.kernel_tag in A16W16_KID_DISPATCH_TAGS: inst_extra_param = ",\n std::optional,\n int" else: inst_extra_param = "" @@ -1272,7 +1281,7 @@ def gen_flatmm_instance( instance_impl_preamble, instance_impl_host_tu_split, record_one_instantiation, - A16W16_TUNE_HOST_EXTRA, + A16W16_LAUNCH_HOST_EXTRA, **_unused, ): """gfx950 a16w16_flatmm launcher emit.""" @@ -1364,7 +1373,7 @@ def gen_flatmm_instance( #endif // launcher only on regular host pass """ Path(os.path.join(cg.impl_path, f"{k.name}.cuh")).write_text(INSTANCE_IMPL) - record_one_instantiation(cg, k, kernel_func, kargs_name, A16W16_TUNE_HOST_EXTRA) + record_one_instantiation(cg, k, kernel_func, kargs_name, A16W16_LAUNCH_HOST_EXTRA) def gen_flatmm_splitk_instance( @@ -1381,11 +1390,19 @@ def gen_flatmm_splitk_instance( instance_impl_preamble, instance_impl_host_tu_split, record_one_instantiation, - A16W16_TUNE_HOST_EXTRA, + A16W16_WORKSPACE_LAUNCH_HOST_EXTRA, BIAS_HOST_VALIDATE, **_unused, ): - """gfx950 a16w16_flatmm_splitk launcher emit (uses ws_handle + reduce kernel call).""" + """Emit a gfx950 split-K launcher using a caller-owned typed workspace.""" + workspace_dtype, _workspace_ptr_type, workspace_aiter_dtype = splitk_workspace_type( + k + ) + if workspace_dtype != "fp32_t": + raise ValueError( + f"gfx950 kid {k.name} declares {workspace_dtype} workspace, but " + "the current gfx950 main/reduce kernels support fp32_t only" + ) _kargs_explicit_param, fwd_decl_kargs_tpl, fwd_decl_kargs_fnarg = ( kargs_template_vars(k.kernel_tag, kargs_name) ) @@ -1394,7 +1411,7 @@ def gen_flatmm_splitk_instance( template using {k.name}_Traits = {traits_name}<{k.BLOCK_SIZE}, opus::seq<{k.B_M}, {k.B_N}, {k.B_K}>, - opus::tuple<{da}, {db}, fp32_t, fp32_t, {da}>, + opus::tuple<{da}, {db}, {workspace_dtype}, fp32_t, {da}>, opus::seq<{k.VEC_A}, {k.VEC_B}, {k.VEC_C}>, opus::seq<{k.W_M}, {k.W_N}, {k.W_K}>, {k.WG_PER_CU}, @@ -1402,7 +1419,7 @@ def gen_flatmm_splitk_instance( {has_oob_str}>; """ - preamble = instance_impl_preamble() + preamble = instance_impl_preamble('\n#include "opus_gemm_common.cuh"') host_tu_split = instance_impl_host_tu_split( traits_header, pipeline_header, @@ -1420,12 +1437,12 @@ def gen_flatmm_splitk_instance( aiter_tensor_t &XQ, aiter_tensor_t &WQ, aiter_tensor_t &Y, + aiter_tensor_t &workspace, std::optional bias, int splitK) {{{{ static_assert(std::is_same::value, - "splitk main kernel uses fp32 workspace; D_C template param must be fp32_t " - "(Y can be bf16 or fp32; reduce kernel handles the cast / passthrough)"); + "split_k launcher uses the fp32 launch specialization"); int batch = XQ.size(0); int M = XQ.size(1); @@ -1459,55 +1476,35 @@ def gen_flatmm_splitk_instance( "need total_iters >= pfk*B_K = ", pfk * {k.B_K}, " (pfk=", pfk, ")"); - int num_tiles_m = (M + {k.B_M} - 1) / {k.B_M}; - int num_tiles_n = (N + {k.B_N} - 1) / {k.B_N}; - int padded_M = num_tiles_m * {k.B_M}; - int padded_N = num_tiles_n * {k.B_N}; - - // Per-stream workspace handle (process-global registry, mutex-protected - // in opus_gemm.cu). Replaces the prior `static thread_local` cache -- - // under TBO two CPU threads drive two streams concurrently, and each - // captured graph must bake in its own buffer pointer. Eager: lazy- - // create. Capture: must be pre-warmed via - // aiter.opus_gemm_workspace_init() on the capture stream. - // (opus_splitk_ws_handle is already a complete type at this point via - // the traits header included at the top of this launcher .cuh.) - extern opus_splitk_ws_handle* opus_splitk_ws_get(hipStream_t, bool); - + int num_tiles_m = 1 + (M - 1) / {k.B_M}; + int num_tiles_n = 1 + (N - 1) / {k.B_N}; + const size_t padded_M_size = opus_checked_extent_product( + {{static_cast(num_tiles_m), static_cast({k.B_M})}}, + "{k.name}"); + const size_t padded_N_size = opus_checked_extent_product( + {{static_cast(num_tiles_n), static_cast({k.B_N})}}, + "{k.name}"); + const size_t workspace_slice_numel = opus_checked_extent_product( + {{padded_M_size, padded_N_size}}, "{k.name}"); + AITER_CHECK(padded_M_size <= static_cast(std::numeric_limits::max()) + && padded_N_size <= static_cast(std::numeric_limits::max()) + && workspace_slice_numel <= static_cast(std::numeric_limits::max()), + "{k.name}: padded workspace extents exceed 32-bit kernel stride limits"); + int padded_M = static_cast(padded_M_size); + int padded_N = static_cast(padded_N_size); + + const size_t required_numel = opus_checked_extent_product( + {{static_cast(split_k), static_cast(batch), + workspace_slice_numel}}, + "{k.name}"); + void* workspace_ptr_ = opus_validate_workspace( + workspace, XQ, {workspace_aiter_dtype}, required_numel, 16, "{k.name}"); auto stream = aiter::getCurrentHIPStream(); - hipStreamCaptureStatus capture_status = hipStreamCaptureStatusNone; - HIP_CALL(hipStreamIsCapturing(stream, &capture_status)); - const bool capturing = (capture_status != hipStreamCaptureStatusNone); - auto* ws_handle_ = opus_splitk_ws_get(stream, /*allow_create=*/!capturing); - - size_t ws_bytes = (size_t)split_k * (size_t)batch - * (size_t)padded_M * (size_t)padded_N * sizeof(float); - if (ws_handle_->ptr == nullptr || ws_bytes > ws_handle_->bytes) - {{ - AITER_CHECK(!capturing, - "splitk workspace grow inside HIP graph capture is not " - "supported (hipMalloc / hipFree are stream-capture-illegal). " - "Warm the cache once eagerly with the largest workspace before " - "capturing. Call aiter.opus_gemm_workspace_init() on the capture " - "stream first."); - - void* new_ptr = nullptr; - const size_t kGrowAlign = (size_t)4 * 1024 * 1024; - size_t grow_bytes = ((ws_bytes + kGrowAlign - 1) / kGrowAlign) * kGrowAlign; - HIP_CALL(hipMalloc(&new_ptr, grow_bytes)); - if (ws_handle_->ptr != nullptr) - {{ - HIP_CALL(hipDeviceSynchronize()); - HIP_CALL(hipFree(ws_handle_->ptr)); - }} - ws_handle_->ptr = new_ptr; - ws_handle_->bytes = grow_bytes; - }} {kargs_name} kargs{{{{}}}}; kargs.ptr_a = XQ.data_ptr(); kargs.ptr_b = WQ.data_ptr(); - kargs.ws_handle = ws_handle_; + kargs.ptr_ws = workspace_ptr_; kargs.ptr_c = Y.data_ptr(); kargs.ptr_bias = ptr_bias_; kargs.m = M; kargs.n = N; kargs.k = K; kargs.batch = batch; @@ -1518,7 +1515,7 @@ def gen_flatmm_splitk_instance( kargs.stride_c = N; kargs.stride_a_batch = XQ.stride(0); kargs.stride_b_batch = WQ.stride(0); - kargs.stride_ws_batch = padded_M * padded_N; + kargs.stride_ws_batch = static_cast(workspace_slice_numel); kargs.stride_c_batch = M * N; kargs.stride_bias_batch = stride_bias_batch_; @@ -1536,7 +1533,7 @@ def gen_flatmm_splitk_instance( if (bias.has_value()) {{{{ splitk_reduce_kernel <<>>( - ws_handle_, + workspace_ptr_, reinterpret_cast<__bf16*>(Y.data_ptr()), split_k, M, N, batch, padded_M, padded_N, reinterpret_cast(ptr_bias_), @@ -1544,7 +1541,7 @@ def gen_flatmm_splitk_instance( }}}} else {{{{ splitk_reduce_kernel <<>>( - ws_handle_, + workspace_ptr_, reinterpret_cast<__bf16*>(Y.data_ptr()), split_k, M, N, batch, padded_M, padded_N, nullptr, 0); @@ -1553,7 +1550,7 @@ def gen_flatmm_splitk_instance( if (bias.has_value()) {{{{ splitk_reduce_kernel <<>>( - ws_handle_, + workspace_ptr_, reinterpret_cast(Y.data_ptr()), split_k, M, N, batch, padded_M, padded_N, reinterpret_cast(ptr_bias_), @@ -1561,7 +1558,7 @@ def gen_flatmm_splitk_instance( }}}} else {{{{ splitk_reduce_kernel <<>>( - ws_handle_, + workspace_ptr_, reinterpret_cast(Y.data_ptr()), split_k, M, N, batch, padded_M, padded_N, nullptr, 0); @@ -1572,7 +1569,13 @@ def gen_flatmm_splitk_instance( #endif // launcher only on regular host pass """ Path(os.path.join(cg.impl_path, f"{k.name}.cuh")).write_text(INSTANCE_IMPL) - record_one_instantiation(cg, k, kernel_func, kargs_name, A16W16_TUNE_HOST_EXTRA) + record_one_instantiation( + cg, + k, + kernel_func, + kargs_name, + A16W16_WORKSPACE_LAUNCH_HOST_EXTRA, + ) def _assert_m_align(k, tile_mult): @@ -1615,12 +1618,14 @@ def _assert_m_align(k, tile_mult): aiter_tensor_t &Y, aiter_tensor_t &x_scale, aiter_tensor_t &w_scale, + std::optional workspace, int splitK) { using Traits = @@NAME@@_Traits; constexpr bool DIRECT_ONLY = @@DIRECT@@; constexpr bool PREFETCH_SCALE = @@PREFETCH@@; constexpr bool PRELOAD_SF_LDS = @@PRELOAD@@; + constexpr bool SPLITK_PRELOAD_SF_LDS = @@SPLITK_PRELOAD@@; AITER_CHECK(splitK >= 1, "splitK must be >= 1"); if constexpr (DIRECT_ONLY) { @@ -1654,15 +1659,15 @@ def _assert_m_align(k, tile_mult): const int num_tiles_n = (N + Traits::B_N - 1) / Traits::B_N; const int padded_M = num_tiles_m * Traits::B_M; const int padded_N = num_tiles_n * Traits::B_N; - const size_t partial_bytes = (size_t)split_k * (size_t)batch - * (size_t)padded_M * (size_t)padded_N * sizeof(float); + const size_t required_numel = (size_t)split_k * (size_t)batch + * (size_t)padded_M * (size_t)padded_N; auto stream = aiter::getCurrentHIPStream(); opus_gemm_scale_splitk_kargs_gfx950 kargs{}; kargs.ptr_a = O.data_ptr(); kargs.ptr_b = wo_a.data_ptr(); - kargs.ws_handle = nullptr; + kargs.ptr_ws = nullptr; kargs.m = M; kargs.n = N; kargs.k = K; kargs.batch = batch; kargs.split_k = split_k; kargs.stride_a = (int)O.stride(0); @@ -1681,6 +1686,8 @@ def _assert_m_align(k, tile_mult): dim3 grid_main(num_tiles_m * num_tiles_n * split_k, 1, batch); dim3 block_main(Traits::BLOCK_SIZE); if (no_split_k) { + AITER_CHECK(!workspace.has_value(), + "@@NAME@@ splitK == 1 does not use workspace"); kargs.ptr_c = Y.data_ptr(); kargs.stride_c = (int)Y.stride(0); kargs.stride_c_batch = (int)Y.stride(1); @@ -1695,35 +1702,20 @@ def _assert_m_align(k, tile_mult): } if constexpr (!DIRECT_ONLY) { - extern opus_splitk_ws_handle* opus_splitk_ws_get(hipStream_t, bool); - hipStreamCaptureStatus capture_status = hipStreamCaptureStatusNone; - HIP_CALL(hipStreamIsCapturing(stream, &capture_status)); - const bool capturing = (capture_status != hipStreamCaptureStatusNone); - auto* ws_handle = opus_splitk_ws_get(stream, /*allow_create=*/!capturing); - - const size_t ws_bytes = partial_bytes; - if (ws_handle->ptr == nullptr || ws_bytes > ws_handle->bytes) { - AITER_CHECK(!capturing, - "splitk workspace grow inside HIP graph capture is not supported"); - void* new_ptr = nullptr; - const size_t kGrowAlign = (size_t)4 * 1024 * 1024; - size_t grow_bytes = ((ws_bytes + kGrowAlign - 1) / kGrowAlign) * kGrowAlign; - HIP_CALL(hipMalloc(&new_ptr, grow_bytes)); - if (ws_handle->ptr != nullptr) { - HIP_CALL(hipDeviceSynchronize()); - HIP_CALL(hipFree(ws_handle->ptr)); - } - ws_handle->ptr = new_ptr; - ws_handle->bytes = grow_bytes; - } - kargs.ws_handle = ws_handle; + AITER_CHECK(workspace.has_value(), + "@@NAME@@ splitK > 1 requires workspace"); + void* workspace_ptr = opus_validate_workspace( + workspace.value(), O, AITER_DTYPE_fp32, required_numel, 16, "@@NAME@@"); + kargs.ptr_ws = workspace_ptr; // Pass all 4 template args explicitly (D_OUT=void: the split-K main kernel // writes an fp32 workspace, so its output dtype is irrelevant; the reduce // kernel casts to the runtime Y dtype). The fused host TU only sees a // no-default forward decl of @@KERNEL@@, so relying on the template's // default args here would fail overload resolution ("no matching function"). - @@KERNEL@@ + // The workspace specialization may use a lower-register-pressure scale + // path than this kid's tuned splitK=1 direct-output specialization. + @@KERNEL@@ <<>>(kargs); constexpr int REDUCE_VEC = 8; @@ -1736,13 +1728,13 @@ def _assert_m_align(k, tile_mult): if (Y.dtype() == AITER_DTYPE_bf16) { opus_bmm_splitk_reduce_kernel<__bf16, REDUCE_VEC, REDUCE_BS> <<>>( - ws_handle, reinterpret_cast<__bf16*>(Y.data_ptr()), + workspace_ptr, reinterpret_cast<__bf16*>(Y.data_ptr()), split_k, M, N, batch, padded_M, padded_N, y_stride_c, y_stride_c_batch); } else { opus_bmm_splitk_reduce_kernel <<>>( - ws_handle, reinterpret_cast(Y.data_ptr()), + workspace_ptr, reinterpret_cast(Y.data_ptr()), split_k, M, N, batch, padded_M, padded_N, y_stride_c, y_stride_c_batch); } @@ -1797,7 +1789,7 @@ def gen_bmm_mxscale_flatmm_splitk_instance( {k.WG_PER_CU}>; """ - preamble = instance_impl_preamble() + preamble = instance_impl_preamble('\n#include "opus_gemm_common.cuh"') host_tu_split = instance_impl_host_tu_split( traits_header, pipeline_header, @@ -1811,24 +1803,27 @@ def gen_bmm_mxscale_flatmm_splitk_instance( # header that defines this kernel -- so the launcher body's <<<...>>> call # needs a visible declaration. On the non-fused device pass the pipeline # header (via splitk_reduce_gfx950.cuh) provides a compatible definition, so - # this is just a harmless redeclaration there. opus_splitk_ws_handle is a - # complete type in both passes via the included traits/pipeline header. + # this is just a harmless redeclaration there. reduce_fwd_decl = """ template __global__ void opus_bmm_splitk_reduce_kernel( - const opus_splitk_ws_handle* __restrict__ ws_handle, + const void* __restrict__ workspace, D_OUT* __restrict__ out, int split_k, int M, int N, int batch, int padded_M, int padded_N, int stride_c, int stride_c_batch); """ + workspace_preload_sf = ( + k.preload_sf if k.workspace_preload_sf is None else k.workspace_preload_sf + ) launcher = ( _BMM_MXSCALE_SPLITK_LAUNCHER_BODY.replace("@@NAME@@", k.name) .replace("@@KERNEL@@", kernel_func) .replace("@@DIRECT@@", "true" if k.direct_only else "false") .replace("@@PREFETCH@@", "true" if k.prefetch_scale else "false") .replace("@@PRELOAD@@", "true" if k.preload_sf else "false") + .replace("@@SPLITK_PRELOAD@@", "true" if workspace_preload_sf else "false") ) INSTANCE_IMPL = ( @@ -1841,6 +1836,7 @@ def gen_bmm_mxscale_flatmm_splitk_instance( host_extra = ( ",\n aiter_tensor_t &x_scale," "\n aiter_tensor_t &w_scale," + "\n std::optional workspace," "\n int splitK" ) for dtype in k.output_dtypes: @@ -1861,25 +1857,27 @@ def gen_bmm_mxscale_flatmm_splitk_instance( direct = "true" if k.direct_only else "false" prefetch = "true" if k.prefetch_scale else "false" preload = "true" if k.preload_sf else "false" + splitk_preload = "true" if workspace_preload_sf else "false" - def _dev(dtype_tag, d_out, dir_flag, pfk_flag): + def _dev(dtype_tag, d_out, dir_flag, pfk_flag, preload_flag): decl = ( f"template __global__ void {kernel_func}<\n" - f" {k.name}_Traits, {d_out}, {dir_flag}, {pfk_flag}, {preload}>({kargs_name});\n" + f" {k.name}_Traits, {d_out}, {dir_flag}, {pfk_flag}, " + f"{preload_flag}>({kargs_name});\n" ) cg._device_instantiations.append( {"kid_name": k.name, "dtype": dtype_tag, "device_decl": decl} ) - _dev("bf16", "__bf16", direct, prefetch) - _dev("fp32", "float", direct, prefetch) + _dev("bf16", "__bf16", direct, prefetch, preload) + _dev("fp32", "float", direct, prefetch, preload) if not k.direct_only: # Split-K > 1 workspace path: host launches . DIRECT_ONLY is false here (direct kids - # never take the workspace path), but PREFETCH_SCALE / PRELOAD_SF_LDS must - # match the kid, else the instantiation is missing for - # prefetch/preload kids -> undefined symbol at load. - _dev("void", "void", "false", prefetch) + # PREFETCH_SCALE, SPLITK_PRELOAD_SF_LDS>. DIRECT_ONLY is false here + # (direct kids never take the workspace path). The split-K preload flag + # normally inherits the kid, with a per-instance compiler workaround + # permitted by workspace_preload_sf. + _dev("void", "void", "false", prefetch, splitk_preload) _BMM_MXSCALE_MINTERLEAVE_LAUNCHER_BODY = r""" @@ -1896,11 +1894,14 @@ def _dev(dtype_tag, d_out, dir_flag, pfk_flag): aiter_tensor_t &Y, aiter_tensor_t &x_scale, aiter_tensor_t &w_scale, - int /*splitK*/) + std::optional workspace, + int splitK) { using Traits = @@NAME@@_Traits; constexpr bool SKIP_SCALE_WAIT = @@SKIP@@; constexpr int MI = 2; + AITER_CHECK(splitK == 1, "@@NAME@@ requires splitK == 1"); + AITER_CHECK(!workspace.has_value(), "@@NAME@@ does not use workspace"); const int M = O.size(0); const int batch = O.size(1); @@ -1922,7 +1923,7 @@ def _dev(dtype_tag, d_out, dir_flag, pfk_flag): opus_gemm_scale_splitk_kargs_gfx950 kargs{}; kargs.ptr_a = O.data_ptr(); kargs.ptr_b = wo_a.data_ptr(); - kargs.ws_handle = nullptr; + kargs.ptr_ws = nullptr; kargs.m = M; kargs.n = N; kargs.k = K; kargs.batch = batch; const int num_tiles_m = M / Traits::B_M; const int num_tiles_n = N / Traits::B_N; @@ -2022,6 +2023,7 @@ def gen_bmm_mxscale_minterleave_instance( host_extra = ( ",\n aiter_tensor_t &x_scale," "\n aiter_tensor_t &w_scale," + "\n std::optional workspace," "\n int splitK" ) for dtype in k.output_dtypes: @@ -2100,6 +2102,7 @@ def _emit_bmm_specialized( host_extra = ( ",\n aiter_tensor_t &x_scale," "\n aiter_tensor_t &w_scale," + "\n std::optional workspace," "\n int splitK" ) for dtype in k.output_dtypes: @@ -2139,9 +2142,11 @@ def _emit_bmm_specialized( aiter_tensor_t &Y, aiter_tensor_t &x_scale, aiter_tensor_t &w_scale, + std::optional workspace, int @@SPLITK_ARG@@) { using Traits = @@NAME@@_Traits; + AITER_CHECK(!workspace.has_value(), "@@NAME@@ does not use workspace"); """ _BMM_SPEC_KARGS = r""" @@ -2150,7 +2155,7 @@ def _emit_bmm_specialized( opus_gemm_scale_splitk_kargs_gfx950 kargs{}; kargs.ptr_a = O.data_ptr(); kargs.ptr_b = wo_a.data_ptr(); - kargs.ws_handle = nullptr; + kargs.ptr_ws = nullptr; kargs.m = M; kargs.n = N; kargs.k = K; kargs.batch = batch; kargs.stride_a = (int)O.stride(0); kargs.stride_b = (int)wo_a.stride(1); @@ -2171,8 +2176,9 @@ def _emit_bmm_specialized( # ---- wave8n2 (kid 132) ---- _BMM_WAVE8N2_LAUNCHER_BODY = ( - _BMM_SPEC_SIG.replace("@@SPLITK_ARG@@", "/*splitK*/") - + r""" const int M = O.size(0); + _BMM_SPEC_SIG.replace("@@SPLITK_ARG@@", "splitK") + + r""" AITER_CHECK(splitK == 1, "@@NAME@@ requires splitK == 1"); + const int M = O.size(0); const int batch = O.size(1); const int N = wo_a.size(1); const int K = O.size(2); @@ -2249,8 +2255,9 @@ def _cppbool(v): # ---- wave4m2_selfload (kids 134/142/148) ---- _BMM_WAVE4M2_LAUNCHER_BODY = ( - _BMM_SPEC_SIG.replace("@@SPLITK_ARG@@", "/*splitK*/") - + r""" constexpr bool SKIP_SCALE_WAIT = @@SSW@@; + _BMM_SPEC_SIG.replace("@@SPLITK_ARG@@", "splitK") + + r""" AITER_CHECK(splitK == 1, "@@NAME@@ requires splitK == 1"); + constexpr bool SKIP_SCALE_WAIT = @@SSW@@; constexpr bool PACK_SCALE_ON_DEMAND = @@PSOD@@; const int M = O.size(0); const int batch = O.size(1); @@ -2355,7 +2362,7 @@ def gen_bmm_mxscale_wave4m2_selfload_instance( opus_gemm_scale_splitk_kargs_gfx950 kargs{}; kargs.ptr_a = O.data_ptr(); kargs.ptr_b = wo_a.data_ptr(); - kargs.ws_handle = nullptr; + kargs.ptr_ws = nullptr; kargs.m = M; kargs.n = N; kargs.k = K; kargs.batch = batch; const int num_tiles_m = M / Traits::B_M; const int num_tiles_n = N / Traits::B_N; @@ -2397,7 +2404,8 @@ def gen_bmm_mxscale_wave4m2_selfload_instance( """ _BMM_MOUTER_LAUNCHER_BODY = ( - _BMM_SPEC_SIG.replace("@@SPLITK_ARG@@", "/*splitK*/") + _BMM_SPEC_SIG.replace("@@SPLITK_ARG@@", "splitK") + + ' AITER_CHECK(splitK == 1, "@@NAME@@ requires splitK == 1");\n' + _BMM_MOUTER_CHECKS + _BMM_MOUTER_KARGS + " const int m_per_wg = (num_tiles_m >= 16) ? 2 : 1;\n" @@ -2408,6 +2416,7 @@ def gen_bmm_mxscale_wave4m2_selfload_instance( # num_tiles_m]); reuses the same mouter kernel. _BMM_MOUTER_TUNABLE_LAUNCHER_BODY = ( _BMM_SPEC_SIG.replace("@@SPLITK_ARG@@", "splitK") + + ' AITER_CHECK(splitK >= 1, "@@NAME@@ requires splitK >= 1");\n' + _BMM_MOUTER_CHECKS + _BMM_MOUTER_KARGS + " int m_per_wg = splitK;\n" @@ -2511,10 +2520,13 @@ def gen_bmm_mxscale_mouter_tunable_instance( aiter_tensor_t &Y, aiter_tensor_t &x_scale, aiter_tensor_t &w_scale, - int /*splitK*/) + std::optional workspace, + int splitK) { using Bf16Traits = @@NAME@@_Bf16Traits; using Fp32Traits = @@NAME@@_Fp32Traits; + AITER_CHECK(splitK == 1, "@@NAME@@ requires splitK == 1"); + AITER_CHECK(!workspace.has_value(), "@@NAME@@ does not use workspace"); const int M = O.size(0); const int batch = O.size(1); const int N = wo_a.size(1); @@ -2624,6 +2636,7 @@ def gen_bmm_mxscale_pipeline_instance( host_extra = ( ",\n aiter_tensor_t &x_scale," "\n aiter_tensor_t &w_scale," + "\n std::optional workspace," "\n int splitK" ) for dtype in k.output_dtypes: @@ -2663,6 +2676,7 @@ def gen_bmm_mxscale_pipeline_instance( aiter_tensor_t &Y, aiter_tensor_t &x_scale, aiter_tensor_t &w_scale, + std::optional workspace, int splitK) { using Traits = @@NAME@@_Traits; @@ -2704,7 +2718,7 @@ def gen_bmm_mxscale_pipeline_instance( opus_gemm_scale_splitk_kargs_gfx950 kargs{}; kargs.ptr_a = O.data_ptr(); kargs.ptr_b = wo_a.data_ptr(); - kargs.ws_handle = nullptr; + kargs.ptr_ws = nullptr; kargs.m = M; kargs.n = N; kargs.k = K; kargs.batch = batch; kargs.split_k = split_k; kargs.stride_a = (int)O.stride(0); @@ -2723,6 +2737,8 @@ def gen_bmm_mxscale_pipeline_instance( dim3 grid_main(num_tiles_m * num_tiles_n * split_k, 1, batch); dim3 block_main(Traits::BLOCK_SIZE); if (no_split_k) { + AITER_CHECK(!workspace.has_value(), + "@@NAME@@ splitK == 1 does not use workspace"); kargs.ptr_c = Y.data_ptr(); kargs.stride_c = (int)Y.stride(0); kargs.stride_c_batch = (int)Y.stride(1); @@ -2736,34 +2752,19 @@ def gen_bmm_mxscale_pipeline_instance( return; } - extern opus_splitk_ws_handle* opus_splitk_ws_get(hipStream_t, bool); - hipStreamCaptureStatus capture_status = hipStreamCaptureStatusNone; - HIP_CALL(hipStreamIsCapturing(stream, &capture_status)); - const bool capturing = (capture_status != hipStreamCaptureStatusNone); - auto* ws_handle = opus_splitk_ws_get(stream, /*allow_create=*/!capturing); - const size_t ws_bytes = counter_offset + counter_bytes; - if (ws_handle->ptr == nullptr || ws_bytes > ws_handle->bytes) { - AITER_CHECK(!capturing, - "splitk workspace grow inside HIP graph capture is not supported"); - void* new_ptr = nullptr; - const size_t kGrowAlign = (size_t)4 * 1024 * 1024; - size_t grow_bytes = ((ws_bytes + kGrowAlign - 1) / kGrowAlign) * kGrowAlign; - HIP_CALL(hipMalloc(&new_ptr, grow_bytes)); - if (ws_handle->ptr != nullptr) { - HIP_CALL(hipDeviceSynchronize()); - HIP_CALL(hipFree(ws_handle->ptr)); - } - ws_handle->ptr = new_ptr; - ws_handle->bytes = grow_bytes; - } - kargs.ws_handle = ws_handle; + const size_t required_numel = (ws_bytes + sizeof(float) - 1) / sizeof(float); + AITER_CHECK(workspace.has_value(), + "@@NAME@@ splitK > 1 requires workspace"); + void* workspace_ptr = opus_validate_workspace( + workspace.value(), O, AITER_DTYPE_fp32, required_numel, 16, "@@NAME@@"); + kargs.ptr_ws = workspace_ptr; kargs.ptr_c = Y.data_ptr(); kargs.stride_c = (int)Y.stride(0); kargs.stride_c_batch = (int)Y.stride(1); kargs.counter_offset_bytes = counter_offset; - HIP_CALL(hipMemsetAsync(static_cast(ws_handle->ptr) + counter_offset, + HIP_CALL(hipMemsetAsync(static_cast(workspace_ptr) + counter_offset, 0, counter_bytes, stream)); if (Y.dtype() == AITER_DTYPE_bf16) { @@KERNEL@@ @@ -2807,7 +2808,7 @@ def gen_bmm_mxscale_fused_instance( kargs_name, da, db, - instance_impl_preamble(), + instance_impl_preamble('\n#include "opus_gemm_common.cuh"'), instance_impl_host_tu_split( traits_header, pipeline_header, tpl, kernel_func, fn ), diff --git a/csrc/opus_gemm/gen_instances.py b/csrc/opus_gemm/gen_instances.py index 4ca863fe01..a43a73aea2 100644 --- a/csrc/opus_gemm/gen_instances.py +++ b/csrc/opus_gemm/gen_instances.py @@ -8,15 +8,12 @@ from pathlib import Path import pandas as pd -import torch from codegen import gen_instances_gfx942 as _gfx942 # noqa: F401 -# Import for side-effect: each arch module self-registers into EMIT_REGISTRY -# and ARCH_MAP_REGISTRY at import time. +# Architecture modules register their code emitters at import time. from codegen import gen_instances_gfx950 as _gfx950 # noqa: F401 from codegen import gen_instances_gfx1250 as _gfx1250 # noqa: F401 from codegen.common import ( - _A16W16_CO_TAGS, _A16W16_TAGS, _GFX942_A16W16_TAGS, _NOSPLIT, @@ -27,7 +24,9 @@ kid_arch as _kid_arch_common, ) from opus_gemm_common import ( - HEURISTIC_DEFAULT_KIDS, + BMM_MXSCALE_KIDS, + DEFAULT_COMPILED_KIDS, + OPUS_MANDATORY_A8_KIDS, OpusGemmInstance, a8w8_kernels_list, a8w8_mxscale_bmm_kernel_lists, @@ -36,17 +35,18 @@ a16w16_flatmm_splitk_kernels_list, a16w16_kernels_list, a16w16_mono_tile_kernels_list, - default_kernels_dict, + default_compiled_kids_for_arch, gfx942_a8w8_kernels_list, gfx942_nosplit_kernels_list, gfx942_splitk_kernels_list, - heuristic_kids_for_arch, + gfx1250_4wave_co_kernels_list, + gfx1250_clusterlaunch_kernels_list, + gfx1250_kernels_list, + gfx1250_splitk_fuse_kernels_list, kernels_list, ) -# Cross-arch maps merged from per-arch contributions. Each arch module -# registers its piece into ARCH_MAP_REGISTRY at import; we merge gfx950 first -# (legacy default) then overlay gfx942 entries. +# Merge the codegen maps registered by each architecture. PIPELINE_HEADER_MAP = { **get_arch_map("gfx950", "pipeline_header"), **get_arch_map("gfx942", "pipeline_header"), @@ -75,25 +75,29 @@ "gfx950": { "forward_decl_include": '#include "gfx950/opus_gemm_traits_a16w16_gfx950.cuh"\n', "kernel": "splitk_reduce_kernel", - "ws_arg": "const opus_splitk_ws_handle* ws_handle", - "ws_type": "const opus_splitk_ws_handle*", + "ws_arg": "const void* ws_ptr", + "ws_type": "const void*", "baseline_has_oob": (True, False), }, "gfx942": { "forward_decl_include": '#include "gfx942/a16w16/opus_gemm_traits_a16w16.cuh"\n', "kernel": "splitk_reduce_kernel_fallback", - "ws_arg": "const opus_splitk_ws_handle* ws_handle", - "ws_type": "const opus_splitk_ws_handle*", + "ws_arg": "const void* ws_ptr", + "ws_type": "const void*", "baseline_has_oob": (True,), }, "gfx1250": { - # gfx1250 cluster/TDM split-K: workspace allocated externally (torch.empty) - # and passed as a direct void* pointer (no ws_handle indirection). + # gfx1250 cluster/TDM split-K: exact-kid bf16/fp32 workspace + separate + # compile-time-split reduce kernel. The shared generator emits both + # workspace types; gen_instances_gfx1250.py adds mixed fp32-bias/bf16-Y. + # Distinct kernel NAME (splitk_reduce_kernel_gfx1250) keeps it from + # colliding with gfx950 in a multi-arch build. "forward_decl_include": '#include "gfx1250/opus_gemm_traits_a16w16_gfx1250.cuh"\n', "kernel": "splitk_reduce_kernel_gfx1250", "ws_arg": "const void* ws_ptr", "ws_type": "const void*", "baseline_has_oob": (True, False), + "forward_decl_extra_template_params": ", int SPLIT_K_, typename D_WS_", }, } @@ -102,40 +106,52 @@ def _kid_name_arch(kid_name): - """Resolve a kid's arch from its symbol name. - - Classified by the `opus_gemm__*` prefix; legacy names carry no arch - token (a16w16 flatmm / persistent / mono_tile, and the opus_bmm_* family) - and are gfx950, matching kid_arch's default. - """ - for ap in SPLITK_REDUCE_ARCHES: - if kid_name.startswith(f"opus_gemm_{ap}_"): - return ap + """Resolve a generated symbol's owning architecture.""" + for arch_prefix in SPLITK_REDUCE_ARCHES: + if kid_name.startswith(f"opus_gemm_{arch_prefix}_"): + return arch_prefix return LEGACY_OPUS_ARCH def _own_arch_device_pass_guard(arch): - """Open/close guard admitting only `arch`'s device pass, plus the host pass. - - In a mixed build (GPU_ARCHS=gfx950;gfx1250) hipcc runs every TU through - one device pass per offload arch, so without a guard gfx950 kid instances - get instantiated for gfx1250 as well: the gfx950 traits then compute their - layouts off a 32-wide wave and trip `BLOCK_SIZE == 4 * get_warp_size()`, - and the reverse direction hits gfx1250-only kernel attributes. A kid is - only ever launched on the arch it was generated for, so the foreign device - pass has nothing to contribute -- guarding out the #include as well leaves - it an empty TU that never parses another arch's headers. - - The host pass stays inside the guard: it is what emits the __device_stub__ - symbols the fused host TU's <<<>>> calls link against, and it is arch - independent. - """ + """Admit the host pass and only this kid's owning device pass.""" return ( f"#if !defined(__HIP_DEVICE_COMPILE__) || defined(__{arch}__)\n", f"#endif // host pass or {arch} device pass\n", ) +def _splitk_reduce_baseline_instantiations( + reduce_kernel, + ws_ptr_type, + has_oob, + vec=16, + block=64, + split_ks=(None,), + workspace_types=(None,), +): + has_oob_str = "true" if has_oob else "false" + configs = ( + ("__bf16", "true", "__bf16"), + ("__bf16", "false", "__bf16"), + ("float", "true", "float"), + ("float", "false", "float"), + ) + out = f"// HAS_OOB={has_oob_str} variants\n" + for split_k in split_ks: + for workspace_type in workspace_types: + tail = "" if split_k is None else f", {split_k}, {workspace_type}" + for out_type, has_bias, bias_type in configs: + out += ( + f"template __global__ void {reduce_kernel}<" + f"{vec}, {block}, {out_type}, {has_bias}, {bias_type}, " + f"{has_oob_str}{tail}>(\n" + f" {ws_ptr_type}, {out_type}*, int, int, int, int, int, int,\n" + f" const {bias_type}*, int);\n" + ) + return out + + # Arches that own an opus_gemm_arch_*.cuh dispatch header, i.e. one set of # lookup tables each. Every generated lookup macro is emitted once per arch and # expanded by that arch's header only: the arches disagree on the a16w16 @@ -150,35 +166,6 @@ def _own_arch_device_pass_guard(arch): LOOKUP_MACRO_ARCHES = ("gfx950", "gfx942", "gfx1250") -def _splitk_reduce_baseline_instantiations( - reduce_kernel, ws_ptr_type, has_oob, vec=16, block=64, split_ks=(None,), d_ws=None -): - # gfx1250 tunes the reduce to VEC=8/BLOCK=128 (coalesced dwordx4 bf16 store, - # no cross-lane shuffle), a per-kid partial type (d_ws), and a COMPILE-TIME split_k (SPLIT_K_ template) dispatched per - # value -> split_ks lists every value the launch helper switches on (0 = the - # runtime-`split_k` fallback, 1..16 = fully-unrolled). gfx950/gfx942 keep the - # legacy 6-param VEC=16/BLOCK=64 fp32-workspace form (split_ks=(None,)). - has_oob_str = "true" if has_oob else "false" - out = f"// HAS_OOB={has_oob_str} variants\n" - for sk in split_ks: - tail = "" if sk is None else f", {sk}, {d_ws}" - out += ( - f"template __global__ void {reduce_kernel}<{vec}, {block}, __bf16, true, __bf16, {has_oob_str}{tail}>(\n" - f" {ws_ptr_type}, __bf16*, int, int, int, int, int, int,\n" - f" const __bf16*, int);\n" - f"template __global__ void {reduce_kernel}<{vec}, {block}, __bf16, false, __bf16, {has_oob_str}{tail}>(\n" - f" {ws_ptr_type}, __bf16*, int, int, int, int, int, int,\n" - f" const __bf16*, int);\n" - f"template __global__ void {reduce_kernel}<{vec}, {block}, float, true, float, {has_oob_str}{tail}>(\n" - f" {ws_ptr_type}, float*, int, int, int, int, int, int,\n" - f" const float*, int);\n" - f"template __global__ void {reduce_kernel}<{vec}, {block}, float, false, float, {has_oob_str}{tail}>(\n" - f" {ws_ptr_type}, float*, int, int, int, int, int, int,\n" - f" const float*, int);\n" - ) - return out - - def _pipeline_header_for(k): if getattr(k, "is_4g_safe", False): # 4g_safe is gfx950-only (no gfx942 sibling pipeline exists). @@ -212,48 +199,25 @@ def _kernel_func_for(k): **{tag: ("bf16_t", "bf16_t") for tag in _A16W16_TAGS}, } -# All a16w16 tags share the 4-arg (XQ, WQ, Y, int splitK) lookup-table slot. -A16W16_TUNE_TAGS = set(_A16W16_TAGS) -# ... except the pre-compiled (.co) families, which get their own flat-array -# dispatch table. Their launcher signature has no workspace, so a function -# pointer to one does not fit the arch's OpusA16W16NoscaleKernel type and it -# cannot share a table with the split-K kids. -A16W16_CO_TUNE_TAGS = set(_A16W16_CO_TAGS) -A8W8_TUNE_TAGS = {"a8w8_blockscale_bpreshuffle_singlebuf"} -# NOSCALE: 3-arg launchers (a16w16 family + a8w8 non-scale). -NOSCALE_TAGS = A16W16_TUNE_TAGS | {"a8w8"} - -# SplitK tags live in the dispatch slot; each instance's traits pick -# the actual workspace dtype and the reduce launcher writes the requested Y. -# For the two _ws families that slot is NOT the output dtype at all -- it is -# the split-K PARTIAL type (traits D_C), which the main kernel stores and the -# reduce reads, so it must be instantiated from the kid's own -# splitk_workspace_dtype. Getting this wrong is a page fault, not a wrong -# number: the host sizes the buffer from the same field. +# A16W16 uses separate direct-output and workspace launcher tables. +A16W16_KID_DISPATCH_TAGS = set(_A16W16_TAGS) +A8W8_BPRESHUFFLE_TAGS = {"a8w8_blockscale_bpreshuffle_singlebuf"} +# Three-tensor launchers: A16W16 and A8W8 no-scale. +NOSCALE_TAGS = A16W16_KID_DISPATCH_TAGS | {"a8w8"} + +# Split-K tags live in the workspace dispatch table and use their existing +# host specialization; each instance's traits pick the actual +# workspace dtype. Fused kids write Y in-kernel; the other tags launch a +# standalone reducer. SPLITK_TAGS = { "a16w16_flatmm_splitk", "a16w16_cluster_tdm_splitk_ws", "a16w16_clusterlaunch_tdm_splitk_ws", - # fused single-kernel split-K: lookup still forces (D_C=fp32 traits), - # but its launcher NAME avoids the "_splitk_" substring so the reduce-TU - # detection (:867 / :799) never emits a reduce kernel for it. "a16w16_clusterlaunch_tdm_splitk_fuse", + "a16w16_em3en4_lds1_pgr2_sk", *_SPLITK, } -_WS_PARTIAL_TAGS = { - "a16w16_cluster_tdm_splitk_ws", - "a16w16_clusterlaunch_tdm_splitk_ws", -} - - -def _ws_partial_ctype(k): - """The kid's split-K partial ctype, or None if its slot is a real dtype.""" - if k.kernel_tag not in _WS_PARTIAL_TAGS: - return None - return getattr(k, "splitk_workspace_dtype", "fp32_t") - - TRAITS_NAME_MAP = { **get_arch_map("gfx950", "traits_name"), **get_arch_map("gfx942", "traits_name"), @@ -268,11 +232,6 @@ def _ws_partial_ctype(k): def _kargs_template_vars(kernel_tag, kargs_name): - # a8w8_mxscale BMM flatmm splitK kernel has two extra compile-time booleans - # (DIRECT_ONLY, PREFETCH_SCALE) plus a non-void D_OUT after Traits. The fused - # host TU must forward-declare all four template params so the launcher body - # (which launches gemm_a8w8_mxscale_flatmm_splitk_kernel) compiles without pulling in the device pipeline header. if kernel_tag in ( "a8w8_mxscale_bmm_flatmm_splitk", "a8w8_mxscale_bmm_fused", @@ -282,17 +241,9 @@ def _kargs_template_vars(kernel_tag, kargs_name): ", typename D_OUT, bool DIRECT_ONLY, bool PREFETCH_SCALE, bool PRELOAD_SF_LDS", kargs_name, ) - # BMM M-tile-interleaved kernel: . The - # fused host TU must forward-declare all three template params so the launcher - # body's gemm_a8w8_mxscale_flatmm_minterleave_kernel - # <<<...>>> call compiles without the device pipeline header. if kernel_tag == "a8w8_mxscale_bmm_minterleave": return "", ", typename D_OUT, bool SKIP_SCALE_WAIT", kargs_name - # BMM specialized pipelines: forward-declare the exact kernel template params - # so the fused host TU's <<<...>>> call compiles against only the traits header. if kernel_tag == "a8w8_mxscale_bmm_pipeline": - # scale-pipeline kernels are templated on a single Traits (output dtype is - # baked into the traits tuple) -> no extra template params. return "", "", kargs_name if kernel_tag in ( "a8w8_mxscale_bmm_mouter", @@ -356,11 +307,15 @@ def instance_impl_host_tu_split( ) -# Launcher signature tails after Y. -A16W16_TUNE_HOST_EXTRA = ",\n std::optional,\n int" -A8W8_SCALE_HOST_EXTRA = ( - ",\n std::optional x_scale," - "\n std::optional w_scale" +# Extra parameters appended to each generated launcher signature. +A16W16_LAUNCH_HOST_EXTRA = ",\n std::optional,\n int" +A16W16_WORKSPACE_LAUNCH_HOST_EXTRA = ( + ",\n aiter_tensor_t &workspace," + "\n std::optional," + "\n int" +) +A8W8_BLOCKSCALE_HOST_EXTRA = ( + ",\n aiter_tensor_t &x_scale," "\n aiter_tensor_t &w_scale" ) @@ -374,6 +329,19 @@ def _make_host_decl(kid_name, dtype, host_extra_params): ) +def _make_a8w8_bpreshuffle_host_decl(kid_name, dtype, _host_extra_params): + """Emit the ``XQ,WQ,x_scale,w_scale,Y`` host declaration.""" + return ( + f"template void\n" + f"{kid_name}<{dtype}>(\n" + f" aiter_tensor_t &XQ,\n" + f" aiter_tensor_t &WQ,\n" + f" aiter_tensor_t &x_scale,\n" + f" aiter_tensor_t &w_scale,\n" + f" aiter_tensor_t &Y);\n" + ) + + def _make_device_decl( kid_name, dtype, kernel_func, kargs_name, kargs_explicit_param="" ): @@ -384,23 +352,21 @@ def _make_device_decl( def _record_one_instantiation( - self_obj, k, kernel_func, kargs_name, host_extra, kargs_explicit_param="" + self_obj, + k, + kernel_func, + kargs_name, + host_extra, + kargs_explicit_param="", + host_decl_factory=_make_host_decl, ): - """Record (host_decl, device_decl) for every dtype the kid is referenced with. - - For the _ws split-K families the template slot is the partial type rather - than the output dtype, so instantiate the kid's splitk_workspace_dtype -- - output_dtypes would give the wrong one and the dispatch table's reference - would not link. - """ - ws_ctype = _ws_partial_ctype(k) - dtypes = (ws_ctype,) if ws_ctype is not None else tuple(k.output_dtypes) - for CDtype in dtypes: + """Record (host_decl, device_decl) for every (kid, dtype) in k.output_dtypes.""" + for CDtype in k.output_dtypes: self_obj._host_instantiations.append( { "kid_name": k.name, "dtype": CDtype, - "host_decl": _make_host_decl(k.name, CDtype, host_extra), + "host_decl": host_decl_factory(k.name, CDtype, host_extra), } ) self_obj._device_instantiations.append( @@ -538,9 +504,11 @@ def gen_instance(self, k: OpusGemmInstance): "record_one_instantiation": _record_one_instantiation, "make_host_decl": _make_host_decl, "make_device_decl": _make_device_decl, - "A16W16_TUNE_HOST_EXTRA": A16W16_TUNE_HOST_EXTRA, - "A8W8_SCALE_HOST_EXTRA": A8W8_SCALE_HOST_EXTRA, - "A16W16_TUNE_TAGS": A16W16_TUNE_TAGS, + "A16W16_LAUNCH_HOST_EXTRA": A16W16_LAUNCH_HOST_EXTRA, + "A16W16_WORKSPACE_LAUNCH_HOST_EXTRA": (A16W16_WORKSPACE_LAUNCH_HOST_EXTRA), + "A8W8_BLOCKSCALE_HOST_EXTRA": A8W8_BLOCKSCALE_HOST_EXTRA, + "make_a8w8_bpreshuffle_host_decl": (_make_a8w8_bpreshuffle_host_decl), + "A16W16_KID_DISPATCH_TAGS": A16W16_KID_DISPATCH_TAGS, "BIAS_HOST_VALIDATE": self.BIAS_HOST_VALIDATE, } dispatch_emit(self, k, **emit_kwargs) @@ -576,378 +544,185 @@ def gen_instance(self, k: OpusGemmInstance): }} """ - def gen_lookup_dict(self, kernels_dict): - """Emit opus_gemm_lookup.h with the (M,N,K)->kernel macros. - - One macro per (CTYPE, arch): see LOOKUP_MACRO_ARCHES. - - Tuned-CSV driven lookup consumed by opus_gemm.cu's runtime - `opus_dispatch_a16w16`. The BF16 / FP32 split - mirrors `gen_a16w16_tune_lookup` and exists because splitk kids - (200..210) are only emitted as `` (their traits - static_assert D_C==float, so referencing `splitk` - produces a linker error). - - Outdtype-aware bucketing - ------------------------ - kernels_dict tuple keys carry the outdtype string in slot 3 - ((M, N, K, outdtype_str, arch), produced by get_tune_dict). The BF16 - macro picks up rows whose outdtype is "torch.bfloat16" and the - FP32 macro picks up rows whose outdtype is "torch.float32"; - same-(M,N,K) rows with different outdtypes therefore land in - different macros and the two C++ maps can resolve to different - kernels for the same shape. Legacy CSVs without an outdtype - column are normalized to bf16 by get_tune_dict, so they only - populate the BF16 map -- matching pre-outdtype-split behavior. - - Per-kid template argument rule: - - * a16w16 kid 4..9 -> `` (both bf16/fp32 exist). - * a16w16_flatmm 100..115 -> `` (both exist). - * a16w16_flatmm_splitk -> always ``. Splitk rows - with outdtype=bf16 land in the BF16 map (with forced - template arg) and rows with outdtype=fp32 land in - the FP32 map (also with ). Both work because the - splitk reduce kernel handles the cast / passthrough at - launch time based on the actual Y dtype. - """ - # Sorted flat-array layout (was: {(M,N,K), kernel} initializer list for std::unordered_map). + def gen_a16w16_kid_dispatch(self, kernels_dict): + """Emit per-arch A16W16 direct and workspace launcher tables.""" HEADER = """#pragma once // SPDX-License-Identifier: MIT // Copyright (C) 2025-2026, Advanced Micro Devices, Inc. All rights reserved. // -// Auto-generated. Do not edit. See gen_instances.py:gen_lookup_dict. +// Auto-generated. Do not edit. See gen_instances.py:gen_a16w16_kid_dispatch. // -// Per-(CTYPE, arch) sorted flat arrays for (M,N,K)->kernel runtime dispatch. -// Same (M,N,K) can resolve to different kernels in the BF16 vs FP32 -// tables because get_tune_dict keys winners on (M, N, K, outdtype_str, arch) -// and gen_lookup_dict buckets the rows into per-(CTYPE, arch) macros below. -// splitk kids appear in either table with their dispatch template forced -// to ; their traits pick the workspace dtype and the reduce -// launcher writes the requested Y dtype. -// -// Lookup is std::lower_bound on the lex-ordered (M, N, K) key. See -// opus_gemm_arch_gfx950.cuh for the dispatch wrapper. -""" - - ENTRY_MATCH_CTYPE = """\ - {{ {{{M}, {N}, {K}}}, &{kernel_name} }}, \\ +// Per-arch sorted flat arrays for strict kid dispatch. Non-workspace tables +// contain five-argument OpusA16W16Kernel pointers. Workspace tables contain +// six-argument OpusA16W16WorkspaceKernel pointers. Never combine them with +// the five-argument table. """ - ENTRY_FORCE_FP32 = """\ - {{ {{{M}, {N}, {K}}}, &{kernel_name} }}, \\ + NON_WORKSPACE_ENTRY = """\ + {{ {kid}, &{kernel_name} }}, \\ """ - # _ws families: the template slot is the split-K partial type, per-kid. - ENTRY_WS_PARTIAL = """\ - {{ {{{M}, {N}, {K}}}, &{kernel_name}<{ctype}> }}, \\ + WORKSPACE_ENTRY = """\ + {{ {kid}, &{kernel_name} }}, \\ """ - # Map ctype short name -> CSV outdtype string emitted by the - # tuner's result_to_df. - ctype_to_outdtype = { - "bf16_t": "torch.bfloat16", - "fp32_t": "torch.float32", - } - - def _emit_map(f, macro_name: str, ctype: str, arch: str): - # No body line break between `\` and the first entry; macro continuation requires every line - # that participates in the definition ... - f.write(f"#define {macro_name}(CTYPE) \\\n") - target_outdtype = ctype_to_outdtype.get(ctype) - # Collect all (M, N, K, kernel_name, is_splitk) rows for this - # CTYPE first, so we can sort lex on (M, N, K) before emitting. - rows = [] - for mnk, k in kernels_dict.items(): - if self.istune and isinstance(mnk, int): - # tune mode shouldn't reach here (gen_lookup_dict is - # for the runtime (M,N,K) map). Skip defensively. - continue - if not (isinstance(mnk, tuple) and mnk[0] > 0): - continue - if len(mnk) >= 4: - row_outdtype = str(mnk[3]) - if target_outdtype is not None and row_outdtype != target_outdtype: - continue - # Pre-compiled (.co) kids have their own function-pointer type - # (no workspace argument), so they cannot go in this table -- - # they go in the parallel CO table emitted by _emit_co_map() - # below, which the gfx1250 dispatch consults first. - if k.kernel_tag in A16W16_CO_TUNE_TAGS: - continue - is_splitk = k.kernel_tag in SPLITK_TAGS - if not is_splitk and ctype not in k.output_dtypes: - continue - if _kid_arch_common(k) != arch: - continue - rows.append( - ( - int(mnk[0]), - int(mnk[1]), - int(mnk[2]), - k.name, - is_splitk, - _ws_partial_ctype(k), - ) - ) - - rows.sort(key=lambda r: (r[0], r[1], r[2])) - n = len(rows) - for i, (M, N, K, name, is_splitk, ws_ctype) in enumerate(rows): - if ws_ctype is not None: - line = ENTRY_WS_PARTIAL.format( - M=M, N=N, K=K, kernel_name=name, ctype=ws_ctype - ) - else: - entry = ENTRY_FORCE_FP32 if is_splitk else ENTRY_MATCH_CTYPE - line = entry.format(M=M, N=N, K=K, kernel_name=name) - if i == n - 1: - # Last entry: drop the trailing `\` so the macro - # ends cleanly. Strip the line's continuation. + def _write_rows(f, macro_name, rows, entry, function_like=False): + f.write(f"#define {macro_name}_SIZE {len(rows)}\n") + macro_suffix = "(CTYPE)" if function_like else "" + if not rows: + f.write(f"#define {macro_name}{macro_suffix}\n\n") + return + f.write(f"#define {macro_name}{macro_suffix} \\\n") + for index, (kid, name) in enumerate(rows): + line = entry.format(kid=kid, kernel_name=name) + if index == len(rows) - 1: line = line.rstrip().rstrip("\\").rstrip() + "\n" f.write(line) f.write("\n") - def _emit_co_map(f, macro_name: str, arch: str): - """The (M, N, K) table for the pre-compiled (.co) families. - - Its own macro (and no CTYPE parameter) because these launchers take - no workspace argument, so their function pointers do not fit the - arch's shared entry type. Without this table a tuned CSV row naming - a .co winner was silently dropped from the runtime lookup: the - Python path honoured it while the C++ `opus_gemm` entry fell back to - the heuristic for the same shape. - """ - f.write(f"#define {macro_name}() \\\n") - rows = [] - for mnk, k in kernels_dict.items(): - if not (isinstance(mnk, tuple) and mnk[0] > 0): - continue - if k.kernel_tag not in A16W16_CO_TUNE_TAGS: - continue - if _kid_arch_common(k) != arch: - continue - # One output dtype per co family, and it is the real C dtype (no - # workspace to stand in for it), so rows whose CSV outdtype is - # something else cannot run on this kid. - if len(mnk) >= 4: - want = ctype_to_outdtype.get(k.output_dtypes[0]) - if want is not None and str(mnk[3]) != want: - continue - rows.append( - (int(mnk[0]), int(mnk[1]), int(mnk[2]), k.name, k.output_dtypes[0]) - ) - - rows.sort(key=lambda r: (r[0], r[1], r[2])) - n = len(rows) - for i, (M, N, K, name, ctype) in enumerate(rows): - line = f" {{ {{{M}, {N}, {K}}}, &{name}<{ctype}> }}, \\\n" - if i == n - 1: - line = line.rstrip().rstrip("\\").rstrip() + "\n" - f.write(line) - f.write("\n") - - with open(os.path.join(self.working_path, "opus_gemm_lookup.h"), "w") as f: - f.write(HEADER) - for arch in LOOKUP_MACRO_ARCHES: - suffix = arch.upper() - _emit_map( - f, f"GENERATE_OPUS_LOOKUP_TABLE_BF16_{suffix}", "bf16_t", arch - ) - _emit_map( - f, f"GENERATE_OPUS_LOOKUP_TABLE_FP32_{suffix}", "fp32_t", arch - ) - _emit_co_map(f, f"GENERATE_OPUS_LOOKUP_TABLE_CO_{suffix}", arch) - - def gen_a16w16_tune_lookup(self, kernels_dict): - """Emit opus_gemm_a16w16_tune_lookup.h with int-ID-to-kernel maps for tuning. - - One macro per (CTYPE, arch): see LOOKUP_MACRO_ARCHES. - - Three a16w16-family tags share the 4-arg launcher signature - (XQ, WQ, Y, int splitK): - * a16w16 (split-barrier) - output_dtypes=["fp32_t", "bf16_t"] - * a16w16_flatmm (warp-spec) - output_dtypes=["bf16_t", "fp32_t"] - * a16w16_flatmm_splitk - output_dtypes=["fp32_t"] ONLY - (main kernel writes fp32 workspace; Y=bf16 via reduce kernel. - Traits static_assert D_C=float, so no instantiation - exists for these kids.) - - The bf16 lookup map therefore must NOT reference splitk kids (their - specialization is never instantiated -> linker error). The - dispatcher in opus_gemm.cu forces kid>=200 to the branch - anyway, so having them absent from the bf16 map is correct. - - Emit the macros side by side, gated on each kid's output_dtypes set - and on its arch. - """ - # Same flat-array design as gen_lookup_dict, keyed on int kid instead of (M,N,K). - HEADER = """#pragma once -// SPDX-License-Identifier: MIT -// Copyright (C) 2025-2026, Advanced Micro Devices, Inc. All rights reserved. -// -// Auto-generated. Do not edit. See gen_instances.py:gen_a16w16_tune_lookup. -// -// Per-(CTYPE, arch) sorted flat arrays for kid->kernel tune dispatch. Kids whose -// output_dtypes doesn't include CTYPE are omitted from that CTYPE's table -// (splitk kids only live in the fp32 table). See -// opus_gemm_arch_gfx950.cuh for the dispatch wrapper. -""" - ENTRY = """\ - {{ {kid}, &{kernel_name} }}, \\ -""" - - # Pre-compiled (.co) kids take a launcher signature with no workspace, - # so their function pointers do not fit the arch's shared entry type. - # They are emitted into a separate macro with no CTYPE parameter: each - # co family is instantiated for exactly one output dtype. - CO_ENTRY = """\ - {{ {kid}, &{kernel_name}<{ctype}> }}, \\ -""" - - def _emit_map(f, macro_name, ctype, arch): - f.write(f"#define {macro_name}(CTYPE) \\\n") + def _emit_non_workspace_map(f, arch, ctype): rows = [] for kid, k in kernels_dict.items(): - if not (isinstance(kid, int) and k.kernel_tag in A16W16_TUNE_TAGS): + if not ( + isinstance(kid, int) and k.kernel_tag in A16W16_KID_DISPATCH_TAGS + ): continue - if k.kernel_tag in A16W16_CO_TUNE_TAGS: + if _kid_arch_common(k) != arch or k.kernel_tag in SPLITK_TAGS: continue if ctype not in k.output_dtypes: continue if _kid_arch_common(k) != arch: continue - rows.append((kid, k.name, _ws_partial_ctype(k))) + rows.append((kid, k.name)) rows.sort(key=lambda r: r[0]) - n = len(rows) - for i, (kid, name, ws_ctype) in enumerate(rows): - line = ( - CO_ENTRY.format(kid=kid, kernel_name=name, ctype=ws_ctype) - if ws_ctype is not None - else ENTRY.format(kid=kid, kernel_name=name) - ) - if i == n - 1: - line = line.rstrip().rstrip("\\").rstrip() + "\n" - f.write(line) - f.write("\n") + dtype_suffix = "BF16" if ctype == "bf16_t" else "FP32" + macro_name = ( + "GENERATE_A16W16_NONWORKSPACE_KID_DISPATCH_" + f"{arch.upper()}_{dtype_suffix}" + ) + _write_rows(f, macro_name, rows, NON_WORKSPACE_ENTRY, function_like=True) - def _emit_co_map(f, macro_name, arch): - f.write(f"#define {macro_name}() \\\n") + def _emit_workspace_map(f, arch): rows = [] for kid, k in kernels_dict.items(): - if not (isinstance(kid, int) and k.kernel_tag in A16W16_CO_TUNE_TAGS): + if not ( + isinstance(kid, int) and k.kernel_tag in A16W16_KID_DISPATCH_TAGS + ): continue - if _kid_arch_common(k) != arch: + if _kid_arch_common(k) != arch or k.kernel_tag not in SPLITK_TAGS: continue - # One output dtype per co family (asserted, because a second one - # would silently drop an entry here). - assert len(k.output_dtypes) == 1, ( - f"co kid {kid} ({k.name}) must have exactly one output dtype; " - f"got {k.output_dtypes}" - ) - rows.append((kid, k.name, k.output_dtypes[0])) + if "fp32_t" not in k.output_dtypes: + raise ValueError( + f"workspace kid {kid} ({k.name}) has no fp32_t host " + "specialization" + ) + rows.append((kid, k.name)) rows.sort(key=lambda r: r[0]) - n = len(rows) - for i, (kid, name, ctype) in enumerate(rows): - line = CO_ENTRY.format(kid=kid, kernel_name=name, ctype=ctype) - if i == n - 1: - line = line.rstrip().rstrip("\\").rstrip() + "\n" - f.write(line) - f.write("\n") + macro_name = f"GENERATE_A16W16_WORKSPACE_KID_DISPATCH_{arch.upper()}" + _write_rows(f, macro_name, rows, WORKSPACE_ENTRY) with open( - os.path.join(self.working_path, "opus_gemm_a16w16_tune_lookup.h"), "w" + os.path.join(self.working_path, "opus_gemm_a16w16_kid_dispatch.h"), "w" ) as f: f.write(HEADER) - # Use explicit per-CTYPE macro names; the dispatcher in opus_gemm.cu calls the right one from - # each opus_a16w16_tune_dispatch }}, \\ + {{ {kid}, &{kernel_name}<{ctype}> }}, \\ """ - def _emit_map(f, macro_name, ctype): - f.write(f"#define {macro_name}(CTYPE) \\\n") + def _rows(arch, tags, ctype): rows = [] for kid, k in kernels_dict.items(): - if not (isinstance(kid, int) and k.kernel_tag in A8W8_TUNE_TAGS): + if not isinstance(kid, int) or k.kernel_tag not in tags: continue - if ctype not in k.output_dtypes: + if _kid_arch_common(k) != arch or ctype not in k.output_dtypes: continue rows.append((kid, k.name)) rows.sort(key=lambda row: row[0]) + return rows + + def _emit_map(f, macro_name, rows, ctype): + f.write(f"#define {macro_name}_SIZE {len(rows)}\n") + if not rows: + f.write(f"#define {macro_name}\n\n") + return + f.write(f"#define {macro_name} \\\n") for index, (kid, name) in enumerate(rows): - line = entry.format(kid=kid, kernel_name=name) + line = entry.format(kid=kid, kernel_name=name, ctype=ctype) if index == len(rows) - 1: line = line.rstrip().rstrip("\\").rstrip() + "\n" f.write(line) f.write("\n") with open( - os.path.join(self.working_path, "opus_gemm_a8w8_tune_lookup.h"), "w" + os.path.join(self.working_path, "opus_gemm_a8w8_kid_dispatch.h"), "w" ) as f: f.write(header) - _emit_map(f, "GENERATE_A8W8_TUNE_LOOKUP_BF16", "bf16_t") - - def gen_bmm_mxscale_tune_lookup(self, kernels_dict): - """Emit opus_bmm_mxscale_tune_lookup.h: int-kid -> launcher map for the - a8w8_mxscale BMM flatmm split-K family (gfx950-only). - - Mirrors gen_a8w8_tune_lookup, but the kid->name mapping lives in - a8w8_mxscale_bmm_flatmm_splitk_kernels_list (kid-keyed), NOT in - kernels_dict (which is name-keyed for the BMM family so gen_manifest_head - can dedup identical geometries). We iterate the kid-keyed source directly - so every switch kid keeps its historical number, even when two kids share - one launcher symbol (e.g. 0 and 32 -> same geometry -> same &launcher). - - The launcher templates static_assert D_C == float (Y=bf16 is produced by - the reduce kernel from an fp32 workspace), so only the fp32_t - specialization is instantiated -> emit a single fp32 macro. The dispatch - wrapper in opus_bmm.cu combines this with the hand-written specialized - pipelines (mouter / wave*n* / minterleave / pipeline / fused). - """ + _emit_map( + f, + "GENERATE_A8W8_NOSCALE_KID_DISPATCH_GFX950", + _rows("gfx950", {"a8w8"}, "fp32_t"), + "fp32_t", + ) + _emit_map( + f, + "GENERATE_A8W8_BLOCKSCALE_KID_DISPATCH_GFX950", + _rows("gfx950", {"a8w8_scale"}, "fp32_t"), + "fp32_t", + ) + for arch in SPLITK_REDUCE_ARCHES: + for ctype, dtype_suffix in ( + ("bf16_t", "BF16"), + ("fp32_t", "FP32"), + ): + macro_name = ( + "GENERATE_A8W8_BLOCKSCALE_BPRESHUFFLE_KID_DISPATCH_" + f"{arch.upper()}_{dtype_suffix}" + ) + _emit_map( + f, + macro_name, + _rows(arch, A8W8_BPRESHUFFLE_TAGS, ctype), + ctype, + ) + + def gen_bmm_mxscale_kid_dispatch(self): + """Emit the global exact-kid table for gfx950 MXFP8 BMM launchers.""" header = """#pragma once // SPDX-License-Identifier: MIT // Copyright (C) 2025-2026, Advanced Micro Devices, Inc. All rights reserved. -// -// Auto-generated. Do not edit. See gen_instances.py:gen_bmm_mxscale_tune_lookup. -// -// fp32-workspace flat map for a8w8_mxscale BMM flatmm split-K tuning (gfx950). -// See opus_bmm.cu opus_bmm_a8w8_mxscale_tune_dispatch(). +// Auto-generated. Do not edit. """ entry = """\ {{ {kid}, &{kernel_name} }}, \\ """ - # The specialized-pipeline families (minterleave, ...) share the same - # int-kid -> launcher map and uniform launcher signature; concatenate - # their kid-keyed source lists so every kid keeps its historical number. - def _emit_map(f, macro_name, ctype): - f.write(f"#define {macro_name}(CTYPE) \\\n") - rows = [ - (kid, k.name) - for kernels in a8w8_mxscale_bmm_kernel_lists - for kid, k in kernels.items() - if ctype in k.output_dtypes - ] - rows.sort(key=lambda row: row[0]) + rows = sorted( + (kid, instance.name) + for family in a8w8_mxscale_bmm_kernel_lists + for kid, instance in family.items() + if "fp32_t" in instance.output_dtypes + ) + with open( + os.path.join(self.working_path, "opus_bmm_mxscale_kid_dispatch.h"), + "w", + ) as f: + f.write(header) + f.write(f"#define GENERATE_BMM_MXSCALE_KID_DISPATCH_SIZE {len(rows)}\n") + f.write("#define GENERATE_BMM_MXSCALE_KID_DISPATCH(CTYPE) \\\n") for index, (kid, name) in enumerate(rows): line = entry.format(kid=kid, kernel_name=name) if index == len(rows) - 1: @@ -955,12 +730,6 @@ def _emit_map(f, macro_name, ctype): f.write(line) f.write("\n") - with open( - os.path.join(self.working_path, "opus_bmm_mxscale_tune_lookup.h"), "w" - ) as f: - f.write(header) - _emit_map(f, "GENERATE_BMM_MXSCALE_FLATMM_SPLITK_LOOKUP_FP32", "fp32_t") - def gen_manifest_head(self, kernels_dict): # Forward declarations for every launcher symbol the dispatcher references. MANIFEST_HEAD = """#pragma once @@ -970,18 +739,27 @@ def gen_manifest_head(self, kernels_dict): #include #include """ - MANIFEST_SCALE = """ + MANIFEST_BLOCKSCALE = """ template void {kernel_name}( aiter_tensor_t &XQ, aiter_tensor_t &WQ, aiter_tensor_t &Y, - std::optional x_scale, - std::optional w_scale); + aiter_tensor_t &x_scale, + aiter_tensor_t &w_scale); """ - # a8w8 noscale (3 args, no splitK): stays compatible with - # opus_gemm_lookup.h where a8w8 kids live. + MANIFEST_BLOCKSCALE_BPRESHUFFLE = """ +template +void +{kernel_name}( + aiter_tensor_t &XQ, + aiter_tensor_t &WQ, + aiter_tensor_t &x_scale, + aiter_tensor_t &w_scale, + aiter_tensor_t &Y); +""" + # a8w8 noscale (3 args, no splitK) has its own exact-kid table. MANIFEST_NOSCALE_3ARG = """ template void @@ -990,8 +768,8 @@ def gen_manifest_head(self, kernels_dict): aiter_tensor_t &WQ, aiter_tensor_t &Y); """ - # a16w16 family (5 args with optional bias + splitK): shared signature for tune lookup. - MANIFEST_NOSCALE_4ARG = """ + # Non-workspace a16w16 launchers keep the existing five-argument ABI. + MANIFEST_A16W16 = """ template void {kernel_name}( @@ -1001,8 +779,10 @@ def gen_manifest_head(self, kernels_dict): std::optional bias, int splitK); """ - # gfx1250 a16w16 split-K (6 args: workspace tensor passed from Python). - MANIFEST_NOSCALE_6ARG_WS = """ + # External-workspace launchers receive a caller-owned typed tensor. + # This covers both two-stage reducers and gfx1250 fused in-cluster + # reduction. + MANIFEST_A16W16_WORKSPACE = """ template void {kernel_name}( @@ -1013,52 +793,35 @@ def gen_manifest_head(self, kernels_dict): std::optional bias, int splitK); """ - # a8w8_mxscale BMM flatmm split-K launcher: mmajor layout with two fp8 - # scale tensors + an int splitK, dispatched by the hand-written - # opus_bmm.cu switch (not the (M,N,K) lookup table). - MANIFEST_BMM_MXSCALE_SPLITK = """ + MANIFEST_BMM_MXSCALE = """ template void {kernel_name}( - aiter_tensor_t &O, - aiter_tensor_t &wo_a, + aiter_tensor_t &XQ, + aiter_tensor_t &WQ, aiter_tensor_t &Y, aiter_tensor_t &x_scale, aiter_tensor_t &w_scale, + std::optional workspace, int splitK); """ - # The gfx1250 split-K families take the 6-arg (workspace-carrying) - # launcher signature. The pre-compiled .co families do NOT -- they have - # no workspace to pass -- so they fall through to the ordinary 5-arg - # a16w16 declaration below. - GFX1250_SPLITK_TAGS = { - "a16w16_cluster_tdm_splitk_ws", - "a16w16_clusterlaunch_tdm_splitk_ws", - "a16w16_clusterlaunch_tdm_splitk_fuse", - } with open(os.path.join(self.working_path, "opus_gemm_manifest.h"), "w") as f: f.write(MANIFEST_HEAD) for k in kernels_dict.values(): - if k.kernel_tag in ( - "a8w8_mxscale_bmm_flatmm_splitk", - "a8w8_mxscale_bmm_fused", - "a8w8_mxscale_bmm_minterleave", - "a8w8_mxscale_bmm_mouter", - "a8w8_mxscale_bmm_mouter_tunable", - "a8w8_mxscale_bmm_pipeline", - "a8w8_mxscale_bmm_wave8n2", - "a8w8_mxscale_bmm_wave4m2_selfload", - ): - f.write(MANIFEST_BMM_MXSCALE_SPLITK.format(kernel_name=k.name)) - elif k.kernel_tag in A16W16_TUNE_TAGS: - if k.kernel_tag in GFX1250_SPLITK_TAGS: - f.write(MANIFEST_NOSCALE_6ARG_WS.format(kernel_name=k.name)) - else: - f.write(MANIFEST_NOSCALE_4ARG.format(kernel_name=k.name)) - elif k.kernel_tag in NOSCALE_TAGS: + if k.kernel_tag.startswith("a8w8_mxscale_bmm_"): + f.write(MANIFEST_BMM_MXSCALE.format(kernel_name=k.name)) + elif k.kernel_tag in SPLITK_TAGS: + f.write(MANIFEST_A16W16_WORKSPACE.format(kernel_name=k.name)) + elif k.kernel_tag in A16W16_KID_DISPATCH_TAGS: + f.write(MANIFEST_A16W16.format(kernel_name=k.name)) + elif k.kernel_tag == "a8w8": f.write(MANIFEST_NOSCALE_3ARG.format(kernel_name=k.name)) + elif k.kernel_tag == "a8w8_scale": + f.write(MANIFEST_BLOCKSCALE.format(kernel_name=k.name)) + elif k.kernel_tag in A8W8_BPRESHUFFLE_TAGS: + f.write(MANIFEST_BLOCKSCALE_BPRESHUFFLE.format(kernel_name=k.name)) else: - f.write(MANIFEST_SCALE.format(kernel_name=k.name)) + raise ValueError(f"no manifest ABI for kernel tag {k.kernel_tag!r}") # -- Per-pass TU emission -- Replaces the old "one .cpp per (kid, dtype)" scheme. @@ -1079,6 +842,7 @@ def _emit_fused_host_tu(self): This TU needs no arch guard: it is host-pass only, so a mixed build's device passes already see nothing here. """ + host_by_arch = {} for row in self._host_instantiations: arch = _kid_name_arch(row["kid_name"]) @@ -1090,12 +854,15 @@ def _emit_fused_host_tu(self): reduce_abi = SPLITK_REDUCE_ABI_MAP[arch] extra_reduce = SPLITK_REDUCE_EXTRA_MAP.get(arch, {}) extra_forward_decls = extra_reduce.get("forward_decls", lambda: "")() + extra_template_params = reduce_abi.get( + "forward_decl_extra_template_params", "" + ) forward_decls = ( "// Forward declaration only. Specialisations live in per-arch device TUs.\n" f"{reduce_abi['forward_decl_include']}" "template\n" + f" bool HAS_OOB_{extra_template_params}>\n" f"__global__ void {reduce_abi['kernel']}(\n" f" {reduce_abi['ws_arg']}, D_OUT* c_out,\n" " int split_k, int M, int N, int batch,\n" @@ -1219,20 +986,14 @@ def _emit_splitk_reduce_tu(self): reduce_abi = SPLITK_REDUCE_ABI_MAP[reduce_arch] ws_ptr_type = reduce_abi["ws_type"] reduce_kernel = reduce_abi["kernel"] - # gfx1250 reduce: VEC=8/BLOCK=128 + both partial types (the kid picks, - # so both must exist) + compile-time split_k (SPLIT_K_ = 0..16, matching - # the launch-helper switch). Other arches keep the legacy - # VEC=16/BLOCK=64 fp32-workspace 6-param instantiations. if reduce_arch == "gfx1250": reduce_vec, reduce_block = 8, 128 - reduce_split_ks = tuple(range(17)) # 0 (runtime) + 1..16 (unrolled) - # Both partial types: which one a kid uses is per-instance - # (splitk_workspace_dtype), so both instantiations must exist. - reduce_d_ws = ("__bf16", "float") + reduce_split_ks = tuple(range(17)) + reduce_workspace_types = ("__bf16", "float") else: reduce_vec, reduce_block = 16, 64 reduce_split_ks = (None,) - reduce_d_ws = (None,) + reduce_workspace_types = (None,) guard_open, guard_close = _own_arch_device_pass_guard(reduce_arch) contents = ( "// SPDX-License-Identifier: MIT\n" @@ -1252,9 +1013,8 @@ def _emit_splitk_reduce_tu(self): reduce_vec, reduce_block, reduce_split_ks, - d_ws, + reduce_workspace_types, ) - for d_ws in reduce_d_ws for has_oob in reduce_abi["baseline_has_oob"] ) ) @@ -1268,6 +1028,17 @@ def _emit_splitk_reduce_tu(self): ).write_text(contents) def gen_instances(self, kernels_dict): + """Regenerate launchers, manifests and exact-kid tables.""" + # A rerun in an existing blob directory must not leave removed or + # renamed generated policy headers behind. + for legacy_header in ( + "opus_gemm_lookup.h", + "opus_gemm_a16w16_tune_lookup.h", + "opus_gemm_a8w8_tune_lookup.h", + "opus_bmm_mxscale_tune_lookup.h", + ): + Path(self.working_path, legacy_header).unlink(missing_ok=True) + if os.path.exists(self.impl_path): shutil.rmtree(self.impl_path) os.mkdir(self.impl_path) @@ -1295,79 +1066,14 @@ def gen_instances(self, kernels_dict): if needs_reduce_tu: self._emit_splitk_reduce_tu() - self.gen_lookup_dict(kernels_dict) self.gen_manifest_head(kernels_dict) - self.gen_a16w16_tune_lookup(kernels_dict) - self.gen_a8w8_tune_lookup(kernels_dict) - self.gen_bmm_mxscale_tune_lookup(kernels_dict) - - -def get_tune_dict(tune_dict_csv): - """Load a tuned CSV into the lookup-dict shape consumed by gen_lookup_dict. - - Key layout - ---------- - Tuple keys: (M, N, K, outdtype_str, arch). Promoting outdtype into the - key is what lets a single (M, N, K) shape carry distinct winners for - bf16 vs fp32 output (the underlying main kernel hardware rules differ - enough that the best kid is not always the same; e.g. fp32 output - biases reduce-bound shapes toward larger split-K). gen_lookup_dict - then writes outdtype="torch.bfloat16" rows only into the BF16 (M,N,K) - map and outdtype="torch.float32" rows only into the FP32 (M,N,K) map. - - arch is in the key for the same reason: the (M,N,K) tables are emitted - per arch, so a shape tuned on two arches has one winner per arch. With - arch out of the key, whichever CSV row was read last silently evicted - the other arch's winner and that arch fell back to its heuristic. - - Backwards compat - ---------------- - Legacy CSVs without an `outdtype` column are interpreted as - bf16-output (matches what the tuner used to write). int keys from - default_kernels_dict are passed through untouched -- gen_lookup_dict - skips them via the `isinstance(mnk, tuple) and mnk[0] > 0` guard. - """ - tune_dict = default_kernels_dict - if os.path.exists(tune_dict_csv): - tune_df = pd.read_csv(tune_dict_csv) - cu_num = None - try: - if torch.cuda.is_available(): - gpu = torch.cuda.current_device() - cu_num = torch.cuda.get_device_properties(gpu).multi_processor_count - except Exception: # noqa: BLE001 - # torch device enumeration is broken on some ROCm nightlies - # (device_count()==0 / "Invalid device id"); use rocminfo instead. - cu_num = None - if cu_num is None: - try: - from aiter.jit.utils.chip_info import get_cu_num as _rocminfo_cu_num - - cu_num = _rocminfo_cu_num() - except Exception: # noqa: BLE001 - cu_num = None - if cu_num is not None: - tune_df = tune_df[tune_df["cu_num"] == cu_num].reset_index() - # Accept either the legacy "kernelId" column or the new "solidx" column. - kids = _tune_df_kids(tune_df) - has_outdtype = "outdtype" in tune_df.columns - for i in range(len(tune_df)): - if kids is None or pd.isna(kids.loc[i]): - continue - M = tune_df.loc[i, "M"] - N = tune_df.loc[i, "N"] - K = tune_df.loc[i, "K"] - outdtype = ( - str(tune_df.loc[i, "outdtype"]) if has_outdtype else "torch.bfloat16" - ) - kid = int(kids.loc[i]) - if kid in kernels_list: - inst = kernels_list[kid] - tune_dict[(M, N, K, outdtype, _kid_arch_common(inst))] = inst - return tune_dict + self.gen_a16w16_kid_dispatch(kernels_dict) + self.gen_a8w8_kid_dispatch(kernels_dict) + self.gen_bmm_mxscale_kid_dispatch() def _tune_df_kids(df): + """Read kid values from either supported tuned-CSV column name.""" kids = None for col in ("solidx", "kernelId"): if col not in df.columns: @@ -1395,7 +1101,7 @@ def _tune_df_kids(df): "--tune", action="store_true", default=False, - help="generate all kernel instances for tuning (id-based lookup)", + help="generate all kernel instances for tuning (exact-kid dispatch)", ) parser.add_argument( @@ -1414,13 +1120,10 @@ def _tune_df_kids(df): "GEMM CSVs (e.g. aiter/configs/bf16_tuned_gemm.csv and " "aiter/configs/model_configs/*_bf16_tuned_gemm.csv). Each " "file is filtered by `libtype == 'opus'`; surviving rows " - "contribute their `solidx` to the subset-compile set S and " - "are also baked into opus_gemm_lookup.h via " - "GENERATE_OPUS_LOOKUP_TABLE_*. Without this flag we still " - "generate a working module (only HEURISTIC_DEFAULT_KIDS + " - "sidecar contents), the lookup table stays empty, and the " - "C++ dispatch falls through to the heuristic for every " - "untuned shape." + "contribute their `solidx`/`kernelId` only to the subset-compile " + "set S. Runtime callers provide the final kid explicitly. " + "Without this flag the module is generated from the sidecar, " + "per-arch default compile floor, and mandatory family kids." ), ) @@ -1432,7 +1135,8 @@ def _tune_df_kids(df): "Path to the subset-compile sidecar (JSON list of int kids). " "Defaults to {working_path}/compiled_kids.json. The sidecar " "captures the union of CSV opus rows + previous sidecar " - "contents + extra kids + HEURISTIC_DEFAULT_KIDS. JIT supplies a " + "contents + extra kids + DEFAULT_COMPILED_KIDS and mandatory " + "family kids. JIT supplies a " "staged copy and publishes it after successful compilation. " "Tuners pass new candidates with --extra_kids." ), @@ -1468,6 +1172,19 @@ def _tune_df_kids(df): "gfx942_nosplit": gfx942_nosplit_kernels_list, "gfx942_splitk": gfx942_splitk_kernels_list, "gfx942_a8w8": gfx942_a8w8_kernels_list, + "a16w16_cluster_tdm_splitk_ws": gfx1250_kernels_list, + "a16w16_clusterlaunch_tdm_splitk_ws": gfx1250_clusterlaunch_kernels_list, + "a16w16_clusterlaunch_tdm_splitk_fuse": gfx1250_splitk_fuse_kernels_list, + "a16w16_4wave_co": { + kid: instance + for kid, instance in gfx1250_4wave_co_kernels_list.items() + if instance.kernel_tag == "a16w16_4wave_co" + }, + "a16w16_4wave_wl_co": { + kid: instance + for kid, instance in gfx1250_4wave_co_kernels_list.items() + if instance.kernel_tag == "a16w16_4wave_wl_co" + }, } # --- Compute the subset-compile set S ------------------------------------ S = (CSV opus rows' @@ -1522,10 +1239,13 @@ def _expand_tune_paths(spec): sidecar_kids = set() # The compile set: union, intersected with valid kernels_list entries. + # MXFP8 BMM launchers are emitted as one gfx950 family below and deduplicated + # by generated symbol name, so they never participate in the per-kid subset. valid_kids = set(kernels_list.keys()) S = ( - csv_kids | sidecar_kids | set(args.extra_kids) | set(HEURISTIC_DEFAULT_KIDS) + csv_kids | sidecar_kids | set(args.extra_kids) | set(DEFAULT_COMPILED_KIDS) ) & valid_kids + S -= set(BMM_MXSCALE_KIDS) # Per-arch filter: drop kids whose arch_prefix is not in the target build set. _kid_arch = _kid_arch_common @@ -1575,35 +1295,46 @@ def _expand_tune_paths(spec): f"#define OPUS_BUILD_HAS_{a.upper()} 1\n" for a in archs_for_header ) - # gfx950 a8w8 (kid 1, 2) is only needed when the module is built with - # gfx950 support. gfx942 has its own blockscale bpreshuffle A8W8 tune path. - if target_arches is None or "gfx950" in target_arches: - S |= set(a8w8_scale_kernels_list.keys()) - S |= set(a8w8_kernels_list.keys()) + # Family ABI defaults must be linkable even when no tuned row or sidecar + # mentions them. This set is arch-scoped so single-arch builds never pull + # another architecture's launcher symbol into their host TU. + mandatory_arches = ( + set(OPUS_MANDATORY_A8_KIDS) if target_arches is None else set(target_arches) + ) + mandatory_a8_kids = set().union( + *(OPUS_MANDATORY_A8_KIDS.get(arch, frozenset()) for arch in mandatory_arches) + ) + S |= mandatory_a8_kids & valid_kids # Honor --kernel_tag as a developer override that *further restricts* the set (within the a16w16 # / a8w8 families). if args.kernel_tag: tag_keys = set(TAG_TO_LIST.get(args.kernel_tag, {}).keys()) if tag_keys: - # Restrict to the requested family + heuristic defaults. - S = (S & tag_keys) | set(heuristic_kids_for_arch(target_arches)) - if target_arches is None or "gfx950" in target_arches: - S |= set(a8w8_scale_kernels_list.keys()) - S |= set(a8w8_kernels_list.keys()) - - # Heuristic-fallback invariant (single source of truth: opus_gemm_common.py). - required_heuristic = set(heuristic_kids_for_arch(target_arches)) - missing_heuristic = required_heuristic - S - assert not missing_heuristic, ( - f"Subset-compile error: heuristic-fallback kids " - f"{sorted(missing_heuristic)} are missing from the compile set S; " - f"opus_a16w16_heuristic_kid_gfx950() would return an unbakeable " - f"kid. Add them to the compile set or update HEURISTIC_DEFAULT_KIDS " + # Restrict to the requested family + default compile floor. + S = (S & tag_keys) | set(default_compiled_kids_for_arch(target_arches)) + S |= mandatory_a8_kids & valid_kids + + # Default exact-id compile-floor invariant (single source of truth: + # opus_gemm_common.py). C++ and Python both perform exact-kid routing. + required_default = set(default_compiled_kids_for_arch(target_arches)) + missing_default = required_default - S + assert not missing_default, ( + f"Subset-compile error: default exact-id kids " + f"{sorted(missing_default)} are missing from the compile set S. " + f"Add them to the compile set or update DEFAULT_COMPILED_KIDS " f"in csrc/opus_gemm/opus_gemm_common.py." ) - missing_requested = set(args.extra_kids) - S + # The sidecar and request validation describe every emitted route, including + # BMM ids whose shared symbols are generated outside the per-kid subset. + bmm_kids = ( + BMM_MXSCALE_KIDS + if target_arches is None or "gfx950" in target_arches + else frozenset() + ) + compiled_kids = S | bmm_kids + missing_requested = set(args.extra_kids) - compiled_kids if missing_requested: parser.error( "cannot compile requested --extra_kids " @@ -1615,102 +1346,27 @@ def _expand_tune_paths(spec): # Build the per-kid dict that drives codegen. kdict = {kid: kernels_list[kid] for kid in sorted(S)} - # a8w8_mxscale BMM flatmm split-K family (gfx950-only). These live in the - # opus_bmm.cu switch's PRIVATE kid namespace (ints 0/32/64/128/...), which - # collides with the global integer kids in kernels_list/S, so we never put - # them in S. Instead merge them into kdict keyed by kernel NAME: gen_instance - # only reads the value (k), gen_manifest_head emits by k.name, and every - # lookup/tune emitter gates on isinstance(key,int|tuple)+tag so the string - # keys are skipped. Name-keying also auto-dedups kids with identical geometry - # (e.g. switch kids 0 and 32 -> one launcher symbol). Always emitted (like - # a8w8_mxscale) so the opus_bmm dispatch never hits a missing symbol. - if target_arches is None or "gfx950" in target_arches: - for _bmm_list in a8w8_mxscale_bmm_kernel_lists: - for _bmm_k in _bmm_list.values(): - kdict[_bmm_k.name] = _bmm_k + # All 45 BMM ids are exact-routable in the canonical registry. Several ids + # intentionally share one device geometry, so key this codegen-only merge by + # symbol name to emit each host/device specialization once. + if bmm_kids: + for family in a8w8_mxscale_bmm_kernel_lists: + for instance in family.values(): + kdict[instance.name] = instance print( f"[opus gen_instances] subset compile: |S|={len(S)} kids " - f"(CSV={len(csv_kids)}, sidecar={len(sidecar_kids)}, heuristic={len(HEURISTIC_DEFAULT_KIDS)})" + f"(sources: CSV={len(S & csv_kids)}, " + f"sidecar={len(S & sidecar_kids)}, " + f"default-compiled={len(S & required_default)}, " + f"mandatory-a8={len(S & mandatory_a8_kids)}, " + f"extra={len(S & set(args.extra_kids))}); " + f"always-emitted-bmm={len(bmm_kids)}" ) codegen = opus_gemm_codegen(args.working_path, args.tune) codegen.gen_instances(kdict) - # Bake the (M, N, K) -> kernel runtime lookup. - if csv_paths: - # Concatenate all opus rows from all matched CSV files (filtered by libtype). - combined_frames = [] - for path in csv_paths: - try: - df = pd.read_csv(path) - except (pd.errors.EmptyDataError, FileNotFoundError): - continue - if "libtype" not in df.columns: - continue - df = df[df["libtype"] == "opus"] - if df.empty: - continue - # Drop off-arch kids: lookup must only reference symbols S actually emitted. - kids = _tune_df_kids(df) - if kids is None: - continue - a16w16_lookup_rows = kids.apply( - lambda kid: ( - not pd.isna(kid) - and int(kid) in kernels_list - and kernels_list[int(kid)].kernel_tag in A16W16_TUNE_TAGS - ) - ) - df = df[kids.isin(S) & a16w16_lookup_rows] - if df.empty: - continue - if "kernelId" in df.columns: - df = df.copy() - if "solidx" in df.columns: - df["solidx"] = df["solidx"].fillna(df["kernelId"]) - df = df.drop(columns=["kernelId"]) - else: - df = df.rename(columns={"kernelId": "solidx"}) - if "solidx" in df.columns: - df["solidx"] = df["solidx"].astype(int) - combined_frames.append(df) - - if combined_frames: - combined = pd.concat(combined_frames, ignore_index=True).drop_duplicates() - tmp_csv = os.path.join(args.working_path, "_combined_opus_tuned.csv") - combined.to_csv(tmp_csv, index=False) - tune_dict = get_tune_dict(tmp_csv) - try: - os.remove(tmp_csv) - except OSError: - pass - # Filter tune_dict entries to those whose kid is in S (defense - # in depth -- valid_kids should have already caught everything). - filtered = {} - for k, v in tune_dict.items(): - if isinstance(k, tuple) and k[0] > 0: - # Find the kid for this entry by reverse-lookup against S. - filtered[k] = v - else: - filtered[k] = v # default_kernels_dict negative-int entries - codegen.gen_lookup_dict(filtered) - n_real = sum(1 for k in filtered if isinstance(k, tuple) and k[0] > 0) - print( - f"[opus gen_instances] baked {n_real} tuned entries from " - f"{len(csv_paths)} CSV file(s) into opus_gemm_lookup.h" - ) - else: - print( - f"[opus gen_instances] no `libtype=='opus'` rows found in " - f"{len(csv_paths)} CSV file(s); using empty lookup" - ) - elif args.tune_files: - print( - f"[opus gen_instances] --tune_files {args.tune_files} matched no " - f"existing files; using empty lookup" - ) - # Write the generated set inside staging. JIT publishes it to bd_dir only # after the binary is installed, so a failed compile cannot advance it. try: @@ -1718,5 +1374,7 @@ def _expand_tune_paths(spec): except OSError: pass with open(sidecar_path, "w") as f: - json.dump(sorted(S), f) - print(f"[opus gen_instances] wrote sidecar with {len(S)} kids: {sidecar_path}") + json.dump(sorted(compiled_kids), f) + print( + f"[opus gen_instances] wrote sidecar with {len(compiled_kids)} kids: {sidecar_path}" + ) diff --git a/csrc/opus_gemm/include/gfx1250/opus_gemm_arch_gfx1250.cuh b/csrc/opus_gemm/include/gfx1250/opus_gemm_arch_gfx1250.cuh index fcd98a5ee0..a1b9b83c55 100644 --- a/csrc/opus_gemm/include/gfx1250/opus_gemm_arch_gfx1250.cuh +++ b/csrc/opus_gemm/include/gfx1250/opus_gemm_arch_gfx1250.cuh @@ -1,200 +1,193 @@ // SPDX-License-Identifier: MIT // Copyright (C) 2025-2026, Advanced Micro Devices, Inc. All rights reserved. // -// opus_gemm_arch_gfx1250.cuh -- gfx1250-specific dispatch. -// -// Wires both call paths for the cluster/TDM split-K (workspace + reduce) kids: -// * opus_a16w16_tune_dispatch_gfx1250 -- id-based (explicit kernelId) -// * opus_dispatch_a16w16_gfx1250 -- tuned (M,N,K) lookup -> heuristic -// -// Every gfx1250 kid is a split-K kid whose main kernel writes an fp32 workspace -// (output_dtypes = ["fp32_t"]); the reduce kernel casts the partials to the -// runtime Y dtype (bf16/fp32) and folds bias. So all dispatch resolves through -// the tune table -- the specializations exist only to satisfy -// the shared arch-router template instantiation and are never invoked for -// gfx1250 (opus_gemm.cu forces for split-K kids). -// -// Included exactly once, by opus_gemm.cu. Self-contained: the shared flat-array -// dispatch types live in opus_gfx1250_detail (opus_gemm_heuristic_dispatch_ -// gfx1250.cuh), so this header does NOT depend on any gfx950 header and a -// gfx1250-only build compiles without OPUS_BUILD_HAS_GFX950. +// Exact-kid launcher tables for gfx1250. #pragma once #include "../opus_gemm_arch.cuh" #include "../opus_gemm_common.cuh" -#include "opus_gemm_heuristic_dispatch_gfx1250.cuh" // OpusA16W16NoscaleKernel + opus_gfx1250_detail::* + opus_a16w16_heuristic_kid_gfx1250 -#include "opus_gemm_lookup.h" // GENERATE_OPUS_LOOKUP_TABLE_FP32_GFX1250 -#include "opus_gemm_a16w16_tune_lookup.h" // GENERATE_A16W16_TUNE_LOOKUP_FP32_GFX1250 -#include "opus_gemm_manifest.h" // launcher symbols -#include "../opus_gemm_utils.cuh" // bf16_t / fp32_t +#include "../opus_gemm_utils.cuh" +#include "opus_gemm_a16w16_kid_dispatch.h" +#include "opus_gemm_a8w8_kid_dispatch.h" +#include "opus_gemm_manifest.h" -#include // std::lower_bound +#include +#include #include +#include -// ── a16w16 tune dispatch (id-based) ───────────────────────────────────────── -// Only the table is populated (gfx1250 kids are fp32-only split-K). -// The specialization is defensive: it never carries a table (which -// would be an empty array in a gfx1250-only build) and is never called. +#ifndef OPUS_A16W16_DISPATCH_KERNEL_TYPES_DEFINED +#define OPUS_A16W16_DISPATCH_KERNEL_TYPES_DEFINED +using OpusA16W16Kernel = void (*)( + aiter_tensor_t&, aiter_tensor_t&, aiter_tensor_t&, + std::optional, int); +using OpusA16W16WorkspaceKernel = void (*)( + aiter_tensor_t&, aiter_tensor_t&, aiter_tensor_t&, + aiter_tensor_t&, std::optional, int); +#endif -template -inline opus_gfx1250_detail::OpusA16W16TuneKernel -opus_a16w16_tune_dispatch_gfx1250(int id); +#ifndef OPUS_A8W8_DISPATCH_KERNEL_TYPES_DEFINED +#define OPUS_A8W8_DISPATCH_KERNEL_TYPES_DEFINED +using OpusA8W8Kernel = void (*)( + aiter_tensor_t&, aiter_tensor_t&, aiter_tensor_t&); +using OpusA8W8BlockscaleKernel = void (*)( + aiter_tensor_t&, aiter_tensor_t&, aiter_tensor_t&, + aiter_tensor_t&, aiter_tensor_t&); +using OpusA8W8BlockscaleBpreshuffleKernel = void (*)( + aiter_tensor_t&, aiter_tensor_t&, aiter_tensor_t&, + aiter_tensor_t&, aiter_tensor_t&); +#endif -template <> -inline opus_gfx1250_detail::OpusA16W16TuneKernel -opus_a16w16_tune_dispatch_gfx1250(int id) +namespace opus_gfx1250_detail +{ +struct OpusA16W16KidEntry +{ + int kid; + OpusA16W16Kernel func; +}; + +struct OpusA16W16WorkspaceKidEntry +{ + int kid; + OpusA16W16WorkspaceKernel func; +}; + +template +struct OpusA8W8KidEntry +{ + int kid; + Kernel func; +}; + +template +inline const Entry* find_kid(const std::array& entries, int kid) { - using namespace opus_gfx1250_detail; - static constexpr OpusA16W16TuneEntry kTune[] = { - GENERATE_A16W16_TUNE_LOOKUP_FP32_GFX1250(fp32_t) - }; - constexpr size_t kSize = sizeof(kTune) / sizeof(kTune[0]); - OpusA16W16TuneEntry needle{id, nullptr}; - auto it = std::lower_bound(kTune, kTune + kSize, needle, - kid_entry_less); - AITER_CHECK(it != kTune + kSize && it->kid == id, - "Kernel id ", id, - " not found in a16w16 fp32 tune lookup table (gfx1250)"); - return it->func; + const auto it = std::lower_bound( + entries.begin(), entries.end(), kid, + [](const Entry& entry, int value) { return entry.kid < value; }); + return it != entries.end() && it->kid == kid ? &*it : nullptr; } -template <> -inline opus_gfx1250_detail::OpusA16W16TuneKernel -opus_a16w16_tune_dispatch_gfx1250(int id) +inline const OpusA16W16WorkspaceKidEntry* workspace_entry(int kid) { - // gfx1250 split-K kids are emitted only; the reduce kernel handles - // bf16 Y output. opus_gemm.cu always routes split-K kids through , - // so this is unreachable -- but it must compile (no empty-array table). - AITER_CHECK(false, - "opus_gemm gfx1250: a16w16 tune dispatch is not used " - "(split-K kids are fp32-workspace; bf16 Y is produced by the " - "reduce kernel). kid=", id); - return nullptr; + static constexpr std::array< + OpusA16W16WorkspaceKidEntry, + GENERATE_A16W16_WORKSPACE_KID_DISPATCH_GFX1250_SIZE> + kWorkspace = {{GENERATE_A16W16_WORKSPACE_KID_DISPATCH_GFX1250}}; + return find_kid(kWorkspace, kid); } -// ── a16w16 pre-compiled .co dispatch ──────────────────────────────────────── -// Two tables, mirroring the split-K pair above: id-based for the explicit -// kernelId / tuner path, (M, N, K) for the tuned-CSV production path. Both hold -// OpusA16W16CoKernel (no workspace argument -- see the type's comment), and both -// are bf16-out only: the kernel writes bf16 C directly, with no reduce kernel in -// the way to cast anything else. -inline opus_gfx1250_detail::OpusA16W16CoKernel -opus_a16w16_co_tune_dispatch_gfx1250(int id) +template +inline const OpusA16W16KidEntry* non_workspace_entry(int kid); + +template <> +inline const OpusA16W16KidEntry* non_workspace_entry(int kid) { - using namespace opus_gfx1250_detail; - static constexpr OpusA16W16CoTuneEntry kTune[] = { - GENERATE_A16W16_CO_TUNE_LOOKUP_GFX1250() - }; - constexpr size_t kSize = sizeof(kTune) / sizeof(kTune[0]); - OpusA16W16CoTuneEntry needle{id, nullptr}; - auto it = std::lower_bound(kTune, kTune + kSize, needle, - kid_entry_less); - AITER_CHECK(it != kTune + kSize && it->kid == id, - "Kernel id ", id, - " not found in the a16w16 pre-compiled (.co) tune lookup table " - "(gfx1250). Either it is not in this build's compile set, or " - "the build saw no gen_co/co_kernels.json / no matching " - "gen_co//*.co and the whole family is empty (the loader " - "drops kids whose image is missing)."); - return it->func; + static constexpr std::array< + OpusA16W16KidEntry, + GENERATE_A16W16_NONWORKSPACE_KID_DISPATCH_GFX1250_BF16_SIZE> + kKids = {{GENERATE_A16W16_NONWORKSPACE_KID_DISPATCH_GFX1250_BF16(bf16_t)}}; + return find_kid(kKids, kid); } -// (M, N, K) -> pre-compiled kernel, nullptr on miss. No heuristic fallback: the -// shape heuristic never returns a .co kid, so a miss just means "this shape has -// no tuned .co winner" and the caller carries on with the split-K path. -inline opus_gfx1250_detail::OpusA16W16CoKernel -opus_a16w16_co_dispatch_gfx1250(int M, int N, int K) +template <> +inline const OpusA16W16KidEntry* non_workspace_entry(int kid) { - using namespace opus_gfx1250_detail; - static constexpr OpusA16W16CoRuntimeEntry kLookup[] = { - GENERATE_OPUS_LOOKUP_TABLE_CO_GFX1250() - }; - constexpr size_t kSize = sizeof(kLookup) / sizeof(kLookup[0]); - OpusA16W16CoRuntimeEntry needle{{M, N, K}, nullptr}; - auto it = std::lower_bound(kLookup, kLookup + kSize, needle, - shape_entry_less); - if (it != kLookup + kSize && shape_entry_eq(*it, needle)) - return it->func; - return nullptr; + static constexpr std::array< + OpusA16W16KidEntry, + GENERATE_A16W16_NONWORKSPACE_KID_DISPATCH_GFX1250_FP32_SIZE> + kKids = {{GENERATE_A16W16_NONWORKSPACE_KID_DISPATCH_GFX1250_FP32(fp32_t)}}; + return find_kid(kKids, kid); } +} // namespace opus_gfx1250_detail -// ── a16w16 runtime dispatch (tuned lookup -> heuristic fallback) ──────────── -// Both dtype specializations route split-K kids through the tune -// table (the launcher's reduce kernel produces the requested Y dtype). +template +inline OpusA8W8BlockscaleBpreshuffleKernel +opus_a8w8_blockscale_bpreshuffle_kid_dispatch_gfx1250(int id); -namespace opus_gfx1250_detail -{ -inline void check_shape_4g(int M, int N, int K, size_t c_elem_bytes) +template <> +inline OpusA8W8BlockscaleBpreshuffleKernel +opus_a8w8_blockscale_bpreshuffle_kid_dispatch_gfx1250(int id) { - // 4 GiB buffer-resource guard: the launcher builds 32-bit-bounded gmem - // descriptors over A/B/C, so >4 GiB tensors wrap num_records -> silent OOB. - constexpr uint64_t U32_MAX_BYTES = (1ULL << 32) - 1; - const uint64_t a_bytes = (uint64_t)M * (uint64_t)K * sizeof(bf16_t); - const uint64_t b_bytes = (uint64_t)N * (uint64_t)K * sizeof(bf16_t); - const uint64_t c_bytes = (uint64_t)M * (uint64_t)N * (uint64_t)c_elem_bytes; - AITER_CHECK(a_bytes <= U32_MAX_BYTES && b_bytes <= U32_MAX_BYTES - && c_bytes <= U32_MAX_BYTES, - "opus_gemm gfx1250: a16w16 heuristic refuses >4 GiB shape (M=", - M, " N=", N, " K=", K, "): launcher gmem descriptors are 32-bit."); + using Entry = opus_gfx1250_detail::OpusA8W8KidEntry< + OpusA8W8BlockscaleBpreshuffleKernel>; + static constexpr std::array< + Entry, + GENERATE_A8W8_BLOCKSCALE_BPRESHUFFLE_KID_DISPATCH_GFX1250_BF16_SIZE> + kKids = {{GENERATE_A8W8_BLOCKSCALE_BPRESHUFFLE_KID_DISPATCH_GFX1250_BF16}}; + AITER_CHECK(!kKids.empty(), + "no registered kernel for OPUS " + "a8w8_blockscale_bpreshuffle on gfx1250 with bf16 Y"); + const auto* entry = opus_gfx1250_detail::find_kid(kKids, id); + AITER_CHECK(entry != nullptr, + "unknown kid ", id, " for OPUS " + "a8w8_blockscale_bpreshuffle on gfx1250 with bf16 Y"); + return entry->func; } -inline void check_shape_reduce_grid(int M) +template <> +inline OpusA8W8BlockscaleBpreshuffleKernel +opus_a8w8_blockscale_bpreshuffle_kid_dispatch_gfx1250(int id) { - // Every kid this heuristic can return is split-K, and their reduce carries - // one row per grid.y block against a 65535 cap. Past it the reduce writes - // garbage rather than failing -- NaN at M=65536, measured. A .co kid has no - // reduce and handles these shapes, but the heuristic never returns one, so - // the only way through is a tuned row. - AITER_CHECK(M <= 65535, - "opus_gemm gfx1250: a16w16 heuristic refuses M=", M, - " (> 65535): every heuristic kid is split-K and its reduce is " - "capped at 65535 rows. Tune this shape -- the winner will be a " - ".co kid, which has no reduce."); + using Entry = opus_gfx1250_detail::OpusA8W8KidEntry< + OpusA8W8BlockscaleBpreshuffleKernel>; + static constexpr std::array< + Entry, + GENERATE_A8W8_BLOCKSCALE_BPRESHUFFLE_KID_DISPATCH_GFX1250_FP32_SIZE> + kKids = {{GENERATE_A8W8_BLOCKSCALE_BPRESHUFFLE_KID_DISPATCH_GFX1250_FP32}}; + AITER_CHECK(!kKids.empty(), + "no registered kernel for OPUS " + "a8w8_blockscale_bpreshuffle on gfx1250 with fp32 Y"); + const auto* entry = opus_gfx1250_detail::find_kid(kKids, id); + AITER_CHECK(entry != nullptr, + "unknown kid ", id, " for OPUS " + "a8w8_blockscale_bpreshuffle on gfx1250 with fp32 Y"); + return entry->func; } -} // namespace opus_gfx1250_detail +// Empty direct-output tables reject unsupported gfx1250 A16W16 calls. template -inline opus_gfx1250_detail::OpusA16W16NoscaleKernel -opus_dispatch_a16w16_gfx1250(int M, int N, int K, int batch, bool has_bias = false); +inline OpusA16W16Kernel opus_a16w16_kid_dispatch_gfx1250(int kid); template <> -inline opus_gfx1250_detail::OpusA16W16NoscaleKernel -opus_dispatch_a16w16_gfx1250(int M, int N, int K, int batch, bool has_bias) +inline OpusA16W16Kernel opus_a16w16_kid_dispatch_gfx1250(int kid) { - using namespace opus_gfx1250_detail; - static constexpr OpusA16W16RuntimeEntry kLookup[] = { - GENERATE_OPUS_LOOKUP_TABLE_BF16_GFX1250(bf16_t) - }; - constexpr size_t kSize = sizeof(kLookup) / sizeof(kLookup[0]); - OpusA16W16RuntimeEntry needle{{M, N, K}, nullptr}; - auto it = std::lower_bound(kLookup, kLookup + kSize, needle, - shape_entry_less); - if (it != kLookup + kSize && shape_entry_eq(*it, needle)) - return it->func; - (void)batch; - opus_gfx1250_detail::check_shape_4g(M, N, K, sizeof(bf16_t)); - opus_gfx1250_detail::check_shape_reduce_grid(M); - const int kid = opus_a16w16_heuristic_kid_gfx1250(M, N, K, has_bias); - return opus_a16w16_tune_dispatch_gfx1250(kid); + const auto* entry = opus_gfx1250_detail::non_workspace_entry(kid); + AITER_CHECK(entry != nullptr, + "unknown kid ", kid, + " for OPUS a16w16 on gfx1250 with bf16 Y in the " + "non-workspace launch table"); + return entry->func; } template <> -inline opus_gfx1250_detail::OpusA16W16NoscaleKernel -opus_dispatch_a16w16_gfx1250(int M, int N, int K, int batch, bool has_bias) +inline OpusA16W16Kernel opus_a16w16_kid_dispatch_gfx1250(int kid) +{ + const auto* entry = opus_gfx1250_detail::non_workspace_entry(kid); + AITER_CHECK(entry != nullptr, + "unknown kid ", kid, + " for OPUS a16w16 on gfx1250 with fp32 Y in the " + "non-workspace launch table"); + return entry->func; +} + +inline bool opus_a16w16_has_non_workspace_kernel_gfx1250(int id) +{ + return opus_gfx1250_detail::non_workspace_entry(id) != nullptr + || opus_gfx1250_detail::non_workspace_entry(id) != nullptr; +} + +inline bool opus_a16w16_has_workspace_kernel_gfx1250(int id) +{ + return opus_gfx1250_detail::workspace_entry(id) != nullptr; +} + +inline OpusA16W16WorkspaceKernel +opus_a16w16_workspace_dispatch_gfx1250(int id) { - using namespace opus_gfx1250_detail; - static constexpr OpusA16W16RuntimeEntry kLookup[] = { - GENERATE_OPUS_LOOKUP_TABLE_FP32_GFX1250(fp32_t) - }; - constexpr size_t kSize = sizeof(kLookup) / sizeof(kLookup[0]); - OpusA16W16RuntimeEntry needle{{M, N, K}, nullptr}; - auto it = std::lower_bound(kLookup, kLookup + kSize, needle, - shape_entry_less); - if (it != kLookup + kSize && shape_entry_eq(*it, needle)) - return it->func; - (void)batch; - opus_gfx1250_detail::check_shape_4g(M, N, K, sizeof(fp32_t)); - opus_gfx1250_detail::check_shape_reduce_grid(M); - const int kid = opus_a16w16_heuristic_kid_gfx1250(M, N, K, has_bias); - return opus_a16w16_tune_dispatch_gfx1250(kid); + const auto* entry = opus_gfx1250_detail::workspace_entry(id); + AITER_CHECK(entry != nullptr, + "unknown kid ", id, + " for OPUS a16w16 on gfx1250 in the workspace launch table"); + return entry->func; } diff --git a/csrc/opus_gemm/include/gfx1250/opus_gemm_heuristic_dispatch_gfx1250.cuh b/csrc/opus_gemm/include/gfx1250/opus_gemm_heuristic_dispatch_gfx1250.cuh deleted file mode 100644 index ff4b83daa5..0000000000 --- a/csrc/opus_gemm/include/gfx1250/opus_gemm_heuristic_dispatch_gfx1250.cuh +++ /dev/null @@ -1,148 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (C) 2025-2026, Advanced Micro Devices, Inc. All rights reserved. -// -// gfx1250 a16w16 shape-heuristic: (M, N, K, has_bias) -> kid. Pure integer -// mapping (no launcher symbols) so it can be included by the dispatcher TU -// without dragging in the lookup macros. -// -// All gfx1250 kids are cluster/TDM split-K (workspace + reduce). The kernel -// requires M % B_M == 0 and N % B_N == 0 (ragged M/N is not supported; ragged -// K is, via the TDM k_extent clamp). The heuristic therefore picks the largest -// tile from the kid set whose B_M divides M and B_N divides N, preferring the -// B_M=16 "tileN" family for small M and the "tileM" family for larger M. -// -// MUST stay in sync with opus_gemm_common.py :: gfx1250_kernels_list and -// HEURISTIC_DEFAULT_KIDS_GFX1250. -#pragma once - -#include - -#include "aiter_tensor.h" // aiter_tensor_t (torch-free) - -// Shared flat-array dispatch POD types + comparators for gfx1250 (mirrors the -// gfx950 set). gen_instances.py emits the tune / (M,N,K) lookup tables as -// arrays of these; std::lower_bound does the O(log N) runtime match. -namespace opus_gfx1250_detail -{ - -// a16w16-family launcher signature for gfx1250: 3 tensors + workspace + -// std::optional + int splitK. Different from gfx950 (no workspace). -using OpusA16W16NoscaleKernel = void (*)( - aiter_tensor_t &, aiter_tensor_t &, - aiter_tensor_t &, aiter_tensor_t &, std::optional, int); - -struct OpusA16W16Shape -{ - int M; - int N; - int K; -}; - -struct OpusA16W16RuntimeEntry -{ - OpusA16W16Shape key; - OpusA16W16NoscaleKernel func; -}; - -// Comparators are templated on the entry type rather than written once per -// table: the entry types differ only in what they carry NEXT to the key (a -// workspace-carrying function pointer, a workspace-free one, ...), and every -// table is keyed the same way. Explicit template args at the call sites, since -// std::lower_bound has no target type to deduce them from. -// -// Lex order on (M, N, K). Used both during sorting (gen_instances.py emits -// entries in lex order) and by std::lower_bound at lookup time. -template -constexpr bool shape_entry_less(const Entry& a, const Entry& b) noexcept -{ - if (a.key.M != b.key.M) return a.key.M < b.key.M; - if (a.key.N != b.key.N) return a.key.N < b.key.N; - return a.key.K < b.key.K; -} - -template -constexpr bool shape_entry_eq(const Entry& a, const Entry& b) noexcept -{ - return a.key.M == b.key.M && a.key.N == b.key.N && a.key.K == b.key.K; -} - -template -constexpr bool kid_entry_less(const Entry& a, const Entry& b) noexcept -{ - return a.kid < b.kid; -} - -// id -> kernel, same flat-array layout. Sorted by id (the codegen -// always emits in ascending id order). -struct OpusA16W16TuneEntry -{ - int kid; - OpusA16W16NoscaleKernel func; -}; - -using OpusA16W16TuneKernel = OpusA16W16NoscaleKernel; - -// ── pre-compiled .co kids (a16w16_4wave_co) ───────────────────────────────── -// A second function-pointer type, and deliberately so: this family has no -// split-K, no partial buffer and no reduce kernel, so it has nothing to put in -// a workspace, and its launcher is the ordinary 5-arg a16w16 signature (the -// same one gfx950/gfx942 use) rather than the 6-arg gfx1250 one above. -// -// The two could be collapsed by making the workspace argument -// std::optional everywhere, but that would ERASE exactly the -// fact worth keeping: with "needs a workspace" in the type, the production -// dispatch in opus_gemm.cu can consult the .co table first and skip the -// hipMalloc / hipDeviceSynchronize / hipFree it has to do for every split-K -// kid. A merged type cannot tell, from a function pointer alone, whether the -// buffer is going to be read. -// -// A split-K .co variant would use OpusA16W16NoscaleKernel and these types go -// unused for it; nothing here has to change. -using OpusA16W16CoKernel = void (*)( - aiter_tensor_t &, aiter_tensor_t &, - aiter_tensor_t &, std::optional, int); - -struct OpusA16W16CoTuneEntry -{ - int kid; - OpusA16W16CoKernel func; -}; - -struct OpusA16W16CoRuntimeEntry -{ - OpusA16W16Shape key; - OpusA16W16CoKernel func; -}; -} // namespace opus_gfx1250_detail - -// Kid map (B_K=128 chosen here; tuner explores B_K 256/512 + the P/wg space). -// Tiles whose per-TDM direct-copy request count (rows*B_K*2/256) hits the 256 -// SIMD-pair limit on some operand are NOT generated (e.g. 32x256x128) so the -// heuristic must not return them. All returned kids are no-cluster prefetch-3. -// tileN (B_M=16): 20000=16x32, 20003=16x64, 20004=16x128 -// tileM (B_M=32): 20005=32x32, 20006=32x64, 20007=32x128 -// (One P=3 kid per tile in the contiguous plain band [20000,20100).) -// MUST stay in sync with opus_gemm_common.py :: gfx1250_kernels_list (the plain -// kids are assigned contiguously from 20000 in _GFX1250_CTDM_TILES order). -inline int opus_a16w16_heuristic_kid_gfx1250(int M, int N, int K, bool has_bias) -{ - (void)K; - (void)has_bias; // bias is folded by the reduce kernel for every kid. - - // M >= 32 (and M % 32 == 0) -> tileM (B_M=32); widest B_N that divides N. - // (32x256 is unavailable -- per-TDM B req = 256 hits the direct-copy limit; - // fall through to the B_M=16 tileN family for N % 256 == 0.) - if (M % 32 == 0) - { - if (N % 128 == 0) return 20007; // 32x128x128 - if (N % 64 == 0) return 20006; // 32x64x128 - if (N % 32 == 0) return 20005; // 32x32x128 - } - - // Small M (or N not tileM-friendly) -> tileN family (B_M=16). Ragged M/N is - // handled by the TDM row/col clamp + padded workspace, so the smallest - // 16x32 tile is always a valid fallback. - if (N % 128 == 0) return 20004; // 16x128x128 - if (N % 64 == 0) return 20003; // 16x64x128 - return 20000; // 16x32x128 -} diff --git a/csrc/opus_gemm/include/gfx1250/opus_gemm_pipeline_a16w16_cluster_tdm_splitk_ws_gfx1250.cuh b/csrc/opus_gemm/include/gfx1250/opus_gemm_pipeline_a16w16_cluster_tdm_splitk_ws_gfx1250.cuh index 8c92406d0c..64a554dabe 100644 --- a/csrc/opus_gemm/include/gfx1250/opus_gemm_pipeline_a16w16_cluster_tdm_splitk_ws_gfx1250.cuh +++ b/csrc/opus_gemm/include/gfx1250/opus_gemm_pipeline_a16w16_cluster_tdm_splitk_ws_gfx1250.cuh @@ -1,8 +1,9 @@ // SPDX-License-Identifier: MIT // Copyright (C) 2025-2026, Advanced Micro Devices, Inc. All rights reserved. // -// gfx1250 bf16 TDM a16w16 GEMM, 4-wave split-K via fp32 workspace + separate -// reduce kernel. C[M,N] = A[M,K] @ B[N,K]^T (+ bias[N], folded in reduce). +// gfx1250 bf16 TDM a16w16 GEMM, 4-wave split-K via a per-kid bf16/fp32 +// workspace + separate reduce kernel. C[M,N] = A[M,K] @ B[N,K]^T (+ bias[N], +// folded in reduce). // // Plain grid (no cluster): grid = (M/B_M, N/B_N, split_k); each WG owns one // B_M x B_N tile and TDM-loads its own A/B. 4 waves: w0=A producer, w1=B @@ -56,6 +57,7 @@ void gemm_a16w16_cluster_tdm_splitk_ws_kernel_gfx1250(opus_gemm_cluster_tdm_ws_k using T = remove_cvref_t; using DataA = typename T::DataA; using DataB = typename T::DataB; + using DataWS = typename T::DataWS; using DataAcc = typename T::DataAcc; DECLARE_NAMED_BARRIERS(); // __nbar_1..__nbar_15 (we use 1..2*kNumSlots <= 6) @@ -376,21 +378,20 @@ void gemm_a16w16_cluster_tdm_splitk_ws_kernel_gfx1250(opus_gemm_cluster_tdm_ws_k // ---- Store the partial into ws[split_idx][padded_m][padded_n]. ---- // bias is folded once by the reduce kernel (not here). - // The partial type is the traits' D_C -- per kid (splitk_workspace_dtype), and - // the reduce is instantiated with the same D_C so the two cannot diverge. - using DataWs = typename T::DataC; - constexpr int kCVec = T::kCVec; // 4 (fp32 dwordx4 / bf16 dwordx2) - DataWs* ws_ptr = reinterpret_cast(kargs.ptr_ws); + // The exact kid declares the physical workspace type in its traits. Keep + // this in lockstep with the Torch workspace dtype and reducer D_WS. + constexpr int kCVec = T::kCVec; + DataWS* ws_ptr = reinterpret_cast(kargs.ptr_ws); const size_t ws_split = (size_t)split_idx * (size_t)kargs.stride_ws_batch; const size_t ws_base = ws_split + (size_t)tile_row * (size_t)kargs.stride_ws + (size_t)tile_col; const unsigned int ws_bytes = (unsigned int)(((size_t)kargs.stride_ws_batch - ((size_t)tile_row * kargs.stride_ws + tile_col)) * - sizeof(DataWs)); - auto g_ws = make_gmem(ws_ptr + ws_base, ws_bytes); + sizeof(DataWS)); + auto g_ws = make_gmem(ws_ptr + ws_base, ws_bytes); auto u_gc = partition_layout_c(mma, opus::make_tuple((int)kargs.stride_ws, 1_I), opus::make_tuple(wave_m, lane_id % mma.grpn_c, wave_n, lane_id / mma.grpn_c)); __builtin_amdgcn_s_barrier(); - auto reg_c_ws = opus::cast(reg_c); + auto reg_c_ws = opus::cast(reg_c); store(g_ws, reg_c_ws, u_gc, 0); // Consumer epilogue: rendezvous with the producers (matches the producer's diff --git a/csrc/opus_gemm/include/gfx1250/opus_gemm_pipeline_a16w16_clusterlaunch_tdm_splitk_ws_gfx1250.cuh b/csrc/opus_gemm/include/gfx1250/opus_gemm_pipeline_a16w16_clusterlaunch_tdm_splitk_ws_gfx1250.cuh index b9ab50d2bf..6a5a6aa9d7 100644 --- a/csrc/opus_gemm/include/gfx1250/opus_gemm_pipeline_a16w16_clusterlaunch_tdm_splitk_ws_gfx1250.cuh +++ b/csrc/opus_gemm/include/gfx1250/opus_gemm_pipeline_a16w16_clusterlaunch_tdm_splitk_ws_gfx1250.cuh @@ -1,8 +1,9 @@ // SPDX-License-Identifier: MIT // Copyright (C) 2025-2026, Advanced Micro Devices, Inc. All rights reserved. // -// gfx1250 bf16 TDM a16w16 GEMM, 4-wave split-K via fp32 workspace + separate -// reduce kernel -- CLUSTER-LAUNCH variant. C[M,N] = A[M,K] @ B[N,K]^T (+ bias). +// gfx1250 bf16 TDM a16w16 GEMM, 4-wave split-K via a per-kid bf16/fp32 +// workspace + separate reduce kernel -- CLUSTER-LAUNCH variant. +// C[M,N] = A[M,K] @ B[N,K]^T (+ bias). // // CLUSTER (kClusterWgM x kClusterWgN x 1) = a CWGM x CWGN grid of workgroups that // co-reside and share TDM loads via CLUSTER_LOAD_ASYNC multicast (MI400 SPG @@ -89,6 +90,7 @@ void gemm_a16w16_clusterlaunch_tdm_splitk_ws_kernel_gfx1250(opus_gemm_cluster_td using T = remove_cvref_t; using DataA = typename T::DataA; using DataB = typename T::DataB; + using DataWS = typename T::DataWS; using DataAcc = typename T::DataAcc; DECLARE_NAMED_BARRIERS(); // __nbar_1..__nbar_15 (we use 1..2*kNumSlots <= 6) @@ -486,17 +488,16 @@ void gemm_a16w16_clusterlaunch_tdm_splitk_ws_kernel_gfx1250(opus_gemm_cluster_td // ---- Store the partial into ws[split_idx][padded_m][padded_n]. ---- // bias is folded once by the reduce kernel (not here). - // The partial type is the traits' D_C -- per kid (splitk_workspace_dtype), and - // the reduce is instantiated with the same D_C so the two cannot diverge. - using DataWs = typename T::DataC; - constexpr int kCVec = T::kCVec; // 4 (fp32 dwordx4 / bf16 dwordx2) - DataWs* ws_ptr = reinterpret_cast(kargs.ptr_ws); + // The exact kid declares the physical workspace type in its traits. Keep + // this in lockstep with the Torch workspace dtype and reducer D_WS. + constexpr int kCVec = T::kCVec; + DataWS* ws_ptr = reinterpret_cast(kargs.ptr_ws); const size_t ws_split = (size_t)split_idx * (size_t)kargs.stride_ws_batch; const size_t ws_base = ws_split + (size_t)tile_row * (size_t)kargs.stride_ws + (size_t)tile_col; const unsigned int ws_bytes = (unsigned int)(((size_t)kargs.stride_ws_batch - ((size_t)tile_row * kargs.stride_ws + tile_col)) * - sizeof(DataWs)); - auto g_ws = make_gmem(ws_ptr + ws_base, ws_bytes); + sizeof(DataWS)); + auto g_ws = make_gmem(ws_ptr + ws_base, ws_bytes); auto u_gc = partition_layout_c(mma, opus::make_tuple((int)kargs.stride_ws, 1_I), opus::make_tuple(wave_m, lane_id % mma.grpn_c, wave_n, lane_id / mma.grpn_c)); // Consumer epilogue: rendezvous with the producers (all 4 waves) BEFORE the @@ -505,7 +506,7 @@ void gemm_a16w16_clusterlaunch_tdm_splitk_ws_kernel_gfx1250(opus_gemm_cluster_td // Unguarded: the tiles the cluster round-up pushed past the padded workspace // (padded_M = ceil(M/B_M)*B_M, and tile_row >= padded_M is the same test as // tile_row >= M) left at tile_oob, so every WG still here owns a real tile. - auto reg_c_ws = opus::cast(reg_c); + auto reg_c_ws = opus::cast(reg_c); store(g_ws, reg_c_ws, u_gc, 0); // Consumer epilogue: rendezvous with the producers (matches the producer's diff --git a/csrc/opus_gemm/include/gfx1250/opus_gemm_traits_a16w16_gfx1250.cuh b/csrc/opus_gemm/include/gfx1250/opus_gemm_traits_a16w16_gfx1250.cuh index d688d08a57..4afb6f6628 100644 --- a/csrc/opus_gemm/include/gfx1250/opus_gemm_traits_a16w16_gfx1250.cuh +++ b/csrc/opus_gemm/include/gfx1250/opus_gemm_traits_a16w16_gfx1250.cuh @@ -2,9 +2,9 @@ // Copyright (C) 2025-2026, Advanced Micro Devices, Inc. All rights reserved. // // Traits + kargs for the gfx1250 a16w16 cluster/TDM split-K pipeline that -// reduces via an fp32 WORKSPACE + a separate REDUCE kernel (no atomic_add, -// no self-clear, no semaphore). The workspace buffer is allocated externally -// (torch.empty on the Python side) and passed as a direct pointer in kargs. +// reduces via a per-kid bf16/fp32 WORKSPACE + a separate REDUCE kernel (no +// atomic_add, no self-clear, no semaphore). The caller-owned typed workspace +// is passed as a direct pointer in kargs. // // This header is the SINGLE source of truth for every compile-time constant // the pipeline needs: the pipeline file @@ -33,31 +33,16 @@ constexpr int kCtdmLayoutTileN = 0; constexpr int kCtdmLayoutTileM = 1; } // namespace opus_gfx1250 - -#ifndef OPUS_GEMM_SPLITK_WS_HANDLE_DEFINED -#define OPUS_GEMM_SPLITK_WS_HANDLE_DEFINED -// Indirection slot for the split-K fp32 workspace pointer. Captured HIP -// graphs hold the slot address (stable), not the workspace ptr, so a -// post-capture grow + hipFree of the old buffer doesn't dangle the graph. -struct opus_splitk_ws_handle { - void* ptr; // current backing workspace; null until first grow - unsigned long bytes; // current capacity in bytes -}; -#endif - #ifndef OPUS_GEMM_CLUSTER_TDM_WS_KARGS_GFX1250_DEFINED #define OPUS_GEMM_CLUSTER_TDM_WS_KARGS_GFX1250_DEFINED // Kernel arguments for the gfx1250 a16w16 cluster/TDM split-K (workspace) -// pipeline. The main kernel writes partial sums (bf16 by default) into -// ptr_ws laid out as [split_k, padded_M, padded_N] (per host launch; -// batch handled by a per-batch host launch with pointer offsets). The reduce -// kernel consumes it, folds bias once, casts to Y dtype, writes C[M, N]. -// The workspace buffer is allocated externally (torch.empty on the Python -// side) and passed in directly -- no indirection through a handle struct. +// pipeline. The main kernel writes D_WS partial sums into ptr_ws laid out as +// [split_k, padded_M, padded_N]. The reduce kernel consumes them, +// re-accumulates in fp32, folds bias once, casts to Y dtype, and writes C[M,N]. struct opus_gemm_cluster_tdm_ws_kargs_gfx1250 { const void* __restrict__ ptr_a; // bf16 [M, K] const void* __restrict__ ptr_b; // bf16 [N, K] (A @ B^T) - void* __restrict__ ptr_ws; // workspace [split_k, padded_M, padded_N] + void* __restrict__ ptr_ws; // D_WS [split_k, padded_M, padded_N] void* __restrict__ ptr_c; // bf16/fp32 [M, N] (filled by reduce kernel) const void* __restrict__ ptr_bias; // consumed by reduce kernel only int m; @@ -111,54 +96,27 @@ struct opus_gemm_splitk_fuse_kargs_gfx1250 #ifndef OPUS_GEMM_4WAVE_COMPUTE_KARGS_GFX1250_DEFINED #define OPUS_GEMM_4WAVE_COMPUTE_KARGS_GFX1250_DEFINED -// Kernel arguments for the gfx1250 a16w16 SYMMETRIC 4-wave compute pipeline -// (kernel_tag a16w16_4wave_co). Shorter than the split-K kargs above because -// this pipeline has no workspace, no split_k and no bias: every wave both -// TDM-loads and runs WMMA, and bf16 C is stored straight out through LDS. -// -// This kernel is compiled by HIP (not hand-written asm), so it uses the plain -// C++ kernarg ABI -- do NOT add `packed` or p2/p3 padding members the way the -// asm .co kargs structs do. Host and device share this one definition, so the -// two cannot drift. -// EXACTLY 64 BYTES, and it must stay that way: built with -D__HIPCC_RTC__ the -// .co has a single by-value kernarg and no hidden/implicit args, so this struct -// IS the kernarg segment. build_info.json records kernarg_segment_size, which -// makes "still 64" the regression signal that the ABI has not drifted. -// -// The batch strides are 64-bit ELEMENT counts: they cross 2^31 at 4 GiB of bf16, -// and int arithmetic there is UB long before the hardware minds. The budget for -// that came from dropping two fields rather than growing the struct: -// * a batch COUNT -- grid.z already carries it, the kernel never read it; -// * a C batch stride -- the kernel derives m * stride_c, which also stays -// correct for a row-padded C, unlike the M * N the host used to pass. +// Exact kernarg ABI baked into every pre-built 4wave CO image. Batch strides +// are 64-bit element counts; grid.z carries the batch count and C's batch +// stride is derived from m * stride_c in the device pipeline. struct opus_gemm_4wave_compute_kargs_gfx1250 { - const void* __restrict__ ptr_a; // 0 bf16 [batch, M, K] - const void* __restrict__ ptr_b; // 8 bf16 [batch, N, K] (C = A * B^T) - void* __restrict__ ptr_c; // 16 bf16 [batch, M, N] - int64_t stride_a_batch; // 24 = M * stride_a - int64_t stride_b_batch; // 32 = N * stride_b + const void* __restrict__ ptr_a; // 0: bf16 [batch, M, K] + const void* __restrict__ ptr_b; // 8: bf16 [batch, N, K] + void* __restrict__ ptr_c; // 16: bf16 [batch, M, N] + int64_t stride_a_batch; // 24 + int64_t stride_b_batch; // 32 int m; // 40 int n; // 44 int k; // 48 - int stride_a; // 52 A row pitch (>= K) - int stride_b; // 56 B row pitch (>= K) - int stride_c; // 60 C row pitch (>= N) + int stride_a; // 52 + int stride_b; // 56 + int stride_c; // 60 }; static_assert(sizeof(opus_gemm_4wave_compute_kargs_gfx1250) == 64, - "the 4wave_co kernarg segment is baked into the pre-built .co at " - "64 bytes -- changing this size means rebuilding gen_co/*/*.co"); + "the pre-built 4wave CO kernarg ABI must remain 64 bytes"); #endif -// ── 4wave_compute user traits ─────────────────────────────────────────────── -// The compile-time config for the symmetric 4-wave pipeline. Member names are -// the UPPER_CASE spellings the pipeline body already reads, kept verbatim from -// the standalone kernel this was ported from: the register pinning in that body -// is fragile enough (see the pipeline header) that renaming through it is not -// worth the risk. k-prefixed aliases are added for the host launcher, which -// follows the aiter convention. -// -// CLUSTER_WG_M / CLUSTER_WG_N are template parameters rather than the hardcoded -// 4/4 of the standalone, so the instance table can sweep cluster geometry. +// Host-visible configuration shared with the offline 4wave device build. template struct opus_a16w16_4wave_compute_traits_gfx1250 { - static constexpr int BLOCK_SIZE = BLOCK_SIZE_; // 128 = 4 waves x 32 - + static constexpr int BLOCK_SIZE = BLOCK_SIZE_; static constexpr int B_M = B_M_; static constexpr int B_N = B_N_; static constexpr int B_K = B_K_; - using D_A = D_A_; - using D_B = D_B_; - using D_C = D_C_; + using D_A = D_A_; + using D_B = D_B_; + using D_C = D_C_; using D_ACC = D_ACC_; + using DataA = D_A; + using DataB = D_B; + using DataC = D_C; + using DataAcc = D_ACC; static_assert(std::is_same::value, "A/B dtype must match"); - static constexpr int VEC_A = 16 / (int)sizeof(D_A); // 8 for bf16 (b128 ds_read) + static constexpr int VEC_A = 16 / (int)sizeof(D_A); static constexpr int VEC_B = 16 / (int)sizeof(D_B); - - // LDS prefetch ring depth. The pipeline keeps 2 TDMs in flight and reuses - // slot g%P three steps later, so P >= 3 keeps g, g+1, g+2 distinct. static constexpr int NUM_SLOTS = NUM_SLOTS_; - // How many slots the K loop may use is a property of the PIPELINE, not of the - // geometry, so each pipeline asserts its own bound. - static_assert(NUM_SLOTS >= 2, "the ring needs at least two slots"); + static_assert(NUM_SLOTS >= 2, "the LDS ring needs at least two slots"); - // Cluster-launch multicast geometry: a CLUSTER_WG_M x CLUSTER_WG_N grid of - // workgroups per cluster. A is multicast to the CLUSTER_WG_N peers sharing - // an M row, B to the CLUSTER_WG_M peers sharing an N column. static constexpr int CLUSTER_WG_M = CLUSTER_WG_M_; static constexpr int CLUSTER_WG_N = CLUSTER_WG_N_; - // TDM multicast fans out to at most 5 WGs per group, and the per-cluster - // workgroup_mask is 16-bit. static_assert(CLUSTER_WG_M >= 1 && CLUSTER_WG_N >= 1 && CLUSTER_WG_M <= 5 && CLUSTER_WG_N <= 5 && CLUSTER_WG_M * CLUSTER_WG_N <= 16, "cluster dims must be 1..5 per side and <= 16 WGs total"); - // TDM/LDS pad: +16B (one PAD_ELEMS group) per B_K row -> bank-conflict-free - // b128 ds_read. Only the ELEMENT geometry lives here, because that is all the - // host needs to size LDS; the D# pad_interval/pad_amount encoding is derived - // by opus::tdm_traits::padding_auto, and the pipeline static_asserts that its - // pitch matches SMEM_PITCH so the two can never drift. - static_assert((B_K & (B_K - 1)) == 0, "B_K must be a power of 2 for a single pad per row"); - static constexpr int PAD_ELEMS = 16 / (int)sizeof(D_A); // 8 bf16 = +16B - static constexpr int SMEM_PITCH = B_K + PAD_ELEMS; - - // One LDS slot holds the full B_M x B_K (A) / B_N x B_K (B) tile. + static_assert((B_K & (B_K - 1)) == 0, + "B_K must be a power of two for the TDM padding scheme"); + static constexpr int PAD_ELEMS = 16 / (int)sizeof(D_A); + static constexpr int SMEM_PITCH = B_K + PAD_ELEMS; static constexpr int SLOT_BYTES_A = B_M * SMEM_PITCH * (int)sizeof(D_A); static constexpr int SLOT_BYTES_B = B_N * SMEM_PITCH * (int)sizeof(D_B); - static constexpr int SEG_BYTES_A = NUM_SLOTS * SLOT_BYTES_A; - static constexpr int SEG_BYTES_B = NUM_SLOTS * SLOT_BYTES_B; + static constexpr int SEG_BYTES_A = NUM_SLOTS * SLOT_BYTES_A; + static constexpr int SEG_BYTES_B = NUM_SLOTS * SLOT_BYTES_B; static constexpr int SEG_BYTES_AB = SEG_BYTES_A + SEG_BYTES_B; - // 1-WG/CU enforcement via LDS padding, the same trick the split-K traits use - // below. This pipeline is only correct at one workgroup per CU: every tile - // whose A/B segments fit twice in the 320 KB budget (<= 160 KB) raced -- - // non-deterministically wrong at large grids, in BOTH this and the - // wave-layout pipeline. Padding past 160 KB so a second workgroup cannot - // co-reside fixes it; the pad tail is never accessed. - // - // That occupancy is the variable (rather than tile size or register - // pressure) is pinned down by a control group: 71 variants whose registers - // would admit two waves per SIMD but whose LDS does not are all correct, - // while the 61 where both admit two are all wrong. Both races behind it -- a - // write-after-read on the ring and a trailing "zero-extent" transfer that - // zero-fills the slot C stages in -- are fixed; the pad stays only because - // dropping it is a per-shape performance trade. See KNOWN_ISSUES.md issue 1. - // -DOPUS_CO_NO_1WG_PAD drops it, back to 2 WG/CU, for re-measuring. - static constexpr int kHalfLds = 160 * 1024; + // Enforce the one-WG/CU occupancy used to validate the shipped images. + static constexpr int kHalfLds = 160 * 1024; #ifdef OPUS_CO_NO_1WG_PAD - static constexpr int LDS_BYTES = SEG_BYTES_AB; + static constexpr int LDS_BYTES = SEG_BYTES_AB; #else - static constexpr int LDS_BYTES = + static constexpr int LDS_BYTES = (SEG_BYTES_AB <= kHalfLds) ? (kHalfLds + 1024) : SEG_BYTES_AB; #endif static_assert(LDS_BYTES <= 320 * 1024, "LDS exceeds the 320KB/CU budget"); - // aiter-convention aliases for the host launcher (which never sees the - // pipeline header, only this one). - static constexpr int kBlockM = B_M; - static constexpr int kBlockN = B_N; - static constexpr int kBlockK = B_K; - static constexpr int kNumSlots = NUM_SLOTS; - static constexpr int kClusterWgM = CLUSTER_WG_M; - static constexpr int kClusterWgN = CLUSTER_WG_N; + static constexpr int kBlockM = B_M; + static constexpr int kBlockN = B_N; + static constexpr int kBlockK = B_K; + static constexpr int kNumSlots = NUM_SLOTS; + static constexpr int kClusterWgM = CLUSTER_WG_M; + static constexpr int kClusterWgN = CLUSTER_WG_N; static constexpr int kLdsTotalBytes = LDS_BYTES; }; -// ── 4wave with a CONFIGURABLE WAVE LAYOUT (a16w16_4wave_wl_co) ────────────── -// Same geometry as the traits above plus TILE_M x TILE_N, the shape in which the -// four waves tile the block. TILE_M * TILE_N == 4 (checked in the pipeline). -// (4, 1) reproduces the fixed layout of the reference traits, so the two agree -// wherever they overlap; the pipeline that reads this one is -// opus_gemm_pipeline_a16w16_4wave_wl_gfx1250.cuh. kargs are shared -- the block -// geometry lives entirely here, so the kernarg struct never changes. template struct opus_a16w16_4wave_wl_traits_gfx1250 - : opus_a16w16_4wave_compute_traits_gfx1250 { + : opus_a16w16_4wave_compute_traits_gfx1250< + BLOCK_SIZE_, B_M_, B_N_, B_K_, NUM_SLOTS_, + D_A_, D_B_, D_C_, D_ACC_, CLUSTER_WG_M_, CLUSTER_WG_N_> { static constexpr int TILE_M = TILE_M_; static constexpr int TILE_N = TILE_N_; }; // ── User-facing traits = the SINGLE compile-time config the pipeline reads ── -// D_A=D_B=bf16, D_ACC=float (WMMA fp32 acc), D_C MUST be float (main kernel -// writes the fp32 workspace; the reduce kernel casts to the final Y dtype). +// D_A=D_B=bf16, D_ACC=float (WMMA fp32 acc), D_WS is bf16 or fp32. +// The main kernel casts its accumulator to D_WS and the reducer reads the +// same physical type before re-accumulating in fp32. template::value, "A/B dtype must match"); - // D_C is the split-K PARTIAL type, not the output dtype -- the reduce picks the - // output. Either width is legal; the kid table chooses (splitk_workspace_dtype). - static_assert(std::is_same::value || std::is_same::value, - "cluster_tdm_splitk_ws partial workspace must be float or __bf16"); + static_assert(std::is_same::value || + std::is_same::value, + "cluster_tdm_splitk_ws D_WS must be bf16 input storage or " + "fp32 accumulator storage"); // Aliases used by the pipeline / layout helpers. using DataA = D_A; using DataB = D_B; - using DataC = D_C; + using DataWS = D_WS; using DataAcc = D_ACC; static constexpr int VEC_A = 16 / (int)sizeof(D_A); // 8 for bf16 @@ -424,8 +347,9 @@ struct opus_cluster_tdm_splitk_ws_traits_gfx1250 { // gfx1250 LDS max ~320KB. static_assert(kLdsTotalBytes <= 320 * 1024, "LDS exceeds 320KB"); - // Workspace plain store: fp32 dwordx4. - static constexpr int kCVec = 16 / (int)sizeof(DataAcc); // 4 (fp32) + // Four workspace elements per issue. This preserves the accumulator + // register partition for both bf16 and fp32 storage. + static constexpr int kCVec = 4; // ── Warp-derived WMMA register-decomposition constants ─────────────────── // (computed from kWarpRt so device/host passes match) diff --git a/csrc/opus_gemm/include/gfx942/a16w16/opus_gemm_helpers_a16w16.cuh b/csrc/opus_gemm/include/gfx942/a16w16/opus_gemm_helpers_a16w16.cuh index 1d1457f078..7528a06c38 100644 --- a/csrc/opus_gemm/include/gfx942/a16w16/opus_gemm_helpers_a16w16.cuh +++ b/csrc/opus_gemm/include/gfx942/a16w16/opus_gemm_helpers_a16w16.cuh @@ -476,7 +476,7 @@ OPUS_D inline void epilogue_store_workspace_sc0nt( Mma& mma, GC& g_c, const Kargs& kargs, VC& v_c, int wave_id_m, int wave_id_n, int lane_id) { - using D_C = typename T::D_C; + using D_WS = typename T::D_WS; using D_ACC = typename T::D_ACC; auto p_coord_c = opus::make_tuple(wave_id_m, lane_id % mma.grpn_c, @@ -488,16 +488,16 @@ OPUS_D inline void epilogue_store_workspace_sc0nt( return half_m * T::HALF_B_M * kargs.stride_ws + half_n * T::HALF_B_N; }; - if constexpr (std::is_same_v) { + if constexpr (std::is_same_v) { opus::store(g_c, v_c[0][0], u_gc, ws_offset(0, 0), opus::number<3>{}); opus::store(g_c, v_c[0][1], u_gc, ws_offset(0, 1), opus::number<3>{}); opus::store(g_c, v_c[1][0], u_gc, ws_offset(1, 0), opus::number<3>{}); opus::store(g_c, v_c[1][1], u_gc, ws_offset(1, 1), opus::number<3>{}); } else { - auto c00 = opus::cast(v_c[0][0]); - auto c01 = opus::cast(v_c[0][1]); - auto c10 = opus::cast(v_c[1][0]); - auto c11 = opus::cast(v_c[1][1]); + auto c00 = opus::cast(v_c[0][0]); + auto c01 = opus::cast(v_c[0][1]); + auto c10 = opus::cast(v_c[1][0]); + auto c11 = opus::cast(v_c[1][1]); opus::store(g_c, c00, u_gc, ws_offset(0, 0), opus::number<3>{}); opus::store(g_c, c01, u_gc, ws_offset(0, 1), opus::number<3>{}); opus::store(g_c, c10, u_gc, ws_offset(1, 0), opus::number<3>{}); diff --git a/csrc/opus_gemm/include/gfx942/a16w16/opus_gemm_pipeline_a16w16_em3en4_lds1_pgr2_sk.cuh b/csrc/opus_gemm/include/gfx942/a16w16/opus_gemm_pipeline_a16w16_em3en4_lds1_pgr2_sk.cuh index 7dc32d09ca..ca07506f91 100644 --- a/csrc/opus_gemm/include/gfx942/a16w16/opus_gemm_pipeline_a16w16_em3en4_lds1_pgr2_sk.cuh +++ b/csrc/opus_gemm/include/gfx942/a16w16/opus_gemm_pipeline_a16w16_em3en4_lds1_pgr2_sk.cuh @@ -543,13 +543,13 @@ void gemm_a16w16_em3en4_lds1_pgr2_sk_kernel(opus_gemm_splitk_kargs kargs) { using T = opus::remove_cvref_t; using D_A = typename T::D_A; using D_B = typename T::D_B; - using D_C = typename T::D_C; + using D_WS = typename T::D_WS; using D_ACC = typename T::D_ACC; static_assert(T::B_M == 96 && T::B_N == 128 && T::B_K == 128); static_assert(T::T_M == 2 && T::T_N == 2 && T::BLOCK_SIZE == 256); static_assert(T::E_M == 3 && T::E_N == 4 && T::E_K == 8); - static_assert(std::is_same_v, + static_assert(std::is_same_v, "EM3EN4 LDS1/PGR2 splitK main kernel writes fp32 workspace"); int wgid_full = opus::block_id_x(); @@ -581,7 +581,7 @@ void gemm_a16w16_em3en4_lds1_pgr2_sk_kernel(opus_gemm_splitk_kargs kargs) { auto g_b = make_gmem(reinterpret_cast(kargs.ptr_a) + batch_id * kargs.stride_a_batch + row * kargs.stride_a + k_start, ((kargs.m - row) * kargs.stride_a - k_start) * sizeof(D_B)); - auto g_c = make_gmem(opus_splitk_ws_ptr(kargs.ws_handle) + auto g_c = make_gmem(opus_gfx942_uniform_ws_ptr(kargs.ptr_ws) + (size_t)split_id * kargs.batch * kargs.stride_ws_batch + (size_t)batch_id * kargs.stride_ws_batch + (size_t)row * kargs.stride_ws diff --git a/csrc/opus_gemm/include/gfx942/a16w16/opus_gemm_pipeline_a16w16_kbuf1.cuh b/csrc/opus_gemm/include/gfx942/a16w16/opus_gemm_pipeline_a16w16_kbuf1.cuh index 65e4422893..df11bd67ae 100644 --- a/csrc/opus_gemm/include/gfx942/a16w16/opus_gemm_pipeline_a16w16_kbuf1.cuh +++ b/csrc/opus_gemm/include/gfx942/a16w16/opus_gemm_pipeline_a16w16_kbuf1.cuh @@ -31,6 +31,7 @@ void gemm_a16w16_kbuf1_kernel(Kargs kargs) { using D_A = typename T::D_A; using D_B = typename T::D_B; using D_C = typename T::D_C; + using D_WS = typename T::D_WS; using D_ACC = typename T::D_ACC; int wgid; @@ -70,7 +71,7 @@ void gemm_a16w16_kbuf1_kernel(Kargs kargs) { auto g_c = [&]() { if constexpr (IS_SPLITK) { - return make_gmem(opus_splitk_ws_ptr(kargs.ws_handle) + return make_gmem(opus_gfx942_uniform_ws_ptr(kargs.ptr_ws) + (size_t)split_id * kargs.batch * kargs.stride_ws_batch + (size_t)batch_id * kargs.stride_ws_batch + (size_t)row * kargs.stride_ws diff --git a/csrc/opus_gemm/include/gfx942/a16w16/opus_gemm_pipeline_a16w16_kbuf1_large_tile.cuh b/csrc/opus_gemm/include/gfx942/a16w16/opus_gemm_pipeline_a16w16_kbuf1_large_tile.cuh index 39242fb8a1..0809e7b90a 100644 --- a/csrc/opus_gemm/include/gfx942/a16w16/opus_gemm_pipeline_a16w16_kbuf1_large_tile.cuh +++ b/csrc/opus_gemm/include/gfx942/a16w16/opus_gemm_pipeline_a16w16_kbuf1_large_tile.cuh @@ -256,7 +256,10 @@ __global__ __launch_bounds__(Traits::BLOCK_SIZE, 2) void gemm_a16w16_kbuf1_large const bool full_tile = (row + T::B_M <= kargs.m) && (col + T::B_N <= kargs.n); - if (full_tile) { + // The LDS coalescing buffers alias the BF16 A/B allocation. They are large + // enough for BF16 output, but not for two FP32 half tiles. FP32 already has + // a naturally coalesced direct-store path through do_store_if. + if (full_tile && !std::is_same_v) { using LT_C = layout_load_traits; constexpr auto r_elem_c = LT_C::r_elem; constexpr index_t acc_chunk = T::VEC_C * vector_traits::size(); diff --git a/csrc/opus_gemm/include/gfx942/a16w16/opus_gemm_pipeline_a16w16_kbuf2v.cuh b/csrc/opus_gemm/include/gfx942/a16w16/opus_gemm_pipeline_a16w16_kbuf2v.cuh index e2ec665685..172deb217e 100644 --- a/csrc/opus_gemm/include/gfx942/a16w16/opus_gemm_pipeline_a16w16_kbuf2v.cuh +++ b/csrc/opus_gemm/include/gfx942/a16w16/opus_gemm_pipeline_a16w16_kbuf2v.cuh @@ -35,6 +35,7 @@ void gemm_a16w16_kbuf2v_kernel(Kargs kargs) { using D_A = typename T::D_A; using D_B = typename T::D_B; using D_C = typename T::D_C; + using D_WS = typename T::D_WS; using D_ACC = typename T::D_ACC; int wgid; @@ -74,7 +75,7 @@ void gemm_a16w16_kbuf2v_kernel(Kargs kargs) { auto g_c = [&]() { if constexpr (IS_SPLITK) { - return make_gmem(opus_splitk_ws_ptr(kargs.ws_handle) + return make_gmem(opus_gfx942_uniform_ws_ptr(kargs.ptr_ws) + (size_t)split_id * kargs.batch * kargs.stride_ws_batch + (size_t)batch_id * kargs.stride_ws_batch + (size_t)row * kargs.stride_ws diff --git a/csrc/opus_gemm/include/gfx942/a16w16/opus_gemm_pipeline_a16w16_kbuf2v_bk128.cuh b/csrc/opus_gemm/include/gfx942/a16w16/opus_gemm_pipeline_a16w16_kbuf2v_bk128.cuh index 536f374b5a..682edd6ce3 100644 --- a/csrc/opus_gemm/include/gfx942/a16w16/opus_gemm_pipeline_a16w16_kbuf2v_bk128.cuh +++ b/csrc/opus_gemm/include/gfx942/a16w16/opus_gemm_pipeline_a16w16_kbuf2v_bk128.cuh @@ -26,6 +26,7 @@ struct bk64_traits_view { using D_A = typename T::D_A; using D_B = typename T::D_B; using D_C = typename T::D_C; + using D_WS = typename T::D_WS; using D_ACC = typename T::D_ACC; static constexpr int T_M = T::T_M; static constexpr int T_N = T::T_N; @@ -72,6 +73,7 @@ void gemm_a16w16_kbuf2v_bk128_kernel(Kargs kargs) { using D_A = typename T::D_A; using D_B = typename T::D_B; using D_C = typename T::D_C; + using D_WS = typename T::D_WS; using D_ACC = typename T::D_ACC; int wgid; @@ -111,7 +113,7 @@ void gemm_a16w16_kbuf2v_bk128_kernel(Kargs kargs) { auto g_c = [&]() { if constexpr (IS_SPLITK) { - return make_gmem(opus_splitk_ws_ptr(kargs.ws_handle) + return make_gmem(opus_gfx942_uniform_ws_ptr(kargs.ptr_ws) + (size_t)split_id * kargs.batch * kargs.stride_ws_batch + (size_t)batch_id * kargs.stride_ws_batch + (size_t)row * kargs.stride_ws diff --git a/csrc/opus_gemm/include/gfx942/a16w16/opus_gemm_pipeline_a16w16_quad_mfma32_kbuf1.cuh b/csrc/opus_gemm/include/gfx942/a16w16/opus_gemm_pipeline_a16w16_quad_mfma32_kbuf1.cuh index a49b3ca384..fef54f42db 100644 --- a/csrc/opus_gemm/include/gfx942/a16w16/opus_gemm_pipeline_a16w16_quad_mfma32_kbuf1.cuh +++ b/csrc/opus_gemm/include/gfx942/a16w16/opus_gemm_pipeline_a16w16_quad_mfma32_kbuf1.cuh @@ -208,6 +208,8 @@ __global__ __launch_bounds__(Traits::BLOCK_SIZE, 1) void gemm_a16w16_quad_mfma32 using D_A = typename T::D_A; using D_B = typename T::D_B; using D_C = typename T::D_C; + using D_WS = typename T::D_WS; + using D_STORE = std::conditional_t; using D_ACC = typename T::D_ACC; static_assert(T::BLOCK_SIZE == 256); @@ -274,7 +276,7 @@ __global__ __launch_bounds__(Traits::BLOCK_SIZE, 1) void gemm_a16w16_quad_mfma32 ((kargs.n - col) * kargs.stride_b - k_start) * sizeof(D_B)); auto g_c = [&]() { if constexpr (IS_SPLITK) { - return make_gmem(opus_splitk_ws_ptr(kargs.ws_handle) + return make_gmem(opus_gfx942_uniform_ws_ptr(kargs.ptr_ws) + (size_t)split_id * kargs.batch * kargs.stride_ws_batch + (size_t)batch_id * kargs.stride_ws_batch + (size_t)row * kargs.stride_ws @@ -300,7 +302,7 @@ __global__ __launch_bounds__(Traits::BLOCK_SIZE, 1) void gemm_a16w16_quad_mfma32 // 136 bf16 columns keeps C-stage rows 16B aligned while avoiding the // slower 128-column LDS drain pattern on this shape. constexpr int C_LDS_STRIDE = T::HALF_B_N + 8; - constexpr int c_stage_byte = T::HALF_B_M * C_LDS_STRIDE * sizeof(D_C); + constexpr int c_stage_byte = T::HALF_B_M * C_LDS_STRIDE * sizeof(D_STORE); constexpr int smem_bytes = ab_stage_byte > c_stage_byte ? ab_stage_byte : c_stage_byte; static_assert(smem_bytes <= 64 * 1024); @@ -568,10 +570,10 @@ __global__ __launch_bounds__(Traits::BLOCK_SIZE, 1) void gemm_a16w16_quad_mfma32 }; auto read_acc_for_store = [](const float16_acc* acc) { - if constexpr (std::is_same_v) { + if constexpr (std::is_same_v) { return agpr_to_bf16_vgpr_trunc<4>(acc); } else { - return cast(agpr_to_vgpr<4>(acc)); + return cast(agpr_to_vgpr<4>(acc)); } }; @@ -591,17 +593,17 @@ __global__ __launch_bounds__(Traits::BLOCK_SIZE, 1) void gemm_a16w16_quad_mfma32 auto do_full_tile_store = [&]() { using LT_C = layout_load_traits; constexpr auto r_elem_c = LT_C::r_elem; - constexpr index_t c_chunk = T::VEC_C * vector_traits::size(); + constexpr index_t c_chunk = T::VEC_C * vector_traits::size(); constexpr int HALF_TILE_ELEMS = T::HALF_B_M * T::HALF_B_N; constexpr int THREAD_TILE_VEC = HALF_TILE_ELEMS / T::BLOCK_SIZE; static_assert(THREAD_TILE_VEC * T::BLOCK_SIZE == HALF_TILE_ELEMS); - constexpr int MAX_STORE_VEC = 16 / sizeof(D_C); + constexpr int MAX_STORE_VEC = 16 / sizeof(D_STORE); constexpr int STORE_VEC = THREAD_TILE_VEC < MAX_STORE_VEC ? THREAD_TILE_VEC : MAX_STORE_VEC; static_assert(THREAD_TILE_VEC % STORE_VEC == 0); constexpr int STORE_ITERS = THREAD_TILE_VEC / STORE_VEC; constexpr int LDS_STRIDE = C_LDS_STRIDE; - smem s_c = make_smem(reinterpret_cast(smem_storage)); + smem s_c = make_smem(reinterpret_cast(smem_storage)); auto u_lds_c = partition_layout_c(mma, opus::make_tuple(opus::number{}, 1_I), p_coord_c); auto offsets_lds = layout_to_offsets(u_lds_c); @@ -626,7 +628,7 @@ __global__ __launch_bounds__(Traits::BLOCK_SIZE, 1) void gemm_a16w16_quad_mfma32 auto store_one_quadrant = [&](auto& vc, int hm, int hn) { #pragma unroll for (index_t i = 0; i < r_elem_c.value; i++) { - vector_t chunk; + vector_t chunk; #pragma unroll for (index_t j = 0; j < c_chunk; j++) { chunk[j] = vc[i * c_chunk + j]; diff --git a/csrc/opus_gemm/include/gfx942/a16w16/opus_gemm_pipeline_a16w16_wave_k_coop.cuh b/csrc/opus_gemm/include/gfx942/a16w16/opus_gemm_pipeline_a16w16_wave_k_coop.cuh index 37eaa8455e..10e5a46225 100644 --- a/csrc/opus_gemm/include/gfx942/a16w16/opus_gemm_pipeline_a16w16_wave_k_coop.cuh +++ b/csrc/opus_gemm/include/gfx942/a16w16/opus_gemm_pipeline_a16w16_wave_k_coop.cuh @@ -633,12 +633,6 @@ void gemm_a16w16_wave_k_coop_accum_kernel(opus_gemm_noscale_kargs kargs) full_tile ? 0xffffffffu : (unsigned int)(((kargs.n - col) * kargs.stride_b) * sizeof(D_B))); - if (split_id == 0) { - zero_output_tile_bf16x2( - ptr_y, row, col, kargs.m, kargs.n, kargs.stride_c, tid); - __builtin_amdgcn_s_barrier(); - } - __shared__ char smem[LDS_BYTES]; char* smem_a = smem; char* smem_b = smem + A_BYTES; diff --git a/csrc/opus_gemm/include/gfx942/a16w16/opus_gemm_traits_a16w16.cuh b/csrc/opus_gemm/include/gfx942/a16w16/opus_gemm_traits_a16w16.cuh index 3c7b4c2f6f..bb1c7f2e18 100644 --- a/csrc/opus_gemm/include/gfx942/a16w16/opus_gemm_traits_a16w16.cuh +++ b/csrc/opus_gemm/include/gfx942/a16w16/opus_gemm_traits_a16w16.cuh @@ -34,6 +34,10 @@ struct opus_gemm_a16w16_traits { using D_A = opus::tuple_element_t<0, DTYPE>; using D_B = opus::tuple_element_t<1, DTYPE>; using D_C = opus::tuple_element_t<2, DTYPE>; + // Tuple slot 2 is the final output type for non-split kernels and the + // caller-owned workspace type for split-K kernels. Split-K pipelines use + // this explicit alias so workspace width never comes from D_ACC/D_OUT. + using D_WS = opus::tuple_element_t<2, DTYPE>; using D_ACC = opus::tuple_element_t<3, DTYPE>; static_assert(std::is_same::value); @@ -101,31 +105,31 @@ struct opus_gemm_noscale_kargs { #ifndef OPUS_GEMM_SPLITK_KARGS_GFX942_DEFINED #define OPUS_GEMM_SPLITK_KARGS_GFX942_DEFINED -#ifndef OPUS_GEMM_SPLITK_WS_HANDLE_DEFINED -#define OPUS_GEMM_SPLITK_WS_HANDLE_DEFINED -struct opus_splitk_ws_handle { - void* ptr; - unsigned long bytes; -}; -#endif #ifdef __HIP_DEVICE_COMPILE__ -template -__device__ __forceinline__ D_WS* opus_splitk_ws_ptr( - const opus_splitk_ws_handle* __restrict__ ws_handle) { +template +__device__ __forceinline__ auto opus_gfx942_uniform_ws_ptr(Ptr ptr_ws) + -> std::conditional_t< + std::is_const_v>, const D_WS*, D_WS*> +{ + static_assert(std::is_pointer_v); + static_assert(std::is_same_v< + std::remove_cv_t>, void>); + using D_WS_PTR = std::conditional_t< + std::is_const_v>, const D_WS*, D_WS*>; #if defined(__gfx942__) __UINTPTR_TYPE__ ptr_bits = 0; if ((opus::thread_id_x() & (opus::get_warp_size() - 1)) == 0) { - ptr_bits = reinterpret_cast<__UINTPTR_TYPE__>(ws_handle->ptr); + ptr_bits = reinterpret_cast<__UINTPTR_TYPE__>(ptr_ws); } const unsigned lo = __builtin_amdgcn_readfirstlane( static_cast(ptr_bits)); const unsigned hi = __builtin_amdgcn_readfirstlane( static_cast(ptr_bits >> 32)); ptr_bits = (static_cast<__UINTPTR_TYPE__>(hi) << 32) | lo; - return reinterpret_cast(ptr_bits); + return reinterpret_cast(ptr_bits); #else - return reinterpret_cast(ws_handle->ptr); + return reinterpret_cast(ptr_ws); #endif } #endif @@ -134,7 +138,7 @@ __device__ __forceinline__ D_WS* opus_splitk_ws_ptr( struct opus_gemm_splitk_kargs { const void* __restrict__ ptr_a; // bf16 [B, M, K] const void* __restrict__ ptr_b; // bf16 [B, N, K] (pre-transposed) - const opus_splitk_ws_handle* __restrict__ ws_handle; // deref at kernel entry + void* __restrict__ ptr_ws; // D_WS [split_k, B, padded_M, padded_N] void* __restrict__ ptr_c; // bf16 [B, M, N] final output (reduce kernel writes) const void* __restrict__ ptr_bias; // unused (reserved) int m; diff --git a/csrc/opus_gemm/include/gfx942/a16w16/splitk_reduce_gfx942.cuh b/csrc/opus_gemm/include/gfx942/a16w16/splitk_reduce_gfx942.cuh index 3ada5b5d80..b668f2d1e7 100644 --- a/csrc/opus_gemm/include/gfx942/a16w16/splitk_reduce_gfx942.cuh +++ b/csrc/opus_gemm/include/gfx942/a16w16/splitk_reduce_gfx942.cuh @@ -5,14 +5,14 @@ #pragma once #include "../../opus_gemm_utils.cuh" -#include "opus_gemm_traits_a16w16.cuh" // opus_splitk_ws_handle +#include "opus_gemm_traits_a16w16.cuh" // opus_gfx942_uniform_ws_ptr #include template __device__ __forceinline__ void splitk_reduce_kernel_fallback_body( - const opus_splitk_ws_handle* __restrict__ ws_handle, + const void* __restrict__ ws_ptr, D_OUT* __restrict__ c_out, int split_k, int M, int N, int batch, int padded_M, int padded_N, @@ -47,7 +47,7 @@ __device__ __forceinline__ void splitk_reduce_kernel_fallback_body( const int b = bm_id / M; const int m = bm_id - b * M; - const D_WS* workspace = opus_splitk_ws_ptr(ws_handle); + const D_WS* workspace = opus_gfx942_uniform_ws_ptr(ws_ptr); opus::vector_t bias_fp32; if constexpr (HAS_BIAS) { @@ -194,7 +194,7 @@ template __global__ void splitk_reduce_kernel_fallback( - const opus_splitk_ws_handle* __restrict__ ws_handle, + const void* __restrict__ ws_ptr, D_OUT* __restrict__ c_out, int split_k, int M, int N, int batch, int padded_M, int padded_N, @@ -202,14 +202,14 @@ __global__ void splitk_reduce_kernel_fallback( int bias_stride_batch) { splitk_reduce_kernel_fallback_body( - ws_handle, c_out, split_k, M, N, batch, padded_M, padded_N, bias, bias_stride_batch); + ws_ptr, c_out, split_k, M, N, batch, padded_M, padded_N, bias, bias_stride_batch); } template __global__ void splitk_reduce_kernel_bf16ws_fallback( - const opus_splitk_ws_handle* __restrict__ ws_handle, + const void* __restrict__ ws_ptr, D_OUT* __restrict__ c_out, int split_k, int M, int N, int batch, int padded_M, int padded_N, @@ -217,7 +217,7 @@ __global__ void splitk_reduce_kernel_bf16ws_fallback( int bias_stride_batch) { splitk_reduce_kernel_fallback_body( - ws_handle, c_out, split_k, M, N, batch, padded_M, padded_N, bias, bias_stride_batch); + ws_ptr, c_out, split_k, M, N, batch, padded_M, padded_N, bias, bias_stride_batch); } // Exact-N row-block fast path: static split_k unroll, guarded M tail. @@ -225,7 +225,7 @@ template __global__ void splitk_reduce_kernel_exact_n_rowblock( - const opus_splitk_ws_handle* __restrict__ ws_handle, + const void* __restrict__ ws_ptr, D_OUT* __restrict__ c_out, int M, int N, int batch, int padded_M, int padded_N, @@ -257,7 +257,7 @@ __global__ void splitk_reduce_kernel_exact_n_rowblock( if (m >= M) { return; } - const D_WS* workspace = opus_splitk_ws_ptr(ws_handle); + const D_WS* workspace = opus_gfx942_uniform_ws_ptr(ws_ptr); opus::vector_t bias_fp32; if constexpr (HAS_BIAS) { diff --git a/csrc/opus_gemm/include/gfx942/opus_gemm_arch_gfx942.cuh b/csrc/opus_gemm/include/gfx942/opus_gemm_arch_gfx942.cuh index ba01e35b93..d56c9631d8 100644 --- a/csrc/opus_gemm/include/gfx942/opus_gemm_arch_gfx942.cuh +++ b/csrc/opus_gemm/include/gfx942/opus_gemm_arch_gfx942.cuh @@ -1,181 +1,193 @@ // SPDX-License-Identifier: MIT // Copyright (C) 2025-2026, Advanced Micro Devices, Inc. All rights reserved. // -// opus_gemm_arch_gfx942.cuh -- gfx942-specific dispatch implementations. +// Exact-kid launcher tables for gfx942. #pragma once #include "../opus_gemm_arch.cuh" #include "../opus_gemm_common.cuh" -#include "opus_gemm_heuristic_dispatch_gfx942.cuh" // OpusA16W16NoscaleKernel + opus_a16w16_heuristic_dispatch_gfx942<> -#include "opus_gemm_lookup.h" // GENERATE_OPUS_LOOKUP_TABLE_{BF16,FP32}_GFX942 -#include "opus_gemm_a16w16_tune_lookup.h" // GENERATE_A16W16_TUNE_LOOKUP_{BF16,FP32}_GFX942 -#include "opus_gemm_a8w8_tune_lookup.h" // GENERATE_A8W8_TUNE_LOOKUP_BF16 -#include "opus_gemm_manifest.h" // launcher symbols referenced by the lookup macros -#include "../opus_gemm_utils.cuh" // bf16_t / fp32_t - -#include // std::lower_bound +#include "../opus_gemm_utils.cuh" +#include "opus_gemm_a16w16_kid_dispatch.h" +#include "opus_gemm_a8w8_kid_dispatch.h" +#include "opus_gemm_manifest.h" + +#include +#include #include #include +#ifndef OPUS_A16W16_DISPATCH_KERNEL_TYPES_DEFINED +#define OPUS_A16W16_DISPATCH_KERNEL_TYPES_DEFINED +using OpusA16W16Kernel = void (*)( + aiter_tensor_t&, aiter_tensor_t&, aiter_tensor_t&, + std::optional, int); +using OpusA16W16WorkspaceKernel = void (*)( + aiter_tensor_t&, aiter_tensor_t&, aiter_tensor_t&, + aiter_tensor_t&, std::optional, int); +#endif + +#ifndef OPUS_A8W8_DISPATCH_KERNEL_TYPES_DEFINED +#define OPUS_A8W8_DISPATCH_KERNEL_TYPES_DEFINED +using OpusA8W8Kernel = void (*)( + aiter_tensor_t&, aiter_tensor_t&, aiter_tensor_t&); +using OpusA8W8BlockscaleKernel = void (*)( + aiter_tensor_t&, aiter_tensor_t&, aiter_tensor_t&, + aiter_tensor_t&, aiter_tensor_t&); +using OpusA8W8BlockscaleBpreshuffleKernel = void (*)( + aiter_tensor_t&, aiter_tensor_t&, aiter_tensor_t&, + aiter_tensor_t&, aiter_tensor_t&); +#endif + namespace opus_gfx942_detail { -struct OpusA16W16Shape +struct OpusA16W16KidEntry { - int M; - int N; - int K; + int kid; + OpusA16W16Kernel func; }; -struct OpusA16W16RuntimeEntry +struct OpusA16W16WorkspaceKidEntry { - OpusA16W16Shape key; - OpusA16W16NoscaleKernel func; + int kid; + OpusA16W16WorkspaceKernel func; }; -constexpr bool entry_less(const OpusA16W16RuntimeEntry& a, - const OpusA16W16RuntimeEntry& b) noexcept -{ - if (a.key.M != b.key.M) return a.key.M < b.key.M; - if (a.key.N != b.key.N) return a.key.N < b.key.N; - return a.key.K < b.key.K; -} - -constexpr bool entry_eq(const OpusA16W16RuntimeEntry& a, - const OpusA16W16RuntimeEntry& b) noexcept -{ - return a.key.M == b.key.M && a.key.N == b.key.N && a.key.K == b.key.K; -} - -struct OpusA16W16TuneEntry +template +struct OpusA8W8KidEntry { int kid; - OpusA16W16NoscaleKernel func; + Kernel func; }; -constexpr bool tune_entry_less(const OpusA16W16TuneEntry& a, - const OpusA16W16TuneEntry& b) noexcept +template +inline const Entry* find_kid(const std::array& entries, int kid) { - return a.kid < b.kid; + const auto it = std::lower_bound( + entries.begin(), entries.end(), kid, + [](const Entry& entry, int value) { return entry.kid < value; }); + return it != entries.end() && it->kid == kid ? &*it : nullptr; } -using OpusA16W16TuneKernel = OpusA16W16NoscaleKernel; +inline const OpusA16W16WorkspaceKidEntry* workspace_entry(int kid) +{ + static constexpr std::array< + OpusA16W16WorkspaceKidEntry, + GENERATE_A16W16_WORKSPACE_KID_DISPATCH_GFX942_SIZE> + kWorkspace = {{GENERATE_A16W16_WORKSPACE_KID_DISPATCH_GFX942}}; + return find_kid(kWorkspace, kid); +} -using OpusA8W8BlockscaleBPreshuffleKernel = void (*)( - aiter_tensor_t&, aiter_tensor_t&, aiter_tensor_t&, - std::optional, std::optional); +template +inline const OpusA16W16KidEntry* non_workspace_entry(int kid); -struct OpusA8W8TuneEntry +template <> +inline const OpusA16W16KidEntry* non_workspace_entry(int kid) { - int kid; - OpusA8W8BlockscaleBPreshuffleKernel func; -}; + static constexpr std::array< + OpusA16W16KidEntry, + GENERATE_A16W16_NONWORKSPACE_KID_DISPATCH_GFX942_BF16_SIZE> + kKids = {{GENERATE_A16W16_NONWORKSPACE_KID_DISPATCH_GFX942_BF16(bf16_t)}}; + return find_kid(kKids, kid); +} -constexpr bool a8w8_tune_entry_less(const OpusA8W8TuneEntry& a, - const OpusA8W8TuneEntry& b) noexcept +template <> +inline const OpusA16W16KidEntry* non_workspace_entry(int kid) { - return a.kid < b.kid; + static constexpr std::array< + OpusA16W16KidEntry, + GENERATE_A16W16_NONWORKSPACE_KID_DISPATCH_GFX942_FP32_SIZE> + kKids = {{GENERATE_A16W16_NONWORKSPACE_KID_DISPATCH_GFX942_FP32(fp32_t)}}; + return find_kid(kKids, kid); } -} // namespace opus_gfx942_detail -// -- a16w16 runtime dispatch (tuned lookup -> heuristic fallback) ------------- +} // namespace opus_gfx942_detail template -inline OpusA16W16NoscaleKernel -opus_dispatch_a16w16_gfx942(int M, int N, int K, int batch, bool has_bias = false); +inline OpusA16W16Kernel opus_a16w16_kid_dispatch_gfx942(int kid); template <> -inline OpusA16W16NoscaleKernel -opus_dispatch_a16w16_gfx942(int M, int N, int K, int batch, bool has_bias) +inline OpusA16W16Kernel opus_a16w16_kid_dispatch_gfx942(int kid) { - using namespace opus_gfx942_detail; - static constexpr OpusA16W16RuntimeEntry kLookup[] = { - GENERATE_OPUS_LOOKUP_TABLE_BF16_GFX942(bf16_t) - }; - constexpr size_t kSize = sizeof(kLookup) / sizeof(kLookup[0]); - OpusA16W16RuntimeEntry needle{{M, N, K}, nullptr}; - auto it = std::lower_bound(kLookup, kLookup + kSize, needle, entry_less); - if (it != kLookup + kSize && entry_eq(*it, needle)) - { - return it->func; - } - return opus_a16w16_heuristic_dispatch_gfx942(M, N, K, batch, has_bias); + const auto* entry = opus_gfx942_detail::non_workspace_entry(kid); + AITER_CHECK(entry != nullptr, + "unknown kid ", kid, + " for OPUS a16w16 on gfx942 with bf16 Y in the " + "non-workspace launch table"); + return entry->func; } template <> -inline OpusA16W16NoscaleKernel -opus_dispatch_a16w16_gfx942(int M, int N, int K, int batch, bool has_bias) +inline OpusA16W16Kernel opus_a16w16_kid_dispatch_gfx942(int kid) { - using namespace opus_gfx942_detail; - static constexpr OpusA16W16RuntimeEntry kLookup[] = { - GENERATE_OPUS_LOOKUP_TABLE_FP32_GFX942(fp32_t) - }; - constexpr size_t kSize = sizeof(kLookup) / sizeof(kLookup[0]); - OpusA16W16RuntimeEntry needle{{M, N, K}, nullptr}; - auto it = std::lower_bound(kLookup, kLookup + kSize, needle, entry_less); - if (it != kLookup + kSize && entry_eq(*it, needle)) - { - return it->func; - } - return opus_a16w16_heuristic_dispatch_gfx942(M, N, K, batch, has_bias); + const auto* entry = opus_gfx942_detail::non_workspace_entry(kid); + AITER_CHECK(entry != nullptr, + "unknown kid ", kid, + " for OPUS a16w16 on gfx942 with fp32 Y in the " + "non-workspace launch table"); + return entry->func; } -// -- a16w16 tune dispatch (id-based, two specializations) -------------------- - -template -inline opus_gfx942_detail::OpusA16W16TuneKernel -opus_a16w16_tune_dispatch_gfx942(int id); +inline bool opus_a16w16_has_non_workspace_kernel_gfx942(int id) +{ + return opus_gfx942_detail::non_workspace_entry(id) != nullptr + || opus_gfx942_detail::non_workspace_entry(id) != nullptr; +} -template <> -inline opus_gfx942_detail::OpusA16W16TuneKernel -opus_a16w16_tune_dispatch_gfx942(int id) +inline bool opus_a16w16_has_workspace_kernel_gfx942(int id) { - using namespace opus_gfx942_detail; - static constexpr OpusA16W16TuneEntry kTune[] = { - GENERATE_A16W16_TUNE_LOOKUP_BF16_GFX942(bf16_t) - }; - constexpr size_t kSize = sizeof(kTune) / sizeof(kTune[0]); - OpusA16W16TuneEntry needle{id, nullptr}; - auto it = std::lower_bound(kTune, kTune + kSize, needle, tune_entry_less); - AITER_CHECK(it != kTune + kSize && it->kid == id, - "Kernel id ", id, - " not found in a16w16 bf16 tune lookup table (gfx942)"); - return it->func; + return opus_gfx942_detail::workspace_entry(id) != nullptr; } -template <> -inline opus_gfx942_detail::OpusA16W16TuneKernel -opus_a16w16_tune_dispatch_gfx942(int id) +inline OpusA16W16WorkspaceKernel +opus_a16w16_workspace_dispatch_gfx942(int id) { - using namespace opus_gfx942_detail; - static constexpr OpusA16W16TuneEntry kTune[] = { - GENERATE_A16W16_TUNE_LOOKUP_FP32_GFX942(fp32_t) - }; - constexpr size_t kSize = sizeof(kTune) / sizeof(kTune[0]); - OpusA16W16TuneEntry needle{id, nullptr}; - auto it = std::lower_bound(kTune, kTune + kSize, needle, tune_entry_less); - AITER_CHECK(it != kTune + kSize && it->kid == id, - "Kernel id ", id, - " not found in a16w16 fp32 tune lookup table (gfx942)"); - return it->func; + const auto* entry = opus_gfx942_detail::workspace_entry(id); + AITER_CHECK(entry != nullptr, + "unknown kid ", id, + " for OPUS a16w16 on gfx942 in the workspace launch table"); + return entry->func; } -// -- a8w8 tune dispatch (id-based, bf16-output explicit tune API only) -------- +template +inline OpusA8W8BlockscaleBpreshuffleKernel +opus_a8w8_blockscale_bpreshuffle_kid_dispatch_gfx942(int id); -inline opus_gfx942_detail::OpusA8W8BlockscaleBPreshuffleKernel -opus_a8w8_tune_dispatch_gfx942(int id); +template <> +inline OpusA8W8BlockscaleBpreshuffleKernel +opus_a8w8_blockscale_bpreshuffle_kid_dispatch_gfx942(int id) +{ + using Entry = opus_gfx942_detail::OpusA8W8KidEntry< + OpusA8W8BlockscaleBpreshuffleKernel>; + static constexpr std::array< + Entry, + GENERATE_A8W8_BLOCKSCALE_BPRESHUFFLE_KID_DISPATCH_GFX942_BF16_SIZE> + kKids = {{GENERATE_A8W8_BLOCKSCALE_BPRESHUFFLE_KID_DISPATCH_GFX942_BF16}}; + AITER_CHECK(!kKids.empty(), + "no registered kernel for OPUS " + "a8w8_blockscale_bpreshuffle on gfx942 with bf16 Y"); + const auto* entry = opus_gfx942_detail::find_kid(kKids, id); + AITER_CHECK(entry != nullptr, + "unknown kid ", id, " for OPUS " + "a8w8_blockscale_bpreshuffle on gfx942 with bf16 Y"); + return entry->func; +} -inline opus_gfx942_detail::OpusA8W8BlockscaleBPreshuffleKernel -opus_a8w8_tune_dispatch_gfx942(int id) +template <> +inline OpusA8W8BlockscaleBpreshuffleKernel +opus_a8w8_blockscale_bpreshuffle_kid_dispatch_gfx942(int id) { - using namespace opus_gfx942_detail; - static constexpr OpusA8W8TuneEntry kTune[] = { - GENERATE_A8W8_TUNE_LOOKUP_BF16(bf16_t) - }; - constexpr size_t kSize = sizeof(kTune) / sizeof(kTune[0]); - OpusA8W8TuneEntry needle{id, nullptr}; - auto it = std::lower_bound(kTune, kTune + kSize, needle, a8w8_tune_entry_less); - AITER_CHECK(it != kTune + kSize && it->kid == id, - "Kernel id ", id, - " not found in a8w8 bf16 tune lookup table (gfx942)"); - return it->func; + using Entry = opus_gfx942_detail::OpusA8W8KidEntry< + OpusA8W8BlockscaleBpreshuffleKernel>; + static constexpr std::array< + Entry, + GENERATE_A8W8_BLOCKSCALE_BPRESHUFFLE_KID_DISPATCH_GFX942_FP32_SIZE> + kKids = {{GENERATE_A8W8_BLOCKSCALE_BPRESHUFFLE_KID_DISPATCH_GFX942_FP32}}; + AITER_CHECK(!kKids.empty(), + "no registered kernel for OPUS " + "a8w8_blockscale_bpreshuffle on gfx942 with fp32 Y"); + const auto* entry = opus_gfx942_detail::find_kid(kKids, id); + AITER_CHECK(entry != nullptr, + "unknown kid ", id, " for OPUS " + "a8w8_blockscale_bpreshuffle on gfx942 with fp32 Y"); + return entry->func; } diff --git a/csrc/opus_gemm/include/gfx942/opus_gemm_heuristic_dispatch_gfx942.cuh b/csrc/opus_gemm/include/gfx942/opus_gemm_heuristic_dispatch_gfx942.cuh deleted file mode 100644 index 7906b9ba0e..0000000000 --- a/csrc/opus_gemm/include/gfx942/opus_gemm_heuristic_dispatch_gfx942.cuh +++ /dev/null @@ -1,157 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (C) 2025-2026, Advanced Micro Devices, Inc. All rights reserved. -// -// Coarse gfx942 a16w16 fallback heuristic. Exact tuned shapes should hit the -// generated CSV lookup first; this path only needs a sane family-level guess. -#pragma once - -#include -#include - -#include "aiter_tensor.h" -#include "../opus_gemm_common.cuh" - -#ifndef OPUS_A16W16_NOSCALE_KERNEL_DEFINED -#define OPUS_A16W16_NOSCALE_KERNEL_DEFINED -using OpusA16W16NoscaleKernel = void (*)( - aiter_tensor_t &, aiter_tensor_t &, - aiter_tensor_t &, std::optional, int); -#endif - -#define OPUS_GFX942_A16W16_DECL(NAME) \ -template \ -void NAME(aiter_tensor_t &, aiter_tensor_t &, aiter_tensor_t &, \ - std::optional, int) - -OPUS_GFX942_A16W16_DECL(opus_gemm_gfx942_512x128x128x64_2x4_16x16x16_0x0x0); -OPUS_GFX942_A16W16_DECL(opus_gemm_gfx942_p1_256x64x64x64_2x2_16x16x16_0x0x0); -OPUS_GFX942_A16W16_DECL(opus_gemm_gfx942_splitk_legacy_512x128x128x64_2x4_16x16x16_0x0x0); -OPUS_GFX942_A16W16_DECL(opus_gemm_gfx942_splitk_p1_256x64x64x64_2x2_16x16x16_0x0x0); -OPUS_GFX942_A16W16_DECL(opus_gemm_gfx942_splitk_p1_bk128_256x64x64x128_2x2_16x16x16_0x0x0); -OPUS_GFX942_A16W16_DECL(opus_gemm_gfx942_splitk_p1_bk128_bf16ws_256x64x64x128_2x2_16x16x16_0x0x0); -OPUS_GFX942_A16W16_DECL(opus_gemm_gfx942_splitk_em3en4_lds1_pgr2_256x128x96x128_2x2_16x16x16_0x0x0); -OPUS_GFX942_A16W16_DECL(opus_gemm_gfx942_splitk_legacy_512x64x128x64_2x4_16x16x16_0x0x0); -OPUS_GFX942_A16W16_DECL(opus_gemm_gfx942_splitk_legacy_bf16ws_512x128x128x64_2x4_16x16x16_0x0x0); -OPUS_GFX942_A16W16_DECL(opus_gemm_gfx942_wkc_512x16x16x64_1x1_16x16x16_0x0x0); -OPUS_GFX942_A16W16_DECL(opus_gemm_gfx942_wkc_512x16x32x32_1x1_16x16x16_0x0x0); -OPUS_GFX942_A16W16_DECL(opus_gemm_gfx942_wkc_512x32x16x64_1x1_16x16x16_0x0x0); -OPUS_GFX942_A16W16_DECL(opus_gemm_gfx942_wkc_256x32x32x64_1x1_16x16x16_0x0x0); -OPUS_GFX942_A16W16_DECL(opus_gemm_gfx942_wkc_512x16x32x64_1x1_16x16x16_0x0x0); - -#undef OPUS_GFX942_A16W16_DECL - -namespace opus_gfx942_heuristic_detail -{ - -inline bool split_barrier_ok(int N, int K) -{ - const int loops = (K + 63) / 64; - return (N % 16 == 0) && (K % 64 == 0) && (loops >= 2) && (loops % 2 == 0); -} - -inline bool bf16ws_band(int M, int N, int K) -{ - return (K >= 4096) && (K % 64 == 0) && (M >= 104) && (M <= 608) && - (N == 256 || (N >= 512 && N <= 2048)); -} - -template -inline OpusA16W16NoscaleKernel dispatch_bf16(int M, int N, int K) -{ - const bool k64_ok = K % 64 == 0; - const bool k32_ok = K % 32 == 0; - const bool wkc_bk64_ok = K >= 4096 && K % 512 == 0; - const bool p1_ok = K % 128 == 0; - const bool sb_ok = split_barrier_ok(N, K); - - // DSV4 bf16 fallback misses that are not present in the generated CSV lookup. - // Keep this exact to K=4096; adjacent K=7168 bands have separate tuning data. - if (K == 4096) - { - if (p1_ok && ((M == 48 || M == 64) && N == 1024)) - return opus_gemm_gfx942_splitk_p1_bk128_bf16ws_256x64x64x128_2x2_16x16x16_0x0x0; - if (p1_ok && ((M == 128 && N == 512) || (M == 256 && N == 256))) - return opus_gemm_gfx942_splitk_p1_bk128_bf16ws_256x64x64x128_2x2_16x16x16_0x0x0; - if (p1_ok && M == 512 && N == 256) - return opus_gemm_gfx942_splitk_p1_bk128_256x64x64x128_2x2_16x16x16_0x0x0; - if ((M == 48 || M == 64) && N >= 1536 && N <= 2048) - return opus_gemm_gfx942_splitk_legacy_512x64x128x64_2x4_16x16x16_0x0x0; - if ((M == 128 && N == 1024) || (M == 256 && N == 512)) - return opus_gemm_gfx942_splitk_legacy_512x64x128x64_2x4_16x16x16_0x0x0; - if ((M == 128 && N >= 1536 && N <= 2048) || - (M == 256 && N == 1024) || (M == 512 && N == 512)) - return opus_gemm_gfx942_splitk_legacy_512x128x128x64_2x4_16x16x16_0x0x0; - } - - if (K >= 1024 && k32_ok && N >= 1536 && M <= 32) - { - if (M <= 4 && N >= 4096) - return opus_gemm_gfx942_wkc_512x16x16x64_1x1_16x16x16_0x0x0; - if (M <= 16) - return wkc_bk64_ok - ? opus_gemm_gfx942_wkc_512x16x32x64_1x1_16x16x16_0x0x0 - : opus_gemm_gfx942_wkc_512x16x32x32_1x1_16x16x16_0x0x0; - return (M == 32 && K == 4096 && wkc_bk64_ok) - ? opus_gemm_gfx942_wkc_512x16x32x64_1x1_16x16x16_0x0x0 - : opus_gemm_gfx942_wkc_256x32x32x64_1x1_16x16x16_0x0x0; - } - - if (K >= 512 && k64_ok && (N <= 64 || (M <= 128 && N <= 1024) || - (M <= 8 && N <= 1536))) - { - if (N <= 64 && M > 128) - return opus_gemm_gfx942_wkc_512x32x16x64_1x1_16x16x16_0x0x0; - if (N <= 256 || M <= 8 || (M <= 16 && N <= 800)) - return opus_gemm_gfx942_wkc_512x16x16x64_1x1_16x16x16_0x0x0; - return opus_gemm_gfx942_wkc_512x32x16x64_1x1_16x16x16_0x0x0; - } - - if (bf16ws_band(M, N, K)) - return opus_gemm_gfx942_splitk_legacy_bf16ws_512x128x128x64_2x4_16x16x16_0x0x0; - - if (N == 384 && K >= 4096) - { - if (M <= 128) - return opus_gemm_gfx942_wkc_512x32x16x64_1x1_16x16x16_0x0x0; - if (M <= 224) - return opus_gemm_gfx942_splitk_p1_256x64x64x64_2x2_16x16x16_0x0x0; - if (M >= 392 && M <= 512) - return opus_gemm_gfx942_splitk_em3en4_lds1_pgr2_256x128x96x128_2x2_16x16x16_0x0x0; - return opus_gemm_gfx942_splitk_legacy_512x128x128x64_2x4_16x16x16_0x0x0; - } - - if (k64_ok && N >= 4096 && K <= 3200) - { - if (K <= 640 && M <= 128) - return opus_gemm_gfx942_p1_256x64x64x64_2x2_16x16x16_0x0x0; - return opus_gemm_gfx942_512x128x128x64_2x4_16x16x16_0x0x0; - } - - if (sb_ok && M >= 128) - return opus_gemm_gfx942_512x128x128x64_2x4_16x16x16_0x0x0; - - if (N <= 256 && p1_ok) - return opus_gemm_gfx942_splitk_p1_256x64x64x64_2x2_16x16x16_0x0x0; - - return opus_gemm_gfx942_splitk_legacy_512x128x128x64_2x4_16x16x16_0x0x0; -} - -} // namespace opus_gfx942_heuristic_detail - -template -inline OpusA16W16NoscaleKernel opus_a16w16_heuristic_dispatch_gfx942( - int M, int N, int K, int /*batch*/, bool has_bias = false) -{ - using namespace opus_gfx942_heuristic_detail; - - if constexpr (std::is_same_v) - { - if (!has_bias) - return dispatch_bf16(M, N, K); - } - - if (N <= 256 && K % 128 == 0) - return opus_gemm_gfx942_splitk_p1_256x64x64x64_2x2_16x16x16_0x0x0; - - return opus_gemm_gfx942_splitk_legacy_512x128x128x64_2x4_16x16x16_0x0x0; -} diff --git a/csrc/opus_gemm/include/gfx950/opus_bmm_launchers_a8w8_mxscale_gfx950.cuh b/csrc/opus_gemm/include/gfx950/opus_bmm_launchers_a8w8_mxscale_gfx950.cuh index afc4a2d979..1dd5f460b2 100644 --- a/csrc/opus_gemm/include/gfx950/opus_bmm_launchers_a8w8_mxscale_gfx950.cuh +++ b/csrc/opus_gemm/include/gfx950/opus_bmm_launchers_a8w8_mxscale_gfx950.cuh @@ -21,16 +21,49 @@ #include static void opus_bmm_a8w8_common_checks(aiter_tensor_t &O, aiter_tensor_t &wo_a, - aiter_tensor_t &Y, const char *who) + aiter_tensor_t &Y, + aiter_tensor_t &x_scale, + aiter_tensor_t &w_scale, + const char *who) { aiter_detail::g_aiter_can_throw = true; - AITER_CHECK(O.dim() == 3 && wo_a.dim() == 3 && Y.dim() == 3, - who, ": O/wo_a/Y must be 3D " - "([M,batch,K] / [batch,N,K] / [M,batch,N])"); + AITER_CHECK(O.dim() == 3 && wo_a.dim() == 3 && Y.dim() == 3 && + x_scale.dim() == 3 && w_scale.dim() == 3, + who, ": O/wo_a/Y/x_scale/w_scale must all be 3D"); AITER_CHECK(O.dtype() == AITER_DTYPE_fp8 && wo_a.dtype() == AITER_DTYPE_fp8, who, ": O and wo_a must be fp8"); AITER_CHECK(Y.dtype() == AITER_DTYPE_fp32 || Y.dtype() == AITER_DTYPE_bf16, who, ": Y must be fp32 or bf16"); + AITER_CHECK((x_scale.dtype() == AITER_DTYPE_u8 || + x_scale.dtype() == AITER_DTYPE_fp8_e8m0) && + (w_scale.dtype() == AITER_DTYPE_u8 || + w_scale.dtype() == AITER_DTYPE_fp8_e8m0), + who, ": x_scale and w_scale must contain one-byte E8M0 values"); + + const int M = O.size(0); + const int batch = O.size(1); + const int K = O.size(2); + const int N = wo_a.size(1); + AITER_CHECK(M > 0 && batch > 0 && N > 0 && K > 0, + who, ": M, batch, N and K must be positive"); + AITER_CHECK(K % 128 == 0 && N % 128 == 0, + who, ": N and K must be multiples of 128; got N=", N, + ", K=", K); + AITER_CHECK(wo_a.size(0) == batch && wo_a.size(2) == K, + who, ": wo_a must have shape [batch,N,K]"); + AITER_CHECK(Y.size(0) == M && Y.size(1) == batch && Y.size(2) == N, + who, ": Y must have shape [M,batch,N]"); + AITER_CHECK(x_scale.size(0) == M && x_scale.size(1) == batch && + x_scale.size(2) == K / 128, + who, ": x_scale must have shape [M,batch,K/128]"); + AITER_CHECK(w_scale.size(0) == batch && w_scale.size(1) == N / 128 && + w_scale.size(2) == K / 128, + who, ": w_scale must have shape [batch,N/128,K/128]"); + + AITER_CHECK(O.device_id == wo_a.device_id && O.device_id == Y.device_id && + O.device_id == x_scale.device_id && + O.device_id == w_scale.device_id, + who, ": all tensors must be on one device"); // The kernels index A/B along K with unit stride (kargs carries only M/N/batch // strides, never a K stride), so K must be the innermost contiguous dim. The // batch axis position is free -- it is fully described by stride_*_batch -- so @@ -43,6 +76,11 @@ static void opus_bmm_a8w8_common_checks(aiter_tensor_t &O, aiter_tensor_t &wo_a, AITER_CHECK(wo_a.stride(2) == 1, who, ": wo_a must be K-contiguous (stride(2)==1); got stride ", (long)wo_a.stride(2)); + AITER_CHECK(Y.stride(2) == 1, who, + ": Y must be N-contiguous (stride(2)==1); got stride ", + (long)Y.stride(2)); + AITER_CHECK(x_scale.stride(2) == 1 && w_scale.stride(2) == 1, who, + ": scale tensors must be contiguous in their final dimension"); } #endif // __HIP_DEVICE_COMPILE__ diff --git a/csrc/opus_gemm/include/gfx950/opus_bmm_pipeline_a8w8_mxscale_gfx950.cuh b/csrc/opus_gemm/include/gfx950/opus_bmm_pipeline_a8w8_mxscale_gfx950.cuh index 03ac8c4141..19ff0218eb 100644 --- a/csrc/opus_gemm/include/gfx950/opus_bmm_pipeline_a8w8_mxscale_gfx950.cuh +++ b/csrc/opus_gemm/include/gfx950/opus_bmm_pipeline_a8w8_mxscale_gfx950.cuh @@ -701,7 +701,7 @@ __global__ __launch_bounds__(Traits::BLOCK_SIZE, 2) void gemm_a8w8_scale_splitk_ // overflows int32 for large-M batch-in-the-middle layouts. auto g_a = make_gmem(reinterpret_cast(kargs.ptr_a) + (size_t)batch_id*kargs.stride_a_batch + (size_t)row*kargs.stride_a + k_start); auto g_b = make_gmem(reinterpret_cast(kargs.ptr_b) + (size_t)batch_id*kargs.stride_b_batch + (size_t)col*kargs.stride_b + k_start); - auto g_c = make_gmem(reinterpret_cast(kargs.ws_handle->ptr) + (size_t)split_id * kargs.batch * kargs.stride_ws_batch + (size_t)batch_id * kargs.stride_ws_batch + (size_t)row * kargs.stride_ws + col); + auto g_c = make_gmem(reinterpret_cast(kargs.ptr_ws) + (size_t)split_id * kargs.batch * kargs.stride_ws_batch + (size_t)batch_id * kargs.stride_ws_batch + (size_t)row * kargs.stride_ws + col); auto g_sfa = make_gmem(reinterpret_cast(kargs.ptr_sfa) + (size_t)batch_id*kargs.stride_sfa_batch + (size_t)(row/T::GROUP_M)*kargs.stride_sfa + sf_start); auto g_sfb = make_gmem(reinterpret_cast(kargs.ptr_sfb) + (size_t)batch_id*kargs.stride_sfb_batch + (size_t)(col/T::GROUP_N)*kargs.stride_sfb + sf_start); diff --git a/csrc/opus_gemm/include/gfx950/opus_gemm_arch_gfx950.cuh b/csrc/opus_gemm/include/gfx950/opus_gemm_arch_gfx950.cuh index b99ab8dcb1..4dc083c3d4 100644 --- a/csrc/opus_gemm/include/gfx950/opus_gemm_arch_gfx950.cuh +++ b/csrc/opus_gemm/include/gfx950/opus_gemm_arch_gfx950.cuh @@ -1,248 +1,225 @@ // SPDX-License-Identifier: MIT // Copyright (C) 2025-2026, Advanced Micro Devices, Inc. All rights reserved. // -// opus_gemm_arch_gfx950.cuh — gfx950-specific dispatch implementations. -// -// Provides: -// * opus_dispatch_a16w16_gfx950 — tuned (M,N,K) lookup → heuristic -// * opus_a16w16_tune_dispatch_gfx950 — id-based tune dispatch -// -// This header is intended to be included exactly once, by opus_gemm.cu, where -// the arch routers in that TU select the per-arch entry. Other TUs (the -// launcher instances) must NOT include it -- they would each pull in the -// generated lookup macros (~70 KiB) for no gain. -// -// To add a new arch (e.g. gfx942): -// 1. Add OpusGfxArch::Gfx942 to opus_gemm_arch.cuh. -// 2. Create opus_gemm_arch_gfx942.cuh mirroring this file's shape; provide -// the per-arch dispatch functions with whatever lookup / heuristic that -// arch needs (it can reuse the same lookup macros if applicable, or -// its own). -// 3. #include "opus_gemm_arch_gfx942.cuh" in opus_gemm.cu and add a -// `case OpusGfxArch::Gfx942: ...` branch to each arch router there. +// Exact-kid launcher tables for gfx950. #pragma once #include "../opus_gemm_arch.cuh" #include "../opus_gemm_common.cuh" -#include "opus_gemm_heuristic_dispatch_gfx950.cuh" // OpusA16W16NoscaleKernel + opus_a16w16_heuristic_kid_gfx950() -#include "opus_gemm_lookup.h" // GENERATE_OPUS_LOOKUP_TABLE_{BF16,FP32}_GFX950 -#include "opus_gemm_a16w16_tune_lookup.h" // GENERATE_A16W16_TUNE_LOOKUP_{BF16,FP32}_GFX950 -#include "opus_gemm_manifest.h" // launcher symbols referenced by the lookup macros -#include "../opus_gemm_utils.cuh" // bf16_t / fp32_t (torch-free; py_itfs_common.h pulls full ) +#include "../opus_gemm_utils.cuh" +#include "opus_gemm_a16w16_kid_dispatch.h" +#include "opus_gemm_a8w8_kid_dispatch.h" +#include "opus_gemm_manifest.h" -#include // std::lower_bound +#include +#include #include +#include + +// Multi-arch builds include this ABI from several headers; define it once. +#ifndef OPUS_A16W16_DISPATCH_KERNEL_TYPES_DEFINED +#define OPUS_A16W16_DISPATCH_KERNEL_TYPES_DEFINED +using OpusA16W16Kernel = void (*)( + aiter_tensor_t&, aiter_tensor_t&, aiter_tensor_t&, + std::optional, int); +using OpusA16W16WorkspaceKernel = void (*)( + aiter_tensor_t&, aiter_tensor_t&, aiter_tensor_t&, + aiter_tensor_t&, std::optional, int); +#endif + +#ifndef OPUS_A8W8_DISPATCH_KERNEL_TYPES_DEFINED +#define OPUS_A8W8_DISPATCH_KERNEL_TYPES_DEFINED +using OpusA8W8Kernel = void (*)( + aiter_tensor_t&, aiter_tensor_t&, aiter_tensor_t&); +using OpusA8W8BlockscaleKernel = void (*)( + aiter_tensor_t&, aiter_tensor_t&, aiter_tensor_t&, + aiter_tensor_t&, aiter_tensor_t&); +using OpusA8W8BlockscaleBpreshuffleKernel = void (*)( + aiter_tensor_t&, aiter_tensor_t&, aiter_tensor_t&, + aiter_tensor_t&, aiter_tensor_t&); +#endif namespace opus_gfx950_detail { -// Sorted flat-array entries for the runtime (M, N, K) -> kernel lookup -// (was: std::unordered_map, -// OpusA16W16NoscaleKernel, IntTupleHash>). The unordered_map version -// added ~1s of frontend / template instantiation per dispatcher TU -// because of the heavyweight std::function-valued hashtable templates; -// a flat array of POD entries plus std::lower_bound costs essentially -// nothing at parse time and matches the lookup at runtime in O(log N) -// over 339 entries. -// Nested {shape, func} aggregate matches the `{ {M, N, K}, &kernel }` -// initializer the codegen emits. Splitting shape into its own struct -// keeps the comparators small and gives gen_instances.py a stable -// brace pattern to target. -struct OpusA16W16Shape -{ - int M; - int N; - int K; +struct OpusA16W16KidEntry +{ + int kid; + OpusA16W16Kernel func; }; -struct OpusA16W16RuntimeEntry +struct OpusA16W16WorkspaceKidEntry { - OpusA16W16Shape key; - OpusA16W16NoscaleKernel func; + int kid; + OpusA16W16WorkspaceKernel func; }; -// Lex order on (M, N, K). Used both during sorting (gen_instances.py -// emits entries in lex order) and by std::lower_bound at lookup time. -constexpr bool entry_less(const OpusA16W16RuntimeEntry& a, - const OpusA16W16RuntimeEntry& b) noexcept +template +struct OpusA8W8KidEntry { - if (a.key.M != b.key.M) return a.key.M < b.key.M; - if (a.key.N != b.key.N) return a.key.N < b.key.N; - return a.key.K < b.key.K; -} + int kid; + Kernel func; +}; -constexpr bool entry_eq(const OpusA16W16RuntimeEntry& a, - const OpusA16W16RuntimeEntry& b) noexcept +template +inline const Entry* find_kid(const std::array& entries, int kid) { - return a.key.M == b.key.M && a.key.N == b.key.N && a.key.K == b.key.K; + const auto it = std::lower_bound( + entries.begin(), entries.end(), kid, + [](const Entry& entry, int value) { return entry.kid < value; }); + return it != entries.end() && it->kid == kid ? &*it : nullptr; } -// id -> kernel, same flat-array layout. Sorted by id (the -// codegen always emits in ascending id order). -struct OpusA16W16TuneEntry +inline const OpusA16W16WorkspaceKidEntry* workspace_entry(int kid) { - int kid; - OpusA16W16NoscaleKernel func; -}; + static constexpr std::array< + OpusA16W16WorkspaceKidEntry, + GENERATE_A16W16_WORKSPACE_KID_DISPATCH_GFX950_SIZE> + kWorkspace = {{GENERATE_A16W16_WORKSPACE_KID_DISPATCH_GFX950}}; + return find_kid(kWorkspace, kid); +} -constexpr bool tune_entry_less(const OpusA16W16TuneEntry& a, - const OpusA16W16TuneEntry& b) noexcept +template +inline const OpusA16W16KidEntry* non_workspace_entry(int kid); + +template <> +inline const OpusA16W16KidEntry* non_workspace_entry(int kid) { - return a.kid < b.kid; + static constexpr std::array< + OpusA16W16KidEntry, + GENERATE_A16W16_NONWORKSPACE_KID_DISPATCH_GFX950_BF16_SIZE> + kKids = {{GENERATE_A16W16_NONWORKSPACE_KID_DISPATCH_GFX950_BF16(bf16_t)}}; + return find_kid(kKids, kid); } -using OpusA16W16TuneKernel = OpusA16W16NoscaleKernel; -} // namespace opus_gfx950_detail - -// Splitk kid range. Kept in this header (rather than relying on the -// opus_gemm.cu copy in OPUS_SPLITK_KID_MIN/MAX) so the heuristic-fallback -// path below can route splitk kids to tune_dispatch without a -// cross-TU dependency. The numbers must match opus_gemm.cu. -namespace opus_gfx950_detail +template <> +inline const OpusA16W16KidEntry* non_workspace_entry(int kid) { -constexpr int kSplitkKidMin = 200; -constexpr int kSplitkKidMax = 300; -constexpr int kNooobKidOffset = 1000; + static constexpr std::array< + OpusA16W16KidEntry, + GENERATE_A16W16_NONWORKSPACE_KID_DISPATCH_GFX950_FP32_SIZE> + kKids = {{GENERATE_A16W16_NONWORKSPACE_KID_DISPATCH_GFX950_FP32(fp32_t)}}; + return find_kid(kKids, kid); +} +} // namespace opus_gfx950_detail -constexpr bool kid_is_splitk(int kid) noexcept +inline OpusA8W8Kernel opus_a8w8_kid_dispatch_gfx950(int id) { - return (kid >= kSplitkKidMin && kid < kSplitkKidMax) || - (kid >= kSplitkKidMin + kNooobKidOffset && - kid < kSplitkKidMax + kNooobKidOffset); + using Entry = opus_gfx950_detail::OpusA8W8KidEntry; + static constexpr std::array< + Entry, GENERATE_A8W8_NOSCALE_KID_DISPATCH_GFX950_SIZE> + kKids = {{GENERATE_A8W8_NOSCALE_KID_DISPATCH_GFX950}}; + AITER_CHECK(!kKids.empty(), + "no registered kernel for OPUS a8w8 on gfx950"); + const auto* entry = opus_gfx950_detail::find_kid(kKids, id); + AITER_CHECK(entry != nullptr, + "unknown kid ", id, " for OPUS a8w8 on gfx950"); + return entry->func; } -} // namespace opus_gfx950_detail -// ── a16w16 tune dispatch (id-based, two specializations) ──────────────────── -// -// The bf16 table omits splitk kids (their instantiation doesn't -// exist; splitk main kernel hardcodes D_C=float). The fp32 table includes -// all a16w16-family kids; splitk kids appear there with -// hardcoded as well, since the reduce kernel handles fp32 Y output by -// skipping the cast. +inline OpusA8W8BlockscaleKernel +opus_a8w8_blockscale_kid_dispatch_gfx950(int id) +{ + using Entry = + opus_gfx950_detail::OpusA8W8KidEntry; + static constexpr std::array< + Entry, GENERATE_A8W8_BLOCKSCALE_KID_DISPATCH_GFX950_SIZE> + kKids = {{GENERATE_A8W8_BLOCKSCALE_KID_DISPATCH_GFX950}}; + AITER_CHECK(!kKids.empty(), + "no registered kernel for OPUS a8w8_blockscale on gfx950"); + const auto* entry = opus_gfx950_detail::find_kid(kKids, id); + AITER_CHECK(entry != nullptr, + "unknown kid ", id, + " for OPUS a8w8_blockscale on gfx950"); + return entry->func; +} template -inline opus_gfx950_detail::OpusA16W16TuneKernel -opus_a16w16_tune_dispatch_gfx950(int id); +inline OpusA8W8BlockscaleBpreshuffleKernel +opus_a8w8_blockscale_bpreshuffle_kid_dispatch_gfx950(int id); template <> -inline opus_gfx950_detail::OpusA16W16TuneKernel -opus_a16w16_tune_dispatch_gfx950(int id) -{ - using namespace opus_gfx950_detail; - static constexpr OpusA16W16TuneEntry kTune[] = { - GENERATE_A16W16_TUNE_LOOKUP_BF16_GFX950(bf16_t) - }; - constexpr size_t kSize = sizeof(kTune) / sizeof(kTune[0]); - OpusA16W16TuneEntry needle{id, nullptr}; - auto it = std::lower_bound(kTune, kTune + kSize, needle, tune_entry_less); - AITER_CHECK(it != kTune + kSize && it->kid == id, - "Kernel id ", id, - " not found in a16w16 bf16 tune lookup table"); - return it->func; +inline OpusA8W8BlockscaleBpreshuffleKernel +opus_a8w8_blockscale_bpreshuffle_kid_dispatch_gfx950(int id) +{ + using Entry = opus_gfx950_detail::OpusA8W8KidEntry< + OpusA8W8BlockscaleBpreshuffleKernel>; + static constexpr std::array< + Entry, + GENERATE_A8W8_BLOCKSCALE_BPRESHUFFLE_KID_DISPATCH_GFX950_BF16_SIZE> + kKids = {{GENERATE_A8W8_BLOCKSCALE_BPRESHUFFLE_KID_DISPATCH_GFX950_BF16}}; + AITER_CHECK(!kKids.empty(), + "no registered kernel for OPUS " + "a8w8_blockscale_bpreshuffle on gfx950 with bf16 Y"); + const auto* entry = opus_gfx950_detail::find_kid(kKids, id); + AITER_CHECK(entry != nullptr, + "unknown kid ", id, " for OPUS " + "a8w8_blockscale_bpreshuffle on gfx950 with bf16 Y"); + return entry->func; } template <> -inline opus_gfx950_detail::OpusA16W16TuneKernel -opus_a16w16_tune_dispatch_gfx950(int id) -{ - using namespace opus_gfx950_detail; - static constexpr OpusA16W16TuneEntry kTune[] = { - GENERATE_A16W16_TUNE_LOOKUP_FP32_GFX950(fp32_t) - }; - constexpr size_t kSize = sizeof(kTune) / sizeof(kTune[0]); - OpusA16W16TuneEntry needle{id, nullptr}; - auto it = std::lower_bound(kTune, kTune + kSize, needle, tune_entry_less); - AITER_CHECK(it != kTune + kSize && it->kid == id, - "Kernel id ", id, - " not found in a16w16 fp32 tune lookup table"); - return it->func; +inline OpusA8W8BlockscaleBpreshuffleKernel +opus_a8w8_blockscale_bpreshuffle_kid_dispatch_gfx950(int id) +{ + using Entry = opus_gfx950_detail::OpusA8W8KidEntry< + OpusA8W8BlockscaleBpreshuffleKernel>; + static constexpr std::array< + Entry, + GENERATE_A8W8_BLOCKSCALE_BPRESHUFFLE_KID_DISPATCH_GFX950_FP32_SIZE> + kKids = {{GENERATE_A8W8_BLOCKSCALE_BPRESHUFFLE_KID_DISPATCH_GFX950_FP32}}; + AITER_CHECK(!kKids.empty(), + "no registered kernel for OPUS " + "a8w8_blockscale_bpreshuffle on gfx950 with fp32 Y"); + const auto* entry = opus_gfx950_detail::find_kid(kKids, id); + AITER_CHECK(entry != nullptr, + "unknown kid ", id, " for OPUS " + "a8w8_blockscale_bpreshuffle on gfx950 with fp32 Y"); + return entry->func; } -// ── a16w16 runtime dispatch (tuned lookup → heuristic fallback) ───────────── -// -// On miss the heuristic returns an integer kid; we re-dispatch through -// opus_a16w16_tune_dispatch_gfx950<>(). Splitk kids only have a -// instantiation (their traits static_assert D_C=float; the reduce kernel -// templated on Y dtype handles bf16/fp32 output at launch time), so we -// force the branch for those regardless of the dispatcher's -// CDataType template parameter. - +// Workspace kids are absent from the direct-output table. template -inline OpusA16W16NoscaleKernel -opus_dispatch_a16w16_gfx950(int M, int N, int K, int batch, bool has_bias = false); +inline OpusA16W16Kernel opus_a16w16_kid_dispatch_gfx950(int kid); template <> -inline OpusA16W16NoscaleKernel -opus_dispatch_a16w16_gfx950(int M, int N, int K, int batch, bool has_bias) -{ - using namespace opus_gfx950_detail; - static constexpr OpusA16W16RuntimeEntry kLookup[] = { - GENERATE_OPUS_LOOKUP_TABLE_BF16_GFX950(bf16_t) - }; - constexpr size_t kSize = sizeof(kLookup) / sizeof(kLookup[0]); - OpusA16W16RuntimeEntry needle{{M, N, K}, nullptr}; - auto it = std::lower_bound(kLookup, kLookup + kSize, needle, entry_less); - if (it != kLookup + kSize && entry_eq(*it, needle)) - { - return it->func; - } - (void)batch; // heuristic does not currently use batch. - // 4 GiB buffer-resource guard. The heuristic returns one of - // HEURISTIC_DEFAULT_KIDS, all of which are legacy (non-4g_safe) and - // build a single AMDGPU buffer-resource over the whole A/B/C tensors; - // 32-bit num_records wraps when any A/B/C bytes exceed UINT32_MAX, - // producing silent OOB. Refuse fallback for >4 GiB shapes -- the - // caller must register a tuned CSV entry mapping the shape to a - // 4g_safe kid (5000-series / 6000-series). - constexpr uint64_t U32_MAX_BYTES = (1ULL << 32) - 1; - const uint64_t a_bytes = (uint64_t)M * (uint64_t)K * sizeof(bf16_t); - const uint64_t b_bytes = (uint64_t)N * (uint64_t)K * sizeof(bf16_t); - const uint64_t c_bytes = (uint64_t)M * (uint64_t)N * sizeof(bf16_t); - AITER_CHECK(a_bytes <= U32_MAX_BYTES && b_bytes <= U32_MAX_BYTES - && c_bytes <= U32_MAX_BYTES, - "opus a16w16 heuristic fallback refuses >4 GiB shape (M=", - M, " N=", N, " K=", K, - "): legacy kids wrap buffer-resource num_records. " - "Add a tuned CSV entry mapping this shape to a 4g_safe kid " - "(5000-series for split-barrier / 6000-series for mono_tile)."); - // has_bias forces the heuristic to skip persistent kids (which - // do not yet implement HAS_BIAS=true) and return a splitk kid instead. - const int kid = opus_a16w16_heuristic_kid_gfx950(M, N, K, has_bias); - if (kid_is_splitk(kid)) - return opus_a16w16_tune_dispatch_gfx950(kid); - return opus_a16w16_tune_dispatch_gfx950(kid); +inline OpusA16W16Kernel opus_a16w16_kid_dispatch_gfx950(int kid) +{ + const auto* entry = opus_gfx950_detail::non_workspace_entry(kid); + AITER_CHECK(entry != nullptr, + "unknown kid ", kid, + " for OPUS a16w16 on gfx950 with bf16 Y in the " + "non-workspace launch table"); + return entry->func; } template <> -inline OpusA16W16NoscaleKernel -opus_dispatch_a16w16_gfx950(int M, int N, int K, int batch, bool has_bias) -{ - using namespace opus_gfx950_detail; - static constexpr OpusA16W16RuntimeEntry kLookup[] = { - GENERATE_OPUS_LOOKUP_TABLE_FP32_GFX950(fp32_t) - }; - constexpr size_t kSize = sizeof(kLookup) / sizeof(kLookup[0]); - OpusA16W16RuntimeEntry needle{{M, N, K}, nullptr}; - auto it = std::lower_bound(kLookup, kLookup + kSize, needle, entry_less); - if (it != kLookup + kSize && entry_eq(*it, needle)) - { - return it->func; - } - (void)batch; - // 4 GiB buffer-resource guard (see overload for rationale). - // C is fp32 here so the bound is 4 bytes/element. - constexpr uint64_t U32_MAX_BYTES = (1ULL << 32) - 1; - const uint64_t a_bytes = (uint64_t)M * (uint64_t)K * sizeof(bf16_t); - const uint64_t b_bytes = (uint64_t)N * (uint64_t)K * sizeof(bf16_t); - const uint64_t c_bytes = (uint64_t)M * (uint64_t)N * sizeof(fp32_t); - AITER_CHECK(a_bytes <= U32_MAX_BYTES && b_bytes <= U32_MAX_BYTES - && c_bytes <= U32_MAX_BYTES, - "opus a16w16 heuristic fallback refuses >4 GiB shape (M=", - M, " N=", N, " K=", K, - "): legacy kids wrap buffer-resource num_records. " - "Add a tuned CSV entry mapping this shape to a 4g_safe kid " - "(5000-series for split-barrier / 6000-series for mono_tile)."); - const int kid = opus_a16w16_heuristic_kid_gfx950(M, N, K, has_bias); - // splitk kids only have in the tune table; non-splitk kids - // have an entry in the fp32 lookup, so the branch is uniform. - return opus_a16w16_tune_dispatch_gfx950(kid); +inline OpusA16W16Kernel opus_a16w16_kid_dispatch_gfx950(int kid) +{ + const auto* entry = opus_gfx950_detail::non_workspace_entry(kid); + AITER_CHECK(entry != nullptr, + "unknown kid ", kid, + " for OPUS a16w16 on gfx950 with fp32 Y in the " + "non-workspace launch table"); + return entry->func; +} + +inline bool opus_a16w16_has_non_workspace_kernel_gfx950(int id) +{ + return opus_gfx950_detail::non_workspace_entry(id) != nullptr + || opus_gfx950_detail::non_workspace_entry(id) != nullptr; +} + +inline bool opus_a16w16_has_workspace_kernel_gfx950(int id) +{ + return opus_gfx950_detail::workspace_entry(id) != nullptr; +} + +inline OpusA16W16WorkspaceKernel +opus_a16w16_workspace_dispatch_gfx950(int id) +{ + const auto* entry = opus_gfx950_detail::workspace_entry(id); + AITER_CHECK(entry != nullptr, + "unknown kid ", id, + " for OPUS a16w16 on gfx950 in the workspace launch table"); + return entry->func; } diff --git a/csrc/opus_gemm/include/gfx950/opus_gemm_heuristic_dispatch_gfx950.cuh b/csrc/opus_gemm/include/gfx950/opus_gemm_heuristic_dispatch_gfx950.cuh deleted file mode 100644 index 4c8f03542e..0000000000 --- a/csrc/opus_gemm/include/gfx950/opus_gemm_heuristic_dispatch_gfx950.cuh +++ /dev/null @@ -1,129 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (C) 2025-2026, Advanced Micro Devices, Inc. All rights reserved. -// -// a16w16 family heuristic dispatcher (gfx950). -// -// The heuristic is the "no-tuned-CSV fallback" arm of opus_dispatch_a16w16: -// when a runtime (M,N,K) shape has no row in opus_gemm_lookup.h, we still -// need to pick *some* valid a16w16 kernel for it. This file defines that -// pick as a pure ``(M,N,K) -> kid`` mapping; the caller (see -// opus_gemm_arch_gfx950.cuh) then resolves the kid through -// opus_a16w16_tune_dispatch_gfx950<>() against the (gen_instances.py- -// emitted) tune lookup table. -// -// Why kid integers instead of launcher symbol names? -// --------------------------------------------------- -// The previous version of this file returned bare ``&opus_gemm_..._wgpcu1 -// `` symbols. That coupled the heuristic to specific .so symbol -// names, which is a real problem in the subset-compile world: if a build -// excludes the splitk-128 launcher (because the CSV doesn't ask for it -// and the heuristic doesn't either), but the .cuh still references the -// symbol, the link fails. By returning an integer kid here and routing -// through the tune lookup, the only invariant is "every kid this function -// can return must also be in the compiled subset S". That invariant is -// enforced at *codegen* time by csrc/opus_gemm/gen_instances.py -// assert HEURISTIC_DEFAULT_KIDS.issubset(S) -// using the single source of truth in opus_gemm_common.py. -// -// Keep the integer kid returns in opus_a16w16_heuristic_kid_gfx950() below -// in sync with the HEURISTIC_DEFAULT_KIDS frozenset in -// csrc/opus_gemm/opus_gemm_common.py. The two are coupled by intent. -// -// gfx950-specific because the choices below were profiled on MI350's -// 256-CU / 160 KB LDS budget. Future archs will have their own -// opus_gemm_heuristic_dispatch_.cuh next to this one. -#pragma once - -#include - -#include "aiter_tensor.h" // aiter_tensor_t (torch-free) -#include "../opus_gemm_common.cuh" -#include "opus_gemm_manifest.h" - -// a16w16-family launcher signature (split-barrier, flatmm, flatmm_splitk): -// 3 tensors + std::optional + int splitK so all three populate the -// same GENERATE_A16W16_TUNE_LOOKUP_*_GFX950 table. gfx1250's launchers take a -// workspace tensor on top of this, which is why the lookup macros are emitted -// per arch (gen_instances.py :: LOOKUP_MACRO_ARCHES). Non-splitk launchers ignore -// splitK; the splitk launcher treats it as literal KBatch. bias is -// consumed by the split-barrier and splitk launchers; the flatmm launcher -// rejects any non-empty bias up front (HAS_BIAS=false on its warp-spec -// epilogue). -// -// Returns void (in-place on Y); the launchers used to return Y but -// nothing read the return value at any call site, and dropping the -// torch::Tensor return type lets the whole dispatch graph go -// torch-free. -// -// Plain function pointer (was: `std::function`). Every -// callable we ever store in this slot is one of the explicitly -// instantiated `xxx` / `xxx` launcher templates -- -// no captures, no type erasure needed. Switching to a function -// pointer drops a heavyweight `std::function` template instantiation -// from the dispatcher TU's host pass and also avoids the per-call -// virtual-dispatch overhead that std::function pays for the type -// erasure we don't actually use. -using OpusA16W16NoscaleKernel = void (*)( - aiter_tensor_t &, aiter_tensor_t &, - aiter_tensor_t &, std::optional, int); - - -// Pure (M, N, K, has_bias) -> integer kid mapping. No reference to launcher -// symbols here -- the caller resolves the returned kid through -// opus_a16w16_tune_dispatch_gfx950(kid). -// -// IMPORTANT: every kid this function can return MUST also be in -// HEURISTIC_DEFAULT_KIDS in csrc/opus_gemm/opus_gemm_common.py, so -// the subset-compile codegen always includes them in S. -// -// `has_bias` matters because the persistent pipeline does not yet -// implement HAS_BIAS=true; when the user passes a non-empty bias the -// heuristic must stay on the bias-aware splitk family even if the M-bucket -// would otherwise return a persistent kid. Splitk kids 200/206/208 (+ -// nooob mirrors) are all bias-aware (see opus_kid_supports_bias in -// opus_gemm.cu and BIAS_AWARE_KIDS in opus_gemm_common.py). -inline int opus_a16w16_heuristic_kid_gfx950(int M, int N, int K, bool has_bias = false) -{ - const bool split_barrier_ok = - (N % 16 == 0) && (K % 64 == 0) && ((K / 64) % 2 == 0); - - if (M <= 4) - { - // Extremely skinny M: cc recommends (64,64,128) WG=1 for deep K. - // kid 208 (oob) / 1208 (nooob): a16w16_flatmm_splitk_64x64x128_wgpcu1. - if ((M % 64 == 0) && (N % 64 == 0) && (K % 128 == 0)) - return 1208; - return 208; - } - if (M <= 64) - { - // Mid-skinny: cc-recommended medium-M kernel (64,32,128) WG=2. - // kid 206 (oob) / 1206 (nooob): a16w16_flatmm_splitk_64x32x128_wgpcu2. - if ((M % 64 == 0) && (N % 32 == 0) && (K % 128 == 0)) - return 1206; - return 206; - } - if (M <= 128) - { - // Sweet spot: (64,64,64) WG=2. - // kid 200 (oob) / 1200 (nooob): a16w16_flatmm_splitk_64x64x64_wgpcu2. - if ((M % 64 == 0) && (N % 64 == 0) && (K % 64 == 0)) - return 1200; - return 200; - } - // M > 128 - if (split_barrier_ok && !has_bias) - { - // Persistent (256, 256, 64) tile; CDataType-templated by caller. - // kid 300 (oob) / 1300 (nooob). Persistent does not yet support bias -- - // when has_bias is true we fall through to the splitk path below. - if ((M % 256 == 0) && (N % 256 == 0) && (K % 64 == 0)) - return 1300; - return 300; - } - // M > 128 but split-barrier prerequisites failed (or bias requested) -- - // fall back to the splitk sweet-spot tile. Splitk supports bias. - if ((M % 64 == 0) && (N % 64 == 0) && (K % 64 == 0)) - return 1200; - return 200; -} diff --git a/csrc/opus_gemm/include/gfx950/opus_gemm_pipeline_a16w16_flatmm_splitk_gfx950.cuh b/csrc/opus_gemm/include/gfx950/opus_gemm_pipeline_a16w16_flatmm_splitk_gfx950.cuh index e0718b391d..32076ab764 100644 --- a/csrc/opus_gemm/include/gfx950/opus_gemm_pipeline_a16w16_flatmm_splitk_gfx950.cuh +++ b/csrc/opus_gemm/include/gfx950/opus_gemm_pipeline_a16w16_flatmm_splitk_gfx950.cuh @@ -213,7 +213,7 @@ void gemm_a16w16_flatmm_splitk_kernel(opus_gemm_flatmm_splitk_kargs_gfx950 kargs using T = opus::remove_cvref_t; using D_A = typename T::D_A; using D_B = typename T::D_B; - using D_C = typename T::D_C; + using D_WS = typename T::D_WS; using D_ACC = typename T::D_ACC; // grid.x = split_k * num_tiles_m * num_tiles_n (S splits fused to inner axis, @@ -262,8 +262,7 @@ void gemm_a16w16_flatmm_splitk_kernel(opus_gemm_flatmm_splitk_kargs_gfx950 kargs auto g_b = make_gmem(reinterpret_cast(kargs.ptr_b) + batch_id * kargs.stride_b_batch + col * kargs.stride_b + k_start, ((kargs.n - col) * kargs.stride_b - k_start) * sizeof(D_B)); - // Deref the handle slot at entry; survives a post-capture grow. - D_C* ws_ptr = reinterpret_cast(kargs.ws_handle->ptr); + D_WS* ws_ptr = reinterpret_cast(kargs.ptr_ws); auto g_c = make_gmem(ws_ptr + (size_t)split_id * kargs.batch * kargs.stride_ws_batch + (size_t)batch_id * kargs.stride_ws_batch diff --git a/csrc/opus_gemm/include/gfx950/opus_gemm_pipeline_a16w16_mono_tile_4g_safe_gfx950.cuh b/csrc/opus_gemm/include/gfx950/opus_gemm_pipeline_a16w16_mono_tile_4g_safe_gfx950.cuh index e043d824b8..28814ac331 100644 --- a/csrc/opus_gemm/include/gfx950/opus_gemm_pipeline_a16w16_mono_tile_4g_safe_gfx950.cuh +++ b/csrc/opus_gemm/include/gfx950/opus_gemm_pipeline_a16w16_mono_tile_4g_safe_gfx950.cuh @@ -9,8 +9,8 @@ // formulas; M/N tails are absorbed by BR num_records. // // Direct port of the upstream yk_gcn mono-tile BF16 kernel template -// (bf16_gemm/gemm_a16w16_mono_tile_kernel_template.hpp) with two mechanical -// adjustments to fit the aiter codegen contract: +// (bf16_gemm/gemm_a16w16_mono_tile_kernel_template.hpp) with AITER integration +// adjustments: // // (1) `opus_gemm_kargs` (yk_gcn) -> `opus_gemm_mono_tile_kargs_gfx950` // (aiter; defined in opus_gemm_traits_a16w16_gfx950.cuh alongside @@ -18,12 +18,14 @@ // (2) The layout-helper namespace is renamed `gemm_mono_tile` -> // `opus_mono_tile_gfx950` to avoid any ODR clash with a separate // upstream build that ships the original symbol name. +// (3) The output epilogue supports both BF16 and FP32. Its lane shuffle and +// physical stores are derived from D_C; the upstream template only +// instantiates BF16 output. // -// The kernel body is otherwise byte-for-byte identical to the upstream -// reference. Geometry is locked (T_M=2, T_N=4, T_K=1, W_M=W_N=16, W_K=32, -// VEC=8, BLOCK_SIZE=512); tile-divisibility / smem-rep constraints are -// enforced in the traits header static_asserts and re-validated host-side -// by _validate_a16w16_mono_tile in gen_instances.py. +// Geometry remains locked to the upstream reference (T_M=2, T_N=4, T_K=1, +// W_M=W_N=16, W_K=32, VEC=8, BLOCK_SIZE=512); tile-divisibility / smem-rep +// constraints are enforced in the traits header static_asserts and +// re-validated host-side by _validate_a16w16_mono_tile in gen_instances.py. #pragma once #include @@ -425,22 +427,55 @@ void gemm_a16w16_mono_tile_4g_safe_kernel_gfx950(opus_gemm_mono_tile_kargs_gfx95 auto u_gc = make_layout_gc(lane_id, 0, wave_id_n, kargs.stride_c); auto v_c_f16 = cast(v_c); - // For every 8 D_C elements (= 4 u32), swap lane L's upper-half (last 4 - // elems) with lane (L^16)'s lower-half (first 4 elems) using - // v_permlane16_swap_b32. - static_assert(sizeof(D_C) * 8 % sizeof(u32_t) == 0); - constexpr int u32_per_chunk = sizeof(D_C) * 8 / sizeof(u32_t); + // For every 8 D_C elements, swap lane L's upper half (last 4 elements) + // with lane (L^16)'s lower half (first 4 elements). The original BF16 + // template hard-coded the two u32 pairs of an 8xbf16 chunk. FP32 has + // four u32 values per half, so derive the pair count from D_C instead of + // silently leaving elements 4..7 unswapped. + static_assert(sizeof(D_C) * 4 % sizeof(u32_t) == 0); + constexpr int u32_per_half = sizeof(D_C) * 4 / sizeof(u32_t); + constexpr int u32_per_chunk = 2 * u32_per_half; constexpr int num_chunks = sizeof(v_c_f16) / (sizeof(u32_t) * u32_per_chunk); auto* p_u32 = reinterpret_cast(&v_c_f16); static_for([&](auto c) { auto* p = p_u32 + c.value * u32_per_chunk; - auto r0 = __builtin_amdgcn_permlane16_swap(p[0], p[2], false, true); - auto r1 = __builtin_amdgcn_permlane16_swap(p[1], p[3], false, true); - p[0] = r0[0]; p[2] = r0[1]; - p[1] = r1[0]; p[3] = r1[1]; + static_for([&](auto i) { + auto r = __builtin_amdgcn_permlane16_swap( + p[i.value], p[i.value + u32_per_half], false, true); + p[i.value] = r[0]; + p[i.value + u32_per_half] = r[1]; + }); }); - store(g_c, v_c_f16, u_gc, wave_id_m * (T::B_M / T::T_M) * kargs.stride_c + col); + // gmem::_store supports at most one 16-byte b128 transaction. Keep the + // logical 8-element cached layout and split each FP32 issue into two + // explicit 4-element stores. Asking the cached layout itself for vec=4 + // is invalid: its offsets were constructed for vec=8 and would alias the + // second physical store with a different logical issue. + constexpr int c_store_vec = 16 / sizeof(D_C); + constexpr int c_store_groups = T::VEC_C / c_store_vec; + static_assert(T::VEC_C % c_store_vec == 0); + const int c_base = + wave_id_m * (T::B_M / T::T_M) * kargs.stride_c + col; + if constexpr (c_store_groups == 1) { + store(g_c, v_c_f16, u_gc, c_base); + } else { + using c_layout = remove_cvref_t; + constexpr int c_issues = + layout_load_traits::r_elem.value; + auto c_offsets = layout_to_offsets(u_gc); + for (int issue = 0; issue < c_issues; ++issue) { + static_for([&](auto group) { + constexpr int group_offset = group.value * c_store_vec; + const int value_offset = issue * T::VEC_C + group_offset; + auto value = slice( + v_c_f16, int(value_offset), + int(value_offset + c_store_vec)); + g_c.template store( + value, c_offsets[issue] + group_offset, c_base); + }); + } + } #else // Non-gfx950 device pass compiles to an empty stub; host-side arch // routing in opus_gemm.cu prevents any non-gfx950 device from diff --git a/csrc/opus_gemm/include/gfx950/opus_gemm_pipeline_a16w16_mono_tile_gfx950.cuh b/csrc/opus_gemm/include/gfx950/opus_gemm_pipeline_a16w16_mono_tile_gfx950.cuh index acc2b13f93..0034d1ea7e 100644 --- a/csrc/opus_gemm/include/gfx950/opus_gemm_pipeline_a16w16_mono_tile_gfx950.cuh +++ b/csrc/opus_gemm/include/gfx950/opus_gemm_pipeline_a16w16_mono_tile_gfx950.cuh @@ -4,8 +4,8 @@ // Mono-tile BF16 a16w16 pipeline. // // Direct port of the upstream yk_gcn mono-tile BF16 kernel template -// (bf16_gemm/gemm_a16w16_mono_tile_kernel_template.hpp) with two mechanical -// adjustments to fit the aiter codegen contract: +// (bf16_gemm/gemm_a16w16_mono_tile_kernel_template.hpp) with AITER integration +// adjustments: // // (1) `opus_gemm_kargs` (yk_gcn) -> `opus_gemm_mono_tile_kargs_gfx950` // (aiter; defined in opus_gemm_traits_a16w16_gfx950.cuh alongside @@ -13,12 +13,14 @@ // (2) The layout-helper namespace is renamed `gemm_mono_tile` -> // `opus_mono_tile_gfx950` to avoid any ODR clash with a separate // upstream build that ships the original symbol name. +// (3) The output epilogue supports both BF16 and FP32. Its lane shuffle and +// physical stores are derived from D_C; the upstream template only +// instantiates BF16 output. // -// The kernel body is otherwise byte-for-byte identical to the upstream -// reference. Geometry is locked (T_M=2, T_N=4, T_K=1, W_M=W_N=16, W_K=32, -// VEC=8, BLOCK_SIZE=512); tile-divisibility / smem-rep constraints are -// enforced in the traits header static_asserts and re-validated host-side -// by _validate_a16w16_mono_tile in gen_instances.py. +// Geometry remains locked to the upstream reference (T_M=2, T_N=4, T_K=1, +// W_M=W_N=16, W_K=32, VEC=8, BLOCK_SIZE=512); tile-divisibility / smem-rep +// constraints are enforced in the traits header static_asserts and +// re-validated host-side by _validate_a16w16_mono_tile in gen_instances.py. #pragma once #include @@ -396,22 +398,55 @@ void gemm_a16w16_mono_tile_kernel_gfx950(opus_gemm_mono_tile_kargs_gfx950 kargs) auto u_gc = make_layout_gc(lane_id, 0, wave_id_n, kargs.stride_c); auto v_c_f16 = cast(v_c); - // For every 8 D_C elements (= 4 u32), swap lane L's upper-half (last 4 - // elems) with lane (L^16)'s lower-half (first 4 elems) using - // v_permlane16_swap_b32. - static_assert(sizeof(D_C) * 8 % sizeof(u32_t) == 0); - constexpr int u32_per_chunk = sizeof(D_C) * 8 / sizeof(u32_t); + // For every 8 D_C elements, swap lane L's upper half (last 4 elements) + // with lane (L^16)'s lower half (first 4 elements). The original BF16 + // template hard-coded the two u32 pairs of an 8xbf16 chunk. FP32 has + // four u32 values per half, so derive the pair count from D_C instead of + // silently leaving elements 4..7 unswapped. + static_assert(sizeof(D_C) * 4 % sizeof(u32_t) == 0); + constexpr int u32_per_half = sizeof(D_C) * 4 / sizeof(u32_t); + constexpr int u32_per_chunk = 2 * u32_per_half; constexpr int num_chunks = sizeof(v_c_f16) / (sizeof(u32_t) * u32_per_chunk); auto* p_u32 = reinterpret_cast(&v_c_f16); static_for([&](auto c) { auto* p = p_u32 + c.value * u32_per_chunk; - auto r0 = __builtin_amdgcn_permlane16_swap(p[0], p[2], false, true); - auto r1 = __builtin_amdgcn_permlane16_swap(p[1], p[3], false, true); - p[0] = r0[0]; p[2] = r0[1]; - p[1] = r1[0]; p[3] = r1[1]; + static_for([&](auto i) { + auto r = __builtin_amdgcn_permlane16_swap( + p[i.value], p[i.value + u32_per_half], false, true); + p[i.value] = r[0]; + p[i.value + u32_per_half] = r[1]; + }); }); - store(g_c, v_c_f16, u_gc, wave_id_m * (T::B_M / T::T_M) * kargs.stride_c + col); + // gmem::_store supports at most one 16-byte b128 transaction. Keep the + // logical 8-element cached layout and split each FP32 issue into two + // explicit 4-element stores. Asking the cached layout itself for vec=4 + // is invalid: its offsets were constructed for vec=8 and would alias the + // second physical store with a different logical issue. + constexpr int c_store_vec = 16 / sizeof(D_C); + constexpr int c_store_groups = T::VEC_C / c_store_vec; + static_assert(T::VEC_C % c_store_vec == 0); + const int c_base = + wave_id_m * (T::B_M / T::T_M) * kargs.stride_c + col; + if constexpr (c_store_groups == 1) { + store(g_c, v_c_f16, u_gc, c_base); + } else { + using c_layout = remove_cvref_t; + constexpr int c_issues = + layout_load_traits::r_elem.value; + auto c_offsets = layout_to_offsets(u_gc); + for (int issue = 0; issue < c_issues; ++issue) { + static_for([&](auto group) { + constexpr int group_offset = group.value * c_store_vec; + const int value_offset = issue * T::VEC_C + group_offset; + auto value = slice( + v_c_f16, int(value_offset), + int(value_offset + c_store_vec)); + g_c.template store( + value, c_offsets[issue] + group_offset, c_base); + }); + } + } #else // Non-gfx950 device pass compiles to an empty stub; host-side arch // routing in opus_gemm.cu prevents any non-gfx950 device from diff --git a/csrc/opus_gemm/include/gfx950/opus_gemm_pipeline_a8w8_mxscale_flatmm_splitk_gfx950.cuh b/csrc/opus_gemm/include/gfx950/opus_gemm_pipeline_a8w8_mxscale_flatmm_splitk_gfx950.cuh index 9e6f6e465d..d03a3ffd07 100644 --- a/csrc/opus_gemm/include/gfx950/opus_gemm_pipeline_a8w8_mxscale_flatmm_splitk_gfx950.cuh +++ b/csrc/opus_gemm/include/gfx950/opus_gemm_pipeline_a8w8_mxscale_flatmm_splitk_gfx950.cuh @@ -1021,7 +1021,7 @@ void gemm_a8w8_mxscale_flatmm_splitk_kernel(opus_gemm_scale_splitk_kargs_gfx950 (unsigned int)rows_avail * (unsigned int)kargs.stride_c * sizeof(D_OUT)); store_c(g_out); } else { - D_C* ws_c_ptr = reinterpret_cast(kargs.ws_handle->ptr) + D_C* ws_c_ptr = reinterpret_cast(kargs.ptr_ws) + (size_t)split_id * kargs.batch * kargs.stride_ws_batch + (size_t)batch_id * kargs.stride_ws_batch + (size_t)row * kargs.stride_ws @@ -1030,7 +1030,7 @@ void gemm_a8w8_mxscale_flatmm_splitk_kernel(opus_gemm_scale_splitk_kargs_gfx950 store_c(g_c); } } else { - D_C* ws_c_ptr = reinterpret_cast(kargs.ws_handle->ptr) + D_C* ws_c_ptr = reinterpret_cast(kargs.ptr_ws) + (size_t)split_id * kargs.batch * kargs.stride_ws_batch + (size_t)batch_id * kargs.stride_ws_batch + (size_t)row * kargs.stride_ws @@ -1053,7 +1053,7 @@ void gemm_a8w8_mxscale_flatmm_splitk_kernel(opus_gemm_scale_splitk_kargs_gfx950 __builtin_amdgcn_s_barrier(); int* counters = reinterpret_cast( - reinterpret_cast(kargs.ws_handle->ptr) + kargs.counter_offset_bytes); + reinterpret_cast(kargs.ptr_ws) + kargs.counter_offset_bytes); const int num_tiles = num_tiles_m * ceil_div(kargs.n, T::B_N); const int tile_id = batch_id * num_tiles + wgid; if (opus::thread_id_x() == 0) { @@ -1066,7 +1066,7 @@ void gemm_a8w8_mxscale_flatmm_splitk_kernel(opus_gemm_scale_splitk_kargs_gfx950 __builtin_amdgcn_s_barrier(); if (fused_do_reduce) { - const D_C* ws_base = reinterpret_cast(kargs.ws_handle->ptr); + const D_C* ws_base = reinterpret_cast(kargs.ptr_ws); D_OUT* out = reinterpret_cast(kargs.ptr_c); const size_t split_stride = (size_t)kargs.batch * (size_t)kargs.stride_ws_batch; for (int i = int(opus::thread_id_x()); i < T::B_M * T::B_N; i += T::BLOCK_SIZE) { diff --git a/csrc/opus_gemm/include/gfx950/opus_gemm_traits_a16w16_gfx950.cuh b/csrc/opus_gemm/include/gfx950/opus_gemm_traits_a16w16_gfx950.cuh index 6842902c4c..06fe4048bd 100644 --- a/csrc/opus_gemm/include/gfx950/opus_gemm_traits_a16w16_gfx950.cuh +++ b/csrc/opus_gemm/include/gfx950/opus_gemm_traits_a16w16_gfx950.cuh @@ -265,24 +265,7 @@ struct opus_gemm_a16w16_flatmm_traits_gfx950 { // only defined on the device pass; host pass would see 65536 and cause // pfk<3 and break static_asserts. // - // All aiter a16w16 kernels are gfx950-only today. Three-layer enforcement: - // 1. Python: aiter/ops/opus/__init__.py calls _arch._detect_arch({"gfx950"}) - // at import time. On non-gfx950 the import still succeeds (so it - // cannot break the surrounding `from aiter.ops.opus import *` in - // aiter/__init__.py) but gemm_a16w16_opus / opus_gemm_a16w16_tune - // are replaced with stubs that raise RuntimeError on call, plus a - // one-shot RuntimeWarning at import. Helper is reusable for future - // opus submodules with different supported sets. - // 2. Host: opus_dispatch_a16w16 / opus_a16w16_tune_dispatch in - // opus_gemm.cu are arch routers built on opus_get_gfx_arch(). Only - // the gfx950 branch is wired up today (delegates to - // opus_dispatch_a16w16_gfx950); other archs return TORCH_CHECK - // fail with a 'pipeline TBD' message. Future archs are added by - // extending OpusGfxArch + adding a per-arch dispatch function. - // 3. Device: each __global__ kernel body wraps real code in - // #if defined(__gfx950__) so non-gfx950 device passes (in multi-arch - // wheels like GPU_ARCHS='gfx942;gfx950') compile to an empty stub. - // Combined with layer 1/2 the empty stub is unreachable at runtime. + // Python selects the kid; the host router then enters the gfx950 table. static constexpr int WG_PER_CU = WG_PER_CU_; static constexpr int LDS_SIZE_TOTAL = 163840; static constexpr int max_lds_size_per_wg = LDS_SIZE_TOTAL / WG_PER_CU_; @@ -339,13 +322,13 @@ struct opus_gemm_flatmm_kargs_gfx950 { // ============================================================================ // // 7 template parameters match opus_gemm_a16w16_flatmm_traits_gfx950, with -// additional static_assert D_C=float (splitk main kernel writes fp32 +// additional static_assert D_WS=float (splitk main kernel writes fp32 // partial sums to workspace). Ported from // gcnasm/opus_fmm/flatmm_a16w16_4wave_wasp_splitk.cc lines 34-143. template - typename DTYPE_, // opus::tuple, D_C MUST be float + typename DTYPE_, // opus::tuple, D_WS MUST be float typename VEC_, // opus::seq typename MFMA_, // opus::seq int WG_PER_CU_, @@ -364,13 +347,13 @@ struct opus_flatmm_splitk_traits_gfx950 { using D_A = opus::tuple_element_t<0, DTYPE>; using D_B = opus::tuple_element_t<1, DTYPE>; - using D_C = opus::tuple_element_t<2, DTYPE>; + using D_WS = opus::tuple_element_t<2, DTYPE>; using D_ACC = opus::tuple_element_t<3, DTYPE>; using D_BIAS = opus::tuple_element_t<4, DTYPE>; // Split-K writes fp32 partial sums; reduce kernel later casts to bf16. - static_assert(std::is_same_v, - "splitk kernel requires D_C = float for fp32 workspace"); + static_assert(std::is_same_v, + "splitk kernel requires D_WS = float for fp32 workspace"); // Warp-specialized 4-wave layout: 2 producer + 2 consumer. T_K=1 locked. static constexpr int T_M = 2; @@ -585,22 +568,11 @@ struct opus_gemm_persistent_kargs_gfx950 { }; #endif -#ifndef OPUS_GEMM_SPLITK_WS_HANDLE_DEFINED -#define OPUS_GEMM_SPLITK_WS_HANDLE_DEFINED -// Indirection slot for the split-K fp32 workspace pointer. Captured HIP -// graphs hold the slot address (stable), not the workspace ptr, so a -// post-capture grow + hipFree of the old buffer doesn't dangle the graph. -struct opus_splitk_ws_handle { - void* ptr; // current backing workspace; null until first grow - unsigned long bytes; // current capacity in bytes -}; -#endif - #ifndef OPUS_GEMM_FLATMM_SPLITK_KARGS_DEFINED #define OPUS_GEMM_FLATMM_SPLITK_KARGS_DEFINED // Kernel arguments for the a16w16 flatmm split-K pipeline. // -// Main kernel writes fp32 partial results to *ws_handle->ptr, laid out as +// Main kernel writes fp32 partial results to ptr_ws, laid out as // [split_k, B, padded_M, padded_N] (tile-aligned, no per-thread pred on // store). Reduce kernel consumes it and writes C[B, M, N]. // @@ -608,7 +580,7 @@ struct opus_splitk_ws_handle { struct opus_gemm_flatmm_splitk_kargs_gfx950 { const void* __restrict__ ptr_a; // bf16 [B, M, K] const void* __restrict__ ptr_b; // bf16 [B, N, K] (pre-transposed) - const opus_splitk_ws_handle* __restrict__ ws_handle; // deref at kernel entry + void* __restrict__ ptr_ws; // D_WS [split_k, B, padded_M, padded_N] void* __restrict__ ptr_c; // bf16 [B, M, N] (filled by reduce kernel) // bias is consumed only by the reduce kernel (main kernel ignores it). // ptr_bias = nullptr when HAS_BIAS=false; dtype matches D_BIAS (== D_C @@ -643,7 +615,8 @@ struct opus_gemm_flatmm_splitk_kargs_gfx950 { // Locked geometry, derived in the kernel itself: // * T_M = 2, T_N = 4, T_K = 1 -> 8 waves / WG -> BLOCK_SIZE = 8 * 64 = 512. // * W_M = 16, W_N = 16, W_K = 32 (MFMA 16x16x32 BF16). -// * VEC_A = VEC_B = VEC_C = 8. +// * VEC_A = VEC_B = VEC_C = 8 logical elements. The pipeline splits an +// FP32 logical C vector into 4-element physical stores. // // Constraints (mirror the kernel-internal static_asserts in // gemm_a16w16_mono_tile_kernel_template.hpp; static_asserts here surface diff --git a/csrc/opus_gemm/include/gfx950/opus_gemm_traits_a8w8_scale_gfx950.cuh b/csrc/opus_gemm/include/gfx950/opus_gemm_traits_a8w8_scale_gfx950.cuh index c341e04a80..c9b3ac42a5 100644 --- a/csrc/opus_gemm/include/gfx950/opus_gemm_traits_a8w8_scale_gfx950.cuh +++ b/csrc/opus_gemm/include/gfx950/opus_gemm_traits_a8w8_scale_gfx950.cuh @@ -6,7 +6,7 @@ #pragma once #include "../opus_gemm_utils.cuh" -#include "opus_gemm_traits_a16w16_gfx950.cuh" // opus_splitk_ws_handle +#include "opus_gemm_traits_a16w16_gfx950.cuh" template // uint16_t / uint32_t used by the bias-fold and bf16 store paths template __global__ void splitk_reduce_kernel( - const opus_splitk_ws_handle* __restrict__ ws_handle, + const void* __restrict__ ws_ptr, D_OUT* __restrict__ c_out, int split_k, int M, int N, int batch, int padded_M, int padded_N, @@ -81,9 +80,9 @@ __global__ void splitk_reduce_kernel( // gfx950-only kernel body. See opus_gemm_pipeline_a16w16_gfx950.cuh for the // multi-arch wheel rationale. // - // Deref the handle slot at entry; survives a post-capture grow. - const float* __restrict__ workspace = - reinterpret_cast(ws_handle->ptr); + using D_WS = float; + const D_WS* __restrict__ workspace = + reinterpret_cast(ws_ptr); constexpr int VEC = VEC_; constexpr int BLOCK = BLOCK_; constexpr bool HAS_BIAS = HAS_BIAS_; @@ -137,7 +136,7 @@ __global__ void splitk_reduce_kernel( const long split_stride = (long)batch * padded_M * padded_N; auto g_ws = opus::make_gmem(workspace, - (unsigned int)(split_stride * split_k * sizeof(float))); + (unsigned int)(split_stride * split_k * sizeof(D_WS))); opus::vector_t acc; #pragma unroll @@ -270,7 +269,7 @@ __global__ void splitk_reduce_kernel( // splitk_reduce_extra hook). template __global__ void opus_bmm_splitk_reduce_kernel( - const opus_splitk_ws_handle* __restrict__ ws_handle, + const void* __restrict__ workspace_ptr, D_OUT* __restrict__ out, int split_k, int M, int N, int batch, int padded_M, int padded_N, @@ -289,7 +288,7 @@ __global__ void opus_bmm_splitk_reduce_kernel( const int m = bm_id - b * M; const float* __restrict__ workspace = - reinterpret_cast(ws_handle->ptr); + reinterpret_cast(workspace_ptr); const long split_stride = (long)batch * padded_M * padded_N; const int base = b * padded_M * padded_N + m * padded_N + n_base; auto g_ws = make_gmem(workspace, (unsigned int)(split_stride * split_k * sizeof(float))); diff --git a/csrc/opus_gemm/include/opus_bmm.h b/csrc/opus_gemm/include/opus_bmm.h index ba9c7a5514..44ff431bac 100644 --- a/csrc/opus_gemm/include/opus_bmm.h +++ b/csrc/opus_gemm/include/opus_bmm.h @@ -3,6 +3,7 @@ #pragma once #include "aiter_tensor.h" +#include // Opus BMM public C++ API. These frontends use BMM/grouped layouts (for example // DSV4 wo_a) while reusing the shared opus GEMM backend kernels. @@ -10,11 +11,13 @@ // fp8 e8m0 mxscale (block-scale) BMM (zero-copy DSV4 wo_a): O/Y are [M, batch, // *], wo_a/w_scale batch-major. Y dtype in {fp32, bf16}. dim0=M, dim1=batch (K // contiguous); the batch axis memory position is otherwise free (see host -// stride checks). kid-dispatched; driven by bmm_a8w8_mxscale_opus (Python). -void opus_bmm_a8w8_mxscale(aiter_tensor_t& O, - aiter_tensor_t& wo_a, - aiter_tensor_t& Y, - aiter_tensor_t& x_scale, - aiter_tensor_t& w_scale, - int splitK, - int kernelId); +// stride checks). The global kid is exact: no fallback or redirect is allowed. +void opus_gemm_a8w8_mxscale_bmm_launch( + aiter_tensor_t& XQ, + aiter_tensor_t& WQ, + aiter_tensor_t& Y, + aiter_tensor_t& x_scale, + aiter_tensor_t& w_scale, + std::optional workspace, + int kid, + int split_k); diff --git a/csrc/opus_gemm/include/opus_gemm.h b/csrc/opus_gemm/include/opus_gemm.h index 5f5d2cd0f3..f1ce7c78dd 100644 --- a/csrc/opus_gemm/include/opus_gemm.h +++ b/csrc/opus_gemm/include/opus_gemm.h @@ -2,45 +2,35 @@ // Copyright (C) 2025-2026, Advanced Micro Devices, Inc. All rights reserved. #pragma once -// Top-level opus_gemm entry points. Uses aiter_tensor_t (POD, -// torch-free) instead of torch::Tensor so this header costs ~200 -// preprocessed lines instead of the ~50K that + -// drag in. Mirrors the refactor in PR #2932 -// (csrc/include/quant.h). The pybind layer -// (csrc/pybind/opus_gemm_pybind.cu) registers aiter_tensor_t as a -// pybind11 class via AITER_CORE_PYBIND, and Python callers are -// converted with aiter.utility.dtypes.torch_to_aiter_pybind. +// Exact-kid OPUS entry points. aiter_tensor_t keeps this header torch-free. #include "aiter_tensor.h" #include -void opus_gemm(aiter_tensor_t& XQ, - aiter_tensor_t& WQ, - aiter_tensor_t& Y, - std::optional group_layout, - std::optional x_scale, - std::optional w_scale, - std::optional bias); +void opus_gemm_a16w16_launch(aiter_tensor_t& XQ, + aiter_tensor_t& WQ, + aiter_tensor_t& Y, + std::optional bias, + std::optional workspace, + int kid, + int split_k); -void opus_gemm_a16w16_tune(aiter_tensor_t& XQ, +void opus_gemm_a8w8_launch(aiter_tensor_t& XQ, aiter_tensor_t& WQ, aiter_tensor_t& Y, - std::optional bias, - std::optional workspace, - int kernelId, - int splitK); + int kid); -void opus_gemm_a8w8_blockscale_bpreshuffle_tune(aiter_tensor_t& XQ, - aiter_tensor_t& WQ, - std::optional x_scale, - std::optional w_scale, - aiter_tensor_t& Y, - int kernelId); +// Blockscale interfaces require both scale tensors. +void opus_gemm_a8w8_blockscale_launch(aiter_tensor_t& XQ, + aiter_tensor_t& WQ, + aiter_tensor_t& Y, + aiter_tensor_t& x_scale, + aiter_tensor_t& w_scale, + int kid); -// Per-stream splitk workspace init. See opus_gemm.cu for rationale. -void opus_gemm_workspace_init(); - -// Release the per-stream splitk workspace (buffer + handles + registry entry). -// `_release` targets the current stream; `_release_all` tears down every -// registered stream. Both must be called in eager mode (not during capture). -void opus_gemm_workspace_release(); -void opus_gemm_workspace_release_all(); +void opus_gemm_a8w8_blockscale_bpreshuffle_launch( + aiter_tensor_t& XQ, + aiter_tensor_t& WQ, + aiter_tensor_t& x_scale, + aiter_tensor_t& w_scale, + aiter_tensor_t& Y, + int kid); diff --git a/csrc/opus_gemm/include/opus_gemm_common.cuh b/csrc/opus_gemm/include/opus_gemm_common.cuh index 1a6849ad38..7c379211af 100644 --- a/csrc/opus_gemm/include/opus_gemm_common.cuh +++ b/csrc/opus_gemm/include/opus_gemm_common.cuh @@ -15,11 +15,85 @@ // same TU. #pragma once +#if !defined(__HIP_DEVICE_COMPILE__) && !defined(__HIPCC_RTC__) +#include "aiter_tensor.h" + +#include +#include +#include + +// Host-only, family-neutral helpers for caller-owned typed workspaces. Shape, +// architecture, kernel-id, and tile policy stay in the family launchers; this +// layer only enforces the common physical buffer contract. +inline size_t opus_checked_extent_product(std::initializer_list extents, + const char* label) +{ + size_t product = 1; + for(const size_t extent : extents) + { + AITER_CHECK(extent > 0, label, ": workspace extents must be positive"); + AITER_CHECK(product <= std::numeric_limits::max() / extent, + label, + ": workspace extent product overflows size_t"); + product *= extent; + } + return product; +} + +inline void* opus_validate_workspace(aiter_tensor_t& workspace, + const aiter_tensor_t& reference, + AiterDtype expected_dtype, + size_t required_numel, + size_t alignment, + const char* label) +{ + AITER_CHECK(required_numel > 0, + label, + ": required workspace element count must be positive"); + AITER_CHECK(alignment > 0 && (alignment & (alignment - 1)) == 0, + label, + ": workspace alignment must be a non-zero power of two"); + AITER_CHECK(workspace.device_id == reference.device_id, + label, + ": workspace device ", + workspace.device_id, + " must match input device ", + reference.device_id); + AITER_CHECK(workspace.dtype() == expected_dtype, + label, + ": workspace dtype must be ", + AiterDtype_to_str(expected_dtype), + ", got ", + AiterDtype_to_str(workspace.dtype())); + AITER_CHECK(workspace.is_contiguous(), label, ": workspace must be contiguous"); + AITER_CHECK(workspace.numel() >= required_numel, + label, + ": workspace capacity is ", + workspace.numel(), + " elements, but ", + required_numel, + " are required"); + + void* ptr = workspace.data_ptr(); + AITER_CHECK(ptr != nullptr, label, ": workspace data pointer must be non-null"); + AITER_CHECK(reinterpret_cast(ptr) % alignment == 0, + label, + ": workspace address must be aligned to ", + alignment, + " bytes"); + + // The workspace is typed, so ensure its required byte span is representable + // as well as its logical element count. + (void)opus_checked_extent_product( + {required_numel, workspace.element_size()}, label); + return ptr; +} +#endif + #include "gfx950/opus_gemm_traits_a8w8_scale_gfx950.cuh" #include "gfx950/opus_gemm_traits_a8w8_noscale_gfx950.cuh" // Both opus_gemm_a16w16_traits_gfx950 (split-barrier) and // opus_gemm_a16w16_flatmm_traits_gfx950 (warp-spec) live in this one header. #include "gfx950/opus_gemm_traits_a16w16_gfx950.cuh" -// gfx1250 cluster/TDM split-K (workspace + reduce) traits + kargs + -// opus_splitk_ws_handle (guarded; shared with gfx950). +// gfx1250 cluster/TDM split-K (workspace + reduce) traits + direct-pointer kargs. #include "gfx1250/opus_gemm_traits_a16w16_gfx1250.cuh" diff --git a/csrc/opus_gemm/opus_bmm.cu b/csrc/opus_gemm/opus_bmm.cu index 5b94ac0bf0..7fd1393b84 100644 --- a/csrc/opus_gemm/opus_bmm.cu +++ b/csrc/opus_gemm/opus_bmm.cu @@ -16,7 +16,7 @@ #include "opus_gemm_arch.cuh" #include "opus_build_archs.h" #include "opus_gemm_manifest.h" -#include "opus_bmm_mxscale_tune_lookup.h" // GENERATE_BMM_MXSCALE_FLATMM_SPLITK_LOOKUP_FP32 +#include "opus_bmm_mxscale_kid_dispatch.h" #include "opus_gemm_utils.cuh" // bf16_t / fp32_t #include "aiter_stream.h" #include "gfx950/opus_bmm_launchers_a8w8_mxscale_gfx950.cuh" @@ -28,57 +28,57 @@ namespace opus_bmm_detail { // Uniform kid->launcher fn-pointer type. Every kid is codegen'd (no hand-written // adapters), so this namespace only holds the shared type. -using OpusBmmMxscaleFlatmmSplitkKernel = void (*)( +using OpusBmmMxscaleKernel = void (*)( aiter_tensor_t &, aiter_tensor_t &, aiter_tensor_t &, - aiter_tensor_t &, aiter_tensor_t &, int /*splitK*/); + aiter_tensor_t &, aiter_tensor_t &, + std::optional, int /*split_k*/); } // namespace opus_bmm_detail -// Table-driven kid -> launcher dispatch. Launchers come from the generated -// GENERATE_BMM_MXSCALE_FLATMM_SPLITK_LOOKUP_FP32 macro; unknown / untuned kids -// fall back to the 32x128x128 wg2 baseline. -static opus_bmm_detail::OpusBmmMxscaleFlatmmSplitkKernel -opus_bmm_a8w8_mxscale_tune_dispatch(int id) +// Table-driven global exact-kid dispatch. Every registry BMM kid is generated; +// an unknown id is an error rather than a fallback to another kernel. +static opus_bmm_detail::OpusBmmMxscaleKernel +opus_bmm_a8w8_mxscale_exact_dispatch(int kid) { using namespace opus_bmm_detail; - static const std::unordered_map kTune = { - GENERATE_BMM_MXSCALE_FLATMM_SPLITK_LOOKUP_FP32(fp32_t) + static const std::unordered_map kDispatch = { + GENERATE_BMM_MXSCALE_KID_DISPATCH(fp32_t) }; - auto it = kTune.find(id); - if (it != kTune.end()) - return it->second; - return &opus_bmm_a8w8_mxscale_flatmm_splitk_256x32x128x128_2x1_16x16x128_1x128x128_wgpcu2; + auto it = kDispatch.find(kid); + AITER_CHECK(it != kDispatch.end(), + "unknown exact OPUS a8w8_mxscale_bmm kid ", kid); + return it->second; } #endif // OPUS_BUILD_HAS_GFX950 -void opus_bmm_a8w8_mxscale( +void opus_gemm_a8w8_mxscale_bmm_launch( aiter_tensor_t &O, aiter_tensor_t &wo_a, aiter_tensor_t &Y, aiter_tensor_t &x_scale, aiter_tensor_t &w_scale, - int splitK, - int kernelId) + std::optional workspace, + int kid, + int split_k) { // Common dtype/shape validation + arch gate, done once here so the codegen'd // launchers (which omit these to stay lean) and the fused kid 100 wrapper share // one check. The _impl still re-checks internally (idempotent). - opus_bmm_a8w8_common_checks(O, wo_a, Y, - "opus_bmm_a8w8_mxscale"); + opus_bmm_a8w8_common_checks(O, wo_a, Y, x_scale, w_scale, + "opus_gemm_a8w8_mxscale_bmm_launch"); #ifndef OPUS_BUILD_HAS_GFX950 AITER_CHECK(false, - "opus_bmm_a8w8_mxscale requires " + "opus_gemm_a8w8_mxscale_bmm_launch requires " "OPUS_BUILD_HAS_GFX950"); #else { const auto &arch_info = opus_get_arch_info(); AITER_CHECK(arch_info.arch == OpusGfxArch::Gfx950, - "opus_bmm_a8w8_mxscale is gfx950-only; " + "opus_gemm_a8w8_mxscale_bmm_launch is gfx950-only; " "current device ", arch_info.dev, " has gcnArchName='", arch_info.name, "'"); } - // Single table lookup instead of a ~40-case switch (see opus_gemm.cu). - opus_bmm_a8w8_mxscale_tune_dispatch(kernelId)( - O, wo_a, Y, x_scale, w_scale, splitK); + opus_bmm_a8w8_mxscale_exact_dispatch(kid)( + O, wo_a, Y, x_scale, w_scale, workspace, split_k); #endif // OPUS_BUILD_HAS_GFX950 } diff --git a/csrc/opus_gemm/opus_bmm_mxscale_tune.py b/csrc/opus_gemm/opus_bmm_mxscale_tune.py index 0f6388df29..2fe1eb997b 100644 --- a/csrc/opus_gemm/opus_bmm_mxscale_tune.py +++ b/csrc/opus_gemm/opus_bmm_mxscale_tune.py @@ -17,9 +17,14 @@ gfx,b,m,n,k,libtype,kernelId,splitK,us,kernelName,tflops,bw,errRatio ``aiter/ops/batched_gemm_op_a8w8.py:lookup_mxscale_bmm_config`` indexes on ``["gfx","b","m","n","k"]``, dispatches to a backend on the winning row's -``libtype``, and ``bmm_op.py`` reads ``kernelId`` / ``splitK`` off that row, so +``libtype``, and the existing A8W8 caller passes ``kernelId`` / ``splitK`` to +the batch-first ``opus_bmm`` entry, so those columns must match exactly. +The shipped schema has no output-dtype key. This tuner deliberately measures +the production BF16 route; FP32 production calls can execute the selected kid, +but do not have an independently tuned winner in this CSV format. + Verification (the part that catches column-transpose / scale defects): * inputs are *signed* and have *per-128-K-block varied magnitude* (``randn * 2**randint(-4,4)`` per block) so the e8m0 128-block scales span @@ -27,6 +32,10 @@ column permutation (kid312/313 measured ~0.007 there but ~0.7-1.0 on real signed data) -- see the opus_bmm.md root-cause note. * reference is a dequantized fp32 einsum. + * output is allocated as contiguous ``[M,G,N]`` exactly like production and + passed to the batch-first public API through a transpose view. The existing + batch-first activation/scale storage is retained, matching the canonical + production transpose-view inputs. * gate: ``mp_tuner`` runs ``checkAllclose(rtol=1e-2, atol=1e-2)`` and ``post_process`` keeps the fastest candidate whose mismatch fraction is ``<= --errRatio`` (default 0.02). A still-broken tileN COM_REP_N>1 kernel @@ -37,16 +46,18 @@ cd && PYTHONPATH=$PWD \\ python3 csrc/opus_gemm/opus_bmm_mxscale_tune.py -g 16 -m 1,16,64 -n 1024 -k 4096 - # re-tune every shape already in the shipped CSV, write a diffable copy: + # tune shipped shapes not already present in the diffable output copy; + # add --all to force every shipped shape to be measured again: ... opus_bmm_mxscale_tune.py - # overwrite the shipped tuned CSV in place: + # overwrite the shipped tuned CSV in place (--apply implies --all): ... opus_bmm_mxscale_tune.py --apply # from an untuned CSV (columns: b,m,n,k -- or g,m,n,k), 8-way parallel: ... opus_bmm_mxscale_tune.py -i my_untuned.csv -o /tmp/out.csv --mp 8 """ +import math import os import sys from typing import Any, ClassVar @@ -55,7 +66,7 @@ import torch from aiter import dtypes, logger -from aiter.ops.opus.bmm_op import _opus_bmm_a8w8_mxscale_raw +from aiter.ops.opus import opus_bmm from aiter.utility.base_tuner import GemmCommonTuner, TunerCommon from aiter.utility.mp_tuner import mp_tuner @@ -71,7 +82,10 @@ # opus_gemm_common is pure python (stdlib only), so importing the codegen kid # table here does not pull in the build. -from opus_gemm_common import a8w8_mxscale_bmm_kernel_lists +from opus_gemm_common import ( + a8w8_mxscale_bmm_kernel_lists, + bmm_mxscale_global_kid, +) from test_opus_a8w8_bmm import ( GROUP, _quant_block_e8m0, @@ -100,58 +114,58 @@ # Tuning policy: kid -> splitK list. The ONLY hand-maintained per-kid metadata -- # it decides which kids to sweep and with which split-K factors, not their # geometry and not their M alignment. Tile shape, kernelName and m_align all come -# from the codegen instance, so this cannot drift from what compiles. kid 0 (the +# from the codegen instance, so this cannot drift from what compiles. kid 8000 (the # heuristic default) is intentionally not tuned. _TUNE_POLICY = { # flatmm_splitk family: the M=16/32 last-mile tiles, the mid-M SFA/SFB-preload # tiles and the 64x* tiles. All are split-K capable via the fused reduce tail, # except kid646 whose persistent DIRECT_ONLY schedule requires splitK == 1. - 32: _SK, - 64: _SK, - 138: _SK, - 139: _SK, - 256: _SK, - 311: _SK, - 312: _SK, - 313: _SK, - 314: _SK, - 316: _SK, - 317: _SK, - 318: _SK, - 319: _SK, - 320: _SK, - 321: _SK, - 322: _SK, - 323: _SK, - 324: _SK, - 326: _SK, - 327: _SK, - 640: _SK, - 642: _SK, - 646: [1], - 650: _SK, - 653: _SK, + 8032: _SK, + 8064: _SK, + 8138: _SK, + 8139: _SK, + 8256: _SK, + 8311: _SK, + 8312: _SK, + 8313: _SK, + 8314: _SK, + 8316: _SK, + 8317: _SK, + 8318: _SK, + 8319: _SK, + 8320: _SK, + 8321: _SK, + 8322: _SK, + 8323: _SK, + 8324: _SK, + 8326: _SK, + 8327: _SK, + 8640: _SK, + 8642: _SK, + 8646: [1], + 8650: _SK, + 8653: _SK, # fused single-tile launcher. - 100: [1], + 8100: [1], # pipeline family; kid158 preloads both the per-token SFA and the block SFB # panel into LDS. - 149: [1], - 150: [1], - 151: [1], - 152: [1], - 158: [1], + 8149: [1], + 8150: [1], + 8151: [1], + 8152: [1], + 8158: [1], # monolithic mouter / wave pipelines. - 131: [1], - 132: [1], - 134: [1], - 142: [1], - 144: [1], - 148: [1], - 160: [1], - 161: [1], + 8131: [1], + 8132: [1], + 8134: [1], + 8142: [1], + 8144: [1], + 8148: [1], + 8160: [1], + 8161: [1], # minterleave only exists in split-K form. - 162: [2, 4, 8], - 163: [2, 4, 8], + 8162: [1], + 8163: [1], # 128x128x128 tiles, splitK=1 only and deliberately so. They are the largest # BMM tile (COM_REP_M=4 x COM_REP_N=8 -> 32 C fragments, 128 fp32 C values per # lane) at 512 VGPRs / occupancy 1. At splitK=1 they run the Cbf16 @@ -160,22 +174,20 @@ # never won (g2/m256: best kid325 split-K is 23.6us against the 14.5us winner), # and which is also where the clang-22 gfx950 greedy-VGPR miscompile lives (one # C-fragment dword left unmaterialized under --amdgpu-mfma-vgpr-form). - 128: [1], - 137: [1], - 325: [1], + 8128: [1], + 8137: [1], + 8325: [1], } -# Only the flatmm_splitk (non-direct) and minterleave launchers honor splitK>1. +# Only non-direct flatmm split-K launchers are swept with splitK>1. # Any other family sweeping it is a policy bug, so fail loudly at import. for _kid, _sks in _TUNE_POLICY.items(): if any(s > 1 for s in _sks): _tag = _CODEGEN_BMM[_kid].kernel_tag assert ( _tag == "a8w8_mxscale_bmm_flatmm_splitk" - and not getattr(_CODEGEN_BMM[_kid], "direct_only", False) - ) or _tag == "a8w8_mxscale_bmm_minterleave", ( - f"kid {_kid} ({_tag}) is not split-K capable but sweeps {_sks}" - ) + and not _CODEGEN_BMM[_kid].direct_only + ), f"kid {_kid} ({_tag}) is not split-K capable but sweeps {_sks}" def _applicable(kid, g, m, n, k): @@ -196,6 +208,59 @@ def _applicable(kid, g, m, n, k): DEFAULT_OUT = os.path.join(_REPO, "dsv4_bmm_mxscale_retuned.csv") +def _read_shape_csv(path): + """Read ``b/g,m,n,k`` shape rows with a clear schema error.""" + try: + df = pd.read_csv(path) + except FileNotFoundError as exc: + raise FileNotFoundError(f"MXFP8 BMM shape CSV does not exist: {path}") from exc + + df.columns = [str(column).strip().lower() for column in df.columns] + bcol = "b" if "b" in df.columns else "g" if "g" in df.columns else None + required = {"m", "n", "k"} + missing = sorted(required.difference(df.columns)) + if bcol is None or missing: + expected = "b,m,n,k (or g,m,n,k)" + raise ValueError( + f"MXFP8 BMM shape CSV {path!r} must contain {expected}; " + f"got columns {list(df.columns)}" + ) + return [ + (int(row[bcol]), int(row["m"]), int(row["n"]), int(row["k"])) + for _, row in df.iterrows() + ] + + +def _validate_tune_shapes(shapes): + """Normalize, deduplicate and enforce the global MXFP8 BMM contract.""" + valid = [] + seen = set() + for raw_shape in shapes: + try: + g, m, n, k = map(int, raw_shape) + except (TypeError, ValueError) as exc: + raise ValueError( + f"MXFP8 BMM shape must be (G,M,N,K), got {raw_shape!r}" + ) from exc + shape = (g, m, n, k) + if min(shape) <= 0: + raise ValueError( + "MXFP8 BMM requires positive G, M, N and K; " + f"got G={g}, M={m}, N={n}, K={k}" + ) + if n % GROUP or k % GROUP: + raise ValueError( + f"MXFP8 BMM requires N and K to be multiples of {GROUP}; " + f"got G={g}, M={m}, N={n}, K={k}" + ) + if shape not in seen: + seen.add(shape) + valid.append(shape) + if not valid: + raise ValueError("no MXFP8 BMM shapes were provided for tuning") + return valid + + # --------------------------------------------------------------------------- # mp_tuner hooks (module-level so the spawn workers can import them by name). # --------------------------------------------------------------------------- @@ -206,36 +271,73 @@ def _gen_varied(shape, k, device): return (x * amp.repeat_interleave(GROUP)).to(dtypes.bf16) -def gen_bmm_mxscale_data(batch, m, n, k, seed, out_dtype, device="cuda"): - """Return the 6-tuple mp_tuner indexes into: - - 0 O_in [m,g,k] fp8 (mmajor transposed view, K contiguous) +def _workspace_numel(kernel_id, split_k, batch, m, n): + instance = _CODEGEN_BMM[int(kernel_id)] + if split_k <= 1 or instance.kernel_tag not in { + "a8w8_mxscale_bmm_flatmm_splitk", + "a8w8_mxscale_bmm_fused", + }: + return 0 + tiles_m = (m + instance.B_M - 1) // instance.B_M + tiles_n = (n + instance.B_N - 1) // instance.B_N + partial_numel = split_k * batch * tiles_m * instance.B_M * tiles_n * instance.B_N + if instance.kernel_tag != "a8w8_mxscale_bmm_fused": + return partial_numel + counter_offset = (partial_numel * 4 + 255) & ~255 + counter_bytes = batch * tiles_m * tiles_n * 4 + return (counter_offset + counter_bytes + 3) // 4 + + +def gen_bmm_mxscale_data( + batch, m, n, k, seed, out_dtype, kernel_id, split_k, device="cuda" +): + """Return the 7-tuple mp_tuner indexes into: + + 0 O_mx [g,m,k] fp8 batch-first contiguous input 1 W_mx [g,n,k] fp8 (batch-major) - 2 Y [m,g,n] out_dtype output buffer - 3 xs_in [m,g,k/128] uint8 e8m0 per-token scale (mmajor view) + 2 Y [m,g,n] contiguous production-layout output buffer + 3 xs_mx [g,m,k/128] uint8 e8m0 batch-first contiguous scale 4 ws_mx [g,n/128,k/128] uint8 e8m0 128x128-block scale - 5 ref [m,g,n] out_dtype dequant fp32 einsum reference + 5 workspace optional caller-owned FP32 split-K buffer + 6 ref [m,g,n] out_dtype dequant fp32 einsum reference """ torch.manual_seed(seed) O_bf16 = _gen_varied((batch, m, k), k, device) W_bf16 = _gen_varied((batch, n, k), k, device) O_mx, xs_mx, xs_fp32 = _quant_per_token_e8m0(O_bf16) W_mx, ws_mx, ws_fp32 = _quant_block_e8m0(W_bf16) - O_in = O_mx.transpose(0, 1) # [m,g,k] - xs_in = xs_mx.transpose(0, 1) # [m,g,k/128] Y = torch.empty((m, batch, n), dtype=out_dtype, device=device) + + workspace_numel = _workspace_numel(kernel_id, split_k, batch, m, n) + workspace = ( + torch.empty(workspace_numel, dtype=torch.float32, device=device) + if workspace_numel + else None + ) ref = run_torch(O_mx, W_mx, xs_fp32, ws_fp32).transpose(0, 1).to(out_dtype) - return (O_in, W_mx, Y, xs_in, ws_mx, ref) + return (O_mx, W_mx, Y, xs_mx, ws_mx, workspace, ref) -def run_bmm_mxscale_bench(O_in, W_mx, Y, xs_in, ws_mx, kernelId, splitK): +def run_bmm_mxscale_bench(O_mx, W_mx, Y, xs_mx, ws_mx, workspace, kernelId, splitK): """Tuner bench func: run the kid in-place, return Y for checkAllclose.""" - _opus_bmm_a8w8_mxscale_raw(O_in, W_mx, Y, xs_in, ws_mx, splitK, kernelId) + # Production owns contiguous token-major Y. Expose its batch-first view to + # the public API; the adapter transposes it back before the raw launch. + opus_bmm( + O_mx, + W_mx, + Y.transpose(0, 1), + kid=int(kernelId), + layout="mxscale_bmm", + x_scale=xs_mx, + w_scale=ws_mx, + split_k=int(splitK), + workspace=workspace, + ) return Y def _bmm_ref_passthrough(ref): - """ref_func: the fp32 reference is precomputed in gen_data (slot 5).""" + """ref_func: the fp32 reference is precomputed in gen_data (slot 6).""" return ref @@ -251,6 +353,7 @@ class OpusBmmMxscaleTuner(GemmCommonTuner): # sit at the ~1e-4 fp8 e8m0 quant floor; a column-transposed kid is ~0.5. "errRatio": 0.02, "batch": 100, + "config_env_name": "AITER_CONFIG_BATCHED_GEMM_A8W8_BLOCKSCALE_MXSCALE", } KEYS: ClassVar[list[str]] = ["gfx", "b", "m", "n", "k"] @@ -365,25 +468,30 @@ def _intlist(s): "--apply", action="store_true", default=False, - help="overwrite the shipped tuned CSV in place", + help="overwrite the shipped tuned CSV in place (implies --all)", ) # --- shape sourcing ----------------------------------------------------- def _shapes_from_shipped(self): - try: - df = pd.read_csv(SHIPPED_CSV) - except FileNotFoundError: - return [] - return sorted( - {(int(r.b), int(r.m), int(r.n), int(r.k)) for _, r in df.iterrows()} - ) + return sorted(set(_read_shape_csv(SHIPPED_CSV))) def pre_process(self, args): if args.apply: args.tune_file = SHIPPED_CSV + # Reading and writing the shipped CSV otherwise makes every source + # shape look already tuned and silently produces zero tasks. + args.all = True gfx = self.get_gfx() - if args.batch_g and args.M: + if gfx != "gfx950": + raise RuntimeError(f"MXFP8 BMM tuning is gfx950-only; detected {gfx!r}") + + manual_g = args.batch_g is not None + manual_m = args.M is not None + if manual_g != manual_m: + raise ValueError("-g/--batch_g and -m/--M must be provided together") + + if manual_g: shapes = [ (g, m, n, k) for g in args.batch_g @@ -391,19 +499,14 @@ def pre_process(self, args): for n in args.N for k in args.K ] - elif args.untune_file and os.path.exists(args.untune_file): - df = pd.read_csv(args.untune_file) - df.columns = [c.strip().lower() for c in df.columns] - bcol = "b" if "b" in df.columns else "g" - shapes = [ - (int(r[bcol]), int(r["m"]), int(r["n"]), int(r["k"])) - for _, r in df.iterrows() - ] + elif args.untune_file: + shapes = _read_shape_csv(args.untune_file) else: logger.info( "no -g/-m and no untune_file; re-tuning shapes from %s", SHIPPED_CSV ) shapes = self._shapes_from_shipped() + shapes = _validate_tune_shapes(shapes) self.untunedf = pd.DataFrame( [{"gfx": gfx, "b": g, "m": m, "n": n, "k": k} for (g, m, n, k) in shapes], @@ -422,6 +525,148 @@ def pre_process(self, args): logger.info("skipping %d already-tuned shapes", int(mask.sum())) self.untunedf = self.untunedf[~mask].reset_index(drop=True) + # --- saved exact-kid benchmark ------------------------------------------ + def _clear_op_caches(self): + from aiter.ops import batched_gemm_op_a8w8 + from aiter.ops.opus import policy + + policy._load_mxscale_bmm_tuned.cache_clear() + policy.lookup_mxscale_bmm_config.cache_clear() + batched_gemm_op_a8w8._get_mxscale_bmm_launch_plan.cache_clear() + + def run_config(self, args): + from aiter.test_common import checkAllclose, run_perftest + + required = {"libtype", "kernelId", "splitK"} + missing = required.difference(self.untunedf.columns) + if missing: + if missing == required: + return self._run_default_config(args) + raise ValueError( + f"--run_config requires a tuned CSV with {sorted(missing)}" + ) + + results = [] + for seed, (_, row) in enumerate(self.untunedf.iterrows(), start=1): + b, m, n, k = (int(row[name]) for name in ("b", "m", "n", "k")) + if str(row["libtype"]).strip().lower() != "opus": + raise ValueError( + "MXFP8 BMM --run_config only supports libtype=opus; " + f"got {row['libtype']!r} for B={b}, M={m}, N={n}, K={k}" + ) + + saved_kid = int(row["kernelId"]) + kernel_id = saved_kid + if kernel_id not in _CODEGEN_BMM: + legacy_global_kid = bmm_mxscale_global_kid(saved_kid) + if legacy_global_kid in _CODEGEN_BMM: + kernel_id = legacy_global_kid + if kernel_id not in _CODEGEN_BMM: + raise ValueError( + f"saved MXFP8 BMM kid {saved_kid} is not registered on gfx950" + ) + + split_k = int(row["splitK"]) + if split_k not in _applicable(kernel_id, b, m, n, k): + raise ValueError( + f"saved MXFP8 BMM kid {saved_kid} (global {kernel_id}) with " + f"splitK={split_k} is incompatible with " + f"B={b}, M={m}, N={n}, K={k}" + ) + + shape_str = f"B={b},M={m},N={n},K={k},kid={kernel_id},splitK={split_k}" + allowed, allowed_desc = self._get_run_config_err_ratio_limit(row, args) + data = gen_bmm_mxscale_data( + b, + m, + n, + k, + seed, + dtypes.bf16, + kernel_id, + split_k, + ) + data[2].fill_(float("nan")) + out, us = run_perftest( + run_bmm_mxscale_bench, + *data[:6], + kernel_id, + split_k, + num_warmup=args.warmup, + num_iters=args.iters, + ) + err_ratio = checkAllclose( + out, + data[6], + rtol=1e-2, + atol=1e-2, + tol_err_ratio=allowed, + msg=f"run_config {shape_str}", + printLog=args.verbose, + ) + if ( + not math.isfinite(us) + or us <= 0 + or not math.isfinite(err_ratio) + or err_ratio > allowed + ): + raise RuntimeError( + f"saved MXFP8 BMM kid {kernel_id} failed: " + f"us={us}, errRatio={err_ratio} (>{allowed_desc})" + ) + results.append({"shape": shape_str, "e2e_us": us, "status": "ok"}) + return results + + def _run_default_config(self, args): + """Keep shape-only ``--run_config``/``--compare`` on production policy.""" + from aiter.ops.batched_gemm_op_a8w8 import batched_gemm_a8w8_mxscale + from aiter.test_common import checkAllclose, run_perftest + + results = [] + for seed, (_, row) in enumerate(self.untunedf.iterrows(), start=1): + b, m, n, k = (int(row[name]) for name in ("b", "m", "n", "k")) + shape_str = f"({b}, {m}, {n}, {k})" + allowed, allowed_desc = self._get_run_config_err_ratio_limit(row, args) + try: + O_mx, W_mx, _Y, xs_mx, ws_mx, _workspace, ref = gen_bmm_mxscale_data( + b, + m, + n, + k, + seed, + dtypes.bf16, + 8000, + 1, + ) + out, us = run_perftest( + batched_gemm_a8w8_mxscale, + O_mx.transpose(0, 1), + W_mx, + xs_mx.transpose(0, 1), + ws_mx, + dtype=dtypes.bf16, + num_warmup=args.warmup, + num_iters=args.iters, + ) + err_ratio = checkAllclose( + out, + ref, + rtol=1e-2, + atol=1e-2, + msg=f"run_config {shape_str}", + ) + status = ( + "ok" + if err_ratio <= allowed + else f"mismatch:err_ratio={err_ratio:.6g}(>{allowed_desc})" + ) + results.append({"shape": shape_str, "e2e_us": us, "status": status}) + except Exception as exc: # noqa: BLE001 + results.append( + {"shape": shape_str, "e2e_us": -1, "status": f"error:{exc}"} + ) + return results + # --- tuning ------------------------------------------------------------- def tune(self, untunedf, tunedf, args): gfx = self.get_gfx() @@ -445,12 +690,12 @@ def tune(self, untunedf, tunedf, args): ( info, gen_bmm_mxscale_data, - (b, m, n, k, seed, out_dtype), + (b, m, n, k, seed, out_dtype, kid, sk), run_bmm_mxscale_bench, - ([0, 1, 2, 3, 4], kid, sk), + ([0, 1, 2, 3, 4, 5], kid, sk), perf_kwargs, _bmm_ref_passthrough, - ([5],), + ([6],), {}, None, 1e-2, # rtol diff --git a/csrc/opus_gemm/opus_gemm.cu b/csrc/opus_gemm/opus_gemm.cu index bd622c6a6e..599492cf8c 100644 --- a/csrc/opus_gemm/opus_gemm.cu +++ b/csrc/opus_gemm/opus_gemm.cu @@ -1,685 +1,490 @@ // SPDX-License-Identifier: MIT // Copyright (C) 2025-2026, Advanced Micro Devices, Inc. All rights reserved. -// Host-side dispatcher (lookup table + heuristic). +// Host-side family routers and strict exact-kid dispatch. #ifndef __HIP_DEVICE_COMPILE__ #include "opus_gemm_arch.cuh" // OpusGfxArch + opus_get_arch_info / opus_get_gfx_arch #include "opus_build_archs.h" // OPUS_BUILD_HAS_GFX942 / OPUS_BUILD_HAS_GFX950 #ifdef OPUS_BUILD_HAS_GFX950 -#include "gfx950/opus_gemm_arch_gfx950.cuh" // opus_dispatch_a16w16_gfx950 / opus_a16w16_tune_dispatch_gfx950 +#include "gfx950/opus_gemm_arch_gfx950.cuh" // generated gfx950 a16w16 kid dispatch #endif #ifdef OPUS_BUILD_HAS_GFX942 -#include "gfx942/opus_gemm_arch_gfx942.cuh" // opus_dispatch_a16w16_gfx942 / opus_a16w16_tune_dispatch_gfx942 +#include "gfx942/opus_gemm_arch_gfx942.cuh" // generated gfx942 a16w16 kid dispatch #endif #ifdef OPUS_BUILD_HAS_GFX1250 -#include "gfx1250/opus_gemm_arch_gfx1250.cuh" // opus_a16w16_tune_dispatch_gfx1250 (tune-id entry only) +#include "gfx1250/opus_gemm_arch_gfx1250.cuh" // generated gfx1250 exact-kid dispatch #endif #include "opus_gemm_common.cuh" -#ifdef OPUS_BUILD_HAS_GFX950 -#include "gfx950/opus_gemm_heuristic_dispatch_gfx950.cuh" // OpusA16W16NoscaleKernel -#endif -#ifdef OPUS_BUILD_HAS_GFX942 -#include "gfx942/opus_gemm_heuristic_dispatch_gfx942.cuh" -#endif #include "opus_gemm_manifest.h" // a8w8 launcher symbols #include "opus_gemm_utils.cuh" // bf16_t / fp32_t -#include "aiter_stream.h" // aiter::getCurrentHIPStream -#include #include -#include -// a8w8 / a8w8_scale: single hardcoded launcher per dtype (no tuned table). -// Plain fn ptrs; std::function's type-erasure is pure waste here. -using OpusScaleKernel = void (*)( - aiter_tensor_t &, aiter_tensor_t &, - aiter_tensor_t &, - std::optional, std::optional); - -using OpusNoscaleKernel = void (*)( - aiter_tensor_t &, aiter_tensor_t &, - aiter_tensor_t &); +#ifndef OPUS_A8W8_DISPATCH_KERNEL_TYPES_DEFINED +#define OPUS_A8W8_DISPATCH_KERNEL_TYPES_DEFINED +using OpusA8W8Kernel = void (*)( + aiter_tensor_t&, aiter_tensor_t&, aiter_tensor_t&); +using OpusA8W8BlockscaleKernel = void (*)( + aiter_tensor_t&, aiter_tensor_t&, aiter_tensor_t&, + aiter_tensor_t&, aiter_tensor_t&); +using OpusA8W8BlockscaleBpreshuffleKernel = void (*)( + aiter_tensor_t&, aiter_tensor_t&, aiter_tensor_t&, + aiter_tensor_t&, aiter_tensor_t&); +#endif -template -OpusScaleKernel opus_dispatch_scale(int M, int N, int K) +static OpusA8W8Kernel opus_a8w8_kid_dispatch(int kid) { + const auto &info = opus_get_arch_info(); + switch (info.arch) + { + case OpusGfxArch::Gfx950: #ifdef OPUS_BUILD_HAS_GFX950 - return opus_gemm_512x256x256x128_4x2_16x16x128_1x128x128; + return opus_a8w8_kid_dispatch_gfx950(kid); +#else + AITER_CHECK(false, + "opus_gemm_a8w8_launch: module was not built with gfx950 ", + "support for current device ", info.dev); +#endif + return nullptr; + case OpusGfxArch::Gfx942: +#ifdef OPUS_BUILD_HAS_GFX942 + AITER_CHECK(false, + "no registered kernel for OPUS a8w8 on gfx942"); #else - (void)M; - (void)N; - (void)K; + AITER_CHECK(false, + "opus_gemm_a8w8_launch: module was not built with gfx942 ", + "support for current device ", info.dev); +#endif + return nullptr; + case OpusGfxArch::Gfx1250: +#ifdef OPUS_BUILD_HAS_GFX1250 + AITER_CHECK(false, + "no registered kernel for OPUS a8w8 on gfx1250"); +#else + AITER_CHECK(false, + "opus_gemm_a8w8_launch: module was not built with gfx1250 ", + "support for current device ", info.dev); +#endif + return nullptr; + default: + AITER_CHECK(false, + "opus_gemm_a8w8_launch: unsupported current device ", + info.dev, " with gcnArchName='", info.name, "'"); + } return nullptr; +} + +static OpusA8W8BlockscaleKernel +opus_a8w8_blockscale_kid_dispatch(int kid) +{ + const auto &info = opus_get_arch_info(); + switch (info.arch) + { + case OpusGfxArch::Gfx950: +#ifdef OPUS_BUILD_HAS_GFX950 + return opus_a8w8_blockscale_kid_dispatch_gfx950(kid); +#else + AITER_CHECK(false, + "opus_gemm_a8w8_blockscale_launch: module was not built ", + "with gfx950 support for current device ", info.dev); +#endif + return nullptr; + case OpusGfxArch::Gfx942: +#ifdef OPUS_BUILD_HAS_GFX942 + AITER_CHECK(false, + "no registered kernel for OPUS a8w8_blockscale on gfx942"); +#else + AITER_CHECK(false, + "opus_gemm_a8w8_blockscale_launch: module was not built ", + "with gfx942 support for current device ", info.dev); +#endif + return nullptr; + case OpusGfxArch::Gfx1250: +#ifdef OPUS_BUILD_HAS_GFX1250 + AITER_CHECK(false, + "no registered kernel for OPUS a8w8_blockscale on gfx1250"); +#else + AITER_CHECK(false, + "opus_gemm_a8w8_blockscale_launch: module was not built ", + "with gfx1250 support for current device ", info.dev); #endif + return nullptr; + default: + AITER_CHECK(false, + "opus_gemm_a8w8_blockscale_launch: unsupported current device ", + info.dev, " with gcnArchName='", info.name, "'"); + } + return nullptr; } template -OpusNoscaleKernel opus_dispatch_a8w8(int M, int N, int K) +static OpusA8W8BlockscaleBpreshuffleKernel +opus_a8w8_blockscale_bpreshuffle_kid_dispatch(int kid) { + const auto &info = opus_get_arch_info(); + switch (info.arch) + { + case OpusGfxArch::Gfx950: #ifdef OPUS_BUILD_HAS_GFX950 - return opus_gemm_512x256x256x128_2x4_16x16x128_0x0x0; + return opus_a8w8_blockscale_bpreshuffle_kid_dispatch_gfx950(kid); #else - (void)M; - (void)N; - (void)K; - return nullptr; + AITER_CHECK(false, + "opus_gemm_a8w8_blockscale_bpreshuffle_launch: module was ", + "not built with gfx950 support for current device ", info.dev); +#endif + return nullptr; + case OpusGfxArch::Gfx942: +#ifdef OPUS_BUILD_HAS_GFX942 + return opus_a8w8_blockscale_bpreshuffle_kid_dispatch_gfx942(kid); +#else + AITER_CHECK(false, + "opus_gemm_a8w8_blockscale_bpreshuffle_launch: module was ", + "not built with gfx942 support for current device ", info.dev); #endif + return nullptr; + case OpusGfxArch::Gfx1250: +#ifdef OPUS_BUILD_HAS_GFX1250 + return opus_a8w8_blockscale_bpreshuffle_kid_dispatch_gfx1250(kid); +#else + AITER_CHECK(false, + "opus_gemm_a8w8_blockscale_bpreshuffle_launch: module was ", + "not built with gfx1250 support for current device ", info.dev); +#endif + return nullptr; + default: + AITER_CHECK(false, + "opus_gemm_a8w8_blockscale_bpreshuffle_launch: unsupported ", + "current device ", info.dev, " with gcnArchName='", + info.name, "'"); + } + return nullptr; } -// a16w16 arch routers (gfx950/gfx942 only; gfx1250 uses its own dispatch with workspace). -#if defined(OPUS_BUILD_HAS_GFX950) || defined(OPUS_BUILD_HAS_GFX942) template -OpusA16W16NoscaleKernel opus_dispatch_a16w16(int M, int N, int K, int batch, bool has_bias = false) +static OpusA16W16Kernel +opus_a16w16_kid_dispatch(int kid) { switch (opus_get_gfx_arch()) { #ifdef OPUS_BUILD_HAS_GFX950 case OpusGfxArch::Gfx950: - return opus_dispatch_a16w16_gfx950(M, N, K, batch, has_bias); + return opus_a16w16_kid_dispatch_gfx950(kid); #endif #ifdef OPUS_BUILD_HAS_GFX942 case OpusGfxArch::Gfx942: - return opus_dispatch_a16w16_gfx942(M, N, K, batch, has_bias); + return opus_a16w16_kid_dispatch_gfx942(kid); +#endif +#ifdef OPUS_BUILD_HAS_GFX1250 + case OpusGfxArch::Gfx1250: + return opus_a16w16_kid_dispatch_gfx1250(kid); #endif default: { const auto &info = opus_get_arch_info(); AITER_CHECK(false, - "opus_gemm: a16w16 dispatch via this path is only implemented for " - "gfx950/gfx942; gfx1250 uses a separate dispatch with workspace. " + "opus_gemm_a16w16_launch: no non-workspace dispatch table for " "current device ", info.dev, - " has gcnArchName='", info.name, "'"); + " with gcnArchName='", info.name, "'"); + return nullptr; } } } -template -OpusA16W16NoscaleKernel -opus_a16w16_tune_dispatch(int id) +// Query the direct table separately so an unavailable workspace kid is not +// misclassified as direct. +static bool opus_a16w16_has_non_workspace_kernel(int kid) { switch (opus_get_gfx_arch()) { #ifdef OPUS_BUILD_HAS_GFX950 case OpusGfxArch::Gfx950: - return opus_a16w16_tune_dispatch_gfx950(id); + return opus_a16w16_has_non_workspace_kernel_gfx950(kid); #endif #ifdef OPUS_BUILD_HAS_GFX942 case OpusGfxArch::Gfx942: - return opus_a16w16_tune_dispatch_gfx942(id); + return opus_a16w16_has_non_workspace_kernel_gfx942(kid); +#endif +#ifdef OPUS_BUILD_HAS_GFX1250 + case OpusGfxArch::Gfx1250: + return opus_a16w16_has_non_workspace_kernel_gfx1250(kid); #endif default: { const auto &info = opus_get_arch_info(); AITER_CHECK(false, - "opus_gemm_a16w16_tune: dispatch is only implemented for gfx950/gfx942 " - "via this path; gfx1250 uses a separate dispatch with workspace. " + "opus_gemm_a16w16_launch: no non-workspace dispatch table for " "current device ", info.dev, - " has gcnArchName='", info.name, "'"); + " with gcnArchName='", info.name, "'"); + return false; } } } -#endif // OPUS_BUILD_HAS_GFX950 || OPUS_BUILD_HAS_GFX942 - -// ── opus_gemm() — top-level a16w16 / a8w8 entry ───────────────────────────── -void opus_gemm( - aiter_tensor_t &XQ, - aiter_tensor_t &WQ, - aiter_tensor_t &Y, - std::optional group_layout, - std::optional x_scale, - std::optional w_scale, - std::optional bias) +// Query the current architecture's generated workspace table. +static bool opus_a16w16_has_workspace_kernel(int kid) { - aiter_detail::g_aiter_can_throw = true; - AITER_CHECK(XQ.dim() == 3, "XQ must be 3D [batch, M, K]"); - AITER_CHECK(WQ.dim() == 3, "WQ must be 3D [batch, N, K]"); - AITER_CHECK(Y.dim() == 3, "Y must be 3D [batch, M, N]"); - - int M = XQ.size(1); - int N = WQ.size(1); - int K = XQ.size(2); - - bool has_scale = x_scale.has_value() && w_scale.has_value(); - - if (XQ.dtype() == AITER_DTYPE_fp8) + switch (opus_get_gfx_arch()) { - // a8w8 / a8w8_scale launchers are gfx950-only today and don't yet flow through the arch-routed - // dispatcher (they pick a single har... - const auto &arch_info = opus_get_arch_info(); #ifdef OPUS_BUILD_HAS_GFX950 - AITER_CHECK(arch_info.arch == OpusGfxArch::Gfx950, - "opus_gemm: a8w8 path is only implemented for gfx950 today; " - "current device ", arch_info.dev, - " has gcnArchName='", arch_info.name, - "'. Other archs will be added as more pipelines land."); - // a8w8 / a8w8_scale launchers do not consume bias yet; reject up front - // rather than silently dropping it. - AITER_CHECK(!bias.has_value(), - "opus_gemm: bias is not supported on a8w8 / a8w8_scale paths"); - if (has_scale) - { - AITER_CHECK(Y.dtype() == AITER_DTYPE_fp32, - "opus_gemm a8w8_scale only supports fp32 output"); - opus_dispatch_scale(M, N, K)(XQ, WQ, Y, x_scale, w_scale); - } - else - { - AITER_CHECK(Y.dtype() == AITER_DTYPE_fp32, - "opus_gemm a8w8 no-scale only supports fp32 output"); - opus_dispatch_a8w8(M, N, K)(XQ, WQ, Y); - } -#else - AITER_CHECK(false, - "opus_gemm: a8w8 path requires module_deepgemm_opus to be " - "built with OPUS_BUILD_HAS_GFX950; current device ", - arch_info.dev, " has gcnArchName='", arch_info.name, "'"); + case OpusGfxArch::Gfx950: + return opus_a16w16_has_workspace_kernel_gfx950(kid); +#endif +#ifdef OPUS_BUILD_HAS_GFX942 + case OpusGfxArch::Gfx942: + return opus_a16w16_has_workspace_kernel_gfx942(kid); #endif - } - else if (XQ.dtype() == AITER_DTYPE_bf16) - { - // Tuned-lookup-then-heuristic dispatch. splitK=0 = "launcher decides". - int batch = XQ.size(0); - const bool has_bias = bias.has_value(); #ifdef OPUS_BUILD_HAS_GFX1250 - if (opus_get_gfx_arch() == OpusGfxArch::Gfx1250) - { - // A tuned pre-compiled (.co) winner for this shape, if there is one, is - // taken first: that family needs no workspace, so a hit skips the - // hipMalloc / hipDeviceSynchronize / hipFree below entirely. It cannot - // serve bias (no epilogue for it) or a non-bf16 Y (it stores C straight - // out with no reduce kernel to cast), so those shapes fall through to the - // split-K path even when the CSV named a .co kid. - if (!has_bias && Y.dtype() == AITER_DTYPE_bf16) - { - if (auto co_fn = opus_a16w16_co_dispatch_gfx1250(M, N, K)) - { - co_fn(XQ, WQ, Y, bias, 0); - return; - } - } - // Otherwise: every gfx1250 split-K kid needs a workspace. The heuristic - // dispatch returns a 6-arg function pointer (with workspace). We allocate - // a temporary workspace here for the auto/heuristic path. For the tuned - // path, Python allocates via torch.empty. - auto fn = opus_dispatch_a16w16_gfx1250(M, N, K, batch, has_bias); - int padded_M = ((M + 63) / 64) * 64; - int padded_N = ((N + 63) / 64) * 64; - size_t ws_elems = (size_t)16 * padded_M * padded_N; - size_t ws_bytes = ws_elems * sizeof(bf16_t); - void* ws_ptr = nullptr; - HIP_CALL(hipMalloc(&ws_ptr, ws_bytes)); - aiter_tensor_t ws_tensor{}; - ws_tensor.ptr = ws_ptr; - ws_tensor.numel_ = ws_elems; - ws_tensor.ndim = 1; - ws_tensor.shape[0] = (int64_t)ws_elems; - ws_tensor.strides[0] = 1; - ws_tensor.dtype_ = AITER_DTYPE_bf16; - ws_tensor.device_id = 0; - fn(XQ, WQ, Y, ws_tensor, bias, 0); - HIP_CALL(hipDeviceSynchronize()); - HIP_CALL(hipFree(ws_ptr)); - } - else + case OpusGfxArch::Gfx1250: + return opus_a16w16_has_workspace_kernel_gfx1250(kid); #endif + default: { -#if defined(OPUS_BUILD_HAS_GFX950) || defined(OPUS_BUILD_HAS_GFX942) - if (Y.dtype() == AITER_DTYPE_bf16) - { - opus_dispatch_a16w16(M, N, K, batch, has_bias)(XQ, WQ, Y, bias, 0); - } - else if (Y.dtype() == AITER_DTYPE_fp32) - { - opus_dispatch_a16w16(M, N, K, batch, has_bias)(XQ, WQ, Y, bias, 0); - } - else - { - AITER_CHECK(false, "opus_gemm a16w16: unsupported output dtype, expected bf16 or fp32"); - } -#else - AITER_CHECK(false, "opus_gemm: no a16w16 dispatch available for this arch"); -#endif + const auto &info = opus_get_arch_info(); + AITER_CHECK(false, + "opus_gemm_a16w16_launch: no workspace dispatch table for device ", + info.dev, " with gcnArchName='", info.name, "'"); + return false; } } - else - { - AITER_CHECK(false, "opus_gemm: unsupported input dtype, expected fp8 or bf16"); - } } - -// opus_gemm_a16w16_tune() — id-based tune entry. - -// splitk kids: gfx950 [200,300) + nooob [1200,1300); gfx942 [10200, 10300). -static constexpr int OPUS_SPLITK_KID_MIN = 200; -static constexpr int OPUS_SPLITK_KID_MAX = 300; -static constexpr int OPUS_GFX942_KID_OFFSET = 10000; -static constexpr int OPUS_GFX942_SPLITK_KID_MAX = 300; -// gfx1250 split-K kids come in TWO bands, because the pre-compiled .co family -// sits between them: -// [20000, 20100) plain cluster/TDM (fp32 workspace + reduce) -// [20100, 21000) clusterlaunch multicast ws -// [21000, 27000) a16w16_4wave_co / _wl_co -- NOT split-K, see below -// [27000, 30000) FUSED single-kernel family, currently unregistered -// (GFX1250_SPLITK_FUSE_ENABLED in opus_gemm_common.py) while -// its pipeline is being fixed -// The upper band stays covered even while empty so a family landing there gets -// the workspace dispatch. All gfx1250 split-K kids use the lookup ABI -// and fold bias. -// -// KEEP IN LOCKSTEP with _SPLITK_KID_RANGES in aiter/ops/opus/gemm_op_a16w16.py. -static constexpr int OPUS_GFX1250_SPLITK_KID_MIN = 20000; -static constexpr int OPUS_GFX1250_SPLITK_KID_MAX = 21000; -static constexpr int OPUS_GFX1250_SPLITK_FUSE_KID_MIN = 27000; -static constexpr int OPUS_GFX1250_SPLITK_FUSE_KID_MAX = 30000; -// gfx1250 symmetric 4-wave compute, device side loaded from a pre-built .co -// (a16w16_4wave_co). Not a split-K family: no workspace, no split_k, no reduce -// kernel, and bf16 C written straight out. It therefore has its OWN dispatch -// (opus_a16w16_co_tune_dispatch_gfx1250) whose launcher takes the ordinary -// 5-arg a16w16 signature, and is deliberately absent from both -// opus_kid_is_splitk() and opus_kid_is_gfx1250_splitk() -- which is exactly why -// the split-K band above had to be cut in two rather than span this range. A -// split-K .co variant would join those instead of this band. -static constexpr int OPUS_GFX1250_CO_KID_MIN = 21000; -static constexpr int OPUS_GFX1250_CO_KID_MAX = 27000; -// SB a16w16 kids: gfx950 [4,10) + mirrors at +1000/.../+7000. -static constexpr int OPUS_A16W16_SB_KID_MIN = 4; -static constexpr int OPUS_A16W16_SB_KID_MAX = 10; -// Persistent a16w16 kids: compact [300, 316) = 4 tiles × 4 cpol groups. -static constexpr int OPUS_PERSISTENT_KID_MIN = 300; -static constexpr int OPUS_PERSISTENT_KID_MAX = 316; -// Mono-tile a16w16 kids: [1400, 1500). Mono-tile is intrinsically non-OOB -// (no tail handling in the kernel body), so kids land in the >=1000 band -// directly — there is no base/nooob mirror split for this family. See -// opus_gemm_common.py :: a16w16_mono_tile_kernels_list. -static constexpr int OPUS_MONO_TILE_KID_MIN = 1400; -static constexpr int OPUS_MONO_TILE_KID_MAX = 1500; -// non-OOB kid offset -static constexpr int OPUS_NOOOB_KID_OFFSET = 1000; - -// Two disjoint bands with the pre-compiled .co family in between; see the kid -// map above. Defined before opus_kid_is_splitk() because that one calls it. -static inline bool opus_kid_is_gfx1250_splitk(int kid) -{ - return (kid >= OPUS_GFX1250_SPLITK_KID_MIN && kid < OPUS_GFX1250_SPLITK_KID_MAX) || - (kid >= OPUS_GFX1250_SPLITK_FUSE_KID_MIN && - kid < OPUS_GFX1250_SPLITK_FUSE_KID_MAX); -} - -static inline bool opus_kid_is_splitk(int kid) -{ - return (kid >= OPUS_SPLITK_KID_MIN && kid < OPUS_SPLITK_KID_MAX) || - (kid >= OPUS_SPLITK_KID_MIN + OPUS_NOOOB_KID_OFFSET && - kid < OPUS_SPLITK_KID_MAX + OPUS_NOOOB_KID_OFFSET) || - (kid >= OPUS_SPLITK_KID_MIN + OPUS_GFX942_KID_OFFSET && - kid < OPUS_GFX942_SPLITK_KID_MAX + OPUS_GFX942_KID_OFFSET) || - opus_kid_is_gfx1250_splitk(kid); -} - -static inline bool opus_kid_is_gfx1250_co(int kid) +static OpusA16W16WorkspaceKernel +opus_a16w16_workspace_dispatch(int kid) { - return kid >= OPUS_GFX1250_CO_KID_MIN && kid < OPUS_GFX1250_CO_KID_MAX; -} - -static inline bool opus_kid_is_a16w16_sb(int kid) -{ - // SB a16w16 kid bases: 0/1000/2000/.../7000 + [4,10) (cpol mirrors). - for (int base : {0, 1000, 2000, 3000, 4000, 5000, 6000, 7000}) + switch (opus_get_gfx_arch()) { - if (kid >= base + OPUS_A16W16_SB_KID_MIN && kid < base + OPUS_A16W16_SB_KID_MAX) - return true; +#ifdef OPUS_BUILD_HAS_GFX950 + case OpusGfxArch::Gfx950: + return opus_a16w16_workspace_dispatch_gfx950(kid); +#endif +#ifdef OPUS_BUILD_HAS_GFX942 + case OpusGfxArch::Gfx942: + return opus_a16w16_workspace_dispatch_gfx942(kid); +#endif +#ifdef OPUS_BUILD_HAS_GFX1250 + case OpusGfxArch::Gfx1250: + return opus_a16w16_workspace_dispatch_gfx1250(kid); +#endif + default: + { + const auto &info = opus_get_arch_info(); + AITER_CHECK(false, + "opus_gemm_a16w16_launch: no workspace dispatch table for device ", + info.dev, " with gcnArchName='", info.name, "'"); + return nullptr; + } } - return false; -} - -static inline bool opus_kid_is_persistent(int kid) -{ - return (kid >= OPUS_PERSISTENT_KID_MIN && kid < OPUS_PERSISTENT_KID_MAX) || - (kid >= OPUS_PERSISTENT_KID_MIN + OPUS_NOOOB_KID_OFFSET && - kid < OPUS_PERSISTENT_KID_MAX + OPUS_NOOOB_KID_OFFSET); -} - -static inline bool opus_kid_is_mono_tile(int kid) -{ - // Mono-tile lives entirely in the non-OOB band [1400, 1500); no mirror. - return kid >= OPUS_MONO_TILE_KID_MIN && kid < OPUS_MONO_TILE_KID_MAX; -} - -static inline bool opus_kid_is_gfx942_splitk(int kid) -{ - return kid >= OPUS_SPLITK_KID_MIN + OPUS_GFX942_KID_OFFSET && - kid < OPUS_GFX942_SPLITK_KID_MAX + OPUS_GFX942_KID_OFFSET; -} - -static inline bool opus_kid_supports_bias(int kid) -{ - // persistent and mono-tile do not support bias (kargs lacks - // ptr_bias/stride_bias_batch; launchers reject non-empty bias up front). - // gfx942 splitk/SB silently ignored bias; exclude explicitly to surface - // misuse as a clear error. - // gfx1250 cluster_tdm_splitk_ws DOES support bias (the reduce kernel folds - // it once, like gfx950 flatmm_splitk). - return (opus_kid_is_a16w16_sb(kid) || opus_kid_is_splitk(kid)) - && !opus_kid_is_gfx942_splitk(kid); } -void opus_gemm_a16w16_tune( +// Validate A16W16 inputs, then call the matching generated launcher table. +static void opus_gemm_a16w16_launch_impl( aiter_tensor_t &XQ, aiter_tensor_t &WQ, aiter_tensor_t &Y, std::optional bias, std::optional workspace, - int kernelId, - int splitK) + int kid, + int split_k) { aiter_detail::g_aiter_can_throw = true; - AITER_CHECK(XQ.dim() == 3, "XQ must be 3D [batch, M, K]"); - AITER_CHECK(WQ.dim() == 3, "WQ must be 3D [batch, N, K]"); - AITER_CHECK(Y.dim() == 3, "Y must be 3D [batch, M, N]"); + + AITER_CHECK(XQ.is_gpu() && WQ.is_gpu() && Y.is_gpu(), + "opus_gemm_a16w16_launch: XQ, WQ, and Y must be GPU tensors"); + AITER_CHECK(XQ.device_id == WQ.device_id && XQ.device_id == Y.device_id, + "opus_gemm_a16w16_launch: XQ/WQ/Y device ids must match (got ", + XQ.device_id, "/", WQ.device_id, "/", Y.device_id, ")"); + if (bias.has_value()) + { + AITER_CHECK(bias->is_gpu() && bias->device_id == XQ.device_id, + "opus_gemm_a16w16_launch: bias device ", bias->device_id, + " must match XQ device ", XQ.device_id); + } + if (workspace.has_value()) + { + AITER_CHECK(workspace->is_gpu() && workspace->device_id == XQ.device_id, + "opus_gemm_a16w16_launch: workspace device ", + workspace->device_id, " must match input device ", XQ.device_id); + } + AITER_CHECK(XQ.dim() == 3, + "opus_gemm_a16w16_launch: XQ must be 3D [batch, M, K]"); + AITER_CHECK(WQ.dim() == 3, + "opus_gemm_a16w16_launch: WQ must be 3D [batch, N, K]"); + AITER_CHECK(Y.dim() == 3, + "opus_gemm_a16w16_launch: Y must be 3D [batch, M, N]"); AITER_CHECK(XQ.dtype() == WQ.dtype(), - "XQ and WQ should have the same dtype!"); - // Early-gate non-bias-capable kids for a clean error before launcher entry. - AITER_CHECK(!bias.has_value() || opus_kid_supports_bias(kernelId), - "opus_gemm_a16w16_tune: bias is currently only supported on " - "a16w16 split-barrier kids [", OPUS_A16W16_SB_KID_MIN, ", ", - OPUS_A16W16_SB_KID_MAX, ") or a16w16_flatmm_splitk kids [", - OPUS_SPLITK_KID_MIN, ", ", OPUS_SPLITK_KID_MAX, - "); got kid=", kernelId); + "opus_gemm_a16w16_launch: XQ and WQ dtype must match"); if (XQ.dtype() == AITER_DTYPE_bf16) { -#ifdef OPUS_BUILD_HAS_GFX1250 - // gfx1250 pre-compiled (.co) kids. Checked first because they are neither - // split-K (no workspace to pass) nor reachable through the shared - // opus_a16w16_tune_dispatch table (their launcher has its own signature). - if (opus_kid_is_gfx1250_co(kernelId)) + const bool uses_workspace = opus_a16w16_has_workspace_kernel(kid); + const bool has_non_workspace = opus_a16w16_has_non_workspace_kernel(kid); + AITER_CHECK(!(uses_workspace && has_non_workspace), + "opus_gemm_a16w16_launch: kid ", kid, + " appears in both workspace and non-workspace launch tables"); + if (!uses_workspace && !has_non_workspace) { - AITER_CHECK(Y.dtype() == AITER_DTYPE_bf16, - "opus_gemm_a16w16_tune: gfx1250 .co kid writes bf16 C " - "directly (no reduce kernel to cast), so Y must be bf16"); - opus_a16w16_co_tune_dispatch_gfx1250(kernelId)(XQ, WQ, Y, bias, splitK); + const auto &info = opus_get_arch_info(); + AITER_CHECK(false, + "opus_gemm_a16w16_launch: kid ", kid, + " is not compiled in the current OPUS module for device ", + info.dev, " with gcnArchName='", info.name, "'"); } - else -#endif - // All splitk kids (gfx950/gfx942/gfx1250) force : the main kernel - // writes an fp32 workspace and a reduce kernel casts it to Y.dtype() at - // runtime (gfx1250 cluster_tdm_splitk_ws now follows this same pattern). - if (opus_kid_is_splitk(kernelId)) + if (uses_workspace) { + AITER_CHECK(workspace.has_value(), + "opus_gemm_a16w16_launch: workspace kid ", kid, + " requires a workspace tensor"); AITER_CHECK(Y.dtype() == AITER_DTYPE_bf16 || Y.dtype() == AITER_DTYPE_fp32, - "opus_gemm_a16w16_tune splitk kid requires bf16 or fp32 Y " + "opus_gemm_a16w16_launch: workspace kid requires bf16 or fp32 Y " "(reduce kernel writes the correct dtype)"); -#ifdef OPUS_BUILD_HAS_GFX1250 - if (opus_kid_is_gfx1250_splitk(kernelId)) + opus_a16w16_workspace_dispatch(kid)( + XQ, WQ, Y, workspace.value(), bias, split_k); + } + else + { + AITER_CHECK(!workspace.has_value(), + "opus_gemm_a16w16_launch: non-workspace kid ", kid, + " requires workspace=None"); + if (Y.dtype() == AITER_DTYPE_bf16) { - AITER_CHECK(workspace.has_value(), - "gfx1250 split-K kids require a workspace tensor " - "(allocated via torch.empty on the Python side)"); - auto& ws = workspace.value(); - opus_a16w16_tune_dispatch_gfx1250(kernelId)(XQ, WQ, Y, ws, bias, splitK); + opus_a16w16_kid_dispatch(kid)(XQ, WQ, Y, bias, split_k); + } + else if (Y.dtype() == AITER_DTYPE_fp32) + { + opus_a16w16_kid_dispatch(kid)(XQ, WQ, Y, bias, split_k); } else -#endif { -#if defined(OPUS_BUILD_HAS_GFX950) || defined(OPUS_BUILD_HAS_GFX942) - opus_a16w16_tune_dispatch(kernelId)(XQ, WQ, Y, bias, splitK); -#else - AITER_CHECK(false, "opus_gemm_a16w16_tune: non-gfx1250 splitk dispatch unavailable"); -#endif + AITER_CHECK(false, + "opus_gemm_a16w16_launch: unsupported output dtype, expected bf16 or fp32"); } } - else if (Y.dtype() == AITER_DTYPE_bf16) - { -#if defined(OPUS_BUILD_HAS_GFX950) || defined(OPUS_BUILD_HAS_GFX942) - opus_a16w16_tune_dispatch(kernelId)(XQ, WQ, Y, bias, splitK); -#else - AITER_CHECK(false, "opus_gemm_a16w16_tune: non-splitk bf16 dispatch unavailable for this arch"); -#endif - } - else if (Y.dtype() == AITER_DTYPE_fp32) - { -#if defined(OPUS_BUILD_HAS_GFX950) || defined(OPUS_BUILD_HAS_GFX942) - opus_a16w16_tune_dispatch(kernelId)(XQ, WQ, Y, bias, splitK); -#else - AITER_CHECK(false, "opus_gemm_a16w16_tune: non-splitk fp32 dispatch unavailable for this arch"); -#endif - } - else - { - AITER_CHECK(false, - "opus_gemm_a16w16_tune: unsupported output dtype, expected bf16 or fp32"); - } } else { AITER_CHECK(false, - "opus_gemm_a16w16_tune: unsupported input dtype ", + "opus_gemm_a16w16_launch: unsupported input dtype ", AiterDtype_to_str(XQ.dtype()), ", expected bf16"); } } -void opus_gemm_a8w8_blockscale_bpreshuffle_tune( +void opus_gemm_a16w16_launch( aiter_tensor_t &XQ, aiter_tensor_t &WQ, - std::optional x_scale, - std::optional w_scale, aiter_tensor_t &Y, - int kernelId) + std::optional bias, + std::optional workspace, + int kid, + int split_k) { - aiter_detail::g_aiter_can_throw = true; - const auto &arch_info = opus_get_arch_info(); - AITER_CHECK(arch_info.arch == OpusGfxArch::Gfx942, - "opus_gemm_a8w8_blockscale_bpreshuffle_tune is only implemented " - "for gfx942 today; current device ", arch_info.dev, - " has gcnArchName='", arch_info.name, "'"); - AITER_CHECK(XQ.dtype() == AITER_DTYPE_fp8 && WQ.dtype() == AITER_DTYPE_fp8, - "opus_gemm_a8w8_blockscale_bpreshuffle_tune expects fp8 XQ/WQ"); - AITER_CHECK(Y.dtype() == AITER_DTYPE_bf16, - "opus_gemm_a8w8_blockscale_bpreshuffle_tune expects bf16 Y"); - AITER_CHECK(x_scale.has_value() && w_scale.has_value(), - "opus_gemm_a8w8_blockscale_bpreshuffle_tune requires x_scale and w_scale"); - -#ifdef OPUS_BUILD_HAS_GFX942 - opus_a8w8_tune_dispatch_gfx942(kernelId)(XQ, WQ, Y, x_scale, w_scale); -#else - AITER_CHECK(false, - "module_deepgemm_opus was not built with OPUS_BUILD_HAS_GFX942"); -#endif + opus_gemm_a16w16_launch_impl( + XQ, WQ, Y, bias, workspace, kid, split_k); } -// ────────────────────────────────────────────────────────────────────────────── -// Splitk fp32 workspace: per-stream owner. -// -// Each splitk launcher (generated by gen_instances.py) needs a stable -// `opus_splitk_ws_handle*` to feed into both the main kernel and the reduce -// kernel; captured HIP graphs bake in that pointer. Previously this was a -// `static thread_local` slot — one handle per CPU thread — but under -// vLLM/sglang-style TBO two CPU threads drive two streams concurrently, and -// each captured graph needs its own buffer pointer baked in. The TLS form -// also tripped the in-capture grow guard on the second thread. -// -// Now we own the handle by stream: a process-global mutex-protected map -// keyed by hipStream_t. Eager: lazy-create on first lookup. Capture: caller -// must pre-register the handle via opus_gemm_workspace_init(), otherwise the -// lookup throws (cleaner than the prior SIGABRT). The framework calls -// opus_gemm_workspace_init() once per TBO stream eagerly before capture. -// -// Teardown: entries are held for the process lifetime unless explicitly freed -// via opus_gemm_workspace_release() (current stream) or -// opus_gemm_workspace_release_all() (all streams). Both run in eager mode and -// synchronize before freeing. -namespace { -struct SplitkWsRegistry { - std::mutex mu; - struct Owner { - opus_splitk_ws_handle* host; - opus_splitk_ws_handle* device; - }; - std::unordered_map map; -}; -SplitkWsRegistry& splitk_ws_registry() +static void opus_check_a8_family_tensors( + const char* entry, + const aiter_tensor_t &XQ, + const aiter_tensor_t &WQ, + const aiter_tensor_t &Y) { - static SplitkWsRegistry r; - return r; + AITER_CHECK(XQ.is_gpu() && WQ.is_gpu() && Y.is_gpu(), + entry, ": XQ, WQ, and Y must be GPU tensors"); + AITER_CHECK(XQ.device_id == WQ.device_id && XQ.device_id == Y.device_id, + entry, ": XQ/WQ/Y device ids must match (got ", + XQ.device_id, "/", WQ.device_id, "/", Y.device_id, ")"); + int current_device = -1; + HIP_CALL(hipGetDevice(¤t_device)); + AITER_CHECK(current_device == XQ.device_id, + entry, ": current HIP device ", current_device, + " does not match tensor device ", XQ.device_id); + AITER_CHECK(XQ.dtype() == AITER_DTYPE_fp8 && WQ.dtype() == AITER_DTYPE_fp8, + entry, ": expected fp8 XQ/WQ, got ", + AiterDtype_to_str(XQ.dtype()), "/", + AiterDtype_to_str(WQ.dtype())); } -} // anonymous -opus_splitk_ws_handle* opus_splitk_ws_get(hipStream_t s, bool allow_create) +// A8W8 entry points perform common checks before exact-kid table lookup. +static void opus_check_a8_scale_devices( + const char* entry, + const aiter_tensor_t &XQ, + const aiter_tensor_t &x_scale, + const aiter_tensor_t &w_scale) { - auto& R = splitk_ws_registry(); - std::lock_guard g(R.mu); - auto it = R.map.find(s); - if (it != R.map.end()) return it->second->host; - AITER_CHECK(allow_create, - "splitk workspace not initialized for the current CUDA stream. " - "Call aiter.opus_gemm_workspace_init() inside " - "`with torch.cuda.stream(s):` (and warm with the largest " - "expected gemm) before HIP graph capture."); - auto* owner = new SplitkWsRegistry::Owner{}; - opus_splitk_ws_handle* h = nullptr; -#ifdef OPUS_BUILD_HAS_GFX950 - // gfx950 launchers feed the host handle STRAIGHT to the kernel, which - // dereferences ptr/bytes on the device -- so it must be device-visible - // pinned/coherent host memory. - HIP_CALL(hipHostMalloc(reinterpret_cast(&h), - sizeof(opus_splitk_ws_handle), - hipHostMallocCoherent)); - h->ptr = nullptr; - h->bytes = 0; -#else - // gfx942/gfx1250 read a device mirror (opus_splitk_ws_sync_to_device); the - // device never dereferences this host handle. So plain host memory suffices - // and we avoid pinned/coherent allocations entirely -- the OS reclaims plain - // host memory at process exit with no dependency on HIP's pinned-memory - // teardown (which can wedge fragile drivers and hang a subsequent process). - h = new opus_splitk_ws_handle{nullptr, 0}; -#endif - owner->host = h; - owner->device = nullptr; - R.map[s] = owner; - return h; + AITER_CHECK(x_scale.is_gpu() && w_scale.is_gpu(), + entry, ": x_scale and w_scale must be GPU tensors"); + AITER_CHECK(x_scale.device_id == XQ.device_id && + w_scale.device_id == XQ.device_id, + entry, ": scale tensor device ids must match XQ.device_id=", + XQ.device_id, " (got ", x_scale.device_id, "/", + w_scale.device_id, ")"); } -const opus_splitk_ws_handle* opus_splitk_ws_device_handle(hipStream_t s, bool allow_create) +void opus_gemm_a8w8_launch( + aiter_tensor_t &XQ, + aiter_tensor_t &WQ, + aiter_tensor_t &Y, + int kid) { - (void)opus_splitk_ws_get(s, allow_create); - auto& R = splitk_ws_registry(); - std::lock_guard g(R.mu); - auto it = R.map.find(s); - AITER_CHECK(it != R.map.end(), "splitk workspace not initialized for the current CUDA stream."); - if (it->second->device == nullptr) - { - AITER_CHECK(allow_create, - "splitk workspace device handle not initialized for the current CUDA stream. " - "Warm the opus gfx942 splitK launcher eagerly before HIP graph capture."); - HIP_CALL(hipMalloc(reinterpret_cast(&it->second->device), - sizeof(opus_splitk_ws_handle))); - HIP_CALL(hipMemcpy(it->second->device, - it->second->host, - sizeof(opus_splitk_ws_handle), - hipMemcpyHostToDevice)); - } - return it->second->device; + aiter_detail::g_aiter_can_throw = true; + constexpr const char* entry = "opus_gemm_a8w8_launch"; + opus_check_a8_family_tensors(entry, XQ, WQ, Y); + AITER_CHECK(Y.dtype() == AITER_DTYPE_fp32, + entry, ": expected fp32 Y, got ", + AiterDtype_to_str(Y.dtype())); + opus_a8w8_kid_dispatch(kid)(XQ, WQ, Y); } -void opus_splitk_ws_sync_to_device(hipStream_t s) +void opus_gemm_a8w8_blockscale_launch( + aiter_tensor_t &XQ, + aiter_tensor_t &WQ, + aiter_tensor_t &Y, + aiter_tensor_t &x_scale, + aiter_tensor_t &w_scale, + int kid) { - auto& R = splitk_ws_registry(); - std::lock_guard g(R.mu); - auto it = R.map.find(s); - AITER_CHECK(it != R.map.end(), "splitk workspace not initialized for the current CUDA stream."); - if (it->second->device == nullptr) - { - HIP_CALL(hipMalloc(reinterpret_cast(&it->second->device), - sizeof(opus_splitk_ws_handle))); - } - HIP_CALL(hipMemcpy(it->second->device, - it->second->host, - sizeof(opus_splitk_ws_handle), - hipMemcpyHostToDevice)); + aiter_detail::g_aiter_can_throw = true; + constexpr const char* entry = "opus_gemm_a8w8_blockscale_launch"; + opus_check_a8_family_tensors(entry, XQ, WQ, Y); + opus_check_a8_scale_devices(entry, XQ, x_scale, w_scale); + AITER_CHECK(Y.dtype() == AITER_DTYPE_fp32, + entry, ": expected fp32 Y, got ", + AiterDtype_to_str(Y.dtype())); + opus_a8w8_blockscale_kid_dispatch(kid)( + XQ, WQ, Y, x_scale, w_scale); } -void opus_gemm_workspace_init() +void opus_gemm_a8w8_blockscale_bpreshuffle_launch( + aiter_tensor_t &XQ, + aiter_tensor_t &WQ, + aiter_tensor_t &x_scale, + aiter_tensor_t &w_scale, + aiter_tensor_t &Y, + int kid) { - hipStream_t s = aiter::getCurrentHIPStream(); - hipStreamCaptureStatus cap = hipStreamCaptureStatusNone; - HIP_CALL(hipStreamIsCapturing(s, &cap)); - AITER_CHECK(cap == hipStreamCaptureStatusNone, - "opus_gemm_workspace_init must be called in eager mode " - "(not inside HIP graph capture)."); - (void)opus_splitk_ws_get(s, /*allow_create=*/true); -} + aiter_detail::g_aiter_can_throw = true; + constexpr const char* entry = + "opus_gemm_a8w8_blockscale_bpreshuffle_launch"; + opus_check_a8_family_tensors(entry, XQ, WQ, Y); + opus_check_a8_scale_devices(entry, XQ, x_scale, w_scale); -// Free everything a single Owner holds: the GPU workspace data buffer (owned via -// the host handle's `ptr`), the host coherent handle itself, and the device -// mirror. Caller must hold the registry mutex and must have synchronized any -// in-flight work that could still reference the buffer. -static void opus_splitk_ws_free_owner_locked(SplitkWsRegistry::Owner* owner) -{ - if (owner == nullptr) return; - if (owner->host != nullptr) + if (Y.dtype() == AITER_DTYPE_bf16) { - if (owner->host->ptr != nullptr) - { - HIP_CALL(hipFree(owner->host->ptr)); - owner->host->ptr = nullptr; - owner->host->bytes = 0; - } -#ifdef OPUS_BUILD_HAS_GFX950 - HIP_CALL(hipHostFree(owner->host)); // paired with hipHostMalloc above -#else - delete owner->host; // paired with plain `new` for the gfx942/gfx1250 path -#endif - owner->host = nullptr; + opus_a8w8_blockscale_bpreshuffle_kid_dispatch(kid)( + XQ, WQ, x_scale, w_scale, Y); } - if (owner->device != nullptr) + else if (Y.dtype() == AITER_DTYPE_fp32) { - HIP_CALL(hipFree(owner->device)); - owner->device = nullptr; + opus_a8w8_blockscale_bpreshuffle_kid_dispatch(kid)( + XQ, WQ, x_scale, w_scale, Y); } - delete owner; -} - -// Release the splitk workspace (buffer + handles + registry entry) for the -// CURRENT stream. Safe to call when the stream was never registered (no-op). -// Must run in eager mode; frees are stream-capture-illegal. -void opus_gemm_workspace_release() -{ - hipStream_t s = aiter::getCurrentHIPStream(); - hipStreamCaptureStatus cap = hipStreamCaptureStatusNone; - HIP_CALL(hipStreamIsCapturing(s, &cap)); - AITER_CHECK(cap == hipStreamCaptureStatusNone, - "opus_gemm_workspace_release must be called in eager mode " - "(not inside HIP graph capture)."); - // Drain the stream so no in-flight kernel references the buffer being freed. - HIP_CALL(hipStreamSynchronize(s)); - auto& R = splitk_ws_registry(); - std::lock_guard g(R.mu); - auto it = R.map.find(s); - if (it == R.map.end()) return; - opus_splitk_ws_free_owner_locked(it->second); - R.map.erase(it); -} - -// Release the splitk workspace for ALL registered streams and clear the -// registry. Intended for explicit teardown (e.g. before a framework tears down -// its stream pool). Must run in eager mode. -void opus_gemm_workspace_release_all() -{ - auto& R = splitk_ws_registry(); - std::lock_guard g(R.mu); - if (R.map.empty()) return; - // Drain all device work before freeing any buffer (buffers belong to many - // streams; a single device sync covers them all). - HIP_CALL(hipDeviceSynchronize()); - for (auto& kv : R.map) + else { - opus_splitk_ws_free_owner_locked(kv.second); + AITER_CHECK(false, + entry, ": unsupported Y dtype ", + AiterDtype_to_str(Y.dtype()), "; expected bf16 or fp32"); } - R.map.clear(); } #endif // !__HIP_DEVICE_COMPILE__ diff --git a/csrc/opus_gemm/opus_gemm_a8w8_tune.py b/csrc/opus_gemm/opus_gemm_a8w8_tune.py new file mode 100644 index 0000000000..7f6fad8dcc --- /dev/null +++ b/csrc/opus_gemm/opus_gemm_a8w8_tune.py @@ -0,0 +1,358 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. +"""Tune plain OPUS FP8 GEMM; CSV scaleAB selects no-scale or blockscale.""" + +import argparse +import math +from pathlib import Path +from typing import ClassVar + +import pandas as pd +import torch + +from aiter import dtypes +from aiter.ops.opus import opus_gemm +from aiter.utility.base_tuner import GemmCommonTuner +from aiter.utility.mp_tuner import mp_tuner +from csrc.opus_gemm.opus_gemm_common import ( + canonical_output_dtype, + get_kernel_instance, + kernels_list, +) + +_SUPPORTED_GFX = "gfx950" + + +def candidate_kids_for_shape(gfx, m, n, k, scale_ab, outdtype="fp32"): + """Return registered plain A8W8 kids whose launch constraints fit the shape.""" + if min(m, n, k) <= 0 or k % 2: + return [] + family = "a8w8_blockscale" if scale_ab else "a8w8" + candidates = [] + for kid in sorted(kernels_list): + instance = get_kernel_instance(gfx, family, kid, outdtype) + if instance is None: + continue + # These pipelines prime two K tiles and advance in pairs. + loops = (k + instance.B_K - 1) // instance.B_K + if loops < 2 or loops % 2: + continue + if not instance.has_oob and any( + size % tile + for size, tile in zip((m, n, k), (instance.B_M, instance.B_N, instance.B_K)) + ): + continue + if scale_ab and any( + size % group + for size, group in zip( + (m, n, k), (instance.GROUP_M, instance.GROUP_N, instance.GROUP_K) + ) + ): + continue + candidates.append(kid) + return candidates + + +def generate_data(m, n, k, kid, *, device): + instance = kernels_list[kid] + generator = torch.Generator(device=device).manual_seed(0) + x = torch.randn((m, k), device=device, generator=generator).to(torch.float8_e4m3fn) + w = torch.randn((n, k), device=device, generator=generator).to(torch.float8_e4m3fn) + x_scale = w_scale = None + if instance.GROUP_K: + x_scale = 0.5 + torch.rand( + (m // instance.GROUP_M, k // instance.GROUP_K), + device=device, + generator=generator, + ) + w_scale = 0.5 + torch.rand( + (n // instance.GROUP_N, k // instance.GROUP_K), + device=device, + generator=generator, + ) + return { + "x": x, + "w": w, + "out": torch.empty((m, n), device=device, dtype=torch.float32), + "x_scale": x_scale, + "w_scale": w_scale, + } + + +def run_torch(x, w, x_scale, w_scale): + inputs = [] + for tensor, scale in ((x, x_scale), (w, w_scale)): + value = tensor.float() + if scale is not None: + value = value * scale.repeat_interleave( + tensor.shape[0] // scale.shape[0], dim=0 + ).repeat_interleave(tensor.shape[1] // scale.shape[1], dim=1) + inputs.append(value) + return inputs[0] @ inputs[1].T + + +def run_bench(x, w, out, x_scale, w_scale, kid): + opus_gemm(x, w, out, kid=kid, x_scale=x_scale, w_scale=w_scale) + return out + + +def compare_outputs(ref, out, **kwargs): + from aiter.test_common import checkAllclose + + err = checkAllclose(ref, out, rtol=1e-2, atol=1e-2, **kwargs) + # mp_tuner stores four decimals; a rare mismatch must not round down to zero. + return math.ceil(err * 10000) / 10000 + + +_BENCH_KEYS = ("x", "w", "out", "x_scale", "w_scale") +_REF_KEYS = ("x", "w", "x_scale", "w_scale") + + +class OpusA8W8Tuner(GemmCommonTuner): + ARG_DEFAULTS: ClassVar[dict] = { + **GemmCommonTuner.ARG_DEFAULTS, + "tune_file": "/tmp/opus_a8w8_tuned.csv", + "errRatio": 0.0, + } + + def __init__(self): + super().__init__( + "opus_a8w8", + key=["gfx", "cu_num", "M", "N", "K", "dtype", "outdtype", "scaleAB"], + resultList=[ + "libtype", + "kernelId", + "splitK", + "us", + "kernelName", + "tflops", + "bw", + "errRatio", + ], + description="Tune plain OPUS A8W8 GEMM (gfx950, FP8 inputs, FP32 output). " + "CSV scaleAB=False selects no scales; True selects 1x128x128 blockscale.", + ) + + def _setup_specific_arguments(self): + self.parser.add_argument( + "--input_file", + dest="untune_file", + default=argparse.SUPPRESS, + help="Input shape CSV (alias for -i/--untune_file)", + ) + self.parser.add_argument( + "--tuned_file", + dest="tune_file", + default=argparse.SUPPRESS, + help="Output tuned CSV (alias for -o/--tune_file)", + ) + self.parser.add_argument("--libtype", choices=["opus"], default="opus") + for action in self.parser._actions: + if action.dest in { + "splitK", + "compare", + "update_improved", + "min_improvement_pct", + }: + action.help = argparse.SUPPRESS + elif action.dest == "run_config": + action.help = ( + "Benchmark saved kids from TUNED_CSV (defaults to --tuned_file)" + ) + elif action.dest == "errRatio": + action.help = "Maximum mismatch fraction at rtol=atol=1e-2 (default: 0)" + + def _normalize_rows(self, df): + df = df.copy() + missing = {"M", "N", "K"}.difference(df.columns) + if missing: + raise ValueError(f"Shape CSV is missing columns: {sorted(missing)}") + defaults = { + "gfx": self.get_gfx(), + "cu_num": self.get_cu_num(), + "dtype": "fp8", + "outdtype": "fp32", + "scaleAB": False, + } + for column, default in defaults.items(): + if column not in df: + df[column] = default + for column in ("M", "N", "K", "cu_num", "kernelId", "splitK"): + if column not in df: + continue + values = pd.to_numeric(df[column], errors="raise") + minimum = 0 if column in ("kernelId", "splitK") else 1 + if (values.isna() | (values < minimum) | (values % 1 != 0)).any(): + raise ValueError(f"{column} must contain integers >= {minimum}") + df[column] = values.astype("int64") + for column in ("scaleAB", "bias", "bpreshuffle"): + if column in df: + try: + df[column] = df[column].map(lambda v: dtypes.str2bool(str(v))) + except argparse.ArgumentTypeError as exc: + raise ValueError(f"{column} must be True or False") from exc + if column != "scaleAB" and df[column].any(): + raise ValueError( + f"Plain OPUS A8W8 tune does not support {column}=True" + ) + if ( + not df["dtype"] + .isin(["fp8", "float8_e4m3fn", str(torch.float8_e4m3fn)]) + .all() + ): + raise ValueError("Plain OPUS A8W8 tune requires dtype=fp8 (float8_e4m3fn)") + if not df["outdtype"].map(canonical_output_dtype).eq("fp32_t").all(): + raise ValueError("Plain OPUS A8W8 tune requires outdtype=fp32") + if "libtype" in df and not df["libtype"].eq("opus").all(): + raise ValueError("Plain OPUS A8W8 tune requires libtype=opus") + df["dtype"] = str(torch.float8_e4m3fn) + df["outdtype"] = str(torch.float32) + return df + + def get_tuned_gemm_list(self, tuned_gemm_file, columns=None): + df = super().get_tuned_gemm_list(tuned_gemm_file, columns) + return self._normalize_rows(df) if not df.empty else df + + def pre_process(self, args): + gfx = self.get_gfx() + if gfx != _SUPPORTED_GFX: + self.parser.error( + "Plain OPUS A8W8 tuning and --run_config only support " + f"{_SUPPORTED_GFX}; current GPU is {gfx}" + ) + if args.splitK: + self.parser.error("These plain OPUS A8W8 kernels require splitK=0") + if args.compare or args.update_improved: + self.parser.error("Use --run_config TUNED_CSV to benchmark saved OPUS kids") + if args.run_config: + if args.run_config is True: + args.run_config = args.tune_file + if not Path(args.run_config).is_file(): + raise FileNotFoundError(args.run_config) + self.tunedf = self.get_tuned_gemm_list(args.run_config) + self.untunedf = self.tunedf + return + if not args.untune_file: + self.parser.error("--input_file/-i is required for tuning") + df = self._normalize_rows(self.get_untuned_gemm_list(args.untune_file)) + df = df[(df["gfx"] == gfx) & (df["cu_num"] == self.get_cu_num())] + if df.empty: + raise ValueError("No input shapes match the current GPU's gfx/cu_num") + self.untunedf = df[self.keys].drop_duplicates().reset_index(drop=True) + self.tunedf = self.get_tuned_gemm_list(args.tune_file) + if not args.all and not self.tunedf.empty: + tuned_keys = self.tunedf[self.keys].apply(tuple, axis=1) + self.untunedf = self.untunedf[ + ~self.untunedf.apply(tuple, axis=1).isin(tuned_keys) + ].reset_index(drop=True) + + def tune(self, untunedf, tunedf, args): + tasks, tasks_data = [], [] + for row in untunedf.itertuples(index=False): + kids = candidate_kids_for_shape( + row.gfx, row.M, row.N, row.K, row.scaleAB, row.outdtype + ) + if not kids: + raise ValueError(f"No OPUS A8W8 candidates for {row}") + for kid in kids: + tasks.append( + ( + (tuple(row), kid, 0, kernels_list[kid].name), + generate_data, + (row.M, row.N, row.K, kid), + run_bench, + (_BENCH_KEYS, kid), + {"num_warmup": args.warmup, "num_iters": args.iters}, + run_torch, + (_REF_KEYS,), + {}, + None, + 1e-2, + 1e-2, + compare_outputs, + None, + ("out",), + ) + ) + tasks_data.append((len(kids), ())) + return mp_tuner( + tasks, + tasks_data, + mp_num=args.mp, + shape_grouped=args.shape_grouped, + err_ratio=args.errRatio, + timeout=args.timeout, + verbose=args.verbose, + ) + + def getKernelName(self, kernel_id): + return kernels_list[kernel_id].name + + def calculate(self, results, bpes=(1, 1, 4)): + return super().calculate(results, bpes) + + def result_to_df(self, results): + df = super().result_to_df(results) + df["libtype"] = "opus" + return df[self.columns] + + def run_config(self, args): + from aiter.test_common import checkAllclose, run_perftest + + missing = {"kernelId", "splitK", "libtype"}.difference(self.untunedf.columns) + if missing: + raise ValueError( + f"--run_config requires a tuned CSV with {sorted(missing)}" + ) + results = [] + for row in self.untunedf.itertuples(index=False): + kid = row.kernelId + if row.splitK != 0: + raise ValueError(f"OPUS A8W8 kid {kid} requires splitK=0") + if kid not in candidate_kids_for_shape( + row.gfx, row.M, row.N, row.K, row.scaleAB, row.outdtype + ): + raise ValueError(f"Saved kid {kid} is incompatible with {row}") + data = generate_data(row.M, row.N, row.K, kid, device="cuda") + ref = run_torch(*(data[key] for key in _REF_KEYS)) + data["out"].fill_(float("nan")) + out, us = run_perftest( + run_bench, + *(data[key] for key in _BENCH_KEYS), + kid, + num_warmup=args.warmup, + num_iters=args.iters, + ) + err = checkAllclose( + ref, + out, + rtol=1e-2, + atol=1e-2, + tol_err_ratio=args.errRatio, + printLog=args.verbose, + ) + if ( + not math.isfinite(us) + or us <= 0 + or not math.isfinite(err) + or err > args.errRatio + ): + raise RuntimeError(f"Saved kid {kid} failed: {us=}, errRatio={err}") + results.append( + { + "shape": f"M={row.M},N={row.N},K={row.K},scaleAB={row.scaleAB},kid={kid}", + "e2e_us": us, + "errRatio": err, + "status": "ok", + } + ) + return results + + +def main(): + tuner = OpusA8W8Tuner() + tuner.run(tuner.parse_args()) + + +if __name__ == "__main__": + main() diff --git a/csrc/opus_gemm/opus_gemm_common.py b/csrc/opus_gemm/opus_gemm_common.py index 5baac7a492..da1b235705 100644 --- a/csrc/opus_gemm/opus_gemm_common.py +++ b/csrc/opus_gemm/opus_gemm_common.py @@ -1,9 +1,13 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2025-2026, Advanced Micro Devices, Inc. All rights reserved. +"""OPUS kernel registrations shared by selection and code generation.""" + import os import sys from dataclasses import dataclass, field +_A16W16_CO_TAGS = frozenset({"a16w16_4wave_co", "a16w16_4wave_wl_co"}) + # Legacy cache policy = traits default for split-barrier & persistent a16w16 (see # opus_gemm_traits_a16w16_gfx950.cuh). _LEGACY_CACHECTL = (0, 17) @@ -63,18 +67,10 @@ class OpusGemmInstance: arch_prefix: str = "" # Optional generated name tag override for same-pipeline variants. name_tag: str = "" - # SplitK workspace storage dtype; splitK launchers still use fp32 tune dispatch. - # Split-K partial type, per kid. Keep the DEFAULT at fp32: several pipelines - # static_assert on it -- gfx942's em3en4_lds1_pgr2_sk among them -- so a - # bf16 default silently retargets every split-K family on every arch and - # only shows up as a compile error on the arches you did not build. - # The gfx1250 _ws families opt in to bf16 (see _gfx1250_ws_bf16_partial): - # it halves what the reduce reads back and what a sweep allocates, and it - # clears the gate the tuner applies -- at rtol=atol=5e-2, which is what a - # bf16 output gets, err_ratio measures 0.004-0.012 against a 0.05 line, flat - # from split_k=1 to 16. It is the coarser choice: against a 1e-2 tolerance - # the same partials fail from split_k=2 on. - splitk_workspace_dtype: str = "fp32_t" + # Physical workspace storage dtype for this exact kid. External-workspace + # kids must declare it explicitly; non-workspace kids leave it unset. + # The launch-dispatch host specialization remains fp32 independently. + splitk_workspace_dtype: str | None = None # gfx1250 cluster/TDM split-K consumer tiling: "tileN" (split N) or # "tileM" (split M). Only consumed by the a16w16_cluster_tdm_splitk_ws tag. @@ -94,72 +90,47 @@ class OpusGemmInstance: cluster_wg_m: int = 4 cluster_wg_n: int = 4 - # gfx1250 FUSED single-kernel split-K (a16w16_clusterlaunch_tdm_splitk_fuse): - # SplitK and MClusterWg are COMPILE-TIME (cluster dims (SplitK, MClusterWg, 1)), - # so each kid bakes one combo. fuse_ws_dtype = DataWs partial storage - # ("bf16_t" default; "fp32_t" for higher reduce precision). Ignored by every - # other pipeline. fuse_split_k == 0 marks "not a fuse kid". - fuse_split_k: int = 0 - fuse_m_cluster: int = 1 - fuse_ws_dtype: str = "bf16_t" - - # --- a8w8_mxscale BMM flatmm-splitK axes (kernel_tag == - # "a8w8_mxscale_bmm_flatmm_splitk"). The BMM main kernel template is - # gemm_a8w8_mxscale_flatmm_splitk_kernel - # so unlike a16w16 each kid carries two compile-time booleans in addition to - # the tile. direct_only == consumer-self-load direct-store (splitK==1 only); - # prefetch_scale == scale-prefetch variant; fused_reduce == splitK==2 fused - # tail-reduce launch path. These drive both the launcher body and the set of - # device instantiations gen_instances emits for the kid. + # gfx950 a8w8 MXFP8 BMM compile-time axes. BMM instances live in the + # canonical global kid registry, but their generated symbols share the + # ``opus_bmm`` root and a uniform exact-kid launcher signature. direct_only: bool = False prefetch_scale: bool = False fused_reduce: bool = False - # a8w8_mxscale BMM flatmm-splitK only: preload this split's SFA (per-token) + - # SFB (block) scale panels into LDS once, then read scales from LDS in the - # consumer instead of a per-K-tile global buffer_load. Maps to the kernel's - # 5th template bool PRELOAD_SF_LDS. preload_sf: bool = False - # a8w8_mxscale BMM specialized-pipeline axis (minterleave / mouter / - # mouter_tunable / wave4m2_selfload families). Maps to the kernel's trailing - # `bool SKIP_SCALE_WAIT` template param: skip the s_waitcnt on the per-K-tile - # scale load (the scale is issued a tile ahead), trading a correctness margin - # for pipeline overlap. Drives both the launcher body and the device - # instantiation set for the kid. + # Optional override for the D_OUT=void split-K specialization. None keeps + # the direct-output specialization's preload_sf setting; False lets an + # exact kid retain its tuned splitK=1 preload while using the equivalent + # lower-register-pressure path for FP32 workspace partials. + workspace_preload_sf: bool | None = None skip_scale_wait: bool = False - # a8w8_mxscale BMM wave4m2_selfload family extra bool axis (kernel template - # order: ). pack_scale_on_demand: bool = False - # a8w8_mxscale BMM pipeline family (kids 150/151/152): dual bf16/fp32 - # traits + one of the gemm_a8w8_scale_* kernels selected by these flags - # (all-false = plain scale kernel). k1024_only: bool = False k1024_lb1: bool = False - # a8w8_mxscale BMM pipeline family (kid158): preload BOTH SFA (per-token) and - # SFB (block) scale panels into LDS. Maps to the pipeline kernel - # gemm_a8w8_scale_preload_sf_kernel. preload_sf_lds: bool = False - # Symbol root ("opus_gemm" for GEMM, "opus_bmm" for the batched frontends). name_root: str = "opus_gemm" - # --- pre-compiled (.co) families only (kernel_tag in _A16W16_CO_TAGS) ------ - # Device-side entry-point attributes. They never reach the host launcher -- - # it only needs the traits constants -- but they DO reach the symbol name, - # because two entries differing only in these must not collide on one .co. - # co_num_vgpr == 0 omits the amdgpu_num_vgpr attribute entirely. + # gfx1250 fused single-kernel split-K. SplitK and the N-direction cluster + # peer count are compile-time kernel properties. The partial storage dtype + # deliberately uses the shared splitk_workspace_dtype field above so every + # external-workspace kid has one exact-kid dtype source of truth. + fuse_split_k: int = 0 + # Historical field name retained for #4246/tuned-config compatibility; in + # the current N-direction pipeline this is the number of N-tile peers. + fuse_m_cluster: int = 1 + + # Pre-compiled (.co) family metadata. These fields participate in the + # stable symbol/image name and are otherwise consumed only by gfx1250 + # CO code generation. co_num_vgpr: int = 0 co_min_waves_per_eu: int = 1 - # Device-pass flags for this entry, verbatim. A tuple so the dataclass stays - # hashable-ish and no two instances can share a mutable default. co_device_flags: tuple = () - # Free-form name suffix, for entries that differ ONLY in co_device_flags. co_variant: str = "" - # a/b/c/acc spellings; co_traits_args() in the gfx1250 codegen reads these. co_dtypes: tuple = ("bf16_t", "bf16_t", "bf16_t", "fp32_t") - # (TileM, TileN): how the 4 waves tile the block. Only the wave-layout co - # family reads it; (4, 1) is the fixed layout of the original pipeline. co_wave_layout: tuple = (4, 1) + # Optional logical-M limit imposed by this exact kernel's launch geometry. + max_m: int | None = None + @property def name(self) -> str: parts = [ @@ -174,7 +145,6 @@ def name(self) -> str: # tag inserts shift right by one slot when arch_prefix is set tag_at = 1 + (1 if self.arch_prefix else 0) if self.kernel_tag == "a8w8_mxscale_bmm_flatmm_splitk": - # opus_bmm_a8w8_mxscale_flatmm_splitk__wgpcu{N}[_selfload][_scaleprefetch] parts.insert(tag_at, "a8w8_mxscale_flatmm_splitk") parts.append(f"wgpcu{self.WG_PER_CU}") if self.direct_only: @@ -184,7 +154,6 @@ def name(self) -> str: if self.preload_sf: parts.append("sfpreload") elif self.kernel_tag == "a8w8_mxscale_bmm_minterleave": - # opus_bmm_a8w8_mxscale_flatmm_minterleave__wgpcu{N}[_skip_scale_wait] parts.insert(tag_at, "a8w8_mxscale_flatmm_minterleave") parts.append(f"wgpcu{self.WG_PER_CU}") if self.skip_scale_wait: @@ -231,7 +200,7 @@ def name(self) -> str: elif self.kernel_tag == "a16w16_mono_tile": parts.insert(tag_at, "mono_tile") elif self.kernel_tag == "a16w16_cluster_tdm_splitk_ws": - # gfx1250 fp32-workspace split-K with a separate reduce kernel. + # gfx1250 typed-workspace split-K with a separate reduce kernel. # Name it opus_gemm_gfx1250_splitk_* (note the "splitk_" segment) so # the reduce-TU arch detection in gen_instances.py -- which keys on # "opus_gemm__splitk_" -- buckets it like the gfx942 splitk kids. @@ -250,32 +219,19 @@ def name(self) -> str: parts.append(f"c{self.cluster_wg_m}x{self.cluster_wg_n}") parts.append(f"p{self.num_slots}w{self.wg_per_cu}") elif self.kernel_tag == "a16w16_clusterlaunch_tdm_splitk_fuse": - # FUSED single-kernel split-K. The visible segment is "skfuse" (NOT - # "splitk_...") so the reduce-TU detection (keys on "_splitk_" in the - # kernel name, gen_instances.py) never emits a reduce kernel for it. - # It IS still in SPLITK_TAGS (fp32 lookup ABI). m{m}s{split_k}ws{dt} - # + cluster geometry keep each (tile, split_k, m_cluster, ws_dtype) - # symbol unique. + # Keep "splitk_" out of this visible segment: fused kids do not + # need a separate reduce-kernel TU. The historical fuse_m_cluster + # field is an N-peer count, hence the n{} spelling. parts.insert(tag_at, "skfuse") - # fuse_m_cluster now groups N-tile peers (A-multicast); tag as n{}. parts.append(f"n{self.fuse_m_cluster}s{self.fuse_split_k}") - parts.append("wsf32" if self.fuse_ws_dtype == "fp32_t" else "wsbf16") + parts.append( + "wsf32" if self.splitk_workspace_dtype == "fp32_t" else "wsbf16" + ) parts.append(f"p{self.num_slots}w{self.wg_per_cu}") - elif self.kernel_tag.startswith("a16w16_4wave"): - # gfx1250 symmetric 4-wave compute pipeline, loaded from a - # pre-compiled .co (see include/gfx1250/opus_co_launch_gfx1250.cuh). - # This name is used THREE ways and they must stay equal: the host - # launcher symbol, the extern "C" kernel symbol inside the .co, and - # the .co filename. That identity is what lets the codegen emit the - # load call without reading any sidecar. No "splitk_" segment: the - # reduce-TU detection in gen_instances.py keys on that substring and - # this pipeline has no reduce kernel. - # - # The trailing segments carry every axis a .co can differ on: - # cluster dims, ring depth, the pinned VGPR budget and the - # launch_bounds waves-per-EU. Two JSON entries differing in any one - # of them are two distinct symbols and two distinct .co files; if - # they differ ONLY in device_flags, co_variant separates them. + elif self.kernel_tag in _A16W16_CO_TAGS: + # The host launcher symbol, ELF entry point and .co filename must + # remain identical. Encode every device configuration axis that + # can distinguish two pre-built images in this stable name. parts.insert( tag_at, ( @@ -284,8 +240,6 @@ def name(self) -> str: else "4wave_co" ), ) - # Wave layout only shows up for the family that can vary it, so the - # 4wave_co names already on disk are untouched. if self.kernel_tag == "a16w16_4wave_wl_co": parts.append(f"w{self.co_wave_layout[0]}x{self.co_wave_layout[1]}") parts.append(f"c{self.cluster_wg_m}x{self.cluster_wg_n}") @@ -312,40 +266,85 @@ def name(self) -> str: @property def m_align(self) -> int: - """M multiple this kid's generated host guard enforces (1 == any M). - - The launcher family decides it, not the kid: see _BMM_M_ALIGN_TILES and - the AITER_CHECK blocks the matching launcher body in - codegen/gen_instances_gfx950.py emits. Consumers that pick a kid for a - shape (the tuner's candidate filter, the runtime's padded-M lookup) must - read it from here rather than keep their own list -- two hand-maintained - copies is exactly how kid326 ended up excluded from tuning while the - runtime dispatched it anyway. - """ + """M multiple enforced by the generated launcher (1 means tail-safe).""" mult = _BMM_M_ALIGN_TILES.get(self.kernel_tag) if mult is not None: return self.B_M * mult if mult else 1 - # Non-BMM families: has_oob is the codegen flag that says whether the - # tail is masked, and opus_gemm_tune.py already gates on it this way. return 1 if self.has_oob else self.B_M -# a8w8_mxscale BMM launcher family -> the B_M multiple its host guard requires, -# or 0 when the launcher masks a partial M tile and emits no M check at all. -# Mirrors the AITER_CHECK blocks in the launcher bodies of -# codegen/gen_instances_gfx950.py (_BMM_*_LAUNCHER_BODY); gen_instances asserts -# the two agree, so a guard edit that forgets this table fails the build. +def a16w16_flatmm_prefetch_k_iter(instance: OpusGemmInstance) -> int: + """Mirror gfx950 ``Traits::prefetch_k_iter`` for host-side planning. + + The exact launcher and both GEMM/BMM tuning paths must agree on the + minimum number of K tiles a flatmm instance can consume. Keep this + scalar-only calculation next to the canonical instance metadata so the + runtime launch plan and tuner do not drift. + """ + sizeof_da = 2 # BF16 + load_group_m = 64 if instance.W_M >= 32 else 32 + load_group_n = 64 if instance.W_N >= 32 else 32 + load_group_k = instance.W_K * 2 + num_m = instance.B_M // load_group_m + num_n = instance.B_N // load_group_n + num_k = instance.B_K // load_group_k + smem_linear = 64 * 16 // sizeof_da # WARP_SIZE=64 + smem_sub = smem_linear // load_group_k + slots = load_group_m // smem_sub + padding = 16 // sizeof_da if instance.W_M >= 32 else 2 * 16 // sizeof_da + per_group_load = slots * (smem_linear + padding) * sizeof_da + per_iter = (num_m + num_n) * num_k * per_group_load + lds_total = 163840 + return max( + 1, + (lds_total // max(instance.WG_PER_CU, 1)) // max(per_iter, 1), + ) + + +def a8w8_mxscale_flatmm_prefetch_k_iter(instance: OpusGemmInstance) -> int: + """Mirror gfx950 MXFP8 flatmm ``Traits::prefetch_k_iter``.""" + sizeof_da = 1 # FP8 + is_tile_n = instance.B_M == 16 + load_group_m = 16 if is_tile_n else 32 + load_group_n = 16 if is_tile_n else 32 + load_group_k = instance.W_K + num_m = instance.B_M // load_group_m + num_n = instance.B_N // load_group_n + num_k = instance.B_K // load_group_k + smem_linear = 64 * 16 // sizeof_da # WARP_SIZE=64 + smem_sub = smem_linear // load_group_k + slots = load_group_m // smem_sub + padding = 2 * 16 // sizeof_da + per_group_load = slots * (smem_linear + padding) * sizeof_da + per_iter = (num_m + num_n) * num_k * per_group_load + lds_total = 163840 + return max( + 1, + (lds_total // max(instance.WG_PER_CU, 1)) // max(per_iter, 1), + ) + + _BMM_M_ALIGN_TILES = { "a8w8_mxscale_bmm_flatmm_splitk": 0, "a8w8_mxscale_bmm_pipeline": 0, "a8w8_mxscale_bmm_fused": 0, - "a8w8_mxscale_bmm_minterleave": 2, # MI=2 M tiles per WG, baked in - "a8w8_mxscale_bmm_wave4m2_selfload": 2, # LOGICAL_B_M = B_M * 2 + "a8w8_mxscale_bmm_minterleave": 2, + "a8w8_mxscale_bmm_wave4m2_selfload": 2, "a8w8_mxscale_bmm_wave8n2": 1, "a8w8_mxscale_bmm_mouter": 1, "a8w8_mxscale_bmm_mouter_tunable": 1, } +# PR #4320 originally used a private, colliding BMM id namespace. The current +# exact-kid router uses one canonical registry, so gfx950 BMM ids occupy the +# previously empty 8000 band. The low digits intentionally preserve the +# upstream id for tuning/debug correlation. +BMM_MXSCALE_KID_OFFSET = 8000 + + +def bmm_mxscale_global_kid(upstream_kid: int) -> int: + return BMM_MXSCALE_KID_OFFSET + int(upstream_kid) + def _a16w16(bs, bm, bn, bk, tn, wm, wn, wk, has_oob=True, cachectl_a=0, cachectl_b=17): """Factory for a16w16 split-barrier kid instances. @@ -357,8 +356,7 @@ def _a16w16(bs, bm, bn, bk, tn, wm, wn, wk, has_oob=True, cachectl_a=0, cachectl This is the "legacy" policy used by KID 4..9 and 1004..1009 -- the `_LEGACY_CACHECTL` special-case in OpusGemmInstance.name keeps these kids emitting the bare `..._0x0x0` symbol (no `_cA0cB17` suffix) so - the production heuristic dispatcher and the opus tuned CSV stay - bit-compatible. + the Python policy and OPUS tuned CSV stay bit-compatible. """ vec = 16 // 2 # VEC_A = VEC_B = 8 for bf16 inst = OpusGemmInstance( @@ -408,6 +406,7 @@ def _a16w16_flatmm_splitk(bm, bn, bk, wg_per_cu, has_oob=True): ["fp32_t"], wg_per_cu, has_oob=has_oob, + splitk_workspace_dtype="fp32_t", ) @@ -440,54 +439,18 @@ def _a16w16_flatmm(bm, bn, bk, wg_per_cu): # fmt: off # --- per-pipeline kernel instance lists --- a8w8_scale_kernels_list = { - # kid 1 (256x256) is the launcher hardcoded by opus_gemm.cu's - # opus_dispatch_scale (the only a8w8_scale GEMM path). The 128x256 sibling - # kid 720 was removed below. 1: OpusGemmInstance(512, 256, 256, 128, 4, 2, 16, 16, 128, 16, 16, 4, 1, 128, 128, "a8w8_scale", ["fp32_t"]), } -# Dead 128x256 scale GEMM tiles removed (no CSV/dispatch caller): -# - kid 720 (a8w8_scale, fp32 block-scale): only consumer was the removed -# opus_bmm_a8w8_scale mmajor path. -# - kid 710 (a8w8_mxscale, e8m0 block-scale): only consumer was the opus_bmm -# kid 149 hand-written adapter (via the _mmajor sibling), now replaced by -# the BMM-native a8w8_mxscale_bmm_pipeline 128x256 instance. -# Both were the same gemm_a8w8_scale_kernel specialization, differing only in -# scale dtype; opus_dispatch_scale still uses the 256x256 kid 1 above. - def _a8w8_mxscale_bmm_flatmm_splitk( bm, bn, bk, wg_per_cu, direct_only=False, prefetch_scale=False, preload_sf=False ): - """fp8 e8m0 mxscale BATCHED matmul flatmm split-K tile. - - Backs opus_bmm_a8w8_mxscale(); the main kernel - (gemm_a8w8_mxscale_flatmm_splitk_kernel) writes an fp32 workspace and a - shared reduce kernel casts to the Y dtype (bf16/fp32), so output_dtypes is - fp32 workspace here. Locked geometry (matches the hand-written traits in - opus_bmm.cu): BLOCK_SIZE=256 (4 waves), T_M=2/T_N=1, MFMA 16x16x128 (fp8), - VEC=(16,16,4), GROUP=(1,128,128) (per-token M, 128x128 block scale). - direct_only / prefetch_scale are the two kernel compile-time booleans. - """ - # tileN (bm==16): consumers split N (T_M=1, T_N=2). tileM (bm>=32): split M - # (T_M=2, T_N=1). The real T_M/T_N is derived in the C++ traits from B_M; - # these values only drive the generated symbol name, so keep them honest. t_m, t_n = (1, 2) if bm == 16 else (2, 1) inst = OpusGemmInstance( - 256, # BLOCK_SIZE - bm, bn, bk, # BLOCK tile - t_m, t_n, # T_M, T_N (4-wave warp-spec; tileN=1,2 / tileM=2,1) - 16, 16, 128, # W_M, W_N, W_K (MFMA 16x16x128 fp8) -- name only - 16, 16, 4, # VEC_A, VEC_B, VEC_C - 1, 128, 128, # GROUP_M=1 (per-token), GROUP_N=GROUP_K=128 - "a8w8_mxscale_bmm_flatmm_splitk", - # Single host instantiation: the launcher is templated on D_C - # only to satisfy the codegen host-decl machinery; its body ignores D_C - # and branches on Y.dtype() at runtime (native __bf16/float), exactly - # like the hand-written _impl. The fp32 split-K workspace dtype is fixed - # inside the traits, and the reduce kernel casts to the runtime Y dtype. - ["fp32_t"], - wg_per_cu, + 256, bm, bn, bk, t_m, t_n, 16, 16, 128, 16, 16, 4, + 1, 128, 128, "a8w8_mxscale_bmm_flatmm_splitk", ["fp32_t"], + wg_per_cu, splitk_workspace_dtype="fp32_t", ) inst.name_root = "opus_bmm" inst.direct_only = direct_only @@ -496,218 +459,156 @@ def _a8w8_mxscale_bmm_flatmm_splitk( return inst -# fp8 e8m0 mxscale BMM flatmm split-K tiles. kid numbers preserved from the old -# opus_bmm.cu switch so existing tuned CSVs / heuristics keep working. Each kid = -# (B_M, B_N, B_K, WG_PER_CU, direct_only, prefetch_scale). Big-tile pipelines -# (mouter / minterleave / wave*n* / pipeline, kids 131/132/134/140-163/149-152) -# stay monolithic in opus_bmm.cu and are NOT migrated here. _BMM_MXSCALE_SPLITK_TILES = { - # tileN (B_M=16): single 16-row MFMA M-wave so small-M/decode shapes (M<=32) - # don't over-compute a fat B_M tile. Targets the G=2 K=4096 M<=32 gap vs bf16. - 316: (16, 32, 256, 2, False, False), - 317: (16, 32, 256, 2, False, True), # scale prefetch - 318: (16, 32, 128, 2, False, False), - # prefetch-depth sweep: higher WG_PER_CU shrinks per-WG LDS -> shallower - # prefetch_k_iter + more occupancy (small-M/few-tile shapes want this). - 319: (16, 32, 256, 4, False, False), - 314: (16, 32, 512, 2, False, False), # fewer K-iters (8) per WG - # wider-N tileN: larger B_N raises COM_REP_N (more MFMA/iter) to hide - # ds_read+scale latency; WG_PER_CU keeps prefetch_k_iter >= 3. - 313: (16, 64, 256, 2, False, False), # COM_REP_N=2 - 312: (16, 128, 256, 1, False, False), # COM_REP_N=4 - # M=16/32 last-mile (G=2 N=1024 K=4096): 311 = wide-K tileN + scale prefetch; - # 321/323 = 32x32 tileM (exact M=32 fit, no OOB waste, COM_REP_N=2). - 311: (16, 32, 512, 2, False, True), - 321: (32, 32, 256, 2, False, True), - 323: (32, 32, 128, 2, False, True), - # fine tiles (small / mid M) - 320: (64, 32, 256, 2, False, False), - 322: (64, 32, 256, 1, False, False), - # kid324 = kid320 tile + SFA+SFB scale panels preloaded into LDS - # (PRELOAD_SF_LDS; wired via the preload-tiles dict below, not the 6-tuple). - # ATT on kid320 showed ~20% of consumer cycles stalled on vmcnt for the - # per-K-tile global scale load; staging both panels into LDS once (ds_read / - # lgkmcnt) breaks the mid-M valley: G4 K4096 M256 0.93->1.00x, M512 - # 0.94->1.01x, M192 0.91->0.98x vs bf16 (+8-26% TFLOPS over kid320, M128-1024). - # Other attempts (scaleprefetch, B_K=128/512, wg4, 64x64 splitK) all <= kid320. - 640: (32, 64, 256, 2, False, False), - 642: (32, 64, 256, 1, False, False), - 646: (32, 64, 256, 2, True, False), # consumer self-load (splitK==1) - 650: (64, 64, 128, 2, False, False), - 653: (64, 64, 128, 2, False, True), # scale prefetch - # No 64x64x256 kid: mirroring bf16's MT64x64x256 forces wg_per_cu=1 (LDS - # ~198KB), so at M=256 it runs half the WGs and lands 0.77x vs bf16. bf16 only - # wins it via stream-K (refills low tile count), which the flatmm pipeline lacks. + 316: (16, 32, 256, 2, False, False), + 317: (16, 32, 256, 2, False, True), + 318: (16, 32, 128, 2, False, False), + 319: (16, 32, 256, 4, False, False), + 314: (16, 32, 512, 2, False, False), + 313: (16, 64, 256, 2, False, False), + 312: (16, 128, 256, 1, False, False), + 311: (16, 32, 512, 2, False, True), + 321: (32, 32, 256, 2, False, True), + 323: (32, 32, 128, 2, False, True), + 320: (64, 32, 256, 2, False, False), + 322: (64, 32, 256, 1, False, False), + 640: (32, 64, 256, 2, False, False), + 642: (32, 64, 256, 1, False, False), + 646: (32, 64, 256, 2, True, False), + 650: (64, 64, 128, 2, False, False), + 653: (64, 64, 128, 2, False, True), 128: (128, 128, 128, 1, False, False), - 137: (128, 128, 128, 1, False, True), # scale prefetch + 137: (128, 128, 128, 1, False, True), 138: (64, 128, 256, 1, False, False), 139: (128, 64, 256, 1, False, False), - # baseline tiles (guaranteed-runnable fallbacks; kid 0 is the heuristic default) - 256: (32, 256, 128, 1, False, False), - 64: (64, 128, 128, 2, False, False), - 0: (32, 128, 128, 2, False, False), - 32: (32, 128, 128, 2, False, False), + 256: (32, 256, 128, 1, False, False), + 64: (64, 128, 128, 2, False, False), + 0: (32, 128, 128, 2, False, False), + 32: (32, 128, 128, 2, False, False), } -a8w8_mxscale_bmm_flatmm_splitk_kernels_list = { +_bmm_flatmm_local = { kid: _a8w8_mxscale_bmm_flatmm_splitk(bm, bn, bk, wg, direct, prefetch) for kid, (bm, bn, bk, wg, direct, prefetch) in _BMM_MXSCALE_SPLITK_TILES.items() } - -# SFA/SFB-into-LDS preload variants (PRELOAD_SF_LDS). Kept in a separate dict so -# the base 6-tuple stays untouched; each entry is (B_M, B_N, B_K, WG_PER_CU) and -# always sets preload_sf=True (non-direct, non-prefetch). _BMM_MXSCALE_SPLITK_PRELOAD_TILES = { - 324: (64, 32, 256, 2), # = kid320 + SFA/SFB scale panels preloaded to LDS - # mid-M wg1 tiles + SFA/SFB preload (same mechanism as kid324/kid158): staging - # both scale panels into LDS removes the per-K-tile global scale vmcnt load that - # gated the plain/scaleprefetch tiles. On K=4096 M256-2048 this wins +13-17% - # over the old kid137/653/139 picks (kid325 ships G2/M2048, G4/M1024, G8/M512, - # G16/M256; kid326 ships G8/M256). K=1024 gains are ~noise (few K-tiles). kid327 - # kept as a candidate but wins nothing robustly (clock-fragile at cold sclk). - 325: (128, 128, 128, 1), # = kid128/137 tile + preload - 326: (128, 64, 256, 1), # = kid139 tile + preload - 327: (64, 128, 256, 1), # = kid138 tile + preload + 324: (64, 32, 256, 2), + 325: (128, 128, 128, 1), + 326: (128, 64, 256, 1), + 327: (64, 128, 256, 1), } -a8w8_mxscale_bmm_flatmm_splitk_kernels_list.update({ +_bmm_flatmm_local.update({ kid: _a8w8_mxscale_bmm_flatmm_splitk(bm, bn, bk, wg, preload_sf=True) for kid, (bm, bn, bk, wg) in _BMM_MXSCALE_SPLITK_PRELOAD_TILES.items() }) +# ROCm 7.2.4 clang-22 assigns an illegal register class while compiling this +# exact high-pressure PRELOAD_SF_LDS + D_OUT=void specialization after the +# workspace kargs moved to a direct pointer. Its splitK=1 BF16/FP32 kernels +# keep PRELOAD_SF_LDS; only the split-K workspace specialization uses the +# semantically equivalent non-preload implementation (the same geometry as +# local kid 139). +_bmm_flatmm_local[326].workspace_preload_sf = False + def _a8w8_mxscale_bmm_minterleave(bm, bn, bk, wg_per_cu, skip_scale_wait=False): - """fp8 e8m0 mxscale BATCHED matmul M-tile-interleaved tile. - - Backs opus_bmm_a8w8_mxscale() kids 162/163. The main kernel - (gemm_a8w8_mxscale_flatmm_minterleave_kernel) - processes MI=2 consecutive M tiles per WG (baked in the launcher, requires - M % (MI*B_M) == 0); splitK is unused (must be 1). Same locked geometry / - traits as the flatmm split-K family (BLOCK_SIZE=256, T_M=2/T_N=1, MFMA - 16x16x128, VEC=(16,16,4), GROUP=(1,128,128), fp32 workspace tuple slot). - """ t_m, t_n = (1, 2) if bm == 16 else (2, 1) inst = OpusGemmInstance( - 256, # BLOCK_SIZE - bm, bn, bk, # BLOCK tile - t_m, t_n, # T_M, T_N (name only) - 16, 16, 128, # W_M, W_N, W_K (name only) - 16, 16, 4, # VEC_A, VEC_B, VEC_C - 1, 128, 128, # GROUP_M=1 (per-token), GROUP_N=GROUP_K=128 - "a8w8_mxscale_bmm_minterleave", - ["fp32_t"], # single fp32 host stub; body branches on Y.dtype() - wg_per_cu, + 256, bm, bn, bk, t_m, t_n, 16, 16, 128, 16, 16, 4, + 1, 128, 128, "a8w8_mxscale_bmm_minterleave", ["fp32_t"], wg_per_cu, ) inst.name_root = "opus_bmm" inst.skip_scale_wait = skip_scale_wait return inst -# fp8 e8m0 mxscale BMM M-tile-interleaved tiles (kids 162/163). Fixed geometry -# m128n128k128 wg1; the only axis is SKIP_SCALE_WAIT. -_BMM_MXSCALE_MINTERLEAVE_TILES = { - # (B_M, B_N, B_K, WG_PER_CU, skip_scale_wait) - 162: (128, 128, 128, 1, False), - 163: (128, 128, 128, 1, True), # skip per-K-tile scale s_waitcnt -} -a8w8_mxscale_bmm_minterleave_kernels_list = { - kid: _a8w8_mxscale_bmm_minterleave(bm, bn, bk, wg, skip) - for kid, (bm, bn, bk, wg, skip) in _BMM_MXSCALE_MINTERLEAVE_TILES.items() +_bmm_minterleave_local = { + 162: _a8w8_mxscale_bmm_minterleave(128, 128, 128, 1, False), + 163: _a8w8_mxscale_bmm_minterleave(128, 128, 128, 1, True), } def _a8w8_mxscale_bmm_spec(tag, bm, bn, bk, wg_per_cu, **flags): - """Generic fp8 e8m0 mxscale BMM specialized-pipeline tile builder. - - Same locked geometry/traits family as the flatmm split-K kids (BLOCK_SIZE - 256, MFMA 16x16x128, VEC=(16,16,4), GROUP=(1,128,128), fp32 workspace tuple - slot). `tag` selects the kernel family (wave8n2 / wave4m2_selfload); - `flags` sets the family's compile-time axes. - """ t_m, t_n = (1, 2) if bm == 16 else (2, 1) inst = OpusGemmInstance( - 256, bm, bn, bk, t_m, t_n, 16, 16, 128, 16, 16, 4, 1, 128, 128, - tag, ["fp32_t"], wg_per_cu, + 256, bm, bn, bk, t_m, t_n, 16, 16, 128, 16, 16, 4, + 1, 128, 128, tag, ["fp32_t"], wg_per_cu, + splitk_workspace_dtype=("fp32_t" if tag == "a8w8_mxscale_bmm_fused" else None), ) inst.name_root = "opus_bmm" - for key, val in flags.items(): - setattr(inst, key, val) + for key, value in flags.items(): + setattr(inst, key, value) return inst -# fused (kid 100): the only fused-reduce path (splitK counter variant). Same -# 256x32x128x128 wg2 traits as standard kid 0/32, so its device symbols resolve -# to the standard family's TUs -> host-only launcher emit. -a8w8_mxscale_bmm_fused_kernels_list = { +_bmm_fused_local = { 100: _a8w8_mxscale_bmm_spec("a8w8_mxscale_bmm_fused", 32, 128, 128, 2), } -# pipeline (kids 149/150/151/152/158): BLOCK_SIZE 512, m{128,256}n256k128, dual -# bf16/fp32 traits (output dtype baked into the traits tuple), non-splitk scale -# kargs. One of the gemm_a8w8_scale_* kernels selected by flags. The wave -# layout (T_M/T_N/W_*) is derived inside opus_gemm_a8w8_scale_traits_gfx950 from -# BLOCK + , so only B_M/B_N/B_K matter here (the T_M/T_N passed to -# OpusGemmInstance are cosmetic for this tag). + def _a8w8_mxscale_bmm_pipeline(**flags): inst = OpusGemmInstance( - 512, 256, 256, 128, 2, 1, 16, 16, 128, 16, 16, 4, 1, 128, 128, - "a8w8_mxscale_bmm_pipeline", ["fp32_t"], 1, + 512, 256, 256, 128, 2, 1, 16, 16, 128, 16, 16, 4, + 1, 128, 128, "a8w8_mxscale_bmm_pipeline", ["fp32_t"], 1, ) inst.name_root = "opus_bmm" - for key, val in flags.items(): - setattr(inst, key, val) + for key, value in flags.items(): + setattr(inst, key, value) return inst -a8w8_mxscale_bmm_pipeline_kernels_list = { - # kid 149: B_M=128 plain scale pipeline (m128n256k128). Same gemm_a8w8_scale_ - # kernel as kid 150, just half the M tile -> 2x output tiles -> fills more CUs - # on batched wo_a shapes. Was a hand-written cross-module adapter delegating - # to opus_gemm's a8w8_mxscale GEMM launcher; now BMM-native codegen. +_bmm_pipeline_local = { 149: _a8w8_mxscale_bmm_pipeline(B_M=128), 150: _a8w8_mxscale_bmm_pipeline(), 151: _a8w8_mxscale_bmm_pipeline(k1024_only=True), 152: _a8w8_mxscale_bmm_pipeline(k1024_lb1=True), - # kid158: preload BOTH SFA (per-token) and SFB (block) scale panels into LDS. 158: _a8w8_mxscale_bmm_pipeline(preload_sf_lds=True), } - -# mouter (kids 131/144) + mouter_tunable (kids 160/161): wg1 m128n128k128, -# 1 bool axis . Both share gemm_..._mouter_kernel, so the -# tunable variant reuses the mouter device instantiations (host-only emit). -a8w8_mxscale_bmm_mouter_kernels_list = { +_bmm_mouter_local = { 131: _a8w8_mxscale_bmm_spec("a8w8_mxscale_bmm_mouter", 128, 128, 128, 1), - 144: _a8w8_mxscale_bmm_spec("a8w8_mxscale_bmm_mouter", 128, 128, 128, 1, skip_scale_wait=True), + 144: _a8w8_mxscale_bmm_spec( + "a8w8_mxscale_bmm_mouter", 128, 128, 128, 1, skip_scale_wait=True + ), } -a8w8_mxscale_bmm_mouter_tunable_kernels_list = { +_bmm_mouter_tunable_local = { 160: _a8w8_mxscale_bmm_spec("a8w8_mxscale_bmm_mouter_tunable", 128, 128, 128, 1), - 161: _a8w8_mxscale_bmm_spec("a8w8_mxscale_bmm_mouter_tunable", 128, 128, 128, 1, skip_scale_wait=True), + 161: _a8w8_mxscale_bmm_spec( + "a8w8_mxscale_bmm_mouter_tunable", 128, 128, 128, 1, + skip_scale_wait=True, + ), } - -# wave8n2 (kid 132): wg1 m128n128k128, no compile-time flags (logical B_N = 256). -a8w8_mxscale_bmm_wave8n2_kernels_list = { +_bmm_wave8n2_local = { 132: _a8w8_mxscale_bmm_spec("a8w8_mxscale_bmm_wave8n2", 128, 128, 128, 1), } - -# wave4m2_selfload (kids 134/142/148): wg1 m128n128k128, 2 bool axes -# (logical B_M = 128*2 = 256). _BMM_WAVE4M2_TILES = { - # (ssw, psod) 134: (False, False), - 142: (True, False), - 148: (True, True), + 142: (True, False), + 148: (True, True), } -a8w8_mxscale_bmm_wave4m2_selfload_kernels_list = { +_bmm_wave4m2_local = { kid: _a8w8_mxscale_bmm_spec( "a8w8_mxscale_bmm_wave4m2_selfload", 128, 128, 128, 1, - skip_scale_wait=ssw, pack_scale_on_demand=psod, + skip_scale_wait=skip, pack_scale_on_demand=pack, ) - for kid, (ssw, psod) in _BMM_WAVE4M2_TILES.items() + for kid, (skip, pack) in _BMM_WAVE4M2_TILES.items() } -# All name-keyed a8w8_mxscale BMM kernel families (gfx950-only). Kept as a tuple -# of the per-family kid-keyed dicts -- NOT merged into one dict, because int kids -# repeat across families and are deduped downstream by launcher NAME (see -# gen_instances.py). Single source of truth for both consumers there: the codegen -# kdict merge and the BMM int-kid tune-lookup emitter. + +def _globalize_bmm_kids(kernels): + return {bmm_mxscale_global_kid(kid): inst for kid, inst in kernels.items()} + + +a8w8_mxscale_bmm_flatmm_splitk_kernels_list = _globalize_bmm_kids(_bmm_flatmm_local) +a8w8_mxscale_bmm_fused_kernels_list = _globalize_bmm_kids(_bmm_fused_local) +a8w8_mxscale_bmm_minterleave_kernels_list = _globalize_bmm_kids(_bmm_minterleave_local) +a8w8_mxscale_bmm_mouter_kernels_list = _globalize_bmm_kids(_bmm_mouter_local) +a8w8_mxscale_bmm_mouter_tunable_kernels_list = _globalize_bmm_kids( + _bmm_mouter_tunable_local +) +a8w8_mxscale_bmm_pipeline_kernels_list = _globalize_bmm_kids(_bmm_pipeline_local) +a8w8_mxscale_bmm_wave8n2_kernels_list = _globalize_bmm_kids(_bmm_wave8n2_local) +a8w8_mxscale_bmm_wave4m2_selfload_kernels_list = _globalize_bmm_kids( + _bmm_wave4m2_local +) a8w8_mxscale_bmm_kernel_lists = ( a8w8_mxscale_bmm_flatmm_splitk_kernels_list, a8w8_mxscale_bmm_fused_kernels_list, @@ -718,6 +619,10 @@ def _a8w8_mxscale_bmm_pipeline(**flags): a8w8_mxscale_bmm_wave8n2_kernels_list, a8w8_mxscale_bmm_wave4m2_selfload_kernels_list, ) +BMM_MXSCALE_KIDS = frozenset( + kid for family in a8w8_mxscale_bmm_kernel_lists for kid in family +) +assert len(BMM_MXSCALE_KIDS) == sum(map(len, a8w8_mxscale_bmm_kernel_lists)) a8w8_kernels_list = { @@ -998,6 +903,26 @@ def _make_4g_safe(inst: "OpusGemmInstance") -> "OpusGemmInstance": # -- gfx942 kernel lists ------------------------------------------------ Kid offset: gfx942 GFX942_KID_OFFSET = 10000 +# Split-K launch policy is consumed by both the Python Torch-workspace planner +# and the generated host launcher. Keep it beside the canonical instances so +# the two sides cannot silently drift and disagree about workspace capacity. +GFX942_MAX_AUTO_SPLIT_K = 16 +GFX942_MIN_ITERS_PER_SPLIT = 2 +GFX942_QUAD_MFMA32_SPLITK_TAG = "a16w16_quad_mfma32_kbuf1_sk" +GFX942_EVEN_LOOP_SPLITK_TAGS = frozenset( + { + "a16w16_kbuf2v_sk", + "a16w16_kbuf2v_bk128_sk", + GFX942_QUAD_MFMA32_SPLITK_TAG, + } +) + +# gfx942 bf16-workspace launchers can use the exact-N row-block reducer only +# for these output widths. Keep the *set* here beside the instance source so +# runtime selection, tuning, and codegen can consume one value. The detailed +# (VEC, N_VEC, ROWS_PER_BLOCK) reduce configurations remain codegen-owned. +GFX942_BF16WS_EXACT_N = frozenset({64, 128, 256, 384, 512, 1024, 2048}) + def _a16w16_gfx942(bs, bm, bn, bk, tn, wm, wn, wk): """Factory for gfx942 a16w16 kbuf1-large-tile kid instances (kid 10000, @@ -1061,6 +986,7 @@ def _a16w16_splitk_tag_gfx942(bs, bm, bn, bk, tn, wm, wn, wk, tag): tag, ["fp32_t"], arch_prefix="gfx942", + splitk_workspace_dtype="fp32_t", ) @@ -1115,6 +1041,7 @@ def _a16w16_kbuf2v_sk_gfx942(bs, bm, bn, bk, tn, wm, wn, wk): return OpusGemmInstance( bs, bm, bn, bk, 2, tn, wm, wn, wk, vec, vec, 4, 0, 0, 0, "a16w16_kbuf2v_sk", ["fp32_t"], arch_prefix="gfx942", + splitk_workspace_dtype="fp32_t", ) @@ -1124,6 +1051,7 @@ def _a16w16_kbuf2v_bk128_sk_gfx942(bs, bm, bn, bk, tn, wm, wn, wk): return OpusGemmInstance( bs, bm, bn, bk, 2, tn, wm, wn, wk, vec, vec, 4, 0, 0, 0, "a16w16_kbuf2v_bk128_sk", ["fp32_t"], arch_prefix="gfx942", + splitk_workspace_dtype="fp32_t", ) @@ -1151,6 +1079,7 @@ def _a16w16_em3en4_lds1_pgr2_sk_gfx942(bs, bm, bn, bk, tn, wm, wn, wk): return OpusGemmInstance( bs, bm, bn, bk, 2, tn, wm, wn, wk, vec, vec, 4, 0, 0, 0, "a16w16_em3en4_lds1_pgr2_sk", ["fp32_t"], arch_prefix="gfx942", + splitk_workspace_dtype="fp32_t", ) @@ -1227,8 +1156,8 @@ def _a8w8_blockscale_bpreshuffle_singlebuf_gfx942( # -- gfx1250 kernel lists ---------------------------------------------------- # Kid offset: gfx1250 kids live in the 20000+ range, disjoint from gfx950 -# (<10000) and gfx942 (50000+). Today only the cluster/TDM split-K (atomic -# fp32 reduction) pipeline is wired (????:fp32 output, no bias). +# (<10000) and gfx942. Both #4246 families are represented: two-stage +# cluster/TDM kids and fused in-cluster-reduce kids. GFX1250_KID_OFFSET = 20000 @@ -1239,10 +1168,11 @@ def _a16w16_cluster_tdm_splitk_ws_gfx1250(bm, bn, bk, layout, num_slots=3, wg_pe (demon_gcn/wmma_opus_rdna4/gemm_a16w16_cluster_tdm_splitk_reduce_4wave.cc): BLOCK_SIZE=128 (4 waves x 32 = 2 producer + 2 consumer), MFMA 16x16x32, NO-CLUSTER (one WG per B_M x B_N tile). The main kernel WMMA-accumulates in - fp32 and PLAIN-stores each split's partial into an fp32 workspace; a separate - reduce kernel sums the split slices, folds bias, and casts to the Y dtype. - output_dtypes = ["fp32_t"] (only the fp32-workspace main kernel is - instantiated; Y bf16/fp32 is a runtime decision in the reduce kernel). + fp32 and casts each split's partial into the exact kid's typed workspace; a + separate reduce kernel sums the split slices in fp32, folds bias, and casts + to the Y dtype. The two-stage families keep partials in fp32 workspace. The + output_dtypes = ["fp32_t"] token selects the existing host launch-dispatch + specialization; Y bf16/fp32 remains a runtime decision in the reducer. layout: "tileN" (consumers split N; B_N>=32) -> T_M=1, T_N=2; "tileM" (consumers split M; B_M>=32) -> T_M=2, T_N=1. @@ -1259,6 +1189,9 @@ def _a16w16_cluster_tdm_splitk_ws_gfx1250(bm, bn, bk, layout, num_slots=3, wg_pe "a16w16_cluster_tdm_splitk_ws", ["fp32_t"], arch_prefix="gfx1250", + splitk_workspace_dtype="fp32_t", + # The separate reducer places one logical row in each grid.y block. + max_m=65535, ctdm_layout=layout, num_slots=num_slots, wg_per_cu=wg_per_cu, @@ -1323,10 +1256,9 @@ def _ctdm_pick_configs(bm, bn, bk): # per-TDM request count hits the 256 direct-copy limit on some operand (e.g. # 32x256x128, 32x128x256) yield no config and are dropped automatically. _GFX1250_CTDM_TILES = [ - # -- ORIGINAL 11 tiles: KEEP THIS ORDER (indices 0..10) -- the C++ heuristic - # opus_a16w16_heuristic_kid_gfx1250() hardcodes kids 20000/20024/20032 - # (16x32/64/128, idx 0/3/4) and 20040/20048/20056 (32x32/64/128, idx - # 5/6/7), and tuned CSVs reference these numbers. Do NOT reorder/insert. + # -- ORIGINAL 11 tiles: KEEP THIS ORDER (indices 0..10) -- tuned CSVs and + # the Python heuristic reference the stable kids derived from these + # indices. Do NOT reorder/insert. # tileN family (B_M=16) (16, 32, 128, "tileN"), (16, 32, 256, "tileN"), @@ -1366,16 +1298,15 @@ def _ctdm_pick_configs(bm, bn, bk): (128, 128, 128, "tileM"), ] -# Kid numbering (clean, contiguous; heuristic / tuned-CSV back-compat dropped): +# Kid numbering is stable for tuned CSVs and the Python heuristic: # plain (no-cluster) kids occupy [20000, 20100), ONE P=3 kid per tile (P=2 is # dropped -- unvalidated). Tiles the picker rejects (>=256-request TDM # direct-copy, now FIXED) fall back to P=3, 1 WG/CU so every no-spill tile still -# emits a plain kid (LDS(P=3) <= 320 KB for this set). The C++ heuristic -# constants in opus_gemm_heuristic_dispatch_gfx1250.cuh are regenerated to match. +# emits a plain kid (LDS(P=3) <= 320 KB for this set). # # The consumer kExpN stability guard (previously _GFX1250_MAX_KEXPN=8) is removed. gfx1250_kernels_list = {} -GFX1250_PLAIN_KID_OF = {} # (B_M,B_N,B_K) -> kid (P=3; for tuner + heuristic regen) +GFX1250_PLAIN_KID_OF = {} # (B_M,B_N,B_K) -> kid (P=3; tuner + Python heuristic) _GFX1250_KID_BASE = 20000 _p_kid = _GFX1250_KID_BASE for _bm, _bn, _bk, _layout in _GFX1250_CTDM_TILES: @@ -1394,13 +1325,13 @@ def _ctdm_pick_configs(bm, bn, bk): # -- gfx1250 CLUSTER-LAUNCH (multicast) variant ------------------------------ -# Same 4-wave TDM split-K + fp32 workspace + reduce kernel, but launched as a +# Same 4-wave TDM split-K + typed workspace + reduce kernel, but launched as a # (cluster_wg_m x cluster_wg_n x 1) workgroup CLUSTER: peers co-reside and share # A/B TDM loads via CLUSTER_LOAD_ASYNC multicast (named-barrier producer/consumer # handshake, same as the plain base). The host launcher rounds the grid up to the -# cluster dims; the workgroups that round-up adds own no tile and return at their -# cluster-barrier arrival, so no shape needs an exact cluster fill. Distinct kid -# band (20100+) so it never collides with the no-cluster base kids (20000..20099). +# cluster dims; surplus workgroups take the pipeline's uniform tile_oob exit. +# Logical workspace strides use the unrounded tile counts. Clusterlaunch kids +# occupy [20100, 21000), separate from plain kids in [20000, 20100). def _a16w16_clusterlaunch_tdm_splitk_ws_gfx1250( bm, bn, bk, layout, cwm, cwn, num_slots=3, wg_per_cu=2 ): @@ -1470,7 +1401,7 @@ def _gfx1250_valid_cluster_dims(): return dims -# Deterministic kid numbering: 20500 + running index over (tile outer, then +# Deterministic kid numbering: 20100 + running index over (tile outer, then # cluster dim (cwn outer, cwm inner)). Kid numbers are provisional -- a global # renumber is pending. The kExpN stability guard has been removed, so ALL 26 # no-spill tiles are expanded (incl. B_N=256 tileM -> kExpN=16). The multicast @@ -1495,19 +1426,28 @@ def _gfx1250_valid_cluster_dims(): GFX1250_CLUSTERLAUNCH_KIDS = frozenset(gfx1250_clusterlaunch_kernels_list.keys()) -# -- gfx1250 FUSED single-kernel split-K (a16w16_clusterlaunch_tdm_splitk_fuse) -- -# Single kernel: last split WG folds bias + reduces the SplitK-1 partials in-kernel -# (cluster-barrier sync), no separate reduce kernel. SplitK / MClusterWg are -# COMPILE-TIME (cluster dims (SplitK, MClusterWg, 1)); DataWs (bf16/fp32) is a kid -# property. B is TDM-multicast across the MClusterWg M-peers. +# -- gfx1250 FUSED single-kernel split-K ------------------------------- +# The first SplitK-1 WGs publish typed partial tiles to caller-owned storage; +# the last WG consumes those tiles after the cluster barrier and writes Y in +# the same kernel. There is no separate reduce launch, but this remains an +# external-workspace family. Workspace is tile-major: +# [num_tiles_m, num_tiles_n, SplitK-1, B_M, B_N] +# SplitK and the N-peer count are compile-time properties of each exact kid. # -# CURRENTLY UNREGISTERED (see GFX1250_SPLITK_FUSE_ENABLED below): the pipeline is -# still being fixed, so the family contributes no kid. Its band has been moved to -# [27000, 30000) -- the pre-compiled .co family took over the [21000, 27000) head -# while this one is unregistered. +# The final #4246 decision leaves this family unregistered until its pipeline +# is fixed. The factory, emitter, and device source remain available, but no kid +# is visible to exact dispatch/capability queries. Its reserved band starts at +# 27000 because pre-compiled CO kids own [21000, 27000). def _a16w16_splitk_fuse_gfx1250( - bm, bn, bk, layout, split_k, m_cluster, ws_dtype="bf16_t", - num_slots=3, wg_per_cu=2, + bm, + bn, + bk, + layout, + split_k, + n_cluster, + ws_dtype="bf16_t", + num_slots=3, + wg_per_cu=2, ): from dataclasses import replace @@ -1517,77 +1457,44 @@ def _a16w16_splitk_fuse_gfx1250( return replace( inst, kernel_tag="a16w16_clusterlaunch_tdm_splitk_fuse", - # output_dtypes MUST stay ["fp32_t"] (the split-K lookup invariant): the - # host launcher is instantiated ONLY as and opus_gemm.cu forces - # the fuse band to the dispatch slot; the launcher then picks the - # real Y dtype at RUNTIME (if Y.dtype()==bf16 ... else float). Advertising - # bf16_t here would make gen_a16w16_tune_lookup emit &{name} in the - # BF16 tune map -> undefined symbol (that specialization is never built). - # The tuner exempts fuse kids from the output-dtype narrowing separately. + # The host token selects the workspace dispatch ABI. The + # fused launcher chooses the real bf16/fp32 Y type at runtime. output_dtypes=["fp32_t"], + splitk_workspace_dtype=ws_dtype, + # This family reduces in-kernel, without the separate grid.y launch. + max_m=None, fuse_split_k=split_k, - fuse_m_cluster=m_cluster, - fuse_ws_dtype=ws_dtype, + # Historical #4246 field name; physically this is an N-peer count. + fuse_m_cluster=n_cluster, ) -# Registration switch for the whole fused family. False sweeps no (tile, split_k, -# n_cluster, ws) combination at all, so gfx1250_splitk_fuse_kernels_list stays -# empty: no kid to look up, nothing for the tuner to pick, nothing for the codegen -# to emit, and the kid band below is unclaimed. The factory above, the emitter in -# codegen/gen_instances_gfx1250.py and the device pipeline are all still here -- -# flipping this back to True is the only step needed to bring the family back. GFX1250_SPLITK_FUSE_ENABLED = False gfx1250_splitk_fuse_kernels_list = {} -# Kid band the family claims WHEN ENABLED. The current sweep needs 1377 kids, so -# [23000, 30000) has room to spare; while disabled the whole range is free. -# It used to start at 21000 -- the pre-compiled .co family took [21000, 23000) -# while this one was unregistered. GFX1250_SPLITK_FUSE_KID_BASE = 27000 -_sf_kid = GFX1250_SPLITK_FUSE_KID_BASE -# (B_M, B_N, B_K, layout, split_k, m_cluster, ws_dtype) -> kid, for the tuner / -# candidate selection to look a fuse kid up by config. GFX1250_SPLITK_FUSE_KID_OF = {} -# Fuse tiles = the SAME no-spill (B_M, B_N, B_K) set as the clusterlaunch sweep -# (_GFX1250_CLUSTERLAUNCH_TILES), so fuse covers the full tile range. Layout -# follows the base rule (B_M==16 -> tileN, else tileM); wg_per_cu is inherited -# per tile from the clusterlaunch table (the fuse producer shares the same TDM -# request profile, so that wg keeps 2-WG/CU co-residency TDM-budget-safe). -# -# N-DIRECTION MULTICAST: the cluster is (SplitK, n_cluster, 1) where the 2nd dim -# (stored in the fuse_m_cluster field) groups n_cluster N-tile peers that share -# A[M-tile] via TDM multicast (mirrors the clusterlaunch cwn A-multicast that -# wins at small M). n_cluster is swept 1..5 (TDM fan-out <= 5) subject to -# SplitK*n_cluster <= 16 (16-WG cluster budget). -# -# split_k sweep per workspace dtype: -# * bf16 workspace: split_k 2..15 -# * fp32 workspace: split_k 2..8 (kept conservative; the reduce now stages -# partials through a bounded LDS RING (kFuseReduceRing in the pipeline), so -# split_k is NO LONGER LDS-bounded -- this cap could be lifted to 15 too). -# SplitK is capped at 15 (NOT 16): each __cluster_dims__ axis is a 4-bit field. -# SplitK / n_cluster are COMPILE-TIME (cluster dims), so each (tile, split_k, -# n_cluster, ws) is a distinct kid; neither is a runtime knob for fuse. -_FUSE_REDUCE_RING = 3 # must match kFuseReduceRing in the fuse pipeline header +_sf_kid = GFX1250_SPLITK_FUSE_KID_BASE + +# The bounded LDS ring used by the fused reducer must fit inside the traits' +# shared allocation. These constants mirror the fused pipeline. +_FUSE_REDUCE_RING = 3 _FUSE_NUM_SLOTS = 3 def _fuse_ring_lds_ok(bm, bn, bk, wg, ws_bytes): - """Guard: the reduce LDS ring (kFuseReduceRing tiles of B_M*B_N*ws_bytes) - must fit kLdsTotalBytes. Mirrors the traits LDS formula so we never emit a - kid that would fail the pipeline's ring static_assert at compile time.""" pitch = bk + 8 - seg_ab = _FUSE_NUM_SLOTS * (bm + bn) * pitch * 2 # bf16 A/B footprint - lds_total = (160 * 1024 + 1024) if (wg == 1 and seg_ab <= 160 * 1024) else seg_ab + seg_ab = _FUSE_NUM_SLOTS * (bm + bn) * pitch * 2 + lds_total = ( + 160 * 1024 + 1024 if wg == 1 and seg_ab <= 160 * 1024 else seg_ab + ) return _FUSE_REDUCE_RING * bm * bn * ws_bytes <= lds_total -_FUSE_WS_SWEEP = (("bf16_t", 2, 15), ("fp32_t", 4, 8)) # (ws_dtype, elem_bytes, sk_hi) -# N-direction cluster (A-multicast) fan-out: the fuse_m_cluster field holds the -# cluster's 2nd-dim WG count, which for this pipeline groups N-tile peers sharing -# A. TDM multicast fans out to <= 5 WGs, and the cluster (SplitK, n_cluster, 1) -# must satisfy SplitK*n_cluster <= 16 (16-bit workgroup_mask / 16-WG budget). +# BF16 storage covers SplitK 2..15; the conservative FP32 family covers 2..8. +# Cluster dims are (SplitK, n_cluster, 1), so SplitK*n_cluster must fit the +# 16-WG cluster budget and each axis stays within its hardware limit. +_FUSE_WS_SWEEP = (("bf16_t", 2, 15), ("fp32_t", 4, 8)) _FUSE_MAX_NCLUSTER = 5 _fuse_tiles_seen = set() _fuse_tiles = _GFX1250_CLUSTERLAUNCH_TILES if GFX1250_SPLITK_FUSE_ENABLED else () @@ -1598,197 +1505,218 @@ def _fuse_ring_lds_ok(bm, bn, bk, wg, ws_bytes): _layout = "tileN" if _bm == 16 else "tileM" for _ws, _ws_bytes, _sk_hi in _FUSE_WS_SWEEP: if not _fuse_ring_lds_ok(_bm, _bn, _bk, _wg, _ws_bytes): - continue # ring wouldn't fit LDS for this (tile, ws) -- skip + continue for _nc in range(1, _FUSE_MAX_NCLUSTER + 1): for _sk in range(2, _sk_hi + 1): - if _sk * _nc > 16: # SplitK * n_cluster <= 16 (cluster budget) + if _sk * _nc > 16: continue - gfx1250_splitk_fuse_kernels_list[_sf_kid] = _a16w16_splitk_fuse_gfx1250( - _bm, _bn, _bk, _layout, - split_k=_sk, m_cluster=_nc, ws_dtype=_ws, wg_per_cu=_wg, + gfx1250_splitk_fuse_kernels_list[_sf_kid] = ( + _a16w16_splitk_fuse_gfx1250( + _bm, + _bn, + _bk, + _layout, + split_k=_sk, + n_cluster=_nc, + ws_dtype=_ws, + wg_per_cu=_wg, + ) ) GFX1250_SPLITK_FUSE_KID_OF[ (_bm, _bn, _bk, _layout, _sk, _nc, _ws) ] = _sf_kid _sf_kid += 1 -assert _sf_kid <= 30000, f"splitk_fuse gfx1250 kids overflow [27000,30000): {_sf_kid}" + +assert _sf_kid <= 30000, ( + "splitk_fuse gfx1250 kids overflow [27000,30000): " + f"ended at {_sf_kid - 1}" +) GFX1250_SPLITK_FUSE_KIDS = frozenset(gfx1250_splitk_fuse_kernels_list.keys()) +assert bool(GFX1250_SPLITK_FUSE_KIDS) == GFX1250_SPLITK_FUSE_ENABLED -# -- gfx1250 SYMMETRIC 4-WAVE COMPUTE, pre-compiled .co (a16w16_4wave_co) ----- -# No producer/consumer split: all 4 waves issue TDM and run WMMA, each owning -# 32 M rows and the full B_N across the whole K loop (~256 fp32 VGPR/lane of -# accumulators live throughout). No workspace, no split-K, no bias -- bf16 C is -# stored straight out through LDS by one TDM per wave. -# -# THE DEVICE SIDE OF THIS FAMILY IS NOT JIT-COMPILED. The kernel needs -# __builtin_amdgcn_pin_vgpr / amdgpu_num_vgpr / -amdgpu-expert-scheduling-mode, -# none of which a release ROCm toolchain provides, so it is built ahead of time -# by gen_co/build_co.py into gen_co/gfx1250/.co and loaded at runtime. -# gen_co/co_kernels.json is the single source of truth for BOTH the offline -# build and the JIT codegen, which is what keeps the .co filename, the .co's -# extern "C" symbol and the host launcher symbol identical. Adding a variant -# (another VGPR budget, another cluster shape, other device flags) is an edit to -# that file only -- nothing here enumerates configurations. -# A full tile x wave-layout x ring-depth x cluster sweep is a few thousand -# entries, so the co band is 6000 wide. splitk_fuse sits above it at 27000. +# -- gfx1250 symmetric 4-wave compute, pre-compiled .co --------------------- +# JSON is the single source of truth shared by the offline image builder and +# this host-launcher registry. The device kernels need compiler facilities not +# available in the release JIT toolchain, so gen_instances emits host launchers +# only and loads the matching image at runtime. GFX1250_4WAVE_CO_KID_BASE = 21000 GFX1250_4WAVE_CO_KID_END = 27000 - CO_KERNELS_JSON = os.path.join(os.path.dirname(__file__), "gen_co", "co_kernels.json") -# Which kid band each pre-compiled tag is allowed to claim. A tag with no entry -# here is rejected outright rather than landing on some other family's kids. _CO_KID_BANDS = { - "a16w16_4wave_co": (GFX1250_4WAVE_CO_KID_BASE, GFX1250_4WAVE_CO_KID_END), - # Same band: kids are explicit in the JSON, so the two co families just - # continue the numbering rather than each reserving a sub-range. - "a16w16_4wave_wl_co": (GFX1250_4WAVE_CO_KID_BASE, GFX1250_4WAVE_CO_KID_END), + tag: (GFX1250_4WAVE_CO_KID_BASE, GFX1250_4WAVE_CO_KID_END) + for tag in _A16W16_CO_TAGS +} +_CO_DTYPE_BYTES = { + "bf16_t": 2, + "fp16_t": 2, + "fp32_t": 4, + "fp8_t": 1, + "bf8_t": 1, } -_CO_DTYPE_BYTES = {"bf16_t": 2, "fp16_t": 2, "fp32_t": 4, "fp8_t": 1, "bf8_t": 1} - - -def _co_instance_from_json(arch, e): - """One JSON entry -> one OpusGemmInstance. Every field the .co and the host - launcher must agree on comes from this single conversion.""" - tag = e["tag"] - bm, bn, bk = e["tile"] - cwm, cwn = e["cluster"] - lb_threads, lb_waves = e["launch_bounds"] - d = e["dtypes"] - # The host launcher starts the grid with dim3(Traits::BLOCK_SIZE); if these - # two disagree the kernel is launched with a block size its - # __launch_bounds__ never promised, which fails silently rather than loudly. - assert lb_threads == e["block_size"], ( - f"co kid {e['kid']}: launch_bounds[0]={lb_threads} must equal " - f"block_size={e['block_size']}" +def _co_instance_from_json(arch, entry): + """Convert one JSON record to the exact host/device CO configuration.""" + tag = entry["tag"] + bm, bn, bk = entry["tile"] + cwm, cwn = entry["cluster"] + lb_threads, lb_waves = entry["launch_bounds"] + dtypes = entry["dtypes"] + assert lb_threads == entry["block_size"], ( + f"co kid {entry['kid']}: launch_bounds[0]={lb_threads} must equal " + f"block_size={entry['block_size']}" ) - vec_a = 16 // _CO_DTYPE_BYTES[d["a"]] - vec_b = 16 // _CO_DTYPE_BYTES[d["b"]] + vec_a = 16 // _CO_DTYPE_BYTES[dtypes["a"]] + vec_b = 16 // _CO_DTYPE_BYTES[dtypes["b"]] return OpusGemmInstance( - e["block_size"], - bm, bn, bk, - 4, 1, # T_M, T_N: the 4 waves split M, each owns the full N - 16, 16, 32, # WMMA 16x16x32 - vec_a, vec_b, 8, - 0, 0, 0, # GROUP (unused) + entry["block_size"], + bm, + bn, + bk, + 4, + 1, + 16, + 16, + 32, + vec_a, + vec_b, + 8, + 0, + 0, + 0, tag, - # The kernel stores C straight through LDS with no reduce kernel behind - # it to cast anything else, so unlike the gfx1250 split-K kids this is - # the real output dtype, not an ABI placeholder for a workspace slot. - [d["c"]], + [dtypes["c"]], arch_prefix=arch, - num_slots=e["num_slots"], + num_slots=entry["num_slots"], cluster_wg_m=cwm, cluster_wg_n=cwn, - co_num_vgpr=e.get("num_vgpr", 0), + co_num_vgpr=entry.get("num_vgpr", 0), co_min_waves_per_eu=lb_waves, - co_device_flags=tuple(e.get("device_flags", ())), - co_variant=e.get("variant", ""), - co_dtypes=(d["a"], d["b"], d["c"], d["acc"]), - co_wave_layout=tuple(e.get("wave_layout", (4, 1))), + co_device_flags=tuple(entry.get("device_flags", ())), + co_variant=entry.get("variant", ""), + co_dtypes=(dtypes["a"], dtypes["b"], dtypes["c"], dtypes["acc"]), + co_wave_layout=tuple(entry.get("wave_layout", (4, 1))), ) -def co_image_path(path, inst): - """Where the pre-built image for `inst` lives, relative to co_kernels.json. - - The one place the layout `gen_co//.co` is spelled on the Python - side; the C++ loader spells the same thing under $OPUS_GEN_CO_DIR. - """ - return os.path.join(os.path.dirname(path), inst.arch_prefix, f"{inst.name}.co") +def _validate_co_instance(kid, instance): + """Validate the runtime contract baked into one pre-built image.""" + if instance.arch_prefix != "gfx1250": + raise ValueError( + f"CO kid {kid} must target gfx1250, got {instance.arch_prefix!r}" + ) + if instance.output_dtypes != ["bf16_t"]: + raise ValueError(f"CO kid {kid} must expose BF16 output only") + if instance.co_dtypes != ( + "bf16_t", + "bf16_t", + "bf16_t", + "fp32_t", + ): + raise ValueError(f"CO kid {kid} has unsupported dtype contract") + if instance.splitk_workspace_dtype is not None: + raise ValueError(f"CO kid {kid} must not declare workspace storage") + + +def co_image_path(path, instance): + """Return the pre-built image corresponding to an instance JSON file.""" + return os.path.join( + os.path.dirname(path), instance.arch_prefix, f"{instance.name}.co" + ) def _load_co_kernels(path, require_image=True): - """Parse gen_co/co_kernels.json into {kid: OpusGemmInstance}. - - A MISSING FILE IS NOT AN ERROR: the family goes empty, exactly like - GFX1250_SPLITK_FUSE_ENABLED=False. A data file that failed to ship should - cost the co kids, not the whole opus codegen. A malformed one still raises. - - ``require_image`` extends that same degradation to the images themselves: an - entry whose .co is not on disk is DROPPED (with a warning) rather than - registered. Without this, a kid can exist all the way through codegen, the - tuner and the tuned CSV and only fail at its first launch -- which is - precisely what a wheel built from a tree where the .co files were never - committed would ship. build_co.py is the one caller that must see the - unfiltered table: it is what produces the missing images. - """ + """Load pre-built CO instances, safely dropping records with no image.""" if not os.path.exists(path): return {} + import json - with open(path) as f: - doc = json.load(f) + with open(path) as stream: + document = json.load(stream) - out = {} + instances = {} seen_names = {} missing = [] - for arch, entries in doc.items(): - if arch.startswith("_"): # "_comment" + for arch, entries in document.items(): + if arch.startswith("_"): continue - for e in entries: - kid = e["kid"] - lo, hi = _CO_KID_BANDS[e["tag"]] + for entry in entries: + kid = entry["kid"] + tag = entry["tag"] + if tag not in _CO_KID_BANDS: + raise ValueError(f"unsupported CO kernel tag {tag!r} for kid {kid}") + lo, hi = _CO_KID_BANDS[tag] assert lo <= kid < hi, ( - f"co kid {kid} (tag {e['tag']}) outside its band [{lo},{hi})" + f"co kid {kid} (tag {tag}) outside its band [{lo},{hi})" ) - assert kid not in out, f"duplicate co kid {kid} in {path}" - inst = _co_instance_from_json(arch, e) - # The name is the .co filename, the .co's extern "C" symbol AND the - # host launcher symbol. Two entries colliding here would silently - # overwrite one .co with the other. - assert inst.name not in seen_names, ( - f"co kids {seen_names[inst.name]} and {kid} generate the same " - f"symbol {inst.name!r} -- give one of them a distinct " - f'"variant" string' + assert kid not in instances, f"duplicate co kid {kid} in {path}" + instance = _co_instance_from_json(arch, entry) + _validate_co_instance(kid, instance) + assert instance.name not in seen_names, ( + f"co kids {seen_names[instance.name]} and {kid} generate the same " + f"symbol {instance.name!r}" ) - seen_names[inst.name] = kid - if require_image and not os.path.exists(co_image_path(path, inst)): - missing.append((kid, inst.name)) + seen_names[instance.name] = kid + if require_image and not os.path.exists(co_image_path(path, instance)): + missing.append((kid, instance.name)) continue - out[kid] = inst + instances[kid] = instance + if missing: + preview = ", ".join(f"{kid} ({name}.co)" for kid, name in missing[:5]) + suffix = f", ... and {len(missing) - 5} more" if len(missing) > 5 else "" print( f"[opus] {len(missing)} pre-compiled (.co) kid(s) dropped -- image " - f"not found under {os.path.dirname(path)}: " - # A full sweep is thousands of entries; naming every one turns a - # warning into a wall of text on every codegen run. - + ", ".join(f"{kid} ({name}.co)" for kid, name in missing[:5]) - + (f", ... and {len(missing) - 5} more" if len(missing) > 5 else "") - + ". Build them with csrc/opus_gemm/gen_co/build_co.py --llvm-bin " - ".", + f"not found under {os.path.dirname(path)}: {preview}{suffix}. Build " + "them with csrc/opus_gemm/gen_co/build_co.py.", file=sys.stderr, ) - return out + return instances gfx1250_4wave_co_kernels_list = { - kid: inst - for kid, inst in _load_co_kernels(CO_KERNELS_JSON).items() - if inst.kernel_tag.startswith("a16w16_4wave") + kid: instance + for kid, instance in _load_co_kernels(CO_KERNELS_JSON).items() + if instance.kernel_tag in _A16W16_CO_TAGS } -# Every JSON entry, image on disk or not. ONLY for build_co.py (the producer); -# everything that dispatches must use the filtered table above. +# The offline builder must see declared records even before their images exist. gfx1250_4wave_co_kernels_declared = { - kid: inst - for kid, inst in _load_co_kernels(CO_KERNELS_JSON, require_image=False).items() - if inst.kernel_tag.startswith("a16w16_4wave") + kid: instance + for kid, instance in _load_co_kernels( + CO_KERNELS_JSON, require_image=False + ).items() + if instance.kernel_tag in _A16W16_CO_TAGS } -# symbol name -> kid, for candidate selection. Keyed on the name rather than on -# (tile, cluster) because two variants may share a tile and differ only in the -# VGPR budget -- a (tile, cluster) key would silently drop one of them. GFX1250_4WAVE_CO_KID_OF = { - k.name: kid for kid, k in gfx1250_4wave_co_kernels_list.items() + instance.name: kid for kid, instance in gfx1250_4wave_co_kernels_list.items() +} +GFX1250_4WAVE_CO_KIDS = frozenset(gfx1250_4wave_co_kernels_list) + +_GFX1250_PRE_CO_KIDS = ( + frozenset(gfx1250_kernels_list) + | frozenset(gfx1250_clusterlaunch_kernels_list) + | GFX1250_SPLITK_FUSE_KIDS +) +assert not (GFX1250_4WAVE_CO_KIDS & _GFX1250_PRE_CO_KIDS), ( + "gfx1250 CO kids overlap an existing gfx1250 family" +) + +# Flatten the eight BMM launcher tags into the same canonical exact-kid +# registry used by every other OPUS family. +a8w8_mxscale_bmm_kernels_list = { + kid: instance + for family in a8w8_mxscale_bmm_kernel_lists + for kid, instance in family.items() } -GFX1250_4WAVE_CO_KIDS = frozenset(gfx1250_4wave_co_kernels_list.keys()) # combined list (used by production gen_instances / dispatch) kernels_list = { **a8w8_scale_kernels_list, **a8w8_kernels_list, + **a8w8_mxscale_bmm_kernels_list, **a16w16_kernels_list, **a16w16_kernels_list_nooob, **a16w16_kernels_list_cpol, @@ -1813,15 +1741,10 @@ def _load_co_kernels(path, require_image=True): **gfx1250_4wave_co_kernels_list, } -default_kernels_dict = { - (-1): OpusGemmInstance(512, 256, 256, 128, 4, 2, 16, 16, 128, 16, 16, 4, 1, 128, 128, "a8w8_scale", ["fp32_t"]), - (-2): OpusGemmInstance(512, 256, 256, 128, 2, 4, 16, 16, 128, 16, 16, 4, 0, 0, 0, "a8w8", ["fp32_t"]), - (-3): _a16w16(512, 256, 256, 64, 4, 16, 16, 32), # same as a16w16 #9 -} # fmt: on -# Subset-compile kid taxonomy (consumed by gen_instances.py for the `HEURISTIC_DEFAULT_KIDS ? +# Subset-compile kid taxonomy consumed by gen_instances.py. # Splitk kids: a16w16_flatmm_splitk pipeline (kid 200..223 + nooob mirror). SPLITK_KIDS = ( @@ -1830,8 +1753,30 @@ def _load_co_kernels(path, require_image=True): | frozenset(gfx942_splitk_kernels_list.keys()) | frozenset(gfx1250_kernels_list.keys()) | frozenset(gfx1250_clusterlaunch_kernels_list.keys()) + | frozenset(gfx1250_splitk_fuse_kernels_list.keys()) +) + +BMM_MXSCALE_WORKSPACE_TAGS = frozenset( + { + "a8w8_mxscale_bmm_flatmm_splitk", + "a8w8_mxscale_bmm_fused", + } +) +BMM_MXSCALE_WORKSPACE_KIDS = frozenset( + kid + for kid, instance in a8w8_mxscale_bmm_kernels_list.items() + if instance.kernel_tag in BMM_MXSCALE_WORKSPACE_TAGS and not instance.direct_only ) +_SUPPORTED_SPLITK_WORKSPACE_DTYPES = frozenset({"bf16_t", "fp32_t"}) +for _workspace_kid in SPLITK_KIDS: + _workspace_dtype = kernels_list[_workspace_kid].splitk_workspace_dtype + if _workspace_dtype not in _SUPPORTED_SPLITK_WORKSPACE_DTYPES: + raise ValueError( + f"workspace kid {_workspace_kid} must explicitly declare " + f"splitk_workspace_dtype, got {_workspace_dtype!r}" + ) + # Non-splitk a16w16-family kids: split-barrier 4..9 + cpol/nooob mirrors, persistent 300..315 + # cpol/nooob mirrors. NON_SPLITK_KIDS = ( @@ -1845,6 +1790,7 @@ def _load_co_kernels(path, require_image=True): | frozenset(a16w16_persistent_kernels_list_cpol_nooob.keys()) | frozenset(a16w16_mono_tile_kernels_list.keys()) | frozenset(gfx942_nosplit_kernels_list.keys()) + | GFX1250_4WAVE_CO_KIDS ) # 4g_safe kid families. Per-WG-tight BR sizing -- selectable for any shape @@ -1881,8 +1827,10 @@ def _load_co_kernels(path, require_image=True): | SPLITK_KIDS ) -# Heuristic-dispatch fallback kids (gfx950). -HEURISTIC_DEFAULT_KIDS_GFX950 = frozenset( +# Exact-id kernels kept in every default build. The high-level A16 caller-side +# heuristics below the tuned lookup are constrained to these ids; the unified +# public/C++ launch path still receives one already-resolved exact kid. +DEFAULT_COMPILED_KIDS_GFX950 = frozenset( { # splitk fallback (small M / non-aligned big M) 200, @@ -1897,9 +1845,9 @@ def _load_co_kernels(path, require_image=True): } ) -HEURISTIC_DEFAULT_KIDS_GFX942 = frozenset( +DEFAULT_COMPILED_KIDS_GFX942 = frozenset( { - # gfx942 heuristic dispatcher fallbacks. + # Representative exact-id launchers kept in default gfx942 builds. 10000, # gfx942 split-barrier 512x128x128x64 16x16x16 (large problem) 10001, # gfx942 p1 256x64x64x64 10003, # gfx942 p1_bk128 256x64x64x128 @@ -1918,14 +1866,9 @@ def _load_co_kernels(path, require_image=True): } ) -# gfx1250 has no shape-heuristic dispatch yet (tune-id entry only). This set -# is used purely to keep the kid in the subset-compile set S so the tune-id -# path can always reach it. -# Only the kids the C++ heuristic (opus_a16w16_heuristic_kid_gfx1250) can return -# must be force-compiled as the always-available (M,N,K) fallback. Every other -# plain kid and ALL clusterlaunch kids are compiled on demand by the tuner -# (candidate selection + sidecar expansion), so default builds stay small. -HEURISTIC_DEFAULT_KIDS_GFX1250 = ( +# Keep representative plain/clusterlaunch workspace kids and all CO host kids +# in default gfx1250 builds; the tuner compiles other device kids on demand. +DEFAULT_COMPILED_KIDS_GFX1250 = ( frozenset( GFX1250_PLAIN_KID_OF[_t] for _t in ( @@ -1936,59 +1879,210 @@ def _load_co_kernels(path, require_image=True): (32, 64, 128), (32, 128, 128), ) - # The 4wave_co kids are NOT reachable from the C++ shape heuristic; they are - # here only for this set's other job -- membership in the subset-compile set - # S (gen_instances.py), so the explicit-kernelId path can always reach them. - # Their device side is a pre-built .co, so "compiling" one is just the host - # launcher: no per-kid device TU is emitted (see the co emit branch in - # codegen/gen_instances_gfx1250.py). ) + | frozenset({GFX1250_CLUSTERLAUNCH_KID_OF[(16, 32, 128, 2, 1)]}) | GFX1250_4WAVE_CO_KIDS ) -HEURISTIC_DEFAULT_KIDS = ( - HEURISTIC_DEFAULT_KIDS_GFX950 - | HEURISTIC_DEFAULT_KIDS_GFX942 - | HEURISTIC_DEFAULT_KIDS_GFX1250 +DEFAULT_COMPILED_KIDS = ( + DEFAULT_COMPILED_KIDS_GFX950 + | DEFAULT_COMPILED_KIDS_GFX942 + | DEFAULT_COMPILED_KIDS_GFX1250 ) -HEURISTIC_DEFAULT_KIDS_BY_ARCH = { - "gfx950": HEURISTIC_DEFAULT_KIDS_GFX950, - "gfx942": HEURISTIC_DEFAULT_KIDS_GFX942, - "gfx1250": HEURISTIC_DEFAULT_KIDS_GFX1250, +DEFAULT_COMPILED_KIDS_BY_ARCH = { + "gfx950": DEFAULT_COMPILED_KIDS_GFX950, + "gfx942": DEFAULT_COMPILED_KIDS_GFX942, + "gfx1250": DEFAULT_COMPILED_KIDS_GFX1250, } -def heuristic_kids_for_arch(arches): - """Return the heuristic-default kid subset whose arch_prefix matches. - - ``arches`` is an iterable of lowercase arch strings (e.g. ``{"gfx942"}``) - or ``None`` (caller does not know / multi-arch build) -- in the ``None`` - case the full union is returned so the legacy multi-arch behaviour is - preserved. - """ +def default_compiled_kids_for_arch(arches): + """Return the default exact-id compile floor for requested arches.""" if arches is None: - return HEURISTIC_DEFAULT_KIDS + return DEFAULT_COMPILED_KIDS arches = {a.lower() for a in arches} out = frozenset() for arch in arches: - out = out | HEURISTIC_DEFAULT_KIDS_BY_ARCH.get(arch, frozenset()) + out = out | DEFAULT_COMPILED_KIDS_BY_ARCH.get(arch, frozenset()) return out +# Map each architecture and interface to its registered kernel tags. +# An empty set means that the interface exists but has no kernel on that arch. +OPUS_KERNEL_TAGS_BY_ARCH_FAMILY = { + "gfx950": { + "a16w16": frozenset( + { + "a16w16", + "a16w16_flatmm", + "a16w16_flatmm_splitk", + "a16w16_mono_tile", + "a16w16_persistent", + } + ), + "a8w8": frozenset({"a8w8"}), + "a8w8_blockscale": frozenset({"a8w8_scale"}), + "a8w8_mxscale_bmm": frozenset( + { + "a8w8_mxscale_bmm_flatmm_splitk", + "a8w8_mxscale_bmm_fused", + "a8w8_mxscale_bmm_minterleave", + "a8w8_mxscale_bmm_mouter", + "a8w8_mxscale_bmm_mouter_tunable", + "a8w8_mxscale_bmm_pipeline", + "a8w8_mxscale_bmm_wave8n2", + "a8w8_mxscale_bmm_wave4m2_selfload", + } + ), + "a8w8_blockscale_bpreshuffle": frozenset(), + }, + "gfx942": { + "a16w16": frozenset( + { + "a16w16_em3en4_lds1_pgr2_sk", + "a16w16_kbuf1_large_tile", + "a16w16_kbuf1_sk", + "a16w16_kbuf2v", + "a16w16_kbuf2v_bk128", + "a16w16_kbuf2v_bk128_sk", + "a16w16_kbuf2v_sk", + "a16w16_quad_mfma32_kbuf1", + "a16w16_quad_mfma32_kbuf1_sk", + "a16w16_wave_k_coop", + "a16w16_wave_k_coop_accum", + } + ), + "a8w8": frozenset(), + "a8w8_blockscale": frozenset(), + "a8w8_mxscale_bmm": frozenset(), + "a8w8_blockscale_bpreshuffle": frozenset( + {"a8w8_blockscale_bpreshuffle_singlebuf"} + ), + }, + "gfx1250": { + "a16w16": frozenset( + { + "a16w16_cluster_tdm_splitk_ws", + "a16w16_clusterlaunch_tdm_splitk_fuse", + "a16w16_clusterlaunch_tdm_splitk_ws", + "a16w16_4wave_co", + "a16w16_4wave_wl_co", + } + ), + "a8w8": frozenset(), + "a8w8_blockscale": frozenset(), + "a8w8_mxscale_bmm": frozenset(), + "a8w8_blockscale_bpreshuffle": frozenset(), + }, +} + +# Always include these A8W8 kernels in matching-architecture subset builds. +OPUS_MANDATORY_A8_KIDS = { + "gfx950": frozenset({1, 2}), + "gfx942": frozenset({11000}), + "gfx1250": frozenset(), +} + + +def canonical_output_dtype(output_dtype) -> str | None: + """Normalize supported output dtype names for registry lookup.""" + if output_dtype is None: + return None + value = str(output_dtype).strip().lower() + aliases = { + "bf16": "bf16_t", + "bfloat16": "bf16_t", + "bf16_t": "bf16_t", + "torch.bfloat16": "bf16_t", + "fp32": "fp32_t", + "float": "fp32_t", + "float32": "fp32_t", + "fp32_t": "fp32_t", + "torch.float32": "fp32_t", + } + return aliases.get(value, value) + + +def get_kernel_instance( + arch: str, + family: str, + kid: int, + output_dtype=None, +) -> OpusGemmInstance | None: + """Return a kernel registered for ``(arch, interface, kid, Y.dtype)``.""" + arch = str(arch).lower() + family = str(family).lower() + family_tags = OPUS_KERNEL_TAGS_BY_ARCH_FAMILY.get(arch, {}).get(family) + if family_tags is None: + return None + + try: + kid = int(kid) + except (TypeError, ValueError): + return None + + instance = kernels_list.get(kid) + if instance is None or instance.kernel_tag not in family_tags: + return None + instance_arch = (instance.arch_prefix or "gfx950").lower() + if instance_arch != arch: + return None + + dtype = canonical_output_dtype(output_dtype) + if dtype is not None: + # Workspace reducers own the final cast, so their host dispatch + # specialization in ``output_dtypes`` is not the Y dtype contract. + if family == "a16w16" and kid in SPLITK_KIDS: + # The current gfx942 BF16-workspace reducer is exact-N and writes + # BF16 only. Other A16 workspace reducers support BF16/FP32 Y. + allowed = ( + {"bf16_t"} + if arch == "gfx942" and instance.splitk_workspace_dtype == "bf16_t" + else {"bf16_t", "fp32_t"} + ) + output_compatible = dtype in allowed + elif family == "a8w8_mxscale_bmm": + output_compatible = dtype in {"bf16_t", "fp32_t"} + else: + output_compatible = dtype in instance.output_dtypes + if not output_compatible: + return None + return instance + + +def kernel_needs_external_workspace(arch: str, family: str, kid: int) -> bool: + """Return whether a registered kernel requires caller-owned workspace. + + Unknown logical keys are errors rather than ``False``: treating an unknown + kid as a non-workspace kernel would let a caller launch it without the + allocation required for memory safety. Capability comes from the existing + ``SPLITK_KIDS`` registry, never from a numeric kid range or tag substring. + The registry includes all enabled two-stage reducers. If the experimental + gfx1250 fused family is re-enabled, its first SplitK-1 WGs also publish + external partial tiles and therefore enter this same capability set. + """ + instance = get_kernel_instance(arch, family, kid) + if instance is None: + raise KeyError( + f"unknown OPUS kernel (arch={arch!r}, family={family!r}, kid={kid!r})" + ) + return int(kid) in SPLITK_KIDS + + def _opus_sidecar_path(): """Return the on-disk path of the subset-compile sidecar. Lives in ``{bd_dir}/`` (one level above the per-module build dir) so it survives ``aiter.jit.core.clear_build("module_deepgemm_opus")`` -- - which ``build_module()`` calls when ``AITER_REBUILD == 1`` -- and is - therefore seeds the last successfully compiled set into the next codegen. + which ``build_module()`` calls when ``AITER_REBUILD == 1`` -- and + seeds the last successfully compiled set into the next codegen. The tuner passes new candidates through ``--extra_kids``; it does not advance this file before compiling. JIT atomically copies the generated sidecar back here after installing the .so, independently of source-cache publication. Its adjacent ``.receipt`` binds the contents to that binary; - the tuner requires both to match before skipping a rebuild. A plain runtime - dispatch uses the CSV/C++ lookup, not this file. + the tuner requires both to match before skipping a rebuild. Runtime exact + dispatch resolves the caller-selected kid without reading this file. """ # Import lazily to avoid circular import at module load (aiter imports # opus_gemm_common, opus_gemm_common imports aiter.jit.core). diff --git a/csrc/opus_gemm/opus_gemm_tune.py b/csrc/opus_gemm/opus_gemm_tune.py index 8dbeda9da8..9f2cacfeb4 100644 --- a/csrc/opus_gemm/opus_gemm_tune.py +++ b/csrc/opus_gemm/opus_gemm_tune.py @@ -1,38 +1,10 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2025-2026, Advanced Micro Devices, Inc. All rights reserved. -"""[DEBUG-ONLY] Single-shape / single-kid opus a16w16 tuner. - -This script used to be the production opus tuning entry point and wrote -directly into a private CSV under aiter/ops/opus/configs/. Production -tuning has moved to gradlib: - - python3 gradlib/gemm_tuner.py --libtype opus - # or as part of a multi-backend tune: - python3 gradlib/gemm_tuner.py --libtype all - -gradlib writes to aiter/configs/bf16_tuned_gemm.csv (or whatever the -user passes via --tuned_file / GTUNE_TUNED), stamping every opus row -with `libtype=='opus'`. The opus runtime dispatch -(aiter/ops/opus/common.py) reads those rows from the global CSV. - -This file is retained for two reasons only: - - 1. **Single-(M,N,K) smoke / debug**: hand-running a specific kid against - a specific shape (-m M -n N -k K --kid K --splitK S) to compare - against the gradlib winner or to investigate a bug. - 2. **Source of truth for tune-time helpers**: candidate_kids_for_shape, - candidate_splitK, kid_rejects_shape / kid_rejects_bias, and - _ensure_kids_compiled live here. They are imported by: - - gradlib's GemmTuner (gradlib/gradlib/GemmTuner.py) for the - production `--libtype opus` path, - - this script's own `tune()` for the single-shape debug path. - csrc/opus_gemm/opus_gemm_common.py only owns the data constants - (SPLITK_KIDS / NON_SPLITK_KIDS / BIAS_AWARE_KIDS / - HEURISTIC_DEFAULT_KIDS) plus _opus_sidecar_path(); it does NOT - re-export any tune-time helper. - -The default output path is /tmp/opus_debug_tuned.csv so this script can -never accidentally pollute the global aiter/configs/ tree. +"""Debug-only single-shape OPUS A16W16 tuner. + +Production tuning uses ``gradlib/gemm_tuner.py``. This module keeps the +single-kid runner and candidate helpers imported by Gradlib; output defaults +to ``/tmp``. """ import json @@ -86,6 +58,7 @@ def _patch_flaky_hip_device_count(): # opus_gemm_common is a sibling file in csrc/opus_gemm/. from opus_gemm_common import ( BIAS_AWARE_KIDS, + GFX942_BF16WS_EXACT_N, GFX1250_CLUSTERLAUNCH_KID_OF, GFX1250_PLAIN_KID_OF, GFX1250_SPLITK_FUSE_ENABLED, @@ -94,29 +67,12 @@ def _patch_flaky_hip_device_count(): NON_SPLITK_KIDS, SPLITK_KIDS, _opus_sidecar_path, - a16w16_flatmm_kernels_list, - a16w16_flatmm_splitk_kernels_list, - a16w16_flatmm_splitk_kernels_list_nooob, - a16w16_kernels_list, - a16w16_kernels_list_cpol, - a16w16_kernels_list_cpol_nooob, - a16w16_kernels_list_nooob, - a16w16_persistent_kernels_list, - a16w16_persistent_kernels_list_cpol, - a16w16_persistent_kernels_list_cpol_nooob, - a16w16_persistent_kernels_list_nooob, - gfx942_nosplit_kernels_list, - gfx942_splitk_kernels_list, - gfx1250_4wave_co_kernels_list, - gfx1250_clusterlaunch_kernels_list, - gfx1250_kernels_list, - gfx1250_splitk_fuse_kernels_list, + a16w16_flatmm_prefetch_k_iter, + kernels_list, ) from aiter import dtypes, logger -from aiter.ops.opus.gemm_op_a16w16 import ( - opus_gemm_a16w16_tune as _opus_gemm_a16w16_tune, -) +from aiter.ops.opus import opus_bmm as _opus_bmm from aiter.utility.base_tuner import INVALID_TIME, GemmCommonTuner from aiter.utility.mp_tuner import mp_tuner @@ -138,6 +94,9 @@ def _patch_flaky_hip_device_count(): # occupancy fit; this only widens which splitK values get benchmarked on the # selected tiles (so some over-occupancy / higher-split candidates are explored). GFX1250_SPLITK_WINDOW_HI_MULT = 4 +GFX1250_CO_TOP_TILES = 6 +GFX1250_CO_TOP_CLUSTERS = 6 +_A16W16_CO_TAGS = frozenset({"a16w16_4wave_co", "a16w16_4wave_wl_co"}) # Tune-time host helpers (defined here, not in opus_gemm_common.py). @@ -145,17 +104,6 @@ def _patch_flaky_hip_device_count(): OCCUPANCY_TILE_BM = 128 OCCUPANCY_TILE_BN = 128 -BF16WS_EXACT_REDUCE_SHAPES = ( - (64, 8), - (128, 4), - (256, 2), - (512, 1), - (1024, 4), - (1024, 2), - (1024, 1), - (2048, 1), -) - EVEN_LOOP_SPLITK_TAGS = frozenset( ( "a16w16_kbuf2v_sk", @@ -174,35 +122,23 @@ def _round_up(a: int, b: int) -> int: def _kid_uses_bf16_workspace(k_inst): - return getattr(k_inst, "splitk_workspace_dtype", "fp32_t") == "bf16_t" + return k_inst.splitk_workspace_dtype == "bf16_t" def _kid_rejects_outdtype(k_inst, out_dtype): - return _kid_uses_bf16_workspace(k_inst) and out_dtype is not dtypes.bf16 + # gfx942's bf16-workspace launchers currently require bf16 Y. gfx1250's + # #4246 two-stage path may write bf16 partials while reducing to bf16 or + # fp32 Y, so workspace dtype alone must not constrain the output dtype. + return ( + k_inst.arch_prefix == "gfx942" + and _kid_uses_bf16_workspace(k_inst) + and out_dtype is not dtypes.bf16 + ) def _flatmm_splitk_pfk(k) -> int: - """Host-side computation of Traits::prefetch_k_iter for a splitk instance. - - Mirrors opus_flatmm_splitk_traits_gfx950's formula so the host can - pre-compute the per-split iter budget without a device call. Hardcodes - LDS=163840 (gfx950), same convention as the traits struct. - """ - sizeof_da = 2 # bf16 - LOAD_GROUP_M = 64 if k.W_M >= 32 else 32 - LOAD_GROUP_N = 64 if k.W_N >= 32 else 32 - LOAD_GROUP_K = k.W_K * 2 - num_m = k.B_M // LOAD_GROUP_M - num_n = k.B_N // LOAD_GROUP_N - num_k = k.B_K // LOAD_GROUP_K - smem_linear = 64 * 16 // sizeof_da # WARP_SIZE=64 - smem_sub = smem_linear // LOAD_GROUP_K - slots = LOAD_GROUP_M // smem_sub - padding = 16 // sizeof_da if k.W_M >= 32 else 2 * 16 // sizeof_da - per_glsz = slots * (smem_linear + padding) * sizeof_da - per_iter = (num_m + num_n) * num_k * per_glsz - lds_total = 163840 - return max(1, (lds_total // max(k.WG_PER_CU, 1)) // max(per_iter, 1)) + """Backward-compatible name for the canonical metadata calculation.""" + return a16w16_flatmm_prefetch_k_iter(k) def _gfx1250_occ_cost(total_wg: int, cu_num: int) -> float: @@ -248,166 +184,146 @@ def _dist(sk): return [min(nz, key=_dist)] -# Pre-compiled (.co) candidate filter. The family is 12 tiles x up to 16 cluster -# dims x 3 wave layouts = 204 kids, and it used to go into every sweep whole. -# Scored against measured per-kid latency on 18 shapes (square, wide-N, narrow-N, -# M=1, ragged, K-tail), the knee is at 5 cluster dims and 4 tiles; one step of -# headroom on each keeps the measured winner for 17 of the 18 and costs 0.16% on -# the last, while cutting the sweep 3.2x. -GFX1250_CO_TOP_TILES = 6 -GFX1250_CO_TOP_CLUSTERS = 6 - -GFX1250_FUSE_TOP_SPLITK = 3 # best-N split_k (by grid-occupancy fit) per fuse tile +GFX1250_FUSE_TOP_SPLITK = 3 +GFX1250_FUSE_MAX_SPLITK = max((key[4] for key in GFX1250_SPLITK_FUSE_KID_OF), default=0) def _gfx1250_fuse_kids_for_tile(M, N, K, cu_num, bm, bn, bk): - """Bounded fuse-kid candidates for ONE tile at shape (M, N, K). + """Return a bounded exact-kid fused set for one tile. - Fuse writes full-N C tiles (needs N % B_N == 0) and clusters n_cluster N-tile - peers (A-multicast; cluster=(SplitK, n_cluster, 1)). To keep the sweep small we - pick: split_k by grid-occupancy fit (top-N), n_cluster in {1 (baseline), max - feasible (most A-multicast)}, and both workspace dtypes. grid WG = split_k * - ceil(N/B_N) * ceil(M/B_M) (n_cluster only groups existing WGs, no extra WG). + SplitK and N-cluster geometry are compile-time, so candidate selection + chooses exact registry entries rather than sweeping a runtime splitK knob. """ if N % bn != 0: return [] layout = "tileN" if bm == 16 else "tileM" - ntn = N // bn - ntm = _ceil_div(M, bm) - base_wg = ntn * ntm + num_tiles_n = N // bn + num_tiles_m = _ceil_div(M, bm) + base_wg = num_tiles_n * num_tiles_m k_steps = _ceil_div(K, bk) - valid_sk = [] - for sk in range(2, 16): # split_k 2..15 - # Balanced K-tile split: every split WG gets >=1 tile iff sk <= k_steps - # (the K tail is TDM-clamped, not handled by emptying WGs). - if sk > k_steps: - continue - valid_sk.append(sk) - if not valid_sk: - return [] - valid_sk.sort(key=lambda sk: (_gfx1250_occ_cost(base_wg * sk, cu_num), sk)) - sk_sel = sorted(set(valid_sk[:GFX1250_FUSE_TOP_SPLITK])) - nc_valid = [nc for nc in range(1, 6) if ntn % nc == 0] # exact N-fill, <=5 - nc_sel = sorted({1, max(nc_valid)}) if nc_valid else [1] + valid_n_cluster = [nc for nc in range(1, 6) if num_tiles_n % nc == 0] + selected_n_cluster = sorted({1, max(valid_n_cluster)}) if valid_n_cluster else [1] + out = [] - for nc in nc_sel: - for sk in sk_sel: - if sk * nc > 16: # cluster budget SplitK*n_cluster <= 16 - continue - for ws in ("bf16_t", "fp32_t"): - kid = GFX1250_SPLITK_FUSE_KID_OF.get((bm, bn, bk, layout, sk, nc, ws)) + for n_cluster in selected_n_cluster: + for workspace_dtype in ("bf16_t", "fp32_t"): + # Select SplitK independently for each exact registry family. + # BF16 and FP32 have different supported SplitK ranges (currently + # 2..15 versus 2..8), and the cluster budget narrows those ranges + # further. Sharing one top-K list can therefore silently erase a + # valid dtype family when its occupancy optimum is outside that + # family's registry range. + valid_split_k = [ + split_k + for split_k in range(2, min(k_steps, GFX1250_FUSE_MAX_SPLITK) + 1) + if ( + bm, + bn, + bk, + layout, + split_k, + n_cluster, + workspace_dtype, + ) + in GFX1250_SPLITK_FUSE_KID_OF + ] + valid_split_k.sort( + key=lambda sk: (_gfx1250_occ_cost(base_wg * sk, cu_num), sk) + ) + for split_k in sorted(valid_split_k[:GFX1250_FUSE_TOP_SPLITK]): + kid = GFX1250_SPLITK_FUSE_KID_OF.get( + ( + bm, + bn, + bk, + layout, + split_k, + n_cluster, + workspace_dtype, + ) + ) if kid is not None: out.append(kid) return out def _gfx1250_fuse_candidates(M, N, K, cu_num, top_tiles=GFX1250_TOP_TILES): - """Fuse candidate set = top-N fuse tiles (by grid-occupancy fit) x the bounded - per-tile (split_k, n_cluster, ws) selection. Replaces dumping all ~1.4k fuse - kids (which explodes to ~900 candidates for wide-N small-M shapes).""" - fuse_tiles = sorted({(t[0], t[1], t[2]) for t in GFX1250_SPLITK_FUSE_KID_OF}) - fuse_tiles = [t for t in fuse_tiles if N % t[1] == 0] - - def _score(t): - bm, bn, bk = t + """Select fused exact kids when that experimental family is registered.""" + fuse_tiles = sorted( + {(key[0], key[1], key[2]) for key in GFX1250_SPLITK_FUSE_KID_OF} + ) + fuse_tiles = [tile for tile in fuse_tiles if N % tile[1] == 0] + + def _score(tile): + bm, bn, bk = tile base = (N // bn) * _ceil_div(M, bm) - ks = max(1, _ceil_div(K, bk)) + k_steps = max(1, _ceil_div(K, bk)) return min( - (_gfx1250_occ_cost(base * sk, cu_num) for sk in range(1, min(16, ks) + 1)), + ( + _gfx1250_occ_cost(base * split_k, cu_num) + for split_k in range(2, min(GFX1250_FUSE_MAX_SPLITK, k_steps) + 1) + ), default=float("inf"), ) fuse_tiles.sort(key=_score) out = set() - for t in fuse_tiles[:top_tiles]: - out.update(_gfx1250_fuse_kids_for_tile(M, N, K, cu_num, *t)) + for tile in fuse_tiles[:top_tiles]: + out.update(_gfx1250_fuse_kids_for_tile(M, N, K, cu_num, *tile)) return out def _gfx1250_cluster_waste(gx: int, gy: int, cwm: int, cwn: int) -> float: - """Share of the launched workgroups that own no tile, in [0, 1). - - The clusterlaunch launcher rounds the tile grid up to whole (cwm x cwn) - clusters; the surplus workgroups are dispatched, arrive at the cluster - barrier and leave at tile_oob without issuing a TDM. They are cheap - individually but they still occupy a cluster slot, and they shrink the - multicast group the survivors merge with, so the ratio is what the sweep - budgets on. split_k scales the launched grid and the tile count alike, so - the ratio does not depend on it. - """ + """Return the fraction of rounded cluster workgroups with no output tile.""" launched = _round_up(gx, cwm) * _round_up(gy, cwn) return (launched - gx * gy) / launched -def _gfx1250_cluster_dims_for_grid(gx, gy, avail, top_clusters, cu_num=0): - """Rank the (cwm, cwn) worth benchmarking for a gx x gy tile grid. - - ``avail`` maps (cwm, cwn) -> kid for one tile. Returns the selected dims, - best first, at most ``top_clusters`` of them (possibly none). - - A cluster side wider than the grid it rides on is dropped outright: the - peers it would multicast to are exactly the workgroups that leave at - tile_oob, so it buys no extra sharing and only launches dead workgroups - (gx == 1 with cwm == 4 launches 4x the workgroups for a single M-tile - column). What is left must fit the GFX1250_MAX_CLUSTER_WASTE budget; if - nothing does, only the tightest-fitting dims stay, so a shape that cannot - fill any cluster still gets its least-wasteful candidate benchmarked. - - Both of those read the surplus workgroups as cost, which holds only once the - launch is big enough for them to displace real work. Pass ``cu_num`` to lift - them while the whole launch still fits one machine's worth: there the dead - workgroups displace nothing and the wider multicast is free. Measured on the - .co family, 129x257x384 on a 64x64 tile is a 3x5 grid whose fastest kid is - c4x4 -- vetoed without this, for 4.9%. Left off (0) by default so the - clusterlaunch sweep keeps the behaviour it was tuned with. - - Ranking is by waste first (bucketed, so a fit within a bucket does not - outrank a bigger multicast group over a rounding crumb), then by the size of - the multicast group, then by how well the cluster's aspect matches the tile - grid's. Measured on a 128x128x128 tile at K=1024..2048, cluster dims that - fill the grid exactly: a 12x12 grid runs 13.5 us at 1x2 and 10.7 us at 3x3, a - 16x16 grid 27.6 us at 1x2 and 18.2 us at 4x4, so a wider cluster is worth - more than anything else here. Aspect breaks ties between equal-size dims and - is worth as much as a size step on a lopsided grid: on a 4x32 grid 2x4 runs - 11.2 us against 13.3 us for 4x2, and 1x4 (11.4 us) beats both 2x2 (13.9 us) - and 4x1 (20.1 us). - """ +def _gfx1250_cluster_dims_for_grid(gx, gy, available, top_clusters, cu_num=0): + """Rank cluster dims by round-up waste; ``cu_num`` enables CO's CU-fit rule.""" - def _free(cwm, cwn): + def _roundup_is_free(cwm, cwn): return cu_num > 0 and _round_up(gx, cwm) * _round_up(gy, cwn) <= cu_num - feas = [ + feasible = [ (_gfx1250_cluster_waste(gx, gy, cwm, cwn), cwm, cwn) - for (cwm, cwn) in avail - if (cwm <= gx and cwn <= gy) or _free(cwm, cwn) + for cwm, cwn in available + if (cwm <= gx and cwn <= gy) or _roundup_is_free(cwm, cwn) ] - if not feas: + if not feasible: return [] - max_waste = 1.0 if _free(1, 1) else GFX1250_MAX_CLUSTER_WASTE - within = [f for f in feas if f[0] <= max_waste] - if not within: - tightest = min(f[0] for f in feas) - within = [f for f in feas if f[0] <= tightest] + admitted = [ + candidate + for candidate in feasible + if candidate[0] <= GFX1250_MAX_CLUSTER_WASTE + or _roundup_is_free(candidate[1], candidate[2]) + ] + if not admitted: + tightest = min(candidate[0] for candidate in feasible) + admitted = [candidate for candidate in feasible if candidate[0] == tightest] bucket = GFX1250_CLUSTER_WASTE_BUCKET - within.sort( - key=lambda f: ( - f[0] if bucket <= 0 else int(f[0] / bucket), # tile-less workgroups - -(f[1] * f[2]), # widest multicast group - abs(f[1] * gy - f[2] * gx), # 0 when cwm/cwn matches the grid aspect - (f[1], f[2]), # keep the pick independent of the kid-table order + admitted.sort( + key=lambda candidate: ( + candidate[0] if bucket <= 0 else int(candidate[0] / bucket), + -(candidate[1] * candidate[2]), + abs(candidate[1] * gy - candidate[2] * gx), + (candidate[1], candidate[2]), ) ) - return [(cwm, cwn) for _w, cwm, cwn in within[:top_clusters]] - - -# tile -> cluster dims -> [kid]. Built once; the wave layouts (and any future -# VGPR budgets) of one (tile, cluster) stay together because nothing host-side -# can rank them -- which of them wins is exactly what the sweep is for. -_GFX1250_CO_BY_TILE: dict = {} -for _kid, _k in gfx1250_4wave_co_kernels_list.items(): - _GFX1250_CO_BY_TILE.setdefault((_k.B_M, _k.B_N, _k.B_K), {}).setdefault( - (_k.cluster_wg_m, _k.cluster_wg_n), [] - ).append(_kid) + return [(cwm, cwn) for _waste, cwm, cwn in admitted[:top_clusters]] + + +# tile -> cluster dims -> all exact kids. Multiple wave-layout variants for one +# geometry remain together because only measurement can rank them. +_GFX1250_CO_BY_TILE: dict[tuple[int, int, int], dict[tuple[int, int], list[int]]] = {} +for _co_kid, _co_instance in kernels_list.items(): + if _co_instance.kernel_tag not in _A16W16_CO_TAGS: + continue + _GFX1250_CO_BY_TILE.setdefault( + (_co_instance.B_M, _co_instance.B_N, _co_instance.B_K), {} + ).setdefault((_co_instance.cluster_wg_m, _co_instance.cluster_wg_n), []).append( + _co_kid + ) def _gfx1250_co_candidates( @@ -418,53 +334,36 @@ def _gfx1250_co_candidates( top_tiles=GFX1250_CO_TOP_TILES, top_clusters=GFX1250_CO_TOP_CLUSTERS, ): - """Pre-compiled (.co) kids worth benchmarking for this shape. - - Same shape as the plain/clusterlaunch selection -- top-N tiles by - grid-occupancy fit, then top-N cluster dims per tile -- with one difference - that matters: this family has NO split-K, so its grid is exactly - ceil(M/B_M) * ceil(N/B_N) and the tile score has nothing to minimise over. - - Every tile carries a c1x1 entry and c1x1 is feasible for any non-empty grid, - so a selected tile always contributes at least one kid. - """ + """Select a bounded tile/cluster subset of pre-built non-split-K kids.""" - def _tile_score(t): - bm, bn, _bk = t + def _tile_score(tile): + bm, bn, _bk = tile return _gfx1250_occ_cost(_ceil_div(M, bm) * _ceil_div(N, bn), cu_num) - tiles = sorted(_GFX1250_CO_BY_TILE, key=lambda t: (_tile_score(t), t)) - out: set[int] = set() - for t in tiles[:top_tiles]: - bm, bn, _bk = t + selected = set() + tiles = sorted(_GFX1250_CO_BY_TILE, key=lambda tile: (_tile_score(tile), tile)) + for tile in tiles[:top_tiles]: + bm, bn, _bk = tile gx, gy = _ceil_div(M, bm), _ceil_div(N, bn) - avail = _GFX1250_CO_BY_TILE[t] + available = _GFX1250_CO_BY_TILE[tile] for dims in _gfx1250_cluster_dims_for_grid( - gx, gy, avail, top_clusters, cu_num=cu_num + gx, gy, available, top_clusters, cu_num ): - out.update(avail[dims]) - return out + selected.update(available[dims]) + return selected def _gfx1250_select_candidates( - M, N, K, cu_num, top_tiles=GFX1250_TOP_TILES, top_clusters=GFX1250_TOP_CLUSTERS + M, + N, + K, + cu_num, + top_tiles=GFX1250_TOP_TILES, + top_clusters=GFX1250_TOP_CLUSTERS, + *, + include_fused=True, ): - """gfx1250 candidate kid set for shape (M,N,K): top-N tiles x {plain + cluster - dims}. - - 1. Tile (top-8): score each plain tile by its best grid-occupancy fit over - splitK in [1, min(16, k_steps)] (occ cost + tiny splitK bias). Smallest - score wins; take top GFX1250_TOP_TILES. - 2. For each selected tile, always include its plain (P=3) kid. - 3. Cluster: any (cwm, cwn) in [1, GFX1250_MAX_CLUSTER_SIDE]^2 except (1,1) can - run this shape -- the clusterlaunch pipeline rounds the grid up to whole - clusters and the tile-less workgroups leave at their cluster barrier, so - cluster-fill divisibility is no longer a constraint. What the round-up costs - is workgroups, so the sweep is bounded by _gfx1250_cluster_dims_for_grid: - drop a cluster side wider than the grid it rides on, drop anything past the - GFX1250_MAX_CLUSTER_WASTE tile-less budget, then rank by (waste bucket, - widest multicast group, grid-aspect match) and keep top_clusters. - """ + """Select top tiles and launcher-supported rounded cluster grids.""" def _tile_score(bm, bn, bk): gx = _ceil_div(M, bm) @@ -483,31 +382,24 @@ def _tile_score(bm, bn, bk): bm, bn, _bk = t gx = _ceil_div(M, bm) gy = _ceil_div(N, bn) - # The kid table also carries cwn==5 dims (TDM multicast fans out to 5 WGs); - # they are left out of the sweep by the [1, MAX_CLUSTER_SIDE]^2 bound. - avail = { + available = { (cwm, cwn): kid for (tbm, tbn, tbk, cwm, cwn), kid in GFX1250_CLUSTERLAUNCH_KID_OF.items() if (tbm, tbn, tbk) == t and cwm <= GFX1250_MAX_CLUSTER_SIDE and cwn <= GFX1250_MAX_CLUSTER_SIDE } - for dims in _gfx1250_cluster_dims_for_grid(gx, gy, avail, top_clusters): - sel.add(avail[dims]) - - # FUSED single-kernel split-K kids: bounded per-shape selection (top-N fuse - # tiles x occupancy-fit split_k x {baseline, max A-multicast n_cluster} x ws), - # instead of dumping all ~1.4k fuse kids (which explodes to ~900 candidates for - # wide-N small-M shapes). kid_rejects_shape still prunes any residual invalids. - # The family is unregistered while its pipeline is being fixed, which is why - # the sweep is 496 kids (28 plain + 468 clusterlaunch) and not ~1.9k. - if GFX1250_SPLITK_FUSE_ENABLED: + for dims in _gfx1250_cluster_dims_for_grid(gx, gy, available, top_clusters): + sel.add(available[dims]) + + # Fused exact kids: bounded by tile, occupancy-fit compile-time SplitK, + # baseline/max feasible N-cluster, and both workspace dtypes. + if ( + GFX1250_SPLITK_FUSE_ENABLED + and include_fused + and os.environ.get("OPUS_TUNE_NO_FUSE") != "1" + ): sel |= _gfx1250_fuse_candidates(M, N, K, cu_num) - - # Pre-compiled (.co) kids. These used to go in wholesale, on the reading that - # the family was a handful of hand-picked variants with nothing for a tile - # ranking to choose between. It is now a swept 12-tile x 16-cluster x - # 3-layout space, so it gets the same treatment as the rest. sel |= _gfx1250_co_candidates(M, N, K, cu_num) return frozenset(sel) @@ -523,17 +415,21 @@ def candidate_splitK(M: int, N: int, K: int, batch: int, cu_num: int, k_inst): Workspace size cap (added with the >4 GiB reduce-BR fix): each candidate splitK value must keep - split_k * batch * padded_M * padded_N * 4 <= UINT32_MAX + split_k * batch * padded_M * padded_N * sizeof(exact-kid D_WS) + <= UINT32_MAX so the splitk_reduce_kernel's buffer-resource num_records stays in range. We compute the same per-slice budget the host reject uses and silently drop any split_k that would push workspace past 4 GiB. """ - # Pre-compiled (.co) kids have no split-K at all: no workspace, no partials, - # no reduce kernel, and the launcher AITER_CHECKs splitK <= 1. Probing any - # other value would just collect exceptions. - if k_inst.kernel_tag == "a16w16_4wave_co": + if k_inst.kernel_tag in _A16W16_CO_TAGS: return [0] + if k_inst.kernel_tag == "a16w16_clusterlaunch_tdm_splitk_fuse": + # Runtime splitK is ignored by this family. Store the baked value in + # tuning rows so the CSV remains self-describing and never suggests a + # false dynamic choice. + return [int(k_inst.fuse_split_k)] + B_K = k_inst.B_K total_iters = _ceil_div(K, B_K) # gfx1250 cluster/TDM split-K triple-buffers but tolerates any k_steps>=1 @@ -556,9 +452,8 @@ def candidate_splitK(M: int, N: int, K: int, batch: int, cu_num: int, k_inst): # Workspace 4 GiB cap. padded_M = _ceil_div(M, k_inst.B_M) * k_inst.B_M padded_N = _ceil_div(N, k_inst.B_N) * k_inst.B_N - per_slice_bytes = ( - batch * padded_M * padded_N * (2 if _kid_uses_bf16_workspace(k_inst) else 4) - ) + workspace_bytes = 2 if _kid_uses_bf16_workspace(k_inst) else 4 + per_slice_bytes = batch * padded_M * padded_N * workspace_bytes UINT32_MAX_BYTES = (1 << 32) - 1 if per_slice_bytes > 0: ws_cap = UINT32_MAX_BYTES // per_slice_bytes @@ -638,30 +533,11 @@ def kid_rejects_shape(k_inst, M, N, K): splitk main kernel's mask_va_tail cover both edge cases, so splitk is safe for any (M, N, K). """ - # The gfx1250 _ws reduce launches as dim3(ceil(N, VEC*BLOCK), M, 1) and - # grid.y is capped at 65535. Past that it does not fail, it writes garbage: - # at M=65536 the output is NaN, while the same shape on a .co kid (no reduce) - # is exact, and M >= 65537 does not launch at all ("invalid configuration - # argument"). Tuning one tends to take the box down with it. The fuse family - # reduces in-kernel and is unaffected. The generated launcher re-checks, for - # a tuned CSV that predates this. - if M > 65535 and _kid_launches_reduce(k_inst): + if k_inst.max_m is not None and M > k_inst.max_m: return True - # Pre-compiled (.co) kids: answered first, because almost none of the rules - # below apply to them. The pipeline builds NO buffer resource at all (it is - # the only a16w16 pipeline with zero make_gmem -- A, B and C all ride TDM - # descriptors with 64-bit base and stride), so the 4 GiB filter underneath - # is not merely satisfied, it is inapplicable. Every tail is handled by the - # D#'s per-dimension saturating clamp, so no M/N/K alignment is required - # either. What IS required: the batch strides are int64 but m/n/k are int, - # and the traits assert the tile it was compiled for. - # Both .co families, not just the compute one: they share the launcher and - # the no-buffer-resource pipeline, so the same reasoning applies. Naming only - # a16w16_4wave_co let a16w16_4wave_wl_co fall through to the rules below, - # where the gfx942 bf16-workspace whitelist rejected all 140 of its kids at - # any N outside that table (384, 32320, 129280) -- silently, since a missing - # candidate looks the same as a candidate that lost. + # CO pipelines use dimension-clamped TDM descriptors for A/B/C, have no + # split-K buffer, and support M/N/K tails. Their scalar extents remain int. if k_inst.kernel_tag in _A16W16_CO_TAGS: return M < 1 or N < 1 or K < 1 @@ -687,20 +563,13 @@ def kid_rejects_shape(k_inst, M, N, K): B_K = k_inst.B_K loops = _ceil_div(K, B_K) - # BF16WS_EXACT_REDUCE_SHAPES is a gfx942 artifact: that family's bf16-workspace - # reduce was only ever validated on the handful of N in the table, so anything - # else is refused. It keyed on "uses a bf16 workspace" alone, which was - # equivalent to "is a gfx942 bf16-ws kid" only while every gfx1250 _ws kid - # declared an fp32 partial. Once those switched to bf16 the rule started - # rejecting them too, and since the table is all powers of two it wiped out - # the whole _ws family at N=384 / 32320 / 129280 -- at N=384 that left 6 of - # 106 candidates, all one tile. The gfx1250 _ws reduce handles a ragged N - # through its tail path, so scope the rule to the family it was written for. - if _kid_uses_bf16_workspace(k_inst) and k_inst.kernel_tag not in _WS_SPLITK_TAGS: + # The exact-N bf16-workspace restriction belongs only to the gfx942 + # reducer. gfx1250 #4246 has a different bf16 reducer and accepts padded N. + if k_inst.arch_prefix == "gfx942" and _kid_uses_bf16_workspace(k_inst): padded_N = _ceil_div(N, k_inst.B_N) * k_inst.B_N if loops < 2 or K % B_K != 0 or padded_N != N: return True - return not any(N == n_exact for n_exact, _ in BF16WS_EXACT_REDUCE_SHAPES) + return N not in GFX942_BF16WS_EXACT_N if k_inst.kernel_tag in ( "a16w16", @@ -748,40 +617,38 @@ def kid_rejects_shape(k_inst, M, N, K): return per_slice_bytes > UINT32_MAX_BYTES if k_inst.kernel_tag == "a16w16_clusterlaunch_tdm_splitk_fuse": - # FUSED single-kernel split-K. split_k / m_cluster are COMPILE-TIME - # (baked per kid). Constraints: - # K even (a16w16 WMMA pairs). - # N % B_N == 0: the last-split C write uses a NON-predicated `store` - # bounded only by buffer num_records (tensor-end). Ragged N would - # intra-row-spill an interior N-tile's OOB columns into the next - # row, so full N-tile alignment is REQUIRED. (Ragged M is fine: OOB - # rows land past num_records and are dropped -> M % B_M NOT required.) - # ceil(M/B_M) % m_cluster == 0 (cluster.y fill). - # balanced K-split: split_k <= k_steps_tot (else some split WG is empty; - # the K tail is TDM-clamped, not handled by emptying WGs). - if K % 2 != 0: + # Fused SplitK/N-cluster values are compile-time. N must be tiled + # exactly because the final store has no intra-row N-tail predicate; + # ragged M remains safe through the bounded output descriptor. + if K % 2 != 0 or N % k_inst.B_N != 0: return True - if N % k_inst.B_N != 0: + split_k = int(k_inst.fuse_split_k) + n_cluster = int(k_inst.fuse_m_cluster) + if split_k < 2 or split_k * n_cluster > 16: return True - split_k = getattr(k_inst, "fuse_split_k", 2) - # fuse_m_cluster holds the cluster's 2nd-dim WG count; for this pipeline it - # groups N-tile peers (cluster.y, A-multicast). Cluster = (SplitK, n_cluster, - # 1); the product must fit the 16-WG cluster budget (kernel static_assert). - n_cluster = getattr(k_inst, "fuse_m_cluster", 1) - if split_k * n_cluster > 16: + num_tiles_m = _ceil_div(M, k_inst.B_M) + num_tiles_n = N // k_inst.B_N + if num_tiles_n % n_cluster != 0: return True - # N-peer cluster fill: ceil(N/B_N) must be a multiple of n_cluster (every - # named multicast WG present, else the cluster barrier stalls). N%B_N==0 is - # already required above, so ceil(N/B_N) == N/B_N. - if _ceil_div(N, k_inst.B_N) % n_cluster != 0: + if split_k > _ceil_div(K, k_inst.B_K): return True - k_steps_tot = _ceil_div(K, k_inst.B_K) - return split_k > k_steps_tot + # Mirror the exact tile-major capacity calculation. Python/C++ use + # size_t checks, but keep the tuner away from an unrepresentable tensor. + workspace_bytes = 2 if _kid_uses_bf16_workspace(k_inst) else 4 + required_bytes = ( + num_tiles_m + * num_tiles_n + * (split_k - 1) + * k_inst.B_M + * k_inst.B_N + * workspace_bytes + ) + return required_bytes > (2**63 - 1) if k_inst.kernel_tag == "a16w16_cluster_tdm_splitk_ws": # gfx1250 WMMA kernel: ragged M/N ARE supported -- the main kernel # TDM-clamps OOB global reads to the real (M, N) extents (tensor_dim1 = - # m - tile_row / n - tile_col), padded partials land in the padded fp32 + # m - tile_row / n - tile_col), padded partials land in the typed # workspace, and the reduce kernel only touches m in [0, M) / n in # [0, N). So M=49 runs as a padded M=64 tile. Ragged K is handled via # the TDM k_extent clamp. K must be even (a16w16 family). Workspace @@ -791,8 +658,9 @@ def kid_rejects_shape(k_inst, M, N, K): padded_M = _ceil_div(M, k_inst.B_M) * k_inst.B_M padded_N = _ceil_div(N, k_inst.B_N) * k_inst.B_N UINT32_MAX_BYTES = (1 << 32) - 1 - # batch=1 in tune path - return 1 * padded_M * padded_N * 4 > UINT32_MAX_BYTES + # batch=1 in tune path; descriptor size follows exact-kid storage. + workspace_bytes = 2 if _kid_uses_bf16_workspace(k_inst) else 4 + return padded_M * padded_N * workspace_bytes > UINT32_MAX_BYTES if k_inst.kernel_tag == "a16w16_clusterlaunch_tdm_splitk_ws": # Same numeric constraints as the plain cluster_tdm_splitk_ws variant. @@ -801,13 +669,8 @@ def kid_rejects_shape(k_inst, M, N, K): padded_M = _ceil_div(M, k_inst.B_M) * k_inst.B_M padded_N = _ceil_div(N, k_inst.B_N) * k_inst.B_N UINT32_MAX_BYTES = (1 << 32) - 1 - # No cluster-fill constraint: the launcher rounds the tile grid up to whole - # (cwm x cwn) clusters and a workgroup the round-up added returns at its - # cluster-barrier arrival, before issuing any TDM. Ragged M/N therefore run on - # any cluster dims, as does a shape whose whole grid is smaller than one - # cluster. 2D clusters (cwm>1 && cwn>1) are no longer locked out either: what - # they used to hang on was the tile-less workgroup streaming ZERO-EXTENT - # multicast loads at peers whose extents were real. + # Rounded OOB workgroups exit after the cluster barrier, so ragged and + # 2D cluster grids do not require exact fill. # NOTE: large output tiles (B_M*B_N >= 16384, e.g. 128x128 / 64x256) used # to fault at runtime, but the root cause was a clang<=22 (HIP<=7.2) # codegen bug in the bounded-buffer C-store address lowering (it sank the @@ -816,7 +679,7 @@ def kid_rejects_shape(k_inst, M, N, K): # voffset barrier, auto-gated to __clang_major__<=22), so large clusterlaunch # tiles are safe to tune again. No tile-area cap here. # batch=1 in tune path - return 1 * padded_M * padded_N * 4 > UINT32_MAX_BYTES + return padded_M * padded_N * 4 > UINT32_MAX_BYTES # kbuf2v_sk and quad_mfma32 splitK families require loops_per_split # (both full and last) even AND >=2. @@ -889,8 +752,10 @@ def kid_rejects_bias(k_inst, bias): """ if not bias: return False - # gfx1250 cluster_tdm_splitk_ws AND clusterlaunch_tdm_splitk_ws both fold - # bias in the shared reduce kernel (bias-aware); neither narrows on bias. + # gfx1250 two-stage families fold the public fp32/bf16, [N]/[batch,N] + # contract in the shared reducer. #4246 fused round-1 accepts only bf16 + # [N], which the boolean tuned key cannot represent, so exclude it from + # bias-bearing tuning rather than producing a row that is unsafe to replay. if k_inst.kernel_tag in ( "a16w16_cluster_tdm_splitk_ws", "a16w16_clusterlaunch_tdm_splitk_ws", @@ -898,6 +763,8 @@ def kid_rejects_bias(k_inst, bias): "a16w16_clusterlaunch_tdm_splitk_fuse", ): return False + if k_inst.kernel_tag == "a16w16_clusterlaunch_tdm_splitk_fuse": + return True if k_inst.kernel_tag not in ("a16w16", "a16w16_flatmm_splitk"): return True # arch_prefix distinguishes gfx942 (no bias) from gfx950 (has bias) @@ -948,13 +815,14 @@ def candidate_kids_for_shape(M, N, K, bias, cu_num): cu_num = int(cu_num) # gfx1250: dedicated candidate filter (top-N tiles by grid-occupancy fit x - # {plain + top-N square cluster dims}). All gfx1250 kids fold bias in the - # reduce kernel, so bias does not narrow the set. + # {plain + top-N round-up-capable cluster dims}). Bias-bearing tuning retains + # the two-stage kids but omits fused round-1 kids because their bf16 [N]-only + # bias contract cannot be encoded by the tuned CSV's boolean bias key. try: from aiter.jit.utils.chip_info import get_gfx_runtime if get_gfx_runtime().lower() == "gfx1250": - return _gfx1250_select_candidates(M, N, K, cu_num) + return _gfx1250_select_candidates(M, N, K, cu_num, include_fused=not bias) except Exception: # noqa: BLE001,S110 pass @@ -994,41 +862,21 @@ def candidate_kids_for_shape(M, N, K, bias, cu_num): pass # unknown arch -> keep legacy multi-arch behaviour # Step 6: drop known-bad kids permanently. - return cands - _OPUS_PERMA_BAD_KIDS + cands = cands - _OPUS_PERMA_BAD_KIDS + return cands # Kids we never want tuner to probe. _OPUS_PERMA_BAD_KIDS = frozenset() -# The two families that stage a split-K partial the reduce then reads. The fuse -# family reduces in-kernel, so its workspace dtype is its own business. -_WS_SPLITK_TAGS = frozenset( - {"a16w16_cluster_tdm_splitk_ws", "a16w16_clusterlaunch_tdm_splitk_ws"} -) - -# The pre-compiled (.co) families. Mirrors codegen/common.py:_A16W16_CO_TAGS. -_A16W16_CO_TAGS = frozenset({"a16w16_4wave_co", "a16w16_4wave_wl_co"}) - - -def _kid_launches_reduce(k_inst): - """True for the gfx1250 families whose launcher runs splitk_reduce separately. - - The gfx942 _sk families and flatmm_splitk launch the same reduce with the - same grid.y, so the 65535 cap is theirs too -- they are left alone here - only because this change is scoped to gfx1250. The fuse family reduces - in-kernel and launches nothing. - """ - return getattr(k_inst, "kernel_tag", "") in _WS_SPLITK_TAGS - - def _ensure_kids_compiled(candidate_kids): """Make sure every kid in ``candidate_kids`` is compiled into the current module_deepgemm_opus.so. Reads the subset-compile sidecar at ``_opus_sidecar_path()`` (lives in ``$JIT_BUILD/`` so it survives clear_build). If any kid in - ``candidate_kids`` (or in ``HEURISTIC_DEFAULT_KIDS``) is missing, or its + ``candidate_kids`` (or in ``DEFAULT_COMPILED_KIDS``) is missing, or its receipt does not match the installed .so, this function: 1. Passes the new candidates as ``--extra_kids`` to codegen. The last @@ -1042,7 +890,7 @@ def _ensure_kids_compiled(candidate_kids): required kid. Without this synchronous step, children would race against the parent's lazy build and the first to dlopen() would get the stale subset .so and fail with - ``AITER_CHECK: Kernel id X not found in a16w16 tune lookup table``. + ``AITER_CHECK: unknown kid X for OPUS a16w16 in the launch table``. Concurrency model ----------------- @@ -1079,19 +927,19 @@ def _ensure_kids_compiled(candidate_kids): True if a rebuild was triggered, False if every required kid was already compiled. """ - from opus_gemm_common import heuristic_kids_for_arch + from opus_gemm_common import default_compiled_kids_for_arch from aiter.jit import core as _jit_core from aiter.jit.utils.file_baton import FileBaton from aiter.jit.utils.jit_cache import compiled_kids_are_current candidate_kids = frozenset(int(k) for k in candidate_kids) - # Restrict the heuristic-default kid set to the running GPU's arch. + # Restrict the default compile floor to the running GPU's arch. try: from aiter.jit.utils.chip_info import get_gfx_runtime _run_arch = get_gfx_runtime().lower() - _heuristic = heuristic_kids_for_arch({_run_arch}) + _defaults = default_compiled_kids_for_arch({_run_arch}) except Exception: # noqa: BLE001 # A runtime probe can fail in a prebuild environment with explicit # targets. Do not require off-arch defaults in that case. @@ -1100,8 +948,8 @@ def _ensure_kids_compiled(candidate_kids): for arch in os.getenv("GPU_ARCHS", "native").split(";") if arch.strip() and arch.strip().lower() != "native" } - _heuristic = heuristic_kids_for_arch(_target_arches or None) - required = candidate_kids | _heuristic + _defaults = default_compiled_kids_for_arch(_target_arches or None) + required = candidate_kids | _defaults def _read_sidecar(path): if not os.path.exists(path): @@ -1278,26 +1126,13 @@ def _reuse_current_binary(): _AITER_VERBOSE = bool(int(os.environ.get("AITER_VERBOSE", "0"))) -# Merge every a16w16-family kid into one tuner search space: * split-barrier a16w16: 4..9 legacy -# cpol = (0, 17) (traits default) ... +# Derive the tuner search space from the canonical registry. Keeping a second +# hand-maintained merge here previously omitted mono-tile and 4g-safe kids and +# could retain stale metadata when a new family reused an existing numeric id. a16w16_all_kernels = { - **a16w16_kernels_list, - **a16w16_kernels_list_nooob, - **a16w16_kernels_list_cpol, - **a16w16_kernels_list_cpol_nooob, - **a16w16_flatmm_kernels_list, - **a16w16_flatmm_splitk_kernels_list, - **a16w16_flatmm_splitk_kernels_list_nooob, - **a16w16_persistent_kernels_list, - **a16w16_persistent_kernels_list_cpol, - **a16w16_persistent_kernels_list_nooob, - **a16w16_persistent_kernels_list_cpol_nooob, - **gfx942_nosplit_kernels_list, - **gfx942_splitk_kernels_list, - **gfx1250_kernels_list, - **gfx1250_clusterlaunch_kernels_list, - **gfx1250_splitk_fuse_kernels_list, - **gfx1250_4wave_co_kernels_list, + kid: instance + for kid, instance in kernels_list.items() + if instance.kernel_tag.startswith("a16w16") } # Arch-filter the kid enumeration so the tuner only dispatches kids whose pipeline body has a @@ -1427,15 +1262,9 @@ def opus_gemm_ref(XQ, WQ, bias=None, out_dtype=None): def run_opus_gemm(XQ, WQ, Y, bias, kernelId, splitK): - """Eager-path tuner func: runs the kernel AND an on-the-fly max_delta check. - - The check raises RuntimeError when the output is numerically off; mp_tuner's - worker catches it and marks the candidate invalid. Used when --no-graph is - passed (i.e. graph mode disabled) so the per-iter check is safe (no CUDA - graph capture). - """ + """Launch one eager kid and reject excessive numerical error.""" _quiet_aiter_logger_once() - _opus_gemm_a16w16_tune(XQ, WQ, Y, bias, kernelId, splitK) + _opus_bmm(XQ, WQ, Y, kid=kernelId, bias=bias, split_k=splitK) ref = opus_gemm_ref(XQ, WQ, bias, Y.dtype) max_delta = (Y.float() - ref.float()).abs().max().item() max_ref = ref.float().abs().max().item() @@ -1476,33 +1305,9 @@ def _quiet_aiter_logger_once(): def run_opus_gemm_bench(XQ, WQ, Y, bias, kernelId, splitK): - """Tuner bench func with capture-safe stream sync + per-task max_delta - safety check. - - Stream sync rationale - --------------------- - When our custom run_perftest replacement (_opus_run_perftest below) uses - torch.cuda.Event to time the graph replay, the end-event record needs - the kernel in flight. The sync inside the bench func itself is for the - WARMUP phase (outside capture), so that max_delta check and torch.bmm - reference see a stable Y before validating. - - The sync is gated on is_current_stream_capturing() because - cudaStreamSynchronize during CUDA graph capture invalidates the graph - (HIP returns hipErrorStreamCaptureInvalidated). - - Correctness gate - ---------------- - mp_tuner.worker's post-run checkAllclose(ref, Y, rtol, atol) gates on - *fraction* of cells above tolerance, not the max single-cell absolute - delta. We add a per-task max_delta check: - * Runs once per (XQ, WQ, Y, kid, splitK) tuple per subprocess. - * Skipped inside CUDA graph capture (.item() forbidden there). - * Raises RuntimeError on violation; mp_tuner.worker marks the - candidate us=-1, err_ratio=1.0. - """ + """Benchmark one kid with a capture-safe, once-per-task accuracy check.""" _quiet_aiter_logger_once() - _opus_gemm_a16w16_tune(XQ, WQ, Y, bias, kernelId, splitK) + _opus_bmm(XQ, WQ, Y, kid=kernelId, bias=bias, split_k=splitK) capturing = torch.cuda.is_current_stream_capturing() @@ -1526,7 +1331,7 @@ def run_opus_gemm_bench(XQ, WQ, Y, bias, kernelId, splitK): raise RuntimeError( f"maxDelta {max_delta:.1f} > bound {bound:.1f} " f"(max|ref|={max_ref:.1f}, scale={MAX_DELTA_SCALE}) " - f"for kid={kernelId} splitK={splitK} bias={bias is not None}" + f"for kid={kernelId} split_k={splitK} bias={bias is not None}" ) # Capture-safe sync: guarantees the warmup kernel has completed before we read Y for the @@ -2359,8 +2164,9 @@ def tune(self, untunedf, tunedf, args): " (or --libtype all to tune all backends in one pass).\n" " gradlib writes to aiter/configs/bf16_tuned_gemm.csv (or the\n" " path passed via --tuned_file / GTUNE_TUNED) and stamps every\n" - " opus row with libtype='opus' so the opus runtime dispatch\n" - " picks it up via aiter.ops.opus.common.lookup_tuned().\n" + " opus row with libtype='opus' so aiter.tuned_gemm passes its\n" + " resolved solidx directly to aiter.ops.opus.opus_gemm().\n" + " This debug tuner itself uses batch-first opus_bmm() tensors.\n" f" This script writes to {OPUS_DEBUG_TUNED_CSV} by default and\n" " will not pollute the global aiter/configs/ tree.\n" "==============================================================\n" diff --git a/csrc/pybind/opus_gemm_pybind.cu b/csrc/pybind/opus_gemm_pybind.cu index 9c90c0fc20..b9586dacd5 100644 --- a/csrc/pybind/opus_gemm_pybind.cu +++ b/csrc/pybind/opus_gemm_pybind.cu @@ -1,25 +1,22 @@ // SPDX-License-Identifier: MIT // Copyright (C) 2025-2026, Advanced Micro Devices, Inc. All rights reserved. // -// pybind glue is host-only. Skip the entire TU on the device pass so we -// don't pay the libtorch + pybind11 + HIP runtime parse (~15s) for code -// that has no GPU side at all. +// Register the four OPUS launch interfaces on the host pass only. #ifndef __HIP_DEVICE_COMPILE__ #include "rocm_ops.hpp" #include "aiter_stream.h" -#include "opus_gemm.h" #include "opus_bmm.h" +#include "opus_gemm.h" PYBIND11_MODULE(AITER_EXTENSION_NAME, m) { AITER_SET_STREAM_PYBIND - OPUS_GEMM_PYBIND; - OPUS_GEMM_A16W16_TUNE_PYBIND; - OPUS_BMM_A8W8_MXSCALE_PYBIND; - OPUS_GEMM_A8W8_BLOCKSCALE_BPRESHUFFLE_TUNE_PYBIND; - OPUS_GEMM_WORKSPACE_INIT_PYBIND; - OPUS_GEMM_WORKSPACE_RELEASE_PYBIND; + OPUS_GEMM_A16W16_LAUNCH_PYBIND; + OPUS_GEMM_A8W8_LAUNCH_PYBIND; + OPUS_GEMM_A8W8_BLOCKSCALE_LAUNCH_PYBIND; + OPUS_GEMM_A8W8_BLOCKSCALE_BPRESHUFFLE_LAUNCH_PYBIND; + OPUS_GEMM_A8W8_MXSCALE_BMM_LAUNCH_PYBIND; } #endif // !__HIP_DEVICE_COMPILE__ diff --git a/op_tests/test_jit_cache_transaction.py b/op_tests/test_jit_cache_transaction.py index a21a3d7d67..4f8b5573b9 100644 --- a/op_tests/test_jit_cache_transaction.py +++ b/op_tests/test_jit_cache_transaction.py @@ -10,6 +10,7 @@ import multiprocessing import os import re +import runpy import shlex import shutil import socket @@ -1041,7 +1042,7 @@ def build(**kwargs): get_gfx_runtime=lambda: "gfx942" ), "opus_gemm_common": types.SimpleNamespace( - heuristic_kids_for_arch=lambda _arches: {1} + default_compiled_kids_for_arch=lambda _arches: {1} ), } self.tuner_imports = imports @@ -1059,7 +1060,6 @@ def import_dependency(name, *args, **kwargs): "os": os, "sys": sys, "json": json, - "HEURISTIC_DEFAULT_KIDS": {1}, "_opus_sidecar_path": lambda: self.sidecar, }, )["_ensure_kids_compiled"] @@ -1144,13 +1144,13 @@ def test_tuner_runtime_probe_fallback_respects_explicit_build_arches(self): self.tuner_imports["aiter.jit.utils.chip_info"].get_gfx_runtime = mock.Mock( side_effect=RuntimeError("no rocminfo") ) - heuristic = mock.Mock( + defaults = mock.Mock( side_effect=lambda arches: {1} if arches == {"gfx942"} else {1, 200} ) - self.tuner_imports["opus_gemm_common"].heuristic_kids_for_arch = heuristic + self.tuner_imports["opus_gemm_common"].default_compiled_kids_for_arch = defaults with mock.patch.dict(os.environ, {"GPU_ARCHS": "gfx942"}): self.assertFalse(tuner({7})) - heuristic.assert_called_once_with({"gfx942"}) + defaults.assert_called_once_with({"gfx942"}) self.assertEqual(calls, []) def test_tuner_interruption_restores_environment_and_releases_locks(self): @@ -1481,6 +1481,10 @@ def run_cpu_ninja(**kwargs): class TestOpusRequestedKids(unittest.TestCase): def test_real_generator_accepts_valid_requests_and_rejects_filtered_requests(self): generator = JIT_CACHE_PATH.parents[3] / "csrc/opus_gemm/gen_instances.py" + registry = runpy.run_path(str(generator.with_name("opus_gemm_common.py"))) + bmm_kids = sorted(registry["BMM_MXSCALE_KIDS"]) + co_kids = sorted(registry["GFX1250_4WAVE_CO_KIDS"]) + workspace_kid = min(registry["gfx1250_clusterlaunch_kernels_list"]) runner = ( "import os, runpy, sys, types; " "sys.argv = sys.argv[1:]; " @@ -1490,15 +1494,19 @@ def test_real_generator_accepts_valid_requests_and_rejects_filtered_requests(sel "runpy.run_path(sys.argv[0], run_name='__main__')" ) cases = ( - (10006, [], True), - (999999, [], False), - (200, [], False), - (200, ["--kernel_tag", "a16w16"], False), - (10006, ["--kernel_tag", "a8w8"], False), + ("gfx942", [10006], [], True), + ("gfx942", [999999], [], False), + ("gfx942", [200], [], False), + ("gfx942", [200], ["--kernel_tag", "a16w16"], False), + ("gfx942", [10006], ["--kernel_tag", "a8w8"], False), + ("gfx950", bmm_kids, [], True), + ("gfx950", bmm_kids, ["--kernel_tag", "a8w8"], True), + ("gfx942", bmm_kids, [], False), + ("gfx1250", [workspace_kid, *co_kids], [], True), ) - for kid, extra_args, accepted in cases: + for arch, kids, extra_args, accepted in cases: with self.subTest( - kid=kid, extra_args=extra_args + arch=arch, kids=kids, extra_args=extra_args ), tempfile.TemporaryDirectory() as tmp: sidecar = os.path.join(tmp, "compiled_kids.json") _write(sidecar, "[]") @@ -1511,17 +1519,28 @@ def test_real_generator_accepts_valid_requests_and_rejects_filtered_requests(sel "--working_path", tmp, "--extra_kids", - str(kid), + *map(str, kids), *extra_args, ], - env={**os.environ, "GPU_ARCHS": "gfx942"}, + env={**os.environ, "GPU_ARCHS": arch}, capture_output=True, text=True, check=False, ) if accepted: self.assertEqual(result.returncode, 0, result.stderr) - self.assertIn(kid, json.loads(_read(sidecar))) + self.assertTrue(set(kids) <= set(json.loads(_read(sidecar)))) + manifest = _read(os.path.join(tmp, "opus_gemm_manifest.h")) + for name in {registry["kernels_list"][kid].name for kid in kids}: + self.assertEqual(manifest.count(name + "("), 1, name) + if arch == "gfx950": + dispatch = _read( + os.path.join(tmp, "opus_bmm_mxscale_kid_dispatch.h") + ) + emitted = { + int(kid) for kid in re.findall(r"\{\s*(\d+),", dispatch) + } + self.assertEqual(emitted, set(bmm_kids)) else: self.assertNotEqual(result.returncode, 0) self.assertIn( diff --git a/op_tests/test_opus_a16w16_gemm.py b/op_tests/test_opus_a16w16_gemm.py index cef48d9c30..3f30e21d48 100644 --- a/op_tests/test_opus_a16w16_gemm.py +++ b/op_tests/test_opus_a16w16_gemm.py @@ -1,83 +1,28 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2025-2026, Advanced Micro Devices, Inc. All rights reserved. -"""End-to-end regression of gemm_a16w16_opus vs torch.bmm; prints TFLOPs. +"""A16W16 OPUS exact-kid regressions and benchmark coverage. Usage: - python3 op_tests/test_opus_a16w16_gemm.py [-m M -n N -k K -b B] + python3 op_tests/test_opus_a16w16_gemm.py --kid KID [-m M -n N -k K -b B] python3 op_tests/test_opus_a16w16_gemm.py --csv_file - - # opus-only sweep in CUDA-graph mode, golden-checked (default entry): - python3 op_tests/test_opus_a16w16_gemm.py --opus_sweep -n 2048 -k 7168 """ import argparse -import os import sys +import pytest import torch +from aiter.benchmark_data_init import fill + # Skip on unsupported arch via the same probe opus uses at import time. -from aiter.ops.opus._arch import _detect_arch +from aiter.ops.opus._arch import _detect_arch, _device_arch_and_cu +from aiter.ops.opus.launch_plan import _get_cached_a16w16_launch_plan _arch_ok, _detected_gfx = _detect_arch({"gfx950", "gfx942", "gfx1250"}) -if not _arch_ok: - print( - f"[skip] test_opus_a16w16_gemm requires gfx950/gfx942/gfx1250 (detected {_detected_gfx!r})" - ) - sys.exit(0) -from aiter.benchmark_data_init import ( - DATA_DISTS, - add_data_init_args, - fill, - make_generator, -) -from aiter.ops.opus import gemm_a16w16_opus -from aiter.test_common import ( - checkAllclose, - run_perftest, -) - -try: - from aiter.ops.opus import opus_gemm_workspace_init -except Exception: # noqa: BLE001 - opus_gemm_workspace_init = None - - -def _graph_capture_stream(): - """The stream torch.cuda.graph captures on when no `stream=` is passed. - - torch lazily creates a single process-global `default_capture_stream`; we - mirror that here so the opus split-K workspace is registered/grown on the - exact stream a later `with torch.cuda.graph(g):` (as used by run_perftest's - graph mode) will capture on. - """ - g = torch.cuda.graphs.graph - if getattr(g, "default_capture_stream", None) is None: - g.default_capture_stream = torch.cuda.Stream() - return g.default_capture_stream - - -def _prewarm_opus_graph_workspace(A, B, out_dtype): - """Eagerly register + size the opus split-K workspace on the capture stream. - - opus split-K kernels keep a per-stream fp32 workspace backed by raw - hipMalloc; growing it is stream-capture-illegal, so it must be registered - and grown to the shape's size *eagerly* before HIP graph capture. Without - this, capturing an opus split-K shape aborts with "splitk workspace not - initialized for the current CUDA stream". No-op on archs without the - registry (opus_gemm_workspace_init unavailable) or while already capturing. - """ - if opus_gemm_workspace_init is None: - return - if torch.cuda.is_current_stream_capturing(): - return - s = _graph_capture_stream() - with torch.cuda.stream(s): - opus_gemm_workspace_init() - # Warm the exact shape so the workspace buffer reaches its final size. - _ = gemm_a16w16_opus(A, B, None, out_dtype) - s.synchronize() +from aiter.ops.opus import opus_bmm, opus_gemm +from aiter.test_common import checkAllclose, run_perftest def _torch_ref(A: torch.Tensor, B: torch.Tensor, out_dtype): @@ -88,27 +33,8 @@ def _torch_ref(A: torch.Tensor, B: torch.Tensor, out_dtype): return torch.bmm(A.float(), B.float().transpose(-1, -2)).to(out_dtype) -# --------------------------------------------------------------------------- # -# Data initialization (bf16 operands) + seed -# --------------------------------------------------------------------------- # -# The distributions and the seeded generator come from aiter.test_common (the -# shared data-init API); a16w16 only needs bf16 DATA operands, so it wraps -# ``fill`` and reuses ``make_generator`` / ``add_data_init_args`` verbatim. -DATA_INITS = DATA_DISTS - - def _make_tensor(shape, dist="norm", gen=None, const_val=1.0): - """Build a bf16 operand under the requested distribution. - - zero : all zeros - constant : filled with ``const_val`` - uniform : U(-1, 1) - norm : N(0, 1) [default; matches the original torch.randn path] - - Delegates to ``benchmark_data_init.fill`` so the operand init matches every - other op test; ``gen`` seeds the sampled dists (uniform/norm), zero/constant - ignore it. - """ + """Build one reproducible BF16 operand for the benchmark paths.""" return fill( shape, dist, @@ -121,83 +47,96 @@ def _make_tensor(shape, dist="norm", gen=None, const_val=1.0): def _make_b( - batch: int, N: int, K: int, dist: str = "norm", gen=None, const_val: float = 1.0 + batch: int, + N: int, + K: int, + dist: str = "norm", + gen=None, + const_val: float = 1.0, ) -> torch.Tensor: - """Build a B that gemm_a16w16_opus accepts for both batch=1 and batch>1. - - The wrapper rejects 2D B + batch>1 because the opus launcher hardcodes - stride_b_batch == N*K (a broadcast view would silently fault). For the - common "shared weight across batch" case, materialize an explicit - `[batch, N, K]` tensor via the contiguous broadcast pattern. - """ + """Build batch-first physical weights for the exact BMM path.""" B2D = _make_tensor((N, K), dist, gen, const_val) - if batch == 1: - return B2D return B2D.unsqueeze(0).expand(batch, -1, -1).contiguous() -def _make_a( - batch: int, M: int, K: int, dist: str = "norm", gen=None, const_val: float = 1.0 -) -> torch.Tensor: - """Build the [batch, M, K] bf16 activation under the requested dist.""" - return _make_tensor((batch, M, K), dist, gen, const_val) - - -# --------------------------------------------------------------------------- # -# FLOPS + Bandwidth -# --------------------------------------------------------------------------- # -def _tflops(batch, M, N, K, us): - """GEMM TFLOPS (2*b*M*N*K FLOP) from microseconds (None-safe).""" - return (2.0 * batch * M * N * K / us / 1e6) if us else None - +def _run_exact_a16w16( + A: torch.Tensor, + B: torch.Tensor, + Y: torch.Tensor, + *, + kid: int, + split_k: int, + use_graph: bool, +): + kwargs = {"kid": kid, "split_k": split_k} + if not use_graph: + return run_perftest(opus_bmm, A, B, Y, **kwargs) + + arch, cu_num = _device_arch_and_cu(A.device) + plan = _get_cached_a16w16_launch_plan( + arch, + A.shape[1], + B.shape[1], + A.shape[2], + A.shape[0], + cu_num, + False, + A.dtype, + Y.dtype, + kid, + split_k, + ) + workspace = ( + torch.empty( + plan.workspace_spec.shape, + dtype=plan.workspace_spec.dtype, + device=A.device, + ) + if plan.workspace_spec is not None + else None + ) + kwargs["workspace"] = workspace -def _tbs(batch, M, N, K, us, out_bytes=2): - """Operand traffic in TB/s (None-safe). + opus_bmm(A, B, Y, **kwargs) + current = torch.cuda.current_stream(A.device) + side = torch.cuda.Stream(device=A.device) + side.wait_stream(current) + with torch.cuda.stream(side): + opus_bmm(A, B, Y, **kwargs) + side.synchronize() - Bytes moved = A[b,M,K]@bf16 + B@bf16 + out[b,M,N]@out_bytes. B is read once - per batch when materialized (batch>1) and once as a shared weight (batch=1). - bytes / (us*1e-6) / 1e12 == bytes / us / 1e6. - """ - if not us: - return None - a = batch * M * K * 2 - b = (batch * N * K if batch > 1 else N * K) * 2 - o = batch * M * N * out_bytes - return (a + b + o) / us / 1e6 + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=side): + opus_bmm(A, B, Y, **kwargs) + current.wait_stream(side) + _, us = run_perftest(graph.replay, use_cuda_event=True) + return Y, us -def test_a16w16( +def run_a16w16_case( batch: int, M: int, N: int, K: int, - out_dtype=torch.bfloat16, - use_graph=False, *, - dist="norm", - gen=None, - const_val=1.0, - iters=101, - warmup=2, - rotate=0, + kid: int, + split_k: int = 0, + out_dtype=torch.bfloat16, + use_graph: bool = False, ): - # gemm_a16w16_opus accepts either 2D or 3D A; test 3D to exercise the - # batched reshape path. B is 2D when batch==1, 3D contiguous otherwise. - A = _make_a(batch, M, K, dist, gen, const_val) - B = _make_b(batch, N, K, dist, gen, const_val) + A = torch.randn(batch, M, K, device="cuda", dtype=torch.bfloat16) + B = _make_b(batch, N, K) + Y = torch.empty((batch, M, N), device="cuda", dtype=out_dtype) ref = _torch_ref(A, B, out_dtype) - Y, us = run_perftest( - gemm_a16w16_opus, + Y, us = _run_exact_a16w16( A, B, - None, - out_dtype, - testGraph=use_graph, - num_iters=iters, - num_warmup=warmup, - num_rotate_args=rotate, + Y, + kid=kid, + split_k=split_k, + use_graph=use_graph, ) err = checkAllclose( @@ -207,255 +146,768 @@ def test_a16w16( rtol=0.1, atol=0.5, ) - tflops = _tflops(batch, M, N, K, us) - tbs = _tbs(batch, M, N, K, us, out_bytes=Y.element_size()) + flops = 2.0 * batch * M * N * K + tflops = flops / us / 1e6 print( f"[a16w16] batch={batch} M={M} N={N} K={K} dtype={out_dtype} " - f"| {us:.1f}us | {tflops:.2f} TFLOPs | {tbs:.3f} TB/s | err={err}" + f"| {us:.1f}us | {tflops:.2f} TFLOPs | err={err}" ) return err -def load_shapes_from_csv(csv_path): +def load_shapes_from_csv(csv_path, *, default_kid=None, default_split_k=0): import pandas as pd df = pd.read_csv(csv_path) - shapes = list(zip(df["M"].astype(int), df["N"].astype(int), df["K"].astype(int))) - return list(dict.fromkeys(shapes)) - - -def _default_tuned_csv(): - """Locate the shipped dsv4 tuned GEMM CSV inside the aiter package.""" - import aiter - - return os.path.join( - os.path.dirname(aiter.__file__), - "configs", - "model_configs", - "dsv4_bf16_tuned_gemm.csv", + kid_column = next( + (name for name in ("kernelId", "solidx", "kid") if name in df), None ) - - -def load_opus_shapes(csv_path, gfx, N=None, K=None): - """Return the opus_gemm rows (with their tuned reference timing). - - Filters the tuned CSV to rows matching the current arch (``gfx``) and, - optionally, a fixed ``N``/``K``, keeping only rows where ``libtype == - 'opus'`` (i.e. the shapes for which opus_gemm was selected). Returns a - list of dicts sorted by M, each with keys: ``M``, ``csv_us`` (the tuned - latency recorded in the CSV) and ``kernelName``. - """ - import pandas as pd - - df = pd.read_csv(csv_path) - mask = (df["gfx"] == gfx) & (df["libtype"] == "opus") - if N is not None: - mask &= df["N"].astype(int) == N - if K is not None: - mask &= df["K"].astype(int) == K - sub = df.loc[mask].copy() - sub["M"] = sub["M"].astype(int) - # Keep the first row per M (CSV holds one tuned winner per shape). - sub = sub.sort_values("M").drop_duplicates(subset="M", keep="first") + split_column = next((name for name in ("splitK", "split_k") if name in df), None) + if kid_column is None and default_kid is None: + raise ValueError( + "exact-kid CSV sweep needs a kernelId/solidx/kid column or --kid" + ) rows = [] - for _, r in sub.iterrows(): - # splitK may be blank/NaN in some rows; keep it as an int when present. - try: - splitk = int(r["splitK"]) if not pd.isna(r.get("splitK")) else None - except (KeyError, ValueError, TypeError): - splitk = None + for row in df.to_dict("records"): rows.append( - { - "M": int(r["M"]), - "csv_us": float(r["us"]), - "kernelName": str(r.get("kernelName", "")), - "splitK": splitk, - } + ( + int(row["M"]), + int(row["N"]), + int(row["K"]), + int(default_kid if default_kid is not None else row[kid_column]), + ( + int(row[split_column]) + if split_column is not None + else int(default_split_k) + ), + ) ) - return rows + return list(dict.fromkeys(rows)) -def test_opus_shapes_graph( - csv_path, - gfx, - N=2048, - K=7168, - batch=1, +def run_a16w16_csv_sweep( + csv_path: str, + batch: int = 1, + *, + kid: int | None = None, + split_k: int = 0, out_dtype=torch.bfloat16, + use_graph: bool = False, +): + shapes = load_shapes_from_csv(csv_path, default_kid=kid, default_split_k=split_k) + return _run_a16w16_sweep( + shapes, + source=csv_path, + batch=batch, + out_dtype=out_dtype, + use_graph=use_graph, + ) + + +def _run_a16w16_sweep( + shapes, *, - dist="norm", - gen=None, - const_val=1.0, - iters=101, - warmup=2, - rotate=0, + source: str, + batch: int, + out_dtype: torch.dtype, + use_graph: bool, ): - """CUDA-graph-mode opus_gemm sweep with golden check. - - Selects the M values from the tuned CSV where opus_gemm is the winner - (for the given arch / N / K), then times each in CUDA-graph mode and - validates against the torch fp32 reference. - """ - rows = load_opus_shapes(csv_path, gfx, N=N, K=K) - ms = [r["M"] for r in rows] print(f"\n{'=' * 80}") + mode = "graph" if use_graph else "eager" print( - f"opus graph-mode sweep [{gfx}] N={N} K={K} batch={batch}: " - f"{len(rows)} opus shapes -> M={ms}" + f"a16w16 sweep from {source}: {len(shapes)} unique shapes, " + f"batch={batch}, mode={mode}" ) print("=" * 80) - if not rows: - print(f"[skip] no opus_gemm rows in {csv_path} for gfx={gfx} N={N} K={K}") - return True - passed = failed = 0 - perf_rows = [] - out_bytes = torch.empty((), dtype=out_dtype).element_size() - for r in rows: - M = r["M"] - csv_us = r["csv_us"] - kid = r.get("kernelName") or "" - splitk = r.get("splitK") - tag = f"a16w16-graph b={batch} M={M} N={N} K={K}" + for M, N, K, row_kid, row_split_k in shapes: + tag = ( + f"a16w16 b={batch} M={M} N={N} K={K} " + f"kid={row_kid} split_k={row_split_k}" + ) try: - A = _make_a(batch, M, K, dist, gen, const_val) - B = _make_b(batch, N, K, dist, gen, const_val) + A = torch.randn(batch, M, K, device="cuda", dtype=torch.bfloat16) + B = _make_b(batch, N, K) + Y = torch.empty((batch, M, N), device="cuda", dtype=out_dtype) ref = _torch_ref(A, B, out_dtype) - # opus split-K workspace must be grown on the capture stream before - # run_perftest's graph mode captures this shape. - _prewarm_opus_graph_workspace(A, B, out_dtype) - Y, us = run_perftest( - gemm_a16w16_opus, + Y, us = _run_exact_a16w16( A, B, - None, - out_dtype, - testGraph=True, - num_iters=iters, - num_warmup=warmup, - num_rotate_args=rotate, + Y, + kid=row_kid, + split_k=row_split_k, + use_graph=use_graph, ) err = checkAllclose(Y, ref, msg=tag, rtol=0.1, atol=0.5) - tflops = _tflops(batch, M, N, K, us) - tbs = _tbs(batch, M, N, K, us, out_bytes=out_bytes) - # Ratio of measured graph latency to the tuned CSV reference. - ratio = (us / csv_us) if csv_us else float("nan") - splitk_str = "" if splitk in (None, "") else str(splitk) - print( - f"[PASS] {tag} | {us:.1f}us (csv {csv_us:.1f}us, " - f"{ratio:.2f}x) | {tflops:.2f} TFLOPs | {tbs:.3f} TB/s | err={err} " - f"| splitK={splitk_str or '-'} | kid={kid or '-'}" - ) - perf_rows.append((M, us, csv_us, ratio, tflops, tbs, splitk, kid)) + tflops = 2.0 * batch * M * N * K / us / 1e6 + print(f"[PASS] {tag} | {us:.1f}us | {tflops:.2f} TFLOPs | err={err}") passed += 1 except Exception as e: # noqa: BLE001 print(f"[FAIL] {tag} | {type(e).__name__}: {e}") failed += 1 + print(f"\nSummary: {passed} passed, {failed} failed out of {len(shapes)}") + return failed == 0 - if perf_rows: - print(f"\n{'-' * 88}") - print(f"latency vs tuned CSV [{gfx}] N={N} K={K} batch={batch}") - print(f"{'-' * 88}") - print( - f"{'M':>6} | {'graph us':>10} | {'csv us':>10} | {'ratio':>7} | " - f"{'TFLOPs':>9} | {'TB/s':>7} | {'splitK':>6} | note | kernel(kid)" + +def _runtime_arch() -> str | None: + if not torch.cuda.is_available(): + return None + props = torch.cuda.get_device_properties(torch.cuda.current_device()) + return str(props.gcnArchName).split(":", 1)[0].lower() + + +def _assert_matches_golden(actual, A, B, bias=None): + golden = A.float() @ B.float().transpose(-1, -2) + if bias is not None: + golden = golden + bias.float() + # BF16 output has one final rounding; fp32 output is normally much tighter. + atol = 0.5 if actual.dtype == torch.bfloat16 else 0.05 + rtol = 0.03 if actual.dtype == torch.bfloat16 else 1e-3 + torch.testing.assert_close(actual.float(), golden, rtol=rtol, atol=atol) + + +@pytest.mark.parametrize( + ("kid", "M", "N", "K", "split_k"), + ((200, 64, 64, 512, 2), (1400, 192, 256, 128, 0)), +) +def test_gfx950_logical_2d_gemm_matches_torch(kid, M, N, K, split_k): + """The public GEMM adapter adds only a no-copy batch-one raw view.""" + if _runtime_arch() != "gfx950": + pytest.skip("requires gfx950 hardware") + torch.manual_seed(0x2D950 + kid) + A = torch.randn((M, K), device="cuda", dtype=torch.bfloat16) + B = torch.randn((N, K), device="cuda", dtype=torch.bfloat16) + Y = torch.empty((M, N), device="cuda", dtype=torch.bfloat16) + actual = opus_gemm(A, B, Y, kid=kid, split_k=split_k) + torch.cuda.synchronize() + assert actual is Y + _assert_matches_golden(actual, A, B) + + +def test_gfx950_batch_first_bmm_matches_torch(): + """The public BMM contract preserves a real batch dimension.""" + if _runtime_arch() != "gfx950": + pytest.skip("requires gfx950 hardware") + torch.manual_seed(0xB950) + A = torch.randn((2, 192, 128), device="cuda", dtype=torch.bfloat16) + B = torch.randn((2, 256, 128), device="cuda", dtype=torch.bfloat16) + Y = torch.empty((2, 192, 256), device="cuda", dtype=torch.bfloat16) + actual = opus_bmm(A, B, Y, kid=1400) + torch.cuda.synchronize() + assert actual is Y + _assert_matches_golden(actual, A, B) + + +@pytest.mark.parametrize( + ("arch", "kid", "M", "N", "K", "split_k", "out_dtype"), + [ + ("gfx950", 200, 64, 64, 512, 2, torch.bfloat16), + ("gfx950", 200, 64, 64, 512, 2, torch.float32), + ("gfx942", 10200, 128, 128, 512, 2, torch.float32), + ("gfx942", 10210, 128, 128, 512, 2, torch.bfloat16), + ("gfx1250", 20000, 16, 32, 512, 2, torch.bfloat16), + ("gfx1250", 20000, 16, 32, 512, 2, torch.float32), + ], +) +def test_split_k_matches_torch_golden(arch, kid, M, N, K, split_k, out_dtype): + if _runtime_arch() != arch: + pytest.skip(f"requires {arch} hardware") + torch.manual_seed(8192 + kid) + A = torch.randn((1, M, K), device="cuda", dtype=torch.bfloat16) + B = torch.randn((1, N, K), device="cuda", dtype=torch.bfloat16) + Y = torch.empty((1, M, N), device="cuda", dtype=out_dtype) + actual = opus_bmm( + A, + B, + Y, + kid=kid, + split_k=split_k, + ) + torch.cuda.synchronize() + _assert_matches_golden(actual, A, B) + + +@pytest.mark.parametrize("workspace_splits", (None, 1, 16)) +def test_gfx942_short_k_auto_split_matches_torch(workspace_splits): + if _runtime_arch() != "gfx942": + pytest.skip("requires gfx942 hardware") + torch.manual_seed(10201) + A = torch.randn((1, 128), device="cuda", dtype=torch.bfloat16) + B = torch.randn((64, 128), device="cuda", dtype=torch.bfloat16) + Y = torch.empty((1, 64), device="cuda", dtype=torch.bfloat16) + # kid 10201 uses FP32 workspace with 64x64 tiles. The short K launches + # one split; a caller may provide either one slice or a larger workspace. + workspace = ( + None + if workspace_splits is None + else torch.empty( + (workspace_splits, 1, 64, 64), device="cuda", dtype=torch.float32 ) - for M, us, csv_us, ratio, tflops, tbs, splitk, kid in perf_rows: - # >20% slower than the tuned reference is flagged for a closer look. - note = "" if ratio <= 1.20 else "SLOW >1.20x" - splitk_str = "-" if splitk in (None, "") else str(splitk) - print( - f"{M:>6} | {us:>10.2f} | {csv_us:>10.2f} | {ratio:>6.2f}x | " - f"{tflops:>9.2f} | {tbs:>7.3f} | {splitk_str:>6} | " - f"{note or '':<10} | {kid or '-'}" - ) + ) + actual = opus_gemm(A, B, Y, kid=10201, split_k=0, workspace=workspace) + torch.cuda.synchronize() + _assert_matches_golden(actual, A, B) - print(f"\nSummary: {passed} passed, {failed} failed out of {len(rows)}") - return failed == 0 +@pytest.mark.parametrize("kid", (1400, 6400)) +def test_gfx950_mono_fp32_overwrites_poisoned_output(kid): + """Regress the ordinary and 4G-safe mono FP32 physical-store paths.""" + if _runtime_arch() != "gfx950": + pytest.skip("requires gfx950 hardware") -def test_a16w16_csv_sweep( - csv_path: str, - batch: int = 1, - *, - out_dtype=torch.bfloat16, - dist="norm", - gen=None, - const_val=1.0, - iters=101, - warmup=2, - rotate=0, - use_graph=False, + torch.manual_seed(0x950000 + kid) + A = torch.randn((1, 192, 128), device="cuda", dtype=torch.bfloat16) + B = torch.randn((1, 256, 128), device="cuda", dtype=torch.bfloat16) + out = torch.full((1, 192, 256), 12345.0, device="cuda", dtype=torch.float32) + + actual = opus_bmm( + A, + B, + out, + kid=kid, + ) + torch.cuda.synchronize() + + assert actual is out + assert int((actual != 12345.0).sum().item()) == actual.numel() + _assert_matches_golden(actual, A, B) + + +def test_gfx950_bias_dtype_rules_and_numerics(): + if _runtime_arch() != "gfx950": + pytest.skip("requires gfx950 hardware") + A = torch.randn((1, 64, 512), device="cuda", dtype=torch.bfloat16) + B = torch.randn((1, 64, 512), device="cuda", dtype=torch.bfloat16) + bias = torch.randn((64,), device="cuda", dtype=torch.bfloat16) + Y = torch.empty((1, 64, 64), device="cuda", dtype=torch.bfloat16) + actual = opus_bmm( + A, + B, + Y, + kid=200, + bias=bias, + split_k=2, + ) + torch.cuda.synchronize() + _assert_matches_golden(actual, A, B, bias) + + with pytest.raises(RuntimeError, match="bias dtype must match Y dtype"): + opus_bmm( + A, + B, + Y, + kid=200, + bias=bias.float(), + split_k=2, + ) + + +def test_gfx942_workspace_kid_rejects_bias_without_framework_fallback(): + if _runtime_arch() != "gfx942": + pytest.skip("requires gfx942 hardware") + A = torch.randn((1, 128, 4096), device="cuda", dtype=torch.bfloat16) + B = torch.randn((1, 256, 4096), device="cuda", dtype=torch.bfloat16) + Y = torch.empty((1, 128, 256), device="cuda", dtype=torch.bfloat16) + bias = torch.randn((256,), device="cuda", dtype=torch.float32) + + with pytest.raises(ValueError, match="rejects bias on split-K kernels"): + opus_bmm( + A, + B, + Y, + kid=10201, + bias=bias, + split_k=2, + ) + + +def test_gfx1250_bf16_output_accepts_fp32_bias(): + if _runtime_arch() != "gfx1250": + pytest.skip("requires gfx1250 hardware") + A = torch.randn((1, 16, 512), device="cuda", dtype=torch.bfloat16) + B = torch.randn((1, 32, 512), device="cuda", dtype=torch.bfloat16) + Y = torch.empty((1, 16, 32), device="cuda", dtype=torch.bfloat16) + bias = torch.randn((32,), device="cuda", dtype=torch.float32) + actual = opus_bmm( + A, + B, + Y, + kid=20000, + bias=bias, + split_k=2, + ) + torch.cuda.synchronize() + _assert_matches_golden(actual, A, B, bias) + + +@pytest.mark.parametrize( + ("arch", "kid", "K", "split_k", "launch_split_k", "expected_shape"), + ( + ("gfx942", 10201, 128, 0, 1, (1, 1, 64, 64)), + ("gfx942", 10201, 128, 16, 1, (1, 1, 64, 64)), + ("gfx942", 10201, 512, 0, 4, (4, 1, 64, 64)), + ("gfx1250", 20000, 128, 2, 1, None), + ), +) +def test_a16w16_split_k_plan_converges_before_workspace( + arch, kid, K, split_k, launch_split_k, expected_shape ): - shapes = load_shapes_from_csv(csv_path) - print(f"\n{'=' * 80}") - print(f"a16w16 sweep from {csv_path}: {len(shapes)} unique shapes, batch={batch}") - print("=" * 80) - passed = failed = 0 - out_bytes = torch.empty((), dtype=out_dtype).element_size() - for M, N, K in shapes: - tag = f"a16w16 b={batch} M={M} N={N} K={K}" - try: - A = _make_a(batch, M, K, dist, gen, const_val) - B = _make_b(batch, N, K, dist, gen, const_val) - ref = _torch_ref(A, B, out_dtype) - if use_graph: - _prewarm_opus_graph_workspace(A, B, out_dtype) - Y, us = run_perftest( - gemm_a16w16_opus, - A, - B, - None, - out_dtype, - testGraph=use_graph, - num_iters=iters, - num_warmup=warmup, - num_rotate_args=rotate, - ) - err = checkAllclose(Y, ref, msg=tag, rtol=0.1, atol=0.5) - tflops = _tflops(batch, M, N, K, us) - tbs = _tbs(batch, M, N, K, us, out_bytes=out_bytes) - print( - f"[PASS] {tag} | {us:.1f}us | {tflops:.2f} TFLOPs | " - f"{tbs:.3f} TB/s | err={err}" - ) - passed += 1 - except Exception as e: # noqa: BLE001 - print(f"[FAIL] {tag} | {type(e).__name__}: {e}") - failed += 1 - print(f"\nSummary: {passed} passed, {failed} failed out of {len(shapes)}") - return failed == 0 + args = _a16_policy_args(arch, 1, 64, K) + if arch == "gfx942": + args["cu_num"] = 80 + plan = _get_cached_a16w16_launch_plan(**args, kid=kid, split_k=split_k) + + assert plan.resolved_kid == kid + assert plan.workspace_capacity_split_k == launch_split_k + assert plan.abi_split_k == launch_split_k + assert plan.workspace_spec is not None + assert plan.workspace_spec.shape[0] == launch_split_k + if expected_shape is not None: + assert plan.workspace_spec.shape == expected_shape + assert plan.workspace_spec.dtype == torch.float32 + + +@pytest.mark.parametrize( + ("K", "caller_splits", "allocated_splits", "launch_split_k"), + ((128, None, 1, 1), (512, None, 4, 4), (128, 16, 16, 1)), +) +def test_gfx942_workspace_allocation_and_launch_split_k( + monkeypatch, K, caller_splits, allocated_splits, launch_split_k +): + from aiter.ops.opus import gemm_op_a16w16 + + calls = [] + + def capture(_A, _B, _Y, _bias, workspace, kid, split_k): + calls.append((workspace, kid, split_k)) + + monkeypatch.setattr(gemm_op_a16w16, "_device_arch_and_cu", lambda _: ("gfx942", 80)) + monkeypatch.setattr(gemm_op_a16w16, "_opus_gemm_a16w16_launch_raw", capture) + A = torch.empty((1, K), device="meta", dtype=torch.bfloat16) + B = torch.empty((64, K), device="meta", dtype=torch.bfloat16) + Y = torch.empty((1, 64), device="meta", dtype=torch.bfloat16) + workspace = ( + None + if caller_splits is None + else torch.empty((caller_splits, 1, 64, 64), device="meta", dtype=torch.float32) + ) + assert opus_gemm(A, B, Y, kid=10201, split_k=0, workspace=workspace) is Y + assert len(calls) == 1 + actual_workspace, actual_kid, actual_split_k = calls[0] + assert (actual_kid, actual_split_k) == (10201, launch_split_k) + assert actual_workspace.shape == (allocated_splits, 1, 64, 64) + assert actual_workspace.dtype == torch.float32 + if workspace is not None: + assert actual_workspace is workspace + + +@pytest.mark.parametrize( + ("arch", "kid", "K", "split_k", "error"), + ( + ("gfx942", 10201, 64, 0, "too small for gfx942"), + ("gfx942", 10201, 192, 0, "needs even loops per split"), + ("gfx950", 200, 128, 3, "too small for gfx950"), + ), +) +def test_a16w16_launch_plan_preserves_split_k_limits(arch, kid, K, split_k, error): + with pytest.raises(ValueError, match=error): + _get_cached_a16w16_launch_plan( + **_a16_policy_args(arch, 1, 64, K), kid=kid, split_k=split_k + ) -if __name__ == "__main__": - parser = argparse.ArgumentParser( - description="End-to-end test for aiter.ops.opus.gemm_a16w16_opus" + +def test_gfx1250_split_k_reducer_row_limit(): + from csrc.opus_gemm.opus_gemm_common import ( + gfx1250_clusterlaunch_kernels_list, + gfx1250_kernels_list, ) - parser.add_argument( - "-m", - type=int, - default=None, - help="Single-shape M. Passing -m forces the single-shape test.", + + for instances in (gfx1250_kernels_list, gfx1250_clusterlaunch_kernels_list): + assert instances + kid = min(instances) + instance = instances[kid] + args = _a16_policy_args("gfx1250", 65535, instance.B_N, 2 * instance.B_K) + plan = _get_cached_a16w16_launch_plan(**args, kid=kid, split_k=2) + assert plan.workspace_spec.dtype == torch.float32 + + args["M"] = 65536 + with pytest.raises(ValueError, match="requires M <= 65535"): + _get_cached_a16w16_launch_plan(**args, kid=kid, split_k=2) + + +def test_global_a16_stale_opus_row_keeps_framework_fallback(monkeypatch): + import aiter.tuned_gemm as tuned + + warnings = [] + key = ( + "gfx942", + 304, + 32, + 256, + 1024, + False, + str(torch.bfloat16), + str(torch.bfloat16), + False, + False, ) - parser.add_argument( - "-n", - type=int, - default=None, - help="N (default: 2048 for the opus sweep, 512 for single-shape).", + row = {"libtype": "opus", "solidx": 200, "splitK": 2, "kernelName": ""} + monkeypatch.setattr(tuned, "get_GEMM_A16W16_config_", lambda: {key: row}) + monkeypatch.setattr(tuned, "get_gfx", lambda: "gfx942") + monkeypatch.setattr(tuned, "get_cu_num", lambda: 304) + # Revisit the stale key in both padded lookups without the native helper. + monkeypatch.setattr(tuned, "get_padded_m", lambda M, _N, _K, _gl: M) + monkeypatch.setattr(tuned, "_opus_launch", object()) + monkeypatch.setattr( + tuned.logger, + "warning", + lambda message, *args: warnings.append(message % args), ) - parser.add_argument( - "-k", - type=int, - default=None, - help="K (default: 7168 for the opus sweep, 256 for single-shape).", + tuned.get_GEMM_A16W16_config.cache_clear() + try: + config = tuned.get_GEMM_A16W16_config( + 32, + 256, + 1024, + False, + str(torch.bfloat16), + str(torch.bfloat16), + ) + finally: + tuned.get_GEMM_A16W16_config.cache_clear() + + assert (config["libtype"], config["solidx"]) == ("torch", 0) + stale_warnings = [ + message + for message in warnings + if message.startswith("Ignoring invalid OPUS tuned row") + ] + assert len(stale_warnings) == 1 + assert "kid=200, splitK=2" in stale_warnings[0] + + +def _capture_shape_driven_opus_launch(monkeypatch, *, arch, tuned_config): + from aiter.ops import opus + from aiter.ops.opus import gemm_op_a16w16, policy + + calls = [] + + def capture(operation): + def fake_launch(XQ, WQ, Y, **kwargs): + calls.append((operation, XQ, WQ, Y, kwargs)) + return Y + + return fake_launch + + monkeypatch.setattr( + gemm_op_a16w16, + "_device_arch_and_cu", + lambda _device: (arch, {"gfx950": 256, "gfx942": 304, "gfx1250": 80}[arch]), ) - parser.add_argument( - "-b", - "--batch", - type=int, - default=None, - help="Batch size. Defaults to 1 for --opus_sweep, else 8.", + monkeypatch.setattr( + policy, + "lookup_a16w16_opus_config", + lambda **_kwargs: tuned_config, ) + monkeypatch.setattr(gemm_op_a16w16, "_launch_a16w16_gemm", capture("gemm")) + monkeypatch.setattr(gemm_op_a16w16, "_launch_a16w16_bmm", capture("bmm")) + return opus, calls + + +def _a16_policy_args(arch, M, N, K): + return { + "arch": arch, + "M": M, + "N": N, + "K": K, + "batch": 1, + "cu_num": {"gfx950": 256, "gfx942": 304, "gfx1250": 80}[arch], + "has_bias": False, + "input_dtype": torch.bfloat16, + "output_dtype": torch.bfloat16, + } + + +def _mock_a16w16_policy_csv(monkeypatch, read_csv): + from types import SimpleNamespace + + from aiter.ops.opus import policy + + monkeypatch.setattr( + policy, + "AITER_CONFIGS", + SimpleNamespace(AITER_CONFIG_GEMM_BF16_FILE="test.csv"), + ) + monkeypatch.setattr(policy.pd, "read_csv", read_csv) + policy._load_a16w16_opus_tuned.cache_clear() + return policy + + +@pytest.mark.parametrize("case", ("missing", "empty", "partial", "unreadable")) +def test_a16w16_policy_loader_bad_csv_is_a_miss(monkeypatch, case): + import pandas as pd + + warnings = [] + errors = { + "missing": FileNotFoundError(), + "empty": pd.errors.EmptyDataError("empty CSV"), + "unreadable": pd.errors.ParserError("bad CSV"), + } + + def read_csv(_path): + if case in errors: + raise errors[case] + return pd.DataFrame({"libtype": ["opus"], "solidx": [200]}) + + policy = _mock_a16w16_policy_csv(monkeypatch, read_csv) + monkeypatch.setattr( + policy.logger, + "warning", + lambda *args, **_kwargs: warnings.append(args), + ) + try: + assert policy._load_a16w16_opus_tuned() == {} + assert bool(warnings) == (case in ("partial", "unreadable")) + finally: + policy._load_a16w16_opus_tuned.cache_clear() + + +def test_a16w16_policy_loader_skips_malformed_kid_and_splitk_rows(monkeypatch): + import pandas as pd + + from aiter.ops.opus import policy + + key = ( + "gfx950", + 256, + 128, + 64, + 512, + False, + "torch.bfloat16", + "torch.bfloat16", + False, + False, + ) + columns = ( + *policy._A16W16_TUNED_KEY_COLUMNS, + "libtype", + "solidx", + "splitK", + "us", + ) + rows = [ + (*key, "opus", "not-a-kid", 2, 1.0), + (*key, "opus", 200, "not-a-split", 1.5), + (*key, "opus", 200, -1, 1.75), + (*key, "opus", "200", "2", 2.0), + ] + policy = _mock_a16w16_policy_csv( + monkeypatch, lambda _path: pd.DataFrame(rows, columns=columns) + ) + try: + configs = policy._load_a16w16_opus_tuned() + assert [(row["solidx"], row["splitK"]) for row in configs.values()] == [ + (200, 2) + ] + finally: + policy._load_a16w16_opus_tuned.cache_clear() + + +def test_shape_driven_opus_selection_and_rank_route(monkeypatch): + from aiter.ops.opus import gemm_op_a16w16, policy + + tuned = [None] + opus, calls = _capture_shape_driven_opus_launch( + monkeypatch, arch="gfx950", tuned_config=None + ) + monkeypatch.setattr(policy, "lookup_a16w16_opus_config", lambda **_kwargs: tuned[0]) + A = torch.empty((128, 512), device="meta", dtype=torch.bfloat16) + B = torch.empty((64, 512), device="meta", dtype=torch.bfloat16) + + opus.gemm_a16w16_opus(A, B) + tuned[0] = {"solidx": 200, "splitK": 2} + opus.gemm_a16w16_opus(A, B) + opus.gemm_a16w16_opus(A, B, kernelId=206, splitK=3) + tuned[0] = None + opus.gemm_a16w16_opus( + A.unsqueeze(0).expand(2, -1, -1).contiguous(), + B.unsqueeze(0).expand(2, -1, -1).contiguous(), + ) + + assert [(op, args["kid"], args["split_k"]) for op, *_, args in calls] == [ + ("gemm", 1200, 0), + ("gemm", 200, 2), + ("gemm", 206, 3), + ("bmm", 1200, 0), + ] + tuned[0] = {"solidx": -1, "splitK": 0} + warnings = [] + gemm_op_a16w16._warn_invalid_a16w16_tuned_row.cache_clear() + monkeypatch.setattr(gemm_op_a16w16.logger, "warning", warnings.append) + try: + opus.gemm_a16w16_opus(A, B) + opus.gemm_a16w16_opus(A, B) + finally: + gemm_op_a16w16._warn_invalid_a16w16_tuned_row.cache_clear() + + assert [(op, args["kid"], args["split_k"]) for op, *_, args in calls[-2:]] == [ + ("gemm", 1200, 0), + ("gemm", 1200, 0), + ] + assert len(warnings) == 1 + assert "kid=-1, splitK=0" in warnings[0] + + +@pytest.mark.parametrize("kid", (10210, 10213, 10216)) +def test_gfx942_exact_plan_rejects_non_exact_n_bf16_workspace_kid(kid): + with pytest.raises(ValueError, match=rf"gfx942 exact kid {kid} requires N"): + _get_cached_a16w16_launch_plan( + "gfx942", + 256, + 1000, + 4096, + 1, + 304, + False, + torch.bfloat16, + torch.bfloat16, + kid, + 2, + ) + + +@pytest.mark.parametrize( + ("selection", "requested_kid", "resolved_kid"), + (("explicit", 10210, 10200), ("tuned", 10213, 10203)), +) +def test_gfx942_compat_redirects_non_exact_n_bf16_workspace_kid( + monkeypatch, selection, requested_kid, resolved_kid +): + tuned_config = ( + {"solidx": requested_kid, "splitK": 2} if selection == "tuned" else None + ) + opus, calls = _capture_shape_driven_opus_launch( + monkeypatch, arch="gfx942", tuned_config=tuned_config + ) + A = torch.empty((256, 4096), device="meta", dtype=torch.bfloat16) + B = torch.empty((1000, 4096), device="meta", dtype=torch.bfloat16) + kwargs = {"kernelId": requested_kid, "splitK": 2} if selection == "explicit" else {} + + opus.gemm_a16w16_opus(A, B, **kwargs) + + assert [(op, args["kid"], args["split_k"]) for op, *_, args in calls] == [ + ("gemm", resolved_kid, 2) + ] + + +def test_gfx942_compat_rejects_non_exact_n_bf16_workspace_kid_without_sibling( + monkeypatch, +): + from aiter.ops.opus import gemm_op_a16w16, policy + + monkeypatch.setattr( + gemm_op_a16w16, + "_device_arch_and_cu", + lambda _device: ("gfx942", 304), + ) + monkeypatch.setattr(policy, "lookup_a16w16_opus_config", lambda **_kwargs: None) + A = torch.empty((256, 4096), device="meta", dtype=torch.bfloat16) + B = torch.empty((1000, 4096), device="meta", dtype=torch.bfloat16) + + with pytest.raises(ValueError, match="gfx942 exact kid 10216 requires N"): + gemm_op_a16w16.gemm_a16w16_opus(A, B, kernelId=10216, splitK=2) + + +def test_legacy_a16w16_tune_routes_to_family_executor(monkeypatch): + from aiter.ops import deepgemm + from aiter.ops.opus import gemm_op_a16w16 + + calls = [] + + def capture(XQ, WQ, Y, bias=None, **kwargs): + calls.append((XQ, WQ, Y, bias, kwargs)) + return Y + + monkeypatch.setattr(gemm_op_a16w16, "_execute_a16w16", capture) + monkeypatch.setattr( + gemm_op_a16w16, + "_launch_a16w16_gemm", + lambda *_args, **_kwargs: pytest.fail("compatibility used public GEMM route"), + ) + monkeypatch.setattr( + gemm_op_a16w16, + "_launch_a16w16_bmm", + lambda *_args, **_kwargs: pytest.fail("GEMM compatibility used BMM"), + ) + XQ = torch.empty((1, 8, 16), device="meta", dtype=torch.bfloat16) + WQ = torch.empty((1, 32, 16), device="meta", dtype=torch.bfloat16) + Y = torch.empty((1, 8, 32), device="meta", dtype=torch.bfloat16) + + assert gemm_op_a16w16.opus_gemm_a16w16_tune(XQ, WQ, Y, 206, 3) is Y + assert gemm_op_a16w16.opus_gemm_a16w16_tune(XQ, WQ, Y, 207, splitK=4) is Y + with pytest.warns(DeprecationWarning, match="has moved"): + assert deepgemm.opus_gemm_a16w16_tune(XQ, WQ, Y, 208, 5) is Y + assert [ + (tuple(xq.shape), tuple(wq.shape), tuple(y.shape), bias, kwargs) + for xq, wq, y, bias, kwargs in calls + ] == [ + ((1, 8, 16), (1, 32, 16), (1, 8, 32), None, {"kid": 206, "split_k": 3}), + ((1, 8, 16), (1, 32, 16), (1, 8, 32), None, {"kid": 207, "split_k": 4}), + ((1, 8, 16), (1, 32, 16), (1, 8, 32), None, {"kid": 208, "split_k": 5}), + ] + + +@pytest.mark.parametrize("arch", ("gfx950", "gfx1250")) +def test_shape_driven_opus_heuristic_rejects_over_4g_shape(arch): + from aiter.ops.opus.policy import resolve_a16w16_heuristic_candidate + + with pytest.raises(RuntimeError, match="refuses >4 GiB shape"): + resolve_a16w16_heuristic_candidate(**_a16_policy_args(arch, 1, 1, 1 << 31)) + + +def test_gfx950_heuristic_does_not_try_secondary_kids(monkeypatch): + from aiter.ops.opus import policy + + attempted = [] + monkeypatch.setattr( + policy, + "select_a16w16_heuristic_kid", + lambda **_kwargs: 1200, + ) + + def reject_candidate(**kwargs): + attempted.append(kwargs["kid"]) + + monkeypatch.setattr(policy, "_resolve_a16w16_candidate", reject_candidate) + plan = policy.resolve_a16w16_heuristic_candidate( + **_a16_policy_args("gfx950", 128, 64, 512) + ) + + assert plan is None + assert attempted == [1200] + + +if __name__ == "__main__": + if not _arch_ok: + print( + "[skip] test_opus_a16w16_gemm requires " + f"gfx950/gfx942/gfx1250 (detected {_detected_gfx!r})" + ) + sys.exit(0) + + # The standard Aiter runner executes each file directly. + if len(sys.argv) == 1: + sys.exit(pytest.main([__file__])) + + parser = argparse.ArgumentParser(description="A16W16 OPUS exact-kid benchmark") + parser.add_argument("-m", type=int, default=None) + parser.add_argument("-n", type=int, default=None) + parser.add_argument("-k", type=int, default=None) + parser.add_argument("-b", "--batch", type=int, default=None) + parser.add_argument("--kid", type=int, default=None) + parser.add_argument("--split-k", type=int, default=0) parser.add_argument( "-d", "--dtype", @@ -474,119 +926,41 @@ def test_a16w16_csv_sweep( "single-shape test and runs a full sweep instead." ), ) - parser.add_argument( - "--opus_sweep", - action="store_true", - help=( - "Run the CUDA-graph-mode opus_gemm sweep (golden-checked) over " - "the M values whose tuned winner is opus for the given N/K in the " - "tuned CSV (default: dsv4_bf16_tuned_gemm.csv). This is also the " - "DEFAULT action when no -m and no --csv_file is given." - ), - ) - parser.add_argument( - "--tuned_csv", - type=str, - default=None, - metavar="CSV", - help=( - "Tuned GEMM CSV used by --opus_sweep to pick opus shapes. " - "Defaults to the shipped dsv4_bf16_tuned_gemm.csv." - ), - ) parser.add_argument( "--graph", action="store_true", help="Use CUDA-graph mode for the single-shape / --csv_file paths too.", ) - # --- warmup / iteration controls --- - parser.add_argument( - "--iters", - type=int, - default=101, - help="Timed iterations passed to run_perftest (default: 101).", - ) - parser.add_argument( - "--warmup", - type=int, - default=2, - help="Warmup iterations passed to run_perftest (default: 2).", - ) - # --- rotating tensors --- - parser.add_argument( - "--rotate", - type=int, - default=0, - help=( - "num_rotate_args for run_perftest: number of rotated input copies " - "used to defeat L2 caching. 0 (default) lets the framework auto-size " - "the rotation from the L2 cache; 1 disables rotation." - ), - ) - # --- data initialization + seed (shared aiter.test_common API) --- - # add_data_init_args attaches --data-init / --scale-init / --seed; a16w16 - # has no scale operand, so --scale-init is accepted but unused. Default the - # DATA dist to norm to preserve the original torch.randn init. - add_data_init_args(parser, default_dist="norm") - parser.add_argument( - "--const-val", - type=float, - default=1.0, - help="Fill value used by --data-init constant (default: 1.0).", - ) args = parser.parse_args() out_dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float32 - gen = make_generator(args.seed) - if len(args.data_init) != 1: - parser.error( - "--data-init accepts exactly one distribution for the a16w16 benchmark" - ) - data_init = args.data_init[0] - init_kwargs = { - "dist": data_init, - "gen": gen, - "const_val": args.const_val, - "iters": args.iters, - "warmup": args.warmup, - "rotate": args.rotate, - } + if args.csv_file is not None and args.m is not None: + parser.error("--csv_file cannot be combined with -m") - # Default action (no -m and no --csv_file): auto-sweep the opus shapes in - # CUDA-graph mode and print the vs-CSV latency table. So a bare - # `python3 op_tests/test_opus_a16w16_gemm.py` reproduces the table. - run_opus_sweep = args.opus_sweep or (args.m is None and args.csv_file is None) - - if run_opus_sweep: - tuned_csv = args.tuned_csv or _default_tuned_csv() - batch = args.batch if args.batch is not None else 1 - ok = test_opus_shapes_graph( - tuned_csv, - _detected_gfx, - N=args.n if args.n is not None else 2048, - K=args.k if args.k is not None else 7168, - batch=batch, - out_dtype=out_dtype, - **init_kwargs, - ) - sys.exit(0 if ok else 1) - elif args.csv_file is not None: - test_a16w16_csv_sweep( + if args.csv_file is not None: + ok = run_a16w16_csv_sweep( args.csv_file, - batch=(args.batch or 8), + batch=args.batch or 8, + kid=args.kid, + split_k=args.split_k, out_dtype=out_dtype, use_graph=args.graph, - **init_kwargs, ) + sys.exit(0 if ok else 1) else: - # Clamp K>=128 so every kid the heuristic picks has K>=B_K (smallest is 128). - k_eff = max(args.k if args.k is not None else 256, 128) - test_a16w16( - args.batch or 8, - args.m, - args.n if args.n is not None else 512, - k_eff, + if args.kid is None: + parser.error("--kid is required for a single-shape exact run") + M = args.m if args.m is not None else 256 + N = args.n if args.n is not None else 512 + K = max(args.k if args.k is not None else 256, 128) + batch = args.batch or 8 + run_a16w16_case( + batch, + M, + N, + K, + kid=args.kid, + split_k=args.split_k, out_dtype=out_dtype, use_graph=args.graph, - **init_kwargs, ) diff --git a/op_tests/test_opus_a16w16_policy_parity.py b/op_tests/test_opus_a16w16_policy_parity.py new file mode 100644 index 0000000000..3fa74d04c7 --- /dev/null +++ b/op_tests/test_opus_a16w16_policy_parity.py @@ -0,0 +1,411 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2025-2026, Advanced Micro Devices, Inc. All rights reserved. +"""CPU-only parity coverage for the migrated OPUS A16W16 policy.""" + +from __future__ import annotations + +from functools import cache +from itertools import product +from pathlib import Path +from types import SimpleNamespace + +import pandas as pd +import pytest +import torch + +from aiter.ops.opus import policy +from csrc.opus_gemm.opus_gemm_common import get_kernel_instance + + +# Keep the original C++ heuristic reference independent of policy.py. +def _pre_pr_gfx950(M: int, N: int, K: int, has_bias: bool, _output: str) -> int: + split_barrier_ok = N % 16 == 0 and K % 64 == 0 and (K // 64) % 2 == 0 + if M <= 4: + if M % 64 == 0 and N % 64 == 0 and K % 128 == 0: + return 1208 + return 208 + if M <= 64: + if M % 64 == 0 and N % 32 == 0 and K % 128 == 0: + return 1206 + return 206 + if M <= 128: + if M % 64 == 0 and N % 64 == 0 and K % 64 == 0: + return 1200 + return 200 + if split_barrier_ok and not has_bias: + if M % 256 == 0 and N % 256 == 0 and K % 64 == 0: + return 1300 + return 300 + if M % 64 == 0 and N % 64 == 0 and K % 64 == 0: + return 1200 + return 200 + + +def _pre_pr_gfx1250(M: int, N: int, _K: int, _has_bias: bool, _output: str) -> int: + if M % 32 == 0: + if N % 128 == 0: + return 20007 + if N % 64 == 0: + return 20006 + if N % 32 == 0: + return 20005 + if N % 128 == 0: + return 20004 + if N % 64 == 0: + return 20003 + return 20000 + + +def _pre_pr_gfx942_bf16(M: int, N: int, K: int) -> int: + k64_ok = K % 64 == 0 + k32_ok = K % 32 == 0 + wkc_bk64_ok = K >= 4096 and K % 512 == 0 + p1_ok = K % 128 == 0 + loops = (K + 63) // 64 + sb_ok = N % 16 == 0 and K % 64 == 0 and loops >= 2 and loops % 2 == 0 + + if K == 4096: + if p1_ok and M in (48, 64) and N == 1024: + return 10213 + if p1_ok and ((M == 128 and N == 512) or (M == 256 and N == 256)): + return 10213 + if p1_ok and M == 512 and N == 256: + return 10203 + if M in (48, 64) and 1536 <= N <= 2048: + return 10205 + if (M == 128 and N == 1024) or (M == 256 and N == 512): + return 10205 + if ( + (M == 128 and 1536 <= N <= 2048) + or (M == 256 and N == 1024) + or (M == 512 and N == 512) + ): + return 10200 + + if K >= 1024 and k32_ok and N >= 1536 and M <= 32: + if M <= 4 and N >= 4096: + return 10300 + if M <= 16: + return 10305 if wkc_bk64_ok else 10301 + return 10305 if M == 32 and K == 4096 and wkc_bk64_ok else 10303 + + if ( + K >= 512 + and k64_ok + and (N <= 64 or (M <= 128 and N <= 1024) or (M <= 8 and N <= 1536)) + ): + if N <= 64 and M > 128: + return 10302 + if N <= 256 or M <= 8 or (M <= 16 and N <= 800): + return 10300 + return 10302 + + bf16ws_band = ( + K >= 4096 and K % 64 == 0 and 104 <= M <= 608 and (N == 256 or 512 <= N <= 2048) + ) + if bf16ws_band: + return 10210 + + if N == 384 and K >= 4096: + if M <= 128: + return 10302 + if M <= 224: + return 10201 + if 392 <= M <= 512: + return 10204 + return 10200 + + if k64_ok and N >= 4096 and K <= 3200: + if K <= 640 and M <= 128: + return 10001 + return 10000 + if sb_ok and M >= 128: + return 10000 + if N <= 256 and p1_ok: + return 10201 + return 10200 + + +def _pre_pr_gfx942(M: int, N: int, K: int, has_bias: bool, output: str) -> int: + if output == "bf16" and not has_bias: + return _pre_pr_gfx942_bf16(M, N, K) + if N <= 256 and K % 128 == 0: + return 10201 + return 10200 + + +_PRE_PR_HEURISTICS = { + "gfx950": _pre_pr_gfx950, + "gfx942": _pre_pr_gfx942, + "gfx1250": _pre_pr_gfx1250, +} + + +def _around(*boundaries: int) -> tuple[int, ...]: + return tuple( + sorted( + value + for boundary in boundaries + for value in (boundary - 1, boundary, boundary + 1) + if value > 0 + ) + ) + + +_M_SWEEP = _around(4, 8, 16, 32, 48, 64, 104, 128, 224, 256, 392, 512, 608) +_N_SWEEP = _around(16, 32, 64, 128, 256, 384, 512, 768, 800, 1024, 1536, 2048, 4096) +_K_SWEEP = _around(32, 64, 128, 512, 640, 1000, 1024, 2048, 3200, 4096, 5120, 7168) + + +def test_python_heuristics_match_pre_pr_cpp_boundary_sweep(): + mismatches = [] + checked = 0 + for arch, reference in _PRE_PR_HEURISTICS.items(): + for M, N, K, has_bias, output in product( + _M_SWEEP, _N_SWEEP, _K_SWEEP, (False, True), ("bf16", "fp32") + ): + expected = reference(M, N, K, has_bias, output) + actual = policy.select_a16w16_heuristic_kid( + arch=arch, + M=M, + N=N, + K=K, + batch=1, + has_bias=has_bias, + output_dtype=output, + ) + checked += 1 + if actual != expected: + mismatches.append((arch, M, N, K, has_bias, output, expected, actual)) + if len(mismatches) == 20: + break + if mismatches: + break + + assert checked > 500_000 + assert not mismatches, f"Python heuristic differs from C++ reference: {mismatches}" + + +@cache +def _shipped_opus_rows() -> pd.DataFrame: + root = Path(__file__).resolve().parents[1] + paths = [root / "aiter/configs/bf16_tuned_gemm.csv"] + paths.extend( + sorted((root / "aiter/configs/model_configs").glob("*_bf16_tuned_gemm.csv")) + ) + frames = [] + for path in paths: + frame = pd.read_csv(path) + if "libtype" in frame.columns: + frame = frame[frame["libtype"].eq("opus")] + if not frame.empty: + frames.append(frame) + assert frames, "the shipped BF16 tuning files contain no OPUS rows" + return pd.concat(frames, ignore_index=True).drop_duplicates() + + +def _pre_pr_python_tuned_map(rows: pd.DataFrame) -> dict[tuple, dict]: + # aiter/ops/opus/common.py on main used this key (notably without gfx). + columns = policy._A16W16_TUNED_KEY_COLUMNS[1:] + result = {} + for _, row in rows.sort_values("us", kind="stable", na_position="last").iterrows(): + key = tuple(row[column] for column in columns) + result.setdefault(key, row.to_dict()) + return result + + +def _pre_pr_codegen_tuned_map(rows: pd.DataFrame) -> dict[tuple, int]: + # opus_gemm_lookup.h had one arch-specific (M,N,K,outdtype) table. The + # generated entry carried a launcher pointer, represented here by its kid. + result = {} + for _, row in rows.iterrows(): + kid = int(row["solidx"]) + kid_arches = [ + arch + for arch in _PRE_PR_HEURISTICS + if get_kernel_instance(arch, "a16w16", kid) is not None + ] + assert kid_arches == [str(row["gfx"])] + key = ( + kid_arches[0], + int(row["M"]), + int(row["N"]), + int(row["K"]), + str(row["outdtype"]), + ) + result[key] = kid + return result + + +def _dtype(value: object) -> torch.dtype: + return { + "torch.bfloat16": torch.bfloat16, + "torch.float32": torch.float32, + }[str(value)] + + +_EXPECTED_GFX_KEY_FIXES = { + ("gfx950", 4096, 2048, 4096): ((21177, 0), (6401, 0)), + ("gfx950", 8192, 1024, 4096): ((21177, 0), (1401, 0)), + ("gfx950", 16384, 512, 4096): ((21177, 0), (1401, 0)), + ("gfx950", 16384, 512, 7168): ((21177, 0), (1401, 0)), +} + + +def test_shipped_tuned_selection_diff_is_exhaustive( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + rows = _shipped_opus_rows() + merged = tmp_path / "bf16_tuned_gemm.csv" + rows.to_csv(merged, index=False) + monkeypatch.setattr( + policy, + "AITER_CONFIGS", + SimpleNamespace(AITER_CONFIG_GEMM_BF16_FILE=str(merged)), + ) + policy._load_a16w16_opus_tuned.cache_clear() + policy.lookup_a16w16_opus_config.cache_clear() + try: + current = policy._load_a16w16_opus_tuned() + finally: + policy._load_a16w16_opus_tuned.cache_clear() + policy.lookup_a16w16_opus_config.cache_clear() + + pre_pr_python = _pre_pr_python_tuned_map(rows) + pre_pr_codegen = _pre_pr_codegen_tuned_map(rows) + differences = {} + invalid_current_rows = [] + redirected_current_rows = [] + + for key, config in current.items(): + fields = dict(zip(policy._A16W16_TUNED_KEY_COLUMNS, key)) + arch = str(fields["gfx"]) + shape = (arch, int(fields["M"]), int(fields["N"]), int(fields["K"])) + pair = (int(config["solidx"]), int(config["splitK"])) + + # The removed C++ table and the runtime table select the same kid for + # every shipped full-key row. splitK was not stored in the C++ table. + codegen_key = (*shape, str(fields["outdtype"])) + assert pre_pr_codegen[codegen_key] == pair[0] + + old_row = pre_pr_python[key[1:]] + old_pair = (int(old_row["solidx"]), int(old_row["splitK"])) + if old_pair != pair: + assert str(old_row["gfx"]) != arch + assert get_kernel_instance(arch, "a16w16", old_pair[0]) is None + assert shape not in differences + differences[shape] = (old_pair, pair) + + plan = policy.resolve_a16w16_tuned_candidate( + arch=arch, + M=int(fields["M"]), + N=int(fields["N"]), + K=int(fields["K"]), + batch=1, + cu_num=int(fields["cu_num"]), + has_bias=bool(fields["bias"]), + input_dtype=_dtype(fields["dtype"]), + output_dtype=_dtype(fields["outdtype"]), + requested_kid=pair[0], + requested_split_k=pair[1], + ) + if plan is None: + invalid_current_rows.append((*shape, pair)) + elif plan.resolved_kid != pair[0]: + redirected_current_rows.append((*shape, pair, plan.resolved_kid)) + + assert differences == _EXPECTED_GFX_KEY_FIXES + assert invalid_current_rows == [] + assert redirected_current_rows == [] + + +@pytest.mark.parametrize( + ("arch", "cu_num", "M", "N", "K", "has_bias", "output", "expected"), + ( + # Representative launch failures after the independent selector sweep. + ("gfx950", 256, 1, 17, 130, False, torch.bfloat16, None), + ("gfx950", 256, 256, 256, 128, False, torch.bfloat16, (1300, 0)), + ("gfx950", 256, 256, 256, 128, True, torch.bfloat16, None), + ("gfx1250", 256, 31, 127, 4098, False, torch.bfloat16, (20000, 1)), + ("gfx1250", 256, 32, 128, 4098, False, torch.bfloat16, (20007, 1)), + ("gfx942", 80, 256, 768, 7168, False, torch.bfloat16, (10200, 7)), + ("gfx942", 80, 257, 1024, 7168, False, torch.bfloat16, (10210, 4)), + ("gfx942", 80, 32, 1537, 2048, False, torch.bfloat16, (10303, 0)), + ("gfx942", 80, 32, 256, 1024, False, torch.float32, (10201, 8)), + ), +) +def test_untuned_shape_launch_resolution( + arch, cu_num, M, N, K, has_bias, output, expected +): + rows = _shipped_opus_rows() + full_key = ( + arch, + cu_num, + M, + N, + K, + has_bias, + str(torch.bfloat16), + str(output), + False, + False, + ) + current_keys = { + tuple(row[column] for column in policy._A16W16_TUNED_KEY_COLUMNS) + for _, row in rows.iterrows() + } + assert full_key not in current_keys + + plan = policy.resolve_a16w16_heuristic_candidate( + arch=arch, + M=M, + N=N, + K=K, + batch=1, + cu_num=cu_num, + has_bias=has_bias, + input_dtype=torch.bfloat16, + output_dtype=output, + ) + if expected is None: + assert plan is None + return + assert plan is not None + assert (plan.resolved_kid, plan.abi_split_k) == expected + + +@pytest.mark.parametrize( + ("N", "requested", "expected"), + ( + (768, 10210, 10200), + (768, 10213, 10203), + (768, 10216, None), + (768, 10200, 10200), + (768, 10300, 10300), + (64, 10210, 10210), + (1024, 10213, 10213), + (2048, 10216, 10216), + ), +) +def test_gfx942_requested_kid_matches_pre_pr_generated_launcher(N, requested, expected): + # main redirected the two paired BF16-workspace launchers and AITER_CHECKed + # 10216, which has no FP32-workspace sibling. ``None`` is that rejection at + # policy time; unrelated gfx942 kids must remain unchanged. + plan = policy.resolve_a16w16_tuned_candidate( + arch="gfx942", + M=256, + N=N, + K=4096, + batch=1, + cu_num=80, + has_bias=False, + input_dtype=torch.bfloat16, + output_dtype=torch.bfloat16, + requested_kid=requested, + requested_split_k=1, + ) + assert (None if plan is None else plan.resolved_kid) == expected + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/op_tests/test_opus_a8w8_bmm.py b/op_tests/test_opus_a8w8_bmm.py index 1b95b5a8f7..b6cca6574c 100644 --- a/op_tests/test_opus_a8w8_bmm.py +++ b/op_tests/test_opus_a8w8_bmm.py @@ -5,7 +5,7 @@ Covers the mmajor DeepSeek-V4 wo_a path: O/Y are [M, G, *] (transposed views of batch-major [G, M, *]); wo_a + w_scale stay batch-major. Activation scale is per-token e8m0 (GROUP_M=1), weight scale is 128x128-block e8m0. Candidates are -kid 0 (always-runnable baseline) and the public dispatch path; the reference is +kid 8000 (always-runnable baseline) and the public dispatch path; the reference is a dequantized fp32 einsum. Per-kid perf comparison / winner selection lives in ``csrc/opus_gemm/opus_bmm_mxscale_tune.py``. @@ -30,17 +30,39 @@ import aiter from aiter import dtypes from aiter.jit.utils.chip_info import get_gfx -from aiter.ops.batched_gemm_op_a8w8 import lookup_mxscale_bmm_config -from aiter.ops.opus.bmm_op import _opus_bmm_a8w8_mxscale_raw, bmm_a8w8_mxscale_opus +from aiter.ops.opus import opus_bmm +from aiter.ops.opus.policy import lookup_mxscale_bmm_config from aiter.test_common import benchmark, checkAllclose, run_perftest -torch.set_default_device("cuda") - SUPPORTED_GFX = ["gfx950"] # fp8 e8m0 mxscale flatmm is gfx950-only GROUP = 128 # GROUP_N == GROUP_K == 128; GROUP_M == 1 (per-token) _DT = {"fp32": dtypes.fp32, "bf16": dtypes.bf16} +def _run_opus( + x, + weight, + out, + x_scale, + w_scale, + kid, + split_k=1, + workspace=None, +): + opus_bmm( + x.transpose(0, 1), + weight, + out.transpose(0, 1), + kid=int(kid), + layout="mxscale_bmm", + x_scale=x_scale.transpose(0, 1), + w_scale=w_scale, + split_k=int(split_k), + workspace=workspace, + ) + return out + + def _to_e8m0_scale(scale): # Round scale up to a power of two so quantized fp8 values stay in range. e = torch.ceil(torch.log2(scale.to(dtypes.fp32))).to(torch.int32) + 127 @@ -88,7 +110,8 @@ def _block_varied(shape, k): """Signed random tensor whose per-128-K-block magnitude spans several powers of two, so the e8m0 128-block scales cover many exponents. - ``rand()/10`` (non-negative, near-uniform) is what let the shipped kid312/313 + ``rand()/10`` (non-negative, near-uniform) is what let the shipped + global kids 8312/8313 tileN COM_REP_N>1 kernels pass this test at ~0.007 rel while silently transposing output column groups: a pure column permutation over symmetric positive columns barely moves any element, and the collapsed single block @@ -118,19 +141,18 @@ def test_mxscale_bmm(g, m, n, k, dtype): def _call(kid): Y = torch.empty(y_shape, dtype=ydt) - _opus_bmm_a8w8_mxscale_raw(O_in, W_mx, Y, xs_in, ws_mx, 1, kid) + _run_opus(O_in, W_mx, Y, xs_in, ws_mx, kid) return Y - # Correctness-focused: kid 0 (k32 fused) is a fixed baseline with no + # Correctness-focused: global kid 8000 is a fixed baseline with no # tile-alignment requirement (always runnable), plus the public dispatch path # end to end. Per-kid perf comparison / winner selection lives in # csrc/opus_gemm/opus_bmm_mxscale_tune.py, not here. - candidates = {"kid0_k32_fused": (lambda: _call(0), ref)} + candidates = {"kid8000_k32_fused": (lambda: _call(8000), ref)} - # Public backend-neutral entry: no kernelId -> per-(g,m,n,k) tuned-CSV - # lookup + heuristic fallback + libtype backend routing. Exercises the - # whole aiter.batched_gemm_a8w8_mxscale -> bmm_a8w8_mxscale_opus path end - # to end (not the raw binding). + # Public backend-neutral entry: per-(g,m,n,k) tuned-CSV lookup + heuristic + # final-kid resolution + libtype backend routing. Split-one uses the + # checked raw fast path; a future split-K tuned row uses public opus_bmm. candidates["auto (batched_gemm_a8w8_mxscale)"] = ( lambda: aiter.batched_gemm_a8w8_mxscale(O_in, W_mx, xs_in, ws_mx, dtype=ydt), ref, @@ -193,33 +215,30 @@ def _call_raw(kid): # Batch-major output buffer; hand the kernel its [m, g, n] view so the # store lands at Y.stride(1) (batch) = m*n (outermost), N contiguous. Yb = torch.empty((g, m, n), dtype=ydt) - _opus_bmm_a8w8_mxscale_raw(O_in, W_mx, Yb.transpose(0, 1), xs_in, ws_mx, 1, kid) + _run_opus(O_in, W_mx, Yb.transpose(0, 1), xs_in, ws_mx, kid) return Yb # [g, m, n] def _call_auto(): - # Same tuned-CSV lookup the public entry does, but writing into a - # caller-owned batch-major buffer -- which the guarded public entry no - # longer exposes (it returns fresh token-major), so drive the opus - # backend directly with the looked-up kid + the batch-major out= view. + # Resolve the final global kid here, then use public opus_bmm with a + # caller-owned batch-major output view. Yb = torch.empty((g, m, n), dtype=ydt) cfg = lookup_mxscale_bmm_config(g, m, n, k) - bmm_a8w8_mxscale_opus( + _run_opus( O_in, W_mx, + Yb.transpose(0, 1), xs_in, ws_mx, - out=Yb.transpose(0, 1), - dtype=ydt, - kernelId=int(cfg["kernelId"]) if cfg is not None else None, - splitK=int(cfg["splitK"]) if cfg is not None else None, + int(cfg["kernelId"]) if cfg is not None else 8000, + int(cfg["splitK"]) if cfg is not None else 1, ) return Yb - # Correctness-focused: kid 0 (always runnable) as the batch-major baseline, + # Correctness-focused: kid 8000 as the batch-major baseline, # plus the backend dispatch path writing into the batch-major buffer via # out=. Per-kid perf sweep lives in csrc/opus_gemm/opus_bmm_mxscale_tune.py. - candidates = {"kid0_k32_fused": (lambda: _call_raw(0), ref)} - candidates["auto (bmm_a8w8_mxscale_opus)"] = (_call_auto, ref) + candidates = {"kid8000_k32_fused": (lambda: _call_raw(8000), ref)} + candidates["auto (public opus_bmm)"] = (_call_auto, ref) flops = 2.0 * g * m * n * k # fp8 A + fp8 W + e8m0 scales (uint8) + output. @@ -248,17 +267,24 @@ def _call_auto(): return ret +# These two functions are CLI benchmarks whose required arguments come from +# ``main()`` rather than pytest fixtures. Keep their historical names for the +# benchmark report while preventing pytest from invoking them with no shape. +test_mxscale_bmm.__test__ = False +test_mxscale_bmm_batch_first.__test__ = False + + # --- tileN column-map regression guard ------------------------------------ # These COM_REP_N>1 kernels previously transposed output column groups. Keep # them out of the narrow perf table, but always exercise both output layouts # with signed, varied-block-scale data so the bug cannot silently return. -_TILEN_REGRESSION_KIDS = (312, 313) +_TILEN_REGRESSION_KIDS = (8312, 8313) _TILEN_REGRESSION_SHAPE = (2, 16, 128, 1024) # G, M, N, K; accepts both kids _TILEN_REGRESSION_ERR_TOL = 0.003 def check_tilen_column_map(): - """Check kid312/313 column mapping for token- and batch-major output.""" + """Check global kids 8312/8313 for both output memory layouts.""" g, m, n, k = _TILEN_REGRESSION_SHAPE O_mx, xs_mx, xs_fp32 = _quant_per_token_e8m0(_block_varied((g, m, k), k)) W_mx, ws_mx, ws_fp32 = _quant_block_e8m0(_block_varied((g, n, k), k)) @@ -275,7 +301,7 @@ def check_tilen_column_map(): out = torch.full((g, m, n), float("nan"), dtype=dtypes.bf16).transpose( 0, 1 ) - _opus_bmm_a8w8_mxscale_raw(O_in, W_mx, out, xs_in, ws_mx, 1, kid) + _run_opus(O_in, W_mx, out, xs_in, ws_mx, kid) torch.cuda.synchronize() delta = (out.to(dtypes.fp32) - ref).abs() rows = delta.flatten(1).mean(1) / (ref.abs().flatten(1).mean(1) + 1e-9) @@ -290,6 +316,53 @@ def check_tilen_column_map(): return len(_TILEN_REGRESSION_KIDS) * 2 +def check_splitk_workspace(): + """Check baseline and sfpreload-fallback split-K workspace paths.""" + g, n, k = 2, 128, 2048 + checks = 0 + # kid8326 is the ROCm 7.2.4 regression: its splitK=1 path keeps scale + # preload, while D_OUT=void deliberately uses the same non-preload device + # specialization as kid8139. M=128 makes both workspaces tightly sized. + for kid, m in ((8000, 32), (8326, 128)): + O_mx, xs_mx, xs_fp32 = _quant_per_token_e8m0(_block_varied((g, m, k), k)) + W_mx, ws_mx, ws_fp32 = _quant_block_e8m0(_block_varied((g, n, k), k)) + O_in = O_mx.transpose(0, 1) + xs_in = xs_mx.transpose(0, 1) + ref = run_torch(O_mx, W_mx, xs_fp32, ws_fp32).transpose(0, 1) + + auto_out = torch.empty((m, g, n), dtype=dtypes.bf16) + _run_opus(O_in, W_mx, auto_out, xs_in, ws_mx, kid, split_k=2) + + workspace = torch.empty(2 * g * m * n, dtype=torch.float32) + caller_out = torch.empty_like(auto_out) + _run_opus( + O_in, + W_mx, + caller_out, + xs_in, + ws_mx, + kid, + split_k=2, + workspace=workspace, + ) + torch.cuda.synchronize() + for mode, out in (("automatic", auto_out), ("caller-owned", caller_out)): + err = checkAllclose( + ref, + out.to(dtypes.fp32), + rtol=1e-2, + atol=1e-2, + msg=f"MXFP8 BMM kid{kid} split-K {mode} workspace", + ) + assert err <= _TILEN_REGRESSION_ERR_TOL, ( + f"kid{kid} split-K {mode} workspace mismatch ratio {err:.4%} " + f"> {_TILEN_REGRESSION_ERR_TOL:.4%}" + ) + checks += 1 + torch.testing.assert_close(auto_out, caller_out, rtol=0, atol=0) + return checks + + # --- m_align guard --------------------------------------------------------- # Straddles every tile boundary in the family (B_M is 16/32/64/128/256) and every # declared m_align (1 / B_M / 2*B_M), with aligned and unaligned M on both sides. @@ -340,9 +413,9 @@ def _align_run(kid, m, n, k): # plausible value that a mean error would dilute. Y = torch.full((m, _ALIGN_G, n), float("nan"), dtype=dtypes.bf16) try: - _opus_bmm_a8w8_mxscale_raw(O_in, W_mx, Y, xs_in, ws_mx, 1, kid) + _run_opus(O_in, W_mx, Y, xs_in, ws_mx, kid) torch.cuda.synchronize() - except RuntimeError: + except (RuntimeError, ValueError): # The launcher's AITER_CHECK on M surfaces here. Deliberately not a # blanket except: a harness bug must fail loudly, not read as a refusal. return False, 0.0 @@ -366,7 +439,7 @@ def check_m_align(): """Assert OpusGemmInstance.m_align matches what each mxscale BMM kid does. m_align says which M values a kid's launcher accepts (1 == it masks a partial - M tile). Both the runtime's padded-M lookup (aiter/ops/opus/bmm_op.py) and a + M tile). Both the tuned caller's padded-M lookup and a tuner's candidate filter act on it, so a wrong value is not merely cosmetic: too strict hides the fastest kernel from tuning (kid326 lost ~9% at the DSV4 wo_a decode shapes that way, while the runtime dispatched it at those very @@ -414,6 +487,10 @@ def check_m_align(): def main(): + # This module is also imported by pytest discovery. Keep the CLI's + # CUDA-default convenience local to the standalone process so importing + # the helpers cannot change unrelated tests' default tensor device. + torch.set_default_device("cuda") if get_gfx() not in SUPPORTED_GFX: aiter.logger.warning( "opus mxscale flatmm BMM unsupported on %s; skipping", get_gfx() @@ -468,6 +545,10 @@ def main(): aiter.logger.info( "tileN column mapping passed for %d kid/layout combinations", n_tilen_checks ) + n_workspace_checks = check_splitk_workspace() + aiter.logger.info( + "MXFP8 BMM split-K workspace checks passed (%d paths)", n_workspace_checks + ) if args.check_m_align: try: diff --git a/op_tests/test_opus_a8w8_interface.py b/op_tests/test_opus_a8w8_interface.py new file mode 100644 index 0000000000..080a9873b7 --- /dev/null +++ b/op_tests/test_opus_a8w8_interface.py @@ -0,0 +1,290 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2025-2026, Advanced Micro Devices, Inc. All rights reserved. +"""CPU checks for the boundary between general A8W8 APIs and OPUS.""" + +from __future__ import annotations + +import inspect +from types import SimpleNamespace + +import pytest +import torch + +from aiter import dtypes +from aiter.ops import batched_gemm_op_a8w8 as batched_a8w8 +from aiter.ops import gemm_op_a8w8 as general_a8w8 +from aiter.ops.opus import gemm_op_a8w8 as opus_a8w8 +from aiter.ops.opus import opus_gemm + + +def test_general_a8w8_restores_required_scales_and_blockscale_dtypes(): + fake_parameters = inspect.signature(general_a8w8.gemm_a8w8_fake).parameters + assert fake_parameters["x_scale"].default is inspect.Parameter.empty + assert fake_parameters["w_scale"].default is inspect.Parameter.empty + + schema = str(torch.ops.aiter.gemm_a8w8.default._schema) + assert "x_scale=None" not in schema + assert "w_scale=None" not in schema + + XQ = torch.empty((1, 1)) + WQ = torch.empty((1, 1)) + with pytest.raises(RuntimeError, match="x_scale"): + general_a8w8.gemm_a8w8(XQ, WQ) + + scale = torch.empty((1, 1)) + with pytest.raises( + AssertionError, + match="Output dtype=torch.float32 is currently not supported", + ): + general_a8w8.gemm_a8w8_blockscale( + XQ, + WQ, + scale, + scale, + dtype=torch.float32, + ) + + +def test_general_scaled_a8w8_keeps_legacy_backend_route(monkeypatch): + calls = [] + result = torch.empty((2, 3), dtype=torch.bfloat16) + + def fake_ck(XQ, WQ, x_scale, w_scale, bias, dtype, splitK): + calls.append((XQ, WQ, x_scale, w_scale, bias, dtype, splitK)) + return result + + monkeypatch.setattr(general_a8w8, "_ck_a8w8_supported", lambda: True) + monkeypatch.setattr(general_a8w8, "gemm_a8w8_CK", fake_ck) + + XQ = torch.empty((2, 4), dtype=torch.int8) + WQ = torch.empty((3, 4), dtype=torch.int8) + x_scale = torch.empty((2, 1), dtype=torch.float32) + w_scale = torch.empty((1, 3), dtype=torch.float32) + actual = general_a8w8.gemm_a8w8( + XQ, + WQ, + x_scale, + w_scale, + dtype=torch.bfloat16, + splitK=3, + ) + + assert actual.data_ptr() == result.data_ptr() + assert len(calls) == 1 + call = calls[0] + assert call[0] is XQ + assert call[1] is WQ + assert call[2] is x_scale + assert call[3] is w_scale + assert call[4:] == (None, torch.bfloat16, 3) + + +@pytest.mark.parametrize( + ("kid", "raw_name", "layout", "with_scale", "output_dtype"), + [ + pytest.param( + 2, + "_opus_gemm_a8w8_launch_raw", + "plain", + False, + torch.float32, + id="noscale", + ), + pytest.param( + 1, + "_opus_gemm_a8w8_blockscale_launch_raw", + "plain", + True, + torch.float32, + id="blockscale", + ), + pytest.param( + 11000, + "_opus_gemm_a8w8_blockscale_bpreshuffle_launch_raw", + "bpreshuffle", + True, + torch.bfloat16, + id="blockscale-bpreshuffle", + ), + ], +) +def test_a8w8_opus_families_use_explicit_exact_kids( + monkeypatch, + kid, + raw_name, + layout, + with_scale, + output_dtype, +): + calls = [] + raw_names = ( + "_opus_gemm_a8w8_launch_raw", + "_opus_gemm_a8w8_blockscale_launch_raw", + "_opus_gemm_a8w8_blockscale_bpreshuffle_launch_raw", + ) + + for candidate in raw_names: + + def fake(*args, _name=candidate, **kwargs): + calls.append((_name, args, kwargs)) + + monkeypatch.setattr(opus_a8w8, candidate, fake) + + XQ = torch.empty((128, 256), dtype=dtypes.fp8) + WQ = torch.empty((128, 256), dtype=dtypes.fp8) + Y = torch.empty((128, 128), dtype=output_dtype) + x_scale = torch.empty((128, 2), dtype=torch.float32) + w_scale = torch.empty((1, 2), dtype=torch.float32) + kwargs = {"kid": kid, "layout": layout} + if with_scale: + kwargs.update(x_scale=x_scale, w_scale=w_scale) + + assert opus_gemm(XQ, WQ, Y, **kwargs) is Y + assert len(calls) == 1 + + called_name, raw_args, raw_kwargs = calls[0] + assert called_name == raw_name + assert raw_kwargs == {} + assert raw_args[-1] == kid + + raw_xq, raw_wq = raw_args[:2] + assert raw_xq.shape == (1, *XQ.shape) + assert raw_xq.data_ptr() == XQ.data_ptr() + assert raw_wq.shape == (1, *WQ.shape) + assert raw_wq.data_ptr() == WQ.data_ptr() + + if layout == "bpreshuffle": + raw_x_scale, raw_w_scale, raw_y = raw_args[2:5] + elif with_scale: + raw_y, raw_x_scale, raw_w_scale = raw_args[2:5] + else: + assert len(raw_args) == 4 + raw_y = raw_args[2] + + assert raw_y.shape == (1, *Y.shape) + assert raw_y.data_ptr() == Y.data_ptr() + if with_scale: + assert len(raw_args) == 6 + assert raw_x_scale is x_scale + assert raw_w_scale is w_scale + + +def test_bpreshuffle_uses_opus_for_tuned_row(monkeypatch): + opus_calls = [] + ck_calls = [] + config = {"libtype": "opus", "kernelId": 11000} + + def fake_opus_gemm(XQ, WQ, Y, **kwargs): + opus_calls.append(kwargs) + return Y + + def fake_ck(XQ, WQ, x_scale, w_scale, Y, kernelName=""): + ck_calls.append(kernelName) + return Y + + from aiter.ops import opus + + monkeypatch.setattr(general_a8w8, "get_gfx", lambda: "gfx942") + monkeypatch.setattr(general_a8w8, "_hip_blockscale_supported", lambda: True) + monkeypatch.setattr(general_a8w8, "get_CKGEMM_config", lambda *_args: config) + monkeypatch.setattr( + general_a8w8, + "gemm_a8w8_blockscale_bpreshuffle_ck", + fake_ck, + ) + monkeypatch.setattr(opus, "opus_gemm", fake_opus_gemm) + + XQ = torch.empty((2, 128), dtype=dtypes.fp8) + WQ = torch.empty((128, 128), dtype=dtypes.fp8) + x_scale = torch.empty((2, 1), dtype=torch.float32) + w_scale = torch.empty((1, 1), dtype=torch.float32) + + bf16_result = general_a8w8.gemm_a8w8_blockscale_bpreshuffle( + XQ, + WQ, + x_scale, + w_scale, + dtype=torch.bfloat16, + ) + assert bf16_result.dtype == torch.bfloat16 + assert len(opus_calls) == 1 + assert opus_calls[0]["kid"] == 11000 + assert opus_calls[0]["layout"] == "bpreshuffle" + assert opus_calls[0]["x_scale"] is x_scale + assert opus_calls[0]["w_scale"] is w_scale + assert ck_calls == [] + + +def test_mxscale_launch_plan_cache_is_bounded(monkeypatch): + calls = [] + + def resolve(g, m, n, k): + calls.append((g, m, n, k)) + return 8000, 1 + + monkeypatch.setattr( + batched_a8w8, + "_resolve_a8w8_mxscale_bmm_plan", + resolve, + ) + batched_a8w8._get_mxscale_bmm_launch_plan.cache_clear() + try: + for m in range(1025): + assert batched_a8w8._get_mxscale_bmm_launch_plan(2, m, 1024, 4096) == ( + 8000, + 1, + ) + + cache = batched_a8w8._get_mxscale_bmm_launch_plan.cache_info() + assert cache.maxsize == 1024 + assert cache.currsize == 1024 + assert len(calls) == 1025 + finally: + batched_a8w8._get_mxscale_bmm_launch_plan.cache_clear() + + +def test_mxscale_invalid_tuned_kid_warns_and_uses_heuristic( + monkeypatch, + tmp_path, +): + from aiter.ops.opus import policy + + config_path = tmp_path / "mxscale.csv" + config_path.write_text( + "gfx,b,m,n,k,libtype,kernelId,splitK\n" + "gfx950,2,1,1024,4096,opus,8001,1\n" + "gfx950,3,1,1024,4096,other,42,1\n" + ) + warnings = [] + monkeypatch.setattr( + policy, + "AITER_CONFIGS", + SimpleNamespace( + AITER_CONFIG_BATCHED_GEMM_A8W8_BLOCKSCALE_MXSCALE_FILE=str(config_path) + ), + ) + monkeypatch.setattr(policy, "get_gfx", lambda: "gfx950") + monkeypatch.setattr(policy, "get_padded_m", lambda m, _n, _k, _gl: m) + monkeypatch.setattr( + policy.logger, + "warning", + lambda *args, **_kwargs: warnings.append(args), + ) + policy._load_mxscale_bmm_tuned.cache_clear() + policy.lookup_mxscale_bmm_config.cache_clear() + try: + rows = policy._load_mxscale_bmm_tuned(None) + assert rows[("gfx950", 3, 1, 1024, 4096)]["kernelId"] == 42 + assert policy.resolve_a8w8_mxscale_bmm_plan(2, 1, 1024, 4096) == ( + 8640, + 1, + ) + assert len(warnings) == 1 + assert warnings[0][0].startswith("Skipping %d invalid OPUS row") + finally: + policy.lookup_mxscale_bmm_config.cache_clear() + policy._load_mxscale_bmm_tuned.cache_clear() + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/op_tests/test_opus_co_integration.py b/op_tests/test_opus_co_integration.py new file mode 100644 index 0000000000..f4610cb808 --- /dev/null +++ b/op_tests/test_opus_co_integration.py @@ -0,0 +1,378 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2025-2026, Advanced Micro Devices, Inc. All rights reserved. +"""Integration coverage for gfx1250 pre-built OPUS A16W16 kernels.""" + +from __future__ import annotations + +import json +import os +import shutil +import subprocess +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + +from aiter.ops.opus import ( + gemm_a16w16_opus, + gemm_op_a16w16, + opus_bmm, + opus_gemm, +) +from aiter.ops.opus._arch import GFX1250, SUPPORTED_OPUS_ARCHES +from aiter.ops.opus.launch_plan import _get_cached_a16w16_launch_plan +from csrc.opus_gemm.opus_gemm_common import ( + BIAS_AWARE_KIDS, + CO_KERNELS_JSON, + DEFAULT_COMPILED_KIDS_BY_ARCH, + GFX1250_4WAVE_CO_KIDS, + NON_SPLITK_KIDS, + OPUS_KERNEL_TAGS_BY_ARCH_FAMILY, + SPLITK_KIDS, + _load_co_kernels, + co_image_path, + get_kernel_instance, + gfx1250_4wave_co_kernels_list, + kernel_needs_external_workspace, +) + + +def _co_plan(kid: int, **overrides): + arguments = { + "arch": GFX1250, + "M": 128, + "N": 128, + "K": 4096, + "batch": 1, + "cu_num": 80, + "has_bias": False, + "input_dtype": torch.bfloat16, + "output_dtype": torch.bfloat16, + "kid": kid, + "split_k": 0, + } + arguments.update(overrides) + return _get_cached_a16w16_launch_plan(**arguments) + + +def _runtime_arch() -> str | None: + if torch.version.hip is None or not torch.cuda.is_available(): + return None + properties = torch.cuda.get_device_properties(torch.cuda.current_device()) + return str(properties.gcnArchName).split(":", 1)[0].lower() + + +def _representative_kids(kids) -> tuple[int, ...]: + ordered = sorted(kids) + representatives = (ordered[0], ordered[len(ordered) // 2], ordered[-1]) + return tuple(dict.fromkeys(representatives)) + + +def _gfx1250_workspace_representatives() -> tuple[int, ...]: + by_tag = {} + for kid in sorted(SPLITK_KIDS): + instance = get_kernel_instance(GFX1250, "a16w16", kid) + if instance is not None: + by_tag.setdefault(instance.kernel_tag, kid) + return tuple(by_tag.values()) + + +_CO_ROUTE_KIDS = _representative_kids(GFX1250_4WAVE_CO_KIDS) + + +@pytest.fixture(scope="module") +def _gfx1250_device() -> torch.device: + arch = _runtime_arch() + if arch != GFX1250: + pytest.skip(f"requires gfx1250 hardware, got {arch!r}") + return torch.device("cuda", torch.cuda.current_device()) + + +def _run_gfx1250_co_case(kid: int, device: torch.device | None = None) -> None: + instance = gfx1250_4wave_co_kernels_list[kid] + if device is None: + arch = _runtime_arch() + if arch != GFX1250: + raise RuntimeError(f"requires gfx1250 hardware, got {arch!r}") + device = torch.device("cuda", torch.cuda.current_device()) + + M = instance.B_M * instance.cluster_wg_m + 1 + N = instance.B_N * instance.cluster_wg_n + 1 + K = instance.B_K * instance.num_slots + 1 + context = ( + f"kid={kid}, symbol={instance.name}, tag={instance.kernel_tag}, " + f"tile=({instance.B_M},{instance.B_N},{instance.B_K}), " + f"cluster=({instance.cluster_wg_m},{instance.cluster_wg_n}), " + f"shape=({M},{N},{K})" + ) + + generator = torch.Generator(device=device).manual_seed(kid) + A = torch.randn((M, K), device=device, dtype=torch.bfloat16, generator=generator) + B = torch.randn((N, K), device=device, dtype=torch.bfloat16, generator=generator) + Y = torch.full((M, N), float("nan"), device=device, dtype=torch.bfloat16) + + actual = opus_gemm(A, B, Y, kid=kid, split_k=0) + torch.cuda.synchronize(device) + assert actual is Y, context + assert torch.isfinite(actual).all().item(), context + + reference = A.float() @ B.float().T + try: + torch.testing.assert_close(actual.float(), reference, rtol=0.03, atol=0.5) + except AssertionError as error: + raise AssertionError(f"{context}\n{error}") from error + + +def _run_gfx1250_co_override_case(available_kid: int, missing_kid: int) -> None: + override_root = Path(os.environ["OPUS_GEN_CO_DIR"]) + missing_instance = gfx1250_4wave_co_kernels_list[missing_kid] + missing_image = override_root / GFX1250 / f"{missing_instance.name}.co" + try: + _run_gfx1250_co_case(missing_kid) + except RuntimeError as error: + if str(missing_image) not in str(error): + raise AssertionError( + f"override lookup did not report expected path {missing_image}: {error}" + ) from error + else: + raise AssertionError(f"override lookup unexpectedly found {missing_image}") + + _run_gfx1250_co_case(available_kid) + + +def test_co_manifest_uses_supported_architecture_keys(): + document = json.loads(Path(CO_KERNELS_JSON).read_text()) + manifest_arches = {key for key in document if not key.startswith("_")} + + assert GFX1250 in manifest_arches + assert manifest_arches <= SUPPORTED_OPUS_ARCHES + + +def test_gfx1250_co_registry_and_launch_contract(): + assert len(gfx1250_4wave_co_kernels_list) == 219 + assert GFX1250_4WAVE_CO_KIDS == frozenset(gfx1250_4wave_co_kernels_list) + assert min(GFX1250_4WAVE_CO_KIDS) == 21016 + assert max(GFX1250_4WAVE_CO_KIDS) == 21315 + assert GFX1250_4WAVE_CO_KIDS <= NON_SPLITK_KIDS + assert GFX1250_4WAVE_CO_KIDS <= DEFAULT_COMPILED_KIDS_BY_ARCH[GFX1250] + assert GFX1250_4WAVE_CO_KIDS.isdisjoint(SPLITK_KIDS) + assert GFX1250_4WAVE_CO_KIDS.isdisjoint(BIAS_AWARE_KIDS) + assert { + "a16w16_4wave_co", + "a16w16_4wave_wl_co", + } <= OPUS_KERNEL_TAGS_BY_ARCH_FAMILY[GFX1250]["a16w16"] + + kid = min(GFX1250_4WAVE_CO_KIDS) + assert ( + get_kernel_instance(GFX1250, "a16w16", kid, torch.bfloat16) + is gfx1250_4wave_co_kernels_list[kid] + ) + assert get_kernel_instance(GFX1250, "a16w16", kid, torch.float32) is None + assert not kernel_needs_external_workspace(GFX1250, "a16w16", kid) + + for split_k in (0, 1): + plan = _co_plan(kid, split_k=split_k) + assert plan.resolved_kid == kid + assert plan.abi_split_k == split_k + assert plan.workspace_spec is None + + with pytest.raises(ValueError, match="does not support split-K"): + _co_plan(kid, split_k=2) + with pytest.raises(ValueError, match="does not support output dtype"): + _co_plan(kid, output_dtype=torch.float32) + with pytest.raises(ValueError, match="does not support bias"): + _co_plan(kid, has_bias=True) + + +# Registry/artifact checks own exhaustive identity coverage. The route itself is +# kid-independent, so cover both entries/split modes with representative ids. +@pytest.mark.parametrize( + ("kid", "split_k", "entry"), + ( + (_CO_ROUTE_KIDS[0], 0, "opus_bmm"), + (_CO_ROUTE_KIDS[1], 1, "opus_bmm"), + (_CO_ROUTE_KIDS[2], 0, "gemm_a16w16_opus"), + (_CO_ROUTE_KIDS[0], 1, "gemm_a16w16_opus"), + ), +) +def test_gfx1250_co_bmm_reaches_exact_launch(kid, split_k, entry, monkeypatch): + calls = [] + + def capture(XQ, WQ, Y, bias, workspace, launched_kid, launch_split_k): + calls.append((XQ, WQ, Y, bias, workspace, launched_kid, launch_split_k)) + + monkeypatch.setattr(gemm_op_a16w16, "_device_arch_and_cu", lambda _: (GFX1250, 80)) + monkeypatch.setattr(gemm_op_a16w16, "_opus_gemm_a16w16_launch_raw", capture) + A = torch.empty((2, 64, 512), device="meta", dtype=torch.bfloat16) + B = torch.empty((2, 64, 512), device="meta", dtype=torch.bfloat16) + Y = torch.empty((2, 64, 64), device="meta", dtype=torch.bfloat16) + + if entry == "opus_bmm": + actual = opus_bmm(A, B, Y, kid=kid, split_k=split_k) + else: + actual = gemm_a16w16_opus(A, B, kernelId=kid, splitK=split_k, out=Y) + + assert actual is Y + assert len(calls) == 1 + XQ, WQ, output, bias, workspace, launched_kid, launch_split_k = calls[0] + assert XQ is A and WQ is B and output is Y + assert bias is None and workspace is None + assert (launched_kid, launch_split_k) == (kid, split_k) + + +@pytest.mark.parametrize( + "kid", + _gfx1250_workspace_representatives(), +) +def test_gfx1250_workspace_kids_reject_bmm(kid): + with pytest.raises(ValueError, match="workspace kids require batch=1"): + _co_plan(kid, M=1, batch=2) + + +def test_gfx1250_co_assets_and_host_only_codegen(tmp_path, monkeypatch): + symbols_by_kid = { + kid: instance.name for kid, instance in gfx1250_4wave_co_kernels_list.items() + } + assert len(set(symbols_by_kid.values())) == 219 + + for instance in gfx1250_4wave_co_kernels_list.values(): + image = Path(co_image_path(CO_KERNELS_JSON, instance)) + with image.open("rb") as stream: + assert stream.read(4) == b"\x7fELF" + assert instance.splitk_workspace_dtype is None + + image_dir = Path(CO_KERNELS_JSON).parent / GFX1250 + build_info = json.loads((image_dir / "build_info.json").read_text()) + assert len(build_info["kernels"]) == 219 + assert {entry["kernarg_segment_size"] for entry in build_info["kernels"]} == {64} + assert { + entry["kid"]: entry["symbol"] for entry in build_info["kernels"] + } == symbols_by_kid + assert {image.stem for image in image_dir.glob("*.co")} == set( + symbols_by_kid.values() + ) + + monkeypatch.syspath_prepend( + str(Path(__file__).resolve().parents[1] / "csrc" / "opus_gemm") + ) + from codegen.gen_instances_gfx1250 import ( + KARGS_NAME_MAP, + TRAITS_HEADER_MAP, + TRAITS_NAME_MAP, + gen_4wave_co_instance, + ) + + instance = next(iter(gfx1250_4wave_co_kernels_list.values())) + codegen = SimpleNamespace( + impl_path=str(tmp_path), + _host_instantiations=[], + _device_instantiations=[], + ) + gen_4wave_co_instance( + codegen, + instance, + traits_header=TRAITS_HEADER_MAP[instance.kernel_tag], + traits_name=TRAITS_NAME_MAP[instance.kernel_tag], + kargs_name=KARGS_NAME_MAP[instance.kernel_tag], + ) + generated = (tmp_path / f"{instance.name}.cuh").read_text() + assert "opus_co_launch_gfx1250" in generated + assert "aiter_tensor_t &workspace" not in generated + assert "splitK == 0 || splitK == 1" in generated + assert len(codegen._host_instantiations) == 1 + assert codegen._device_instantiations == [] + + +@pytest.mark.parametrize( + "kid", + sorted(gfx1250_4wave_co_kernels_list), + ids=lambda kid: f"kid-{kid}", +) +def test_gfx1250_co_kernel_matches_torch(kid, _gfx1250_device, monkeypatch): + monkeypatch.setenv("OPUS_GEN_CO_DIR", str(Path(CO_KERNELS_JSON).parent)) + _run_gfx1250_co_case(kid, _gfx1250_device) + + +def test_gfx1250_co_dir_override_matches_torch(_gfx1250_device, tmp_path): + kid = max(gfx1250_4wave_co_kernels_list) + missing_kid = min(gfx1250_4wave_co_kernels_list) + instance = gfx1250_4wave_co_kernels_list[kid] + override_root = tmp_path / "co-root" + override_arch_dir = override_root / GFX1250 + override_arch_dir.mkdir(parents=True) + shutil.copy2( + co_image_path(CO_KERNELS_JSON, instance), + override_arch_dir / f"{instance.name}.co", + ) + + environment = os.environ.copy() + environment["OPUS_GEN_CO_DIR"] = str(override_root) + command = ( + "from op_tests.test_opus_co_integration import " + "_run_gfx1250_co_override_case; " + f"_run_gfx1250_co_override_case({kid}, {missing_kid})" + ) + completed = subprocess.run( + [sys.executable, "-c", command], + cwd=Path(__file__).resolve().parents[1], + env=environment, + capture_output=True, + text=True, + timeout=900, + check=False, + ) + assert completed.returncode == 0, completed.stdout + completed.stderr + + +def test_gfx1250_co_loader_handles_missing_assets(tmp_path, capsys): + assert _load_co_kernels(tmp_path / "missing.json") == {} + + document = json.loads(Path(CO_KERNELS_JSON).read_text()) + document[GFX1250] = document[GFX1250][:1] + manifest = tmp_path / "co_kernels.json" + manifest.write_text(json.dumps(document)) + + assert _load_co_kernels(manifest) == {} + assert "1 pre-compiled (.co) kid(s) dropped" in capsys.readouterr().err + assert len(_load_co_kernels(manifest, require_image=False)) == 1 + + +@pytest.mark.parametrize("duplicate", ["kid", "name"]) +def test_gfx1250_co_loader_rejects_duplicate_identity(tmp_path, duplicate): + document = json.loads(Path(CO_KERNELS_JSON).read_text()) + first = document[GFX1250][0] + second = json.loads(json.dumps(first)) + if duplicate == "name": + second["kid"] += 1 + document[GFX1250] = [first, second] + manifest = tmp_path / "co_kernels.json" + manifest.write_text(json.dumps(document)) + + expected = "duplicate co kid" if duplicate == "kid" else "same symbol" + with pytest.raises(AssertionError, match=expected): + _load_co_kernels(manifest, require_image=False) + + +def test_gfx1250_tuner_selects_co_without_split_k(monkeypatch): + monkeypatch.syspath_prepend( + str(Path(__file__).resolve().parents[1] / "csrc" / "opus_gemm") + ) + from opus_gemm_tune import ( + _gfx1250_select_candidates, + candidate_splitK, + kid_rejects_shape, + ) + + selected = _gfx1250_select_candidates(64, 128, 4096, 256) + assert selected & GFX1250_4WAVE_CO_KIDS + assert selected.isdisjoint(range(27000, 30000)) + + instance = gfx1250_4wave_co_kernels_list[min(GFX1250_4WAVE_CO_KIDS)] + assert candidate_splitK(64, 128, 4096, 1, 256, instance) == [0] + assert not kid_rejects_shape(instance, 65, 129, 4097) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/op_tests/tuning_tests/test_config_shape_collision.py b/op_tests/tuning_tests/test_config_shape_collision.py index 53b50b4cf9..4a0750c17b 100644 --- a/op_tests/tuning_tests/test_config_shape_collision.py +++ b/op_tests/tuning_tests/test_config_shape_collision.py @@ -223,6 +223,26 @@ def test_batched_gemm_a8w8_blockscale_mxscale(self): "batched_gemm_a8w8_blockscale_mxscale_tuned", ) + # The generic config registry owns file discovery and merging; the + # dedicated caller policy still owns public-kid normalization. + from aiter.ops.opus import policy + + policy._load_mxscale_bmm_tuned.cache_clear() + rows = policy._load_mxscale_bmm_tuned("opus") + self.assertTrue(rows) + self.assertEqual(len(rows), len(set(rows))) + self.assertEqual( + rows[("gfx950", 2, 1, 1024, 4096)]["kernelId"], + 8311, + "legacy local OPUS kid 311 must become public global kid 8311", + ) + self.assertEqual( + rows[("gfx950", 8, 128, 1024, 4096)]["kernelId"], + 8653, + "legacy local OPUS kid 653 must become public global kid 8653", + ) + policy._load_mxscale_bmm_tuned.cache_clear() + def test_batched_gemm_a8w8_blockscale_mxscale_bpreshuffle(self): self._check_family( "AITER_CONFIG_BATCHED_GEMM_A8W8_BLOCKSCALE_MXSCALE_BPRESHUFFLE",