Add opus fp8 mxscale BMM kernels for gfx950

[yzhou103]

Motivation

Adds the opus fp8 e8m0 mxscale (block-scale) batched GEMM (BMM) path for
gfx950, exposed as aiter.batched_gemm_a8w8_mxscale. The per-kid device
kernels and host launchers are fully code-generated (mirroring the opus_gemm
module), and dispatch is driven by a tuned CSV with padded-M / M-split /
heuristic fallbacks. Also includes a round of dead-kernel and naming cleanup.

# mmajor DSV4 layout: x [M, G, K] fp8, wo_a [G, N, K] fp8 (batch-major, zero-copy),
# x_scale [M, G, K/128], w_scale [G, N/128, K/128], out optional [M, G, N].
aiter.batched_gemm_a8w8_mxscale(x, wo_a, x_scale, w_scale, dtype=torch.bfloat16)

Technical Details

Reuses the shared opus GEMM gfx950 backend; every kid is codegen'd into its own _C{void,bf16,fp32}.device.cu (device symbol) + inlined host launcher, no hand-written adapters.
Table-driven kid -> launcher dispatch (opus_bmm_a8w8_mxscale_tune_dispatch) instead of a large switch.
Kernel families: fused, pipeline, minterleave, mouter / mouter_tunable, wave8n2, wave4m2_selfload, plus OOB-masked sub-tiles (311/313/320/321/324/640/650/653).

Dispatch logic (kernelId=None)
Exact tuned CSV hit → use its kid + splitK.
Padded-M: round M up to the nearest tuned bucket in the same (g,n,k) family; accepted only for OOB-masked sub-tiles (run the real smaller M with no pad/copy).
M-split: for large tile-unaligned M, split into an aligned strong-tile bulk (157/150) + a <256-row OOB-safe tail. Gated on total bulk work g*m >= 8192 (measured break-even), not on g alone.
Heuristic single-kid fallback otherwise.

Test Plan

op_tests/test_opus_a8w8_bmm.py: correctness (fp8 vs bf16 einsum, err 0) across g∈{2,8}, m=1..16384, both m-major and batch-first layouts, and per-family kids.

Test Result

together with #4428

gfx950 x8, DeepSeek-V4-Pro (kv_cache_dtype=fp8, -tp 8, expert parallel).
WOA=1 is the fp8 e8m0 mxscale batched-GEMM path this PR touches, WOA=0 the
existing bf16 dequant + einsum baseline.

Accuracy — gsm8k 5-shot, 1319 problems, flexible-extract / strict-match, stderr ±0.006 on every row:

config	WOA=0	WOA=1
TP8	0.9530 / 0.9538	0.9545 / 0.9538
DP-attention	0.9568 / 0.9568	0.9545 / 0.9553
DP-attention + MTP	0.9500 / 0.9507	0.9530 / 0.9538

The whole spread is 0.9500-0.9568, about one stderr, so the mxscale path is indistinguishable from the baseline.

Throughput — benchmark_serving, random 8192 in / 1024 out, range ratio 0.8, concurrency 128, 1280 prompts, --ignore-eos. Medians over 2-3 rounds per arm:

config	WOA=0 tok/s	WOA=1 tok/s	delta	TPOT delta
TP8	18143.7	18436.1	+1.61%	-1.59%
DP-attention	19814.0	19956.7	+0.72%	-1.28%
DP-attention + MTP	27190.5	27870.9	+2.50%	-1.55%

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[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: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 4320 --add-label <label>

[yzhou103]

[Copilot]

Pull request overview

This PR adds gfx950-only OPUS fp8 (e8m0) mxscale batched GEMM/BMM support as a self-contained kernel + codegen + Python-binding stack, plus tuned configs and a new op_test to validate DeepSeek-V4-style wo_a layouts.

Changes:

• Introduces new gfx950 OPUS BMM frontends (C++ + pybind) and Python wrappers for fp8 e8m0 mxscale flatmm split-K dispatch (tuned CSV + heuristic fallback).
• Extends OPUS codegen/instance plumbing to emit and name the new BMM kernel families and generate the int-kid tuning lookup.
• Adds a new regression/perf op_test and a model-specific tuned CSV; updates the config-collision guard to include the new config family.

Reviewed changes

Copilot reviewed 23 out of 23 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
op_tests/tuning_tests/test_config_shape_collision.py	Adds the new mxscale batched-GEMM config family to the shape-collision guard suite.
op_tests/test_opus_a8w8_bmm.py	New correctness + perf sweep covering mxscale BMM paths (including batch-major round trip).
csrc/pybind/opus_gemm_pybind.cu	Registers the new OPUS BMM pybind entry point.
csrc/opus_gemm/opus_gemm_common.py	Extends instance metadata/naming and adds curated gfx950 mxscale BMM kernel lists for codegen/tuning.
csrc/opus_gemm/opus_bmm.cu	Adds the host-side gfx950 mxscale BMM dispatcher/front-end and split-K reduce explicit instantiations.
csrc/opus_gemm/opus_bmm_bake.cu	Adds a currently-disabled experimental/placeholder TU (needs cleanup before merge).
csrc/opus_gemm/include/opus_bmm.h	New public C++ header declaring the OPUS mxscale BMM entry point.
csrc/opus_gemm/include/gfx950/splitk_reduce_gfx950.cuh	Adds a BMM-specific split-K reduce kernel template for the fp32 workspace → final Y cast/store.
csrc/opus_gemm/include/gfx950/opus_gemm_traits_a8w8_scale_gfx950.cuh	Adds split-K kargs + new mxscale flatmm split-K traits used by the BMM kernel families.
csrc/opus_gemm/include/gfx950/opus_gemm_traits_a16w16_gfx950.cuh	Minor comment/style changes for consistency.
csrc/opus_gemm/include/gfx950/opus_gemm_pipeline_a8w8_scale_gfx950.cuh	Removes the old monolithic kernel body (now provided via the new pipeline header mapping).
csrc/opus_gemm/include/gfx950/opus_bmm_pipeline_a8w8_mxscale_gfx950.cuh	New combined pipeline header providing scaled MFMA helpers + BMM/GEMM scale kernels used by gfx950 codegen.
csrc/opus_gemm/include/gfx950/opus_bmm_launchers_a8w8_mxscale_gfx950.cuh	Adds host-only common checks and forward decls for BMM kernel templates (currently needs signature fixes).
csrc/opus_gemm/gen_instances.py	Extends instance generation to support BMM mxscale tags and emits a BMM kid→launcher tuning map.
csrc/include/rocm_ops.hpp	Adds the pybind macro for opus_bmm_a8w8_mxscale_flatmm_splitk.
aiter/ops/opus/bmm_op.py	New Python raw binding + tuned CSV lookup/heuristic wrapper for mxscale BMM.
aiter/ops/opus/init.py	Exposes the new BMM ops under aiter.ops.opus.
aiter/ops/batched_gemm_op_a8w8.py	Adds backend-neutral public entry batched_gemm_a8w8_mxscale (currently selecting OPUS backend).
aiter/jit/optCompilerConfig.json	Adds opus_bmm.cu to the OPUS JIT module build sources.
aiter/jit/core.py	Registers the new config env var/path and improves tuned→untuned name derivation + dedup key robustness.
aiter/configs/model_configs/dsv4_batched_gemm_a8w8_blockscale_mxscale_tuned.csv	Adds DeepSeek-V4-tuned mxscale batched-GEMM/BMM entries for gfx950.

Comments suppressed due to low confidence (1)

csrc/opus_gemm/include/gfx950/opus_bmm_launchers_a8w8_mxscale_gfx950.cuh:25

• The forward declaration for gemm_a8w8_mxscale_flatmm_splitk_mouter_kernel is missing the SKIP_SCALE_WAIT bool template parameter (the real kernel is template<typename Traits, typename D_OUT, bool SKIP_SCALE_WAIT=false>). Generated launchers for kids that set skip_scale_wait=True will fail to compile against this header.

template<typename Traits, typename D_OUT>
__global__ void gemm_a8w8_mxscale_flatmm_splitk_mouter_kernel(opus_gemm_scale_splitk_kargs_gfx950 kargs);

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[yzhou103]

preload sfa+sfb and vec load opt: