aiter-op-test

# aiter op_test standard How every aiter op test in `op_tests/test_*.py` must be built. The canonical reference in-tree is **`op_tests/test_quant.py`** — match its shape. A test is both a **correctness check** (vs a torch reference) and a **perf sweep** that ends in a **markdown summary table**. Follow this whenever you create a new `test_*.py`, rewrite an old one, or add shapes/candidates to an existing one. ## The hard rules 1. **Mirror `test_quant.py`.** Same imports, same decorator, same table-at-the-end flow. Don't invent a different structure. 2. **`@benchmark()` on the test fn.** It logs the function's call args (the shape params) as table columns automatically and merges the dict you `return`. So the test fn signature *is* the table's left-hand columns — name params accordingly. 3. **Candidates live in a dict; build `ret` in a loop.** Per candidate record raw `us`, plus **`TFLOPS` and `TB/s`**, plus `err` — `ret[f"{name} us"]`, `ret[f"{name} TFLOPS"]`, `ret[f"{name} TB/s"]`, `ret[f"{name} err"]`. Never hand-write ratio columns. (TFLOPS/TB-s section below.) 4. **torch is the reference only** — compute it, compare against it, but do **not** time it and do **not** put it in the table. (A pure-torch candidate is allowed only when torch *is* one of the kernels under test, e.g. `torch.einsum`.) 5. **Time with `run_perftest`, check with `checkAllclose`** — both, for every candidate. Compare in fp32 (`.to(dtypes.fp32)`). 6. **End with a markdown summary table — one per test function.** Sweep the shape lists with `itertools.product`, collect per-shape dicts into a `pd.DataFrame`, print via `aiter.logger.info("... :\n%s", df.to_markdown(index=False))`. A file with several test fns of different arg signatures emits **one table each** — never force-merge them (it scatters NaN columns). Mandatory — a test with no summary table is incomplete. 7. **`__main__` guard.** All argparse + the sweep loop go inside `main()`, called under `if __name__ == "__main__": main()`. The reference (`run_torch`) and the `@benchmark` test fn stay at module top level so other scripts can `import` them for combination testing. 8. **Standard argparse only.** Use `-d/--dtype`, `-b/--batch`, `-s/--mnk` plus *op-specific sweep axes* as needed (e.g. `--layout`, `--modes`, `--mtp`). Those are legitimate data lists. **Do not** add bespoke behavior-toggle flags (no `--dsv4`, no `--only-*`) — every flag is a list the sweep iterates. 9. **Run clean on every supported card.** Gate on `get_gfx()` in `main()` so the test passes on all supported archs; arch-unsupported ops/candidates are filtered out *before* launch. Prefer the kernel's arch-dispatching wrapper over a file-per-arch (full section below). ## Canonical template ```python # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. import argparse import itertools import aiter import pandas as pd import torch import torch.nn.functional as F from aiter import dtypes from aiter.test_common import ( benchmark, checkAllclose, run_perftest, ) from aiter.jit.utils.chip_info import get_gfx # "gfx942", "gfx950", "gfx1250", ... torch.set_default_device("cuda") SUPPORTED_GFX = ["gfx942", "gfx950"] # every card this op is built/validated for def run_torch(x, weight, dtype=dtypes.bf16): # Reference only: fp32 math, cast back. Not timed, not in the table. ... return out.to(dtype) @benchmark() # call args (b, m, n, k, dtype, ...) become the table's left columns def test_op(b, m, n, k, dtype, layout): # build inputs/outputs in the layout the MODEL actually uses (see below); # ref = run_torch(...) candidates = { "triton": lambda: ..., # the path the model really runs "torch_einsum": lambda: ..., # optional torch kernel under test } if <kernel supported on this arch + config>: # e.g. get_gfx() != "gfx1250" candidates["ck"] = lambda: ... # else skip it (see rules below) flops = 2 * b * m * n * k # roofline numerator nbytes = (b * m * k + b * n * k + b * m * n) * x.element_size() ret = {"gfx": get_gfx()} # record the card in the table for name, fn in candidates.items(): out, us = run_perftest(fn) err = checkAllclose(ref.to(dtypes.fp32), out.to(dtypes.fp32), rtol=1e-2, atol=1e-2, msg=f"{name}: <op>") ret[f"{name} us"] = us ret[f"{name} TFLOPS"] = flops / us / 1e6 ret[f"{name} TB/s"] = nbytes / us / 1e6 ret[f"{name} err"] = err return ret def main(): # Whole-op arch gate goes HERE, not inside test_op: @benchmark always returns # the call-args dict, so an in-fn `return` still emits an args-only row. if get_gfx() not in SUPPORTED_GFX: aiter.logger.warning("<op> unsupported on %s; skipping", get_gfx()) return parser = argparse.ArgumentParser( formatter_class=argparse.RawTextHelpFormatter, description="config input of test" ) parser.add_argument("-d", "--dtype", type=dtypes.str2Dtype, nargs="*", default="bf16,", ...) parser.add_argument("-b", "--batch", type=int, nargs="*", default=[...], ...) parser.add_argument("-s", "--mnk", type=dtypes.str2tuple, nargs="*", default=[...], ...) # add -l/--layout ONLY if the op has real layout variants args = parser.parse_args() for dtype in args.dtype: # one table per outer-config element df = [] for layout, b, (m, n, k) in itertools.product( # sweep via itertools.product args.layout, args.batch, args.mnk ): df.append(test_op(b, m, n, k, dtype, layout)) df = pd.DataFrame(df) aiter.logger.info("<op> summary (markdown):\n%s", df.to_markdown(index=False)) if __name__ == "__main__": main() ``` `op_tests/test_batched_gemm_bf16.py` is a complete worked instance of this template (GEMM with triton/CK/einsum candidates, `bmn`/`mbn` layout, V4 shapes). ## Faithful to the real model call (do not test an idealized op) The point of the test is the kernel **as the model invokes it**, not a clean textbook version. Reproduce exactly: - **Preallocated output buffers.** If the model passes `YQ=`/`out=` a buffer it allocated, allocate and pass the same — don't let the kernel allocate its own. - **Real tensor layout, including non-contiguous views.** If the model feeds a transposed view (e.g. `o.transpose(0, 1)`), build the input as a transposed view of a contiguous tensor — not a fresh contiguous tensor of that shape. - **Couple input and output layout.** They are linked in the model. If you sweep an output layout (`bmn` vs `mbn`), the *input* must follow: the `mbn` (model) case is a transposed view of `[m, b, k]` (physically `mbk`) **and** a transposed view of `[m, b, n]` — not contiguous `[b, m, k]`. - **For a torch.einsum candidate, switch the layout by editing the subscript string**, not by transposing afterwards: `->sgr` is physically `[m,b,n]` (mbn), `->gsr` is physically `[b,m,n]` (bmn). Feed it the model's natural contiguous operand (e.g. `o` is contiguous `[s,g,d]`, so `x.transpose(0,1).contiguous()`). A test that quietly uses contiguous inputs when the model uses a transposed view gives the wrong perf **and hides correctness bugs** (see next rule). ## Report TFLOPS and TB/s, not just `us` A bare `us` doesn't say whether a kernel is compute- or memory-bound. Always add both roofline metrics per candidate, derived from the same `us`: ```python flops = 2 * b * m * n * k # GEMM: 2*M*N*K mul-add nbytes = (b * m * k + b * n * k + b * m * n) * x.element_size() # in + weight + out ret[f"{name} TFLOPS"] = flops / us / 1e6 # us -> s is 1e-6, FLOP -> T is 1e-12 ret[f"{name} TB/s"] = nbytes / us / 1e6 ``` Count the FLOPs and bytes the *op* actually does (adjust the formula per op: a quant/norm/attention kernel has its own element-traffic and arithmetic). Use `tensor.element_size()` for dtype width so fp8/bf16/fp16 are handled. Reading the table: small `m` (decode) is memory-bound (high TB/s, low TFLOPS); large `m` is compute-bound (TFLOPS approaches peak). ## Multiple test functions → multiple tables; reference patterns - **One `@benchmark` fn per distinct arg signature, one table each.** A file may hold several (e.g. a main bf16/fp8 sweep + an fp8 nm-asm cross-check with different columns). Give each its own `pd.DataFrame` + `aiter.logger.info(...)` via a tiny `summarize(name, rows)` helper. Forcing them into one table scatters NaN columns and is unreadable. - **Two correctness shapes are common:** - *Multi-candidate* (this template): several kernels vs one torch reference, all timed, dict-loop. - *Single-kernel-vs-reference, in-place output* (e.g. cache-writing kernels): clone the output buffer, run the kernel into one copy and the reference into another, then `checkAllclose`. Still record `us`/TFLOPS/TB-s/err. - **Prefer a shared reference from `aiter.ops.torch_ref`** when one exists; only hand-write `run_torch` when there is none. - **Do all the asserts a case needs; record one representative `err`.** A quant path checks the dequantized output *and* bit-exact scales: use `tol_err_ratio=` for the fraction of allowed element mismatches (fp8/bf16 rounding), and `rtol=0, atol=0` for values that must match to the bit (scales). ## Skip a candidate in configs it does not support Some kernels are only correct for some layouts/dtypes. Running them anyway pollutes the table with wrong-but-fast numbers. **Conditionally add** such a candidate and leave its cells `nan` elsewhere — e.g. `batched_gemm_bf16_CK` returns garbage (`err ≈ 0.99`) on a non-contiguous `mbk` input, so it is only added for `bmn`: ```python if layout == "bmn": candidates["ck"] = lambda: aiter.batched_gemm_bf16_CK(x, weight) ``` When you skip something, say so in a code comment with the reason. The non-zero `err` column is exactly how you discover these — never silently drop a candidate because its error is high; first confirm whether it's a real bug or an unsupported config, then skip with a comment. ## Run on every supported card (arch gating) A test must run **clean on every currently-supported card** — today `gfx942` (MI300) and `gfx950` (MI35x), plus any arch the kernel specifically targets. Never assume one GPU. Detect arch at runtime and filter *before* launching: ```python from aiter.jit.utils.chip_info import get_gfx # "gfx942" / "gfx950" / "gfx1250" / ... ``` - **One test, all archs — drive the arch-dispatching wrapper; do NOT write a file per arch.** Most aiter kernels expose a public wrapper that routes to the wave64/wave32 (or gfx-specific) implementation internally by `get_gfx()` — call that wrapper and the single test covers every arch. Only the *behavioral* arch differences need handling in the test (e.g. gfx1250 uses a linear FP8 layout, so force `preshuffle=False` there; keep the kernel call and the reference in sync). Never import an `*_gfx1250`/arch-suffixed kernel directly to make a parallel test file. - **Op not built/supported on this arch → skip in `main()`** with an allow-list early `return` (skips the whole sweep cleanly — no rows): ```python SUPPORTED_GFX = ["gfx942", "gfx950", "gfx1250"] if get_gfx() not in SUPPORTED_GFX: aiter.logger.warning("<op> unsupported on %s; skipping", get_gfx()) return ``` Do **not** gate by returning from inside the `@benchmark` fn: that wrapper always returns the call-args dict, so an in-fn `return` still emits an args-only NaN row. Prefer a **positive allow-list** (`not in [...]`) over a deny-list so an unknown new card doesn't silently run an unbuilt kernel and crash. - **One candidate / one sub-check unsupported on this arch → drop just it**, with a warning naming arch + reason — e.g. the fp8 nm-asm cross-checks are wave64-only: ```python if get_gfx() != "gfx1250": summarize("hca_fp8", [test_hca_fp8(bs) for bs in args.fp8_bs]) else: aiter.logger.warning("gfx1250: skipping wave64-only fp8 cross-checks") ``` - **Record the card**: put `"gfx": get_gfx()` in the returned dict so one table is self-describing across cards. In-tree precedent: the flydsl `fused_compress_attn` wrappers dispatch wave64/wave32 internally (one `test_flydsl_compress_attn.py` covers all archs); `test_deepgemm.py` / `test_gemm_a4w4.py` allow-list a single arch; `test_gemm_a8w8.py` does per-arch candidate/dtype/shape gating with warnings. ## Deriving model shapes (don't guess) When adding "test op X for model Y" shapes: 1. **Read the real `config.json`** for the actual dims (`grep`/`python -json`), don't assume the dataclass defaults. 2. **Map model semantics → the kernel's `(b, m, n, k)`** and write the mapping in a comment. Worked example — DeepSeek-V4 grouped output LoRA (`atom/models/deepseek_v4.py`, `batched_gemm_bf16(o.transpose(0,1), wo_a, YQ=y)`):