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

[Fyzyukk]

Summary

This PR refactors the OPUS GEMM/BMM around one strict, exact-kernel-id contract from Python through generated C++ dispatch.

• Add public opus_gemm() and opus_bmm() entry points,
• Separate kernel selection from exact execution,
• Use per-call Torch workspaces for split-K; C++ only validates and uses them,
• Keep existing non-OPUS backend behavior intact.

Motivation

The previous path mixed kernel selection, validation, dispatch, and workspace management across Python and C++, resulting in duplicated wrappers, heuristic dispatch, and global workspace state.

This PR defines a strict exact-kid boundary: callers select the kernel and own the output/workspace tensors, while OPUS C++ validates and launches that exact kernel.

Architecture

A16W16

GEMM: TunedGemm / gemm_a16w16()
  -> tuned-row validation and fallback policy
  -> final kid + split_k
  -> opus_gemm()                         [2D]

BMM: caller-selected kid + split_k
  -> opus_bmm()                          [batch-first 3D]

opus_gemm() / opus_bmm()
  -> canonical a16w16 registry lookup
  -> _execute_a16w16()
  -> A16W16LaunchPlan
       - validate exact kid, shape, dtype, bias, and split-K
       - reuse caller workspace or allocate a per-call Torch workspace
  -> _launch_a16w16_backend()
       - eager / graph capture: cached C ABI
       - torch.compile / Meta / FakeTensor: torch.ops boundary
  -> opus_gemm_a16w16_launch()
  -> per-architecture exact-kid table
  -> generated launcher
  -> direct kernel or split-K kernel + reducer

A8W8

aiter/ops/gemm_op_a8w8.py
  -> gemm_a8w8()                         [no-scale GEMM]
  -> gemm_a8w8_blockscale()              [plain block-scale GEMM]
  -> gemm_a8w8_blockscale_bpreshuffle()  [B-preshuffle GEMM]

aiter/ops/batched_gemm_op_a8w8.py
  -> batched_gemm_a8w8_mxscale()         [MXFP8 BMM]
             |
             v
fixed, tuned, or caller-resolved final kid
  -> opus_gemm()                         [three GEMM families]
  -> checked raw fast path               [MXFP8 BMM, split_k <= 1]
  -> opus_bmm(layout="mxscale_bmm")      [MXFP8 BMM, split_k > 1]
             |
             v
canonical A8 family registry lookup and contract validation
  -> family adapter
  -> MXFP8 split-K only: build plan and reuse/allocate FP32 Torch workspace
  -> _launch_a8w8_backend()               [pybind]
  -> family C++ entry
  -> exact-kid table
  -> generated launcher
  -> exact kernel; MXFP8 split-K also launches the reducer

Production paths

Family	High-level or public path	C++ exact entry	Workspace	Current capability
A16W16 GEMM	TunedGemm.mm / gemm_a16w16() -> opus_gemm() -> _execute_a16w16()	opus_gemm_a16w16_launch()	Optional, Torch-owned for workspace kids	gfx942, gfx950, gfx1250
A16W16 BMM	direct opus_bmm() -> _launch_a16w16_bmm() -> _execute_a16w16()	same A16 3D exact launcher	Optional, batch-aware Torch workspace	Exact API only; there is no new high-level BF16 BMM dispatcher
A8W8 no-scale GEMM	gemm_a8w8() -> opus_gemm() -> A8 backend	opus_gemm_a8w8_launch()	None	gfx950 kid 2, FP32 output
A8W8 plain block-scale GEMM	gemm_a8w8_blockscale() -> opus_gemm()	opus_gemm_a8w8_blockscale_launch()	None	gfx950 kid 1, FP32 output
A8W8 block-scale B-preshuffle GEMM	tuned libtype=opus row -> opus_gemm(layout="bpreshuffle")	opus_gemm_a8w8_blockscale_bpreshuffle_launch()	None	gfx942 kid 11000, BF16 output
A8W8 MXFP8 BMM	batched_gemm_a8w8_mxscale() -> split-1 checked raw path or opus_bmm(layout="mxscale_bmm")	opus_gemm_a8w8_mxscale_bmm_launch()	FP32 Torch workspace when split-K requires it	gfx950, 45 global ids, BF16/FP32 output

Workspace ownership and execution

• Split-K workspace is either supplied by the caller or allocated by Torch for that call. Python validates it, and C++ only uses its pointer during the launch.
• Direct kernels reject workspace; workspace kernels require it.
• A16 uses pybind for initial JIT loading. Later eager and graph-capture calls use the cached C ABI, while torch.compile, MetaTensor, and FakeTensor keep the registered torch.ops path.

Current A16 two-stage workspace layouts are:

Architecture	Logical workspace shape	Storage
gfx950	[split_capacity, batch, padded_M, padded_N]	FP32
gfx942	[split_capacity, batch, padded_M, padded_N]	exact BF16/FP32 instance dtype
gfx1250	[split_capacity, padded_M, padded_N]	BF16; current workspace kernels require batch 1

The experimental gfx1250 fused split-K source remains available for repair, but it is not registered and cannot be selected by the public API.

Tuning

Path	OPUS tuning behavior
A16W16 GEMM	Full multi-kid/multi-split tuning through csrc/gemm_a16w16/gemm_a16w16_tune.py --libtype opus
A8W8 no-scale GEMM	Fixed kid 2; no OPUS offline selection
A8W8 plain block-scale GEMM	Fixed kid 1; no OPUS offline selection
A8W8 B-preshuffle GEMM	OPUS candidate generation for gfx942 only; current candidate is kid 11000
A8W8 MXFP8 BMM	Dedicated gfx950 tuner covering applicable kids and split-K values 1/2/4/8

Performance

Comparison on gfx950	Result
A16W16 interface	-24.051%
Executable A8W8 no-scale + FP32 block-scale	-0.062% (flat)
Torch vs implicit stream workspace, all 48 workspace kids, eager	-17.801%
Torch vs implicit stream workspace, all 48 workspace kids, graph replay	-10.618%
Public router vs private family adapter, eager	+1.028%
Public router vs private family adapter, graph replay	+0.005%

All 96 eager and 96 graph cases improved. Matching device show that the gains came from Python/C ABI dispatch, not kernel changes.

Guide

• aiter/ops/opus/__init__.py: strict public GEMM/BMM router and contracts
• aiter/ops/opus/policy.py: high-level candidate and tuned-row policy
• aiter/ops/opus/launch_plan.py: immutable split-K/workspace planning
• aiter/ops/opus/gemm_op_a16w16.py: A16 executor, Torch workspace, and C ABI
• aiter/ops/opus/gemm_op_a8w8.py: A8 family adapters and pybind backend
• csrc/opus_gemm/opus_gemm.cu: A16/A8 GEMM family entries and exact dispatch
• csrc/opus_gemm/opus_bmm.cu: MXFP8 BMM exact dispatch
• csrc/opus_gemm/gen_instances.py and codegen/gen_instances_gfx*.py:
  generated manifests, subset selection, and typed exact-kid tables
• csrc/opus_gemm/opus_gemm_common.py: canonical registry and family metadata

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 4961 --add-label <label>

PR title tags:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf] and op tags like [MLA] are left untouched. Add the no-auto-title label to opt this PR out of title tagging.

[demonsan]

why we need these sha?

[yifehuan]

Could we distinguish the workspace capacity from the effective split-K count in this check?

For example, on gfx942 with BF16, (M, N, K) = (1, 64, 128), kid=10201, and split_k=0, _plan_gfx942_split_k() returns (workspace_capacity=16, abi_split_k=1). This check then compares the capacity 16 against ceil(K / B_K) = 2 and raises ValueError, even though the
effective split-K count has already been reduced to 1. Explicit split_k=1 passes planning.

Reserving more workspace than the launch needs should be safe here; the C++ validation also accepts workspace.numel() >= required_numel. Would it make sense to size the workspace using the converged split-K count, or allow this capacity overestimate while validating
the effective split-K count separately? A regression test covering automatic split-K with a short K would help keep these two steps consistent.

[Fyzyukk]

The K-tile limit was incorrectly applied to workspace capacity. I moved that check into launch planning and size the gfx942 workspace from the converged split-K count. For your example, launch split-K stays 1, and planned capacity now matches it (16 slices → 1).
Verified on gfx942: auto allocation and caller workspaces with 1 or 16 slices all pass the Torch check.

[demonsan]

test opus_a16w16 has already covered splitk test. This file is duplicated. Also, hardcoded kid will easily failed. Remove this test will be better

[Fyzyukk]

Already fixed.

[demonsan]

gfx942 BF16-workspace exact-N guard redirects on one entry point and hard-raises on another — Problem: aiter/ops/opus/policy.py:419-425 (_resolve_a16w16_candidate) redirects kid 10210→10200 and 10213→10203 when N not in GFX942_BF16WS_EXACT_N, but aiter/ops/opus/gemm_op_a16w16.py:419-431 passes the tuned row's solidx and an explicit kernelId verbatim into _execute_a16w16, and aiter/ops/opus/launch_plan.py:285-294 then raises ValueError("gfx942 exact kid 10210 requires N in [64, 128, 256, 384, 512, 1024, 2048]; got N=1000") with no redirect anywhere in between. Impact at runtime: on a gfx942 device, gemm_a16w16_opus(A, B, kernelId=10210, splitK=2) with A: [M,K] bf16 and B: [K,1000] bf16 (N=1000, outside the exact-N set; the guard fires for any splitK) now hard-raises through the compatibility entry the author claims is behavior-preserved, where base's generated launcher (csrc/opus_gemm/codegen/gen_instances_gfx942.py:351-357 at base) silently called the fp32-workspace sibling 10200 and returned a correct Y; the tuned-row path has the same asymmetry via a user AITER_CONFIG_GEMM_BF16_FILE override (the shipped CSV carries no such rows and the base tuner already refused to emit them, so the explicit-kernelId path is the concrete trigger). Action: Author must route gemm_a16w16_opus's tuned-row and explicit-kernelId kids through resolve_a16w16_tuned_candidate (or apply the same 10210/10213 redirect inside _build_a16w16_launch_plan) so all entry points agree. [verified]