[ThruptiRajLakshmanaGowda]

Motivation

ISSUE ID : AICK-2170

A change to the CK-Tile dispatcher bridge or codegen can break numeric correctness,
kernel build, or kernel launch for any of the 13 bridged GEMM operators, and today
nothing catches it until someone runs the tests by hand. Three operators
(stream-K, aquant, abquant) had no GPU correctness test at all, and four more were
registered in ctest but never invoked from CI.

This adds two opt-in Jenkins lanes that fan out across gfx942 and gfx950 and cover
every operator with a bridge on develop, plus the missing tests. Both parameters
default to OFF, so no existing CI run gets slower.

Dependencies

Depends on #11591 — fix(ck-tile): guard the gfx1250 unscaled WMMA builtin behind __gfx1250__.

ck_tile/ops/gemm.hpp does not compile on any non-gfx1250 target on current
develop (19 errors, in both the host and the device pass), which breaks four of the
dispatcher GPU correctness tests this PR adds. That fix is unrelated to this change and
is split out into its own PR. Default CI here is unaffected — both lanes are
defaultValue: false — but Test Plan steps 1, 2 and 4 will not pass until #11591 merges.

Technical Details

1. RUN_DISPATCHER_CORRECTNESS_TESTS — JIT-compiled kernels vs. host reference

Covers all 13 operators rather than a representative sample: a developer touching
the bridge needs signal on all of them, and a lane covering four cannot give that.
Cost is held down by the search budget instead — --budget 64 for the standard
sweep, --budget 16 per variant sweep.

	gfx942	gfx950
shared	gemm_universal sweep + parity, grouped, multi_d (4 layouts), multi_abd, batched_gemm, batched_contraction (fp16+bf16+fp32), stream_k, rowcolquant, tensorquant, aquant, abquant	same
arch-only	—	mx_gemm, bquant C4 (fp8/bf8), C (fp8i4/bf8i4), H3 (mx_*)

Only two things stay gfx950-only, both for ISA reasons rather than policy. mx_gemm
needs the scale-MFMA builtins; it also cannot join the --variant sweep at all,
because mx_gemm_utils exposes only default_fp8_config() / default_fp4_config()
and never plugs into expand_sweep. The bquant H3 (mx_*) configs need the same
builtins. bquant's C4/C paths are held to gfx950 for a softer reason: its config
builders are arch-aware now, so the gate is pending one validating gfx942 run rather
than a known defect. Everything else — including
stream-K and all three quant operators — was widened to run on gfx942 as well over
the course of this branch. Both rationales are recorded inline in ck.groovy.

There is deliberately no gfx1201 stage: gemm_utils._SUPPORTED_ARCHES excludes
it, and widening that guard does not help, because every warp tile in
default_ci_config.json is 64 wide while RDNA4 is warp_size 32, so all 16 strata
enumerate 0 configs and the runner exits 1 on "no configs to run".

One narrowed surface is documented rather than closed: grouped_gemm_{a,ab,b}quant
run only on rcr, because all three config builders hardcode layout="rcr" and the
ctypes libs assert packed rcr strides.

2. RUN_DISPATCHER_PERF_TESTS — ahead-of-time tile_engine instances

Smoke-tier benchmark covering the same 13 operators as the correctness lane, so
the AOT instance path gets the same operator coverage the JIT path gets. The two are
gated separately and neither substitutes for the other: correctness JIT-builds each
config with hipcc and checks it against a host reference, while this lane exercises
tile_engine's codegen, instance builder and benchmark harness.

One dtype, one layout, one problem size, and a fixed 64 AOT instances per operator.

	gfx942	gfx950
fp16 / rcr	gemm_universal, batched_gemm, batched_contraction, gemm_streamk, grouped_gemm	gemm_streamk
fp16 / rcrr	gemm_multi_d, gemm_multi_abd	—
fp8 / rcr	grouped_gemm_rowcolquant, grouped_gemm_tensorquant	gemm_aquant, gemm_bquant, gemm_abquant, mx_gemm

Three details that are load-bearing rather than incidental:

• The per-op cap is explicit, not derived. TILE_ENGINE_SAMPLING_TIER is a
  total budget split evenly across every op with a non-empty <OP>_DATATYPE
  (ops/gemm/CMakeLists.txt:16-23), so adding operators silently shrinks the coverage
  of the ones already there — gfx942 alone would have gone from 125 to 55 instances
  per op. Pinning <OP>_MAX_INSTANCES uses the first-class override path (same file,
  :50-59) and keeps lane cost fixed and independent of how many ops the lane builds.
  The tier is still passed, because each op also forwards it to the sampler as
  --tier to select the sampling strategy.
• One dtype and one layout per op is not just smoke sizing. The per-op budget is
  divided again across (dtype × layout) combos with integer truncation, so a
  multi-combo op can floor a combo to zero instances and build nothing.
• gemm_multi_d and gemm_multi_abd take 4-character layouts (A, B, D, E).
  Passing the rcr the other ops use is a parse error in the instance builder, not a
  narrower selection.

mx_gemm cannot leave the gfx950 branch: its CMakeLists filters targets with a
literal regex, target MATCHES "^gfx950", so on any other arch
benchmark_mx_gemm_all is never created and ninja fails on an unknown target rather
than skipping. It also has no fp16 path.

gemm_streamk builds on both arches. DESIRED_TARGETS was widened to include gfx950
earlier on this branch, but that CMake change is the one part that could not be
exercised locally; this lane is its first real confirmation, and it matches the
correctness lane, which already runs stream-K on both. If the gfx950 AOT build turns
out to be broken, that DESIRED_TARGETS line is revertible on its own without
touching the lane.

3. New search-space runner — dispatcher/tests/test_gemm_search_space.py

One runner covering every variant expand_sweep understands (standard, grouped,
multi_d, multi_abd, stream_k), selected with --variant — they share the same
enumerate/sample/build/report machinery and differ only in which Gpu*Runner to
drive and how the numpy reference is computed.

It enumerates the GemmKernelConfig space against tile_engine's
default_ci_config.json, samples a budget-limited subset with a daily rotating seed
per (dtype, layout) stratum — so even budget 64 spreads ~4 configs across all 16
combinations rather than testing one deeply — compiles via the bridge, runs on GPU,
and reports correctness + TFLOPS as JSON.

This is the dispatcher equivalent of tile_engine's per-op benchmark scripts, with one
deliberate difference: those never verify numerics and always exit 0. This one
validates every kernel against a numpy reference and exits 1 on any mismatch or
build failure.

4. New GPU correctness tests

• test_streamk_gpu_correctness.py — all three reduction strategies. The shared
  sweep only ever builds atomic (default_ci_config.json carries no streamk_config,
  so expand_sweep falls back to ["atomic"]), leaving linear/tree unverified
  through the ctypes bridge. Two shapes: an occupancy-bound 4096³ case, and a
  Stream-K-heavy M=N=128, K=8192 case that only linear/tree can assert tightly,
  since estimate_num_wgs_per_tile is gated on Atomic and returns 1 otherwise.
  Tolerance is CK's own split-K-aware element-wise bound, ported from
  streamk_driver_common.hpp + check_err.hpp.
• test_aquant_gpu_correctness.py — A-side per-(m,k)-group scales; fp8, bf8, and a
  tiled shape spanning a 4×4 output-tile grid that the 1×1-tile cases cannot catch.
• test_abquant_gpu_correctness.py — scales on both operands, completing the quant
  trio; fp8 + bf8 on compv3, plus the gfx950-native eightwaves pipeline (192×256
  tiles, TransposeC=True, bquant_group_n=128), a different MFMA path and B-scale
  granularity than compv3.

All four quant/stream-K bridge tests are now registered in ctest with
SKIP_RETURN_CODE 77. They intentionally take no --gfx from CMake: each
hipcc-compiles with --offload-arch and then runs, so the arch must match the
physical GPU, and the configured GPU_TARGETS may legitimately differ from the card
in the box.

5. Skip semantics — exit 77 instead of exit 0

The script-style tests previously returned 0 when no GPU was present, so a CPU-only
runner reported a green PASS for a test that never touched the GPU. They now return 77,
which ctest maps to skipped via SKIP_RETURN_CODE and the Jenkins lane maps via a
run_ok shell helper. Only 77 is swallowed — any other non-zero still fails the lane.
This matters because execute_cmd is one &&-joined string: without it, a single
unsupported arch would abort the lane and every later operator would silently never run.

--gfx also loses its hardcoded gfx950 default in the quant tests, so "user asked for
gfx950" stays distinguishable from "we are on an unrelated box" and the skip can fire.

6. Fixes and coverage extensions found while wiring this up

• gemm_utils.expand_sweep: skip non-cshuffle epilogues for stream_k — the codegen
  emits only cshuffle for that variant, so those configs surfaced as spurious build
  failures.
• grouped_gemm_aquant_utils._compile_aquant_kernel: add the arch defines
  (CK_USE_OCP_FP8, CK_TILE_USE_OCP_FP8, CK_USE_NATIVE_MX_SUPPORT,
  CK_GFX950_SUPPORT) that CMake normally injects but the standalone hipcc path omitted,
  matching the bquant path. gfx942 is FNUZ, gfx950 is OCP; without this the encoded
  inputs and the compiled kernel disagree and the device returns NaN.
• tile_engine/ops/gemm_streamk/CMakeLists.txt: widen AOT DESIRED_TARGETS to include
  gfx950, and correct the skip warning, which still listed only (gfx90a, gfx942)
  afterwards.
• streamk_gemm_ctypes_lib.cpp: correct a stale comment claiming the lib is Atomic-only.
  The strategy is a codegen-time template constant and GetWorkSpaceSize returns 0 for
  Atomic / non-zero for Linear/Tree, so the one allocate-and-launch path serves all three.
• Widened two existing GPU tests that were narrower than the operator they cover:
  batched-contraction now exercises bf16 and fp32 alongside fp16, and the Multi-D GEMM
  test exercises all four layouts rather than one.
• Dockerfile: add ml_dtypes, required for host-side fp8/bf8 encoding.

Both lanes archive their per-kernel JSON from a finally block — a red lane is exactly
when that artifact is worth having, and buildAndTest throws on failure.

Test Plan

1. gfx942 correctness lane — RUN_DISPATCHER_CORRECTNESS_TESTS=true: the
   gemm_universal sweep at budget 64 across all 16 dtype × layout strata, the parity
   test, the grouped/multi_d/multi_abd/batched_gemm/batched_contraction/rowcolquant/
   tensorquant/aquant/abquant/bquant bridge tests, the grouped/multi_d/multi_abd/
   stream_k variant sweeps at budget 16, and the stream-K registry + bridge tests.
2. gfx950 correctness lane — the same set, plus mx_gemm and the bquant H3
   (mx_*) configs.
3. Both perf lanes — RUN_DISPATCHER_PERF_TESTS=true on gfx942 and gfx950; confirm
   all 9 gfx942 and all 5 gfx950 benchmark_*_all targets build, every script runs, and
   each *_results.json lands as an archived artifact. The gfx950 gemm_streamk result
   is the specific new signal — it is the first exercise of the DESIRED_TARGETS
   widening.
4. Skip path — run the ctest suite on a non-gfx942/gfx950 box and confirm the
   arch-gated tests report Skipped, not Passed and not Failed.
5. Artifact archiving — confirm dispatcher_*_results.json and *_results.json
   appear as build artifacts, including on a deliberately failed run.

Test Result

• Will be run on CI

Submission Checklist

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

[therock-pr-bot]

✅ All Policy Checks Passed

Check	Status	Details
📝 PR Description	✅ Pass	—
⛔ Forbidden Files	✅ Pass	—
🧪 Unit Test	✅ Pass	—
🚫 Draft PR	🔜 To Be Enabled	—
🚩 Feature Flag	🔜 To Be Enabled	—
📊 Code Coverage	🔜 To Be Enabled	—

🎉 All policy checks passed!

📖 Need help? See the Policy FAQ for details on every check and how to fix failures.

🙋 Wish to Override Policy?

[therock-pr-bot]

🎉 All checks passed! This PR is ready for review.

[Copilot]

Pull request overview

Adds CI coverage for the CK-Tile universal GEMM dispatcher by introducing new Jenkins pipeline stages and a Python-based, budget-limited search-space runner that builds and validates sampled dispatcher GEMM configurations.

Changes:

• Extend projects/composablekernel/Jenkinsfile with new boolean parameters and parallel stages to run dispatcher-focused tests.
• Add dispatcher CI entrypoints in groovy/vars/ck.groovy to configure/build dispatcher artifacts and invoke the new runner and correctness tests.
• Introduce dispatcher/tests/test_gemm_search_space.py to enumerate, sample, build, run, and emit JSON results for dispatcher GEMM configs.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 6 comments.

File	Description
projects/composablekernel/Jenkinsfile	Adds CI parameters and stages for running dispatcher GEMM and universal dispatcher tests on specific GPUs.
projects/composablekernel/groovy/vars/ck.groovy	Adds Jenkins shared-library helpers to build dispatcher components and run/archive dispatcher test outputs.
projects/composablekernel/dispatcher/tests/test_gemm_search_space.py	New Python runner for stratified sampling of universal GEMM configs with correctness/perf reporting to JSON.

Suppressed comments (1)

projects/composablekernel/dispatcher/tests/test_gemm_search_space.py:74

• _emulate() currently rounds fp8/bf8 inputs as fp16, which makes the CPU reference diverge from what the GPU kernel actually multiplies. That can hide correctness issues for fp8/bf8. Use the same fp8/bf8 encode/decode path as the existing dispatcher parity tests so the reference matches the kernel’s storage dtype.

def _emulate(x: np.ndarray, dtype: str) -> np.ndarray:
    if dtype == "bf16":
        u32 = np.ascontiguousarray(x, dtype=np.float32).view(np.uint32)
        rounded = (u32 + ((u32 >> 16) & 1) + np.uint32(0x7FFF)) >> 16
        return (rounded.astype(np.uint32) << 16).view(np.float32)
    return x.astype(np.float16).astype(np.float32)

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[ozturkosu]

My request has been resolved.

[brockhargreaves-amd]

Looks good, lgtm after the merge conflicts are fixed.

[assistant-librarian]

TheRock Submodule Bump Activity

Newest first

• 2026-09-18 — Included in TheRock via ROCm/TheRock#8336.