Repository navigation
Conversation
Extend the GEMM bridge so the Dispatcher generates and launches batched GEMM with full parity to the legacy Tile Engine batched_gemm op (fp16 / rcr, adding a batch dimension with per-batch strides). Mirrors the gemm_universal (#8997), grouped (#9000) and stream-K (#9028) bridges. - unified_gemm_codegen.py: new GemmVariant.BATCHED; batched launch branch (ck_tile::BatchedGemmKernel + BatchedGemmHostArgs, 4-arg GridSize, batched CShuffle epilogue); _batched name suffix; GEMM_KEY_BATCHED macro; --variants batched; batched configs share the standard tile/trait sweep. - batched_gemm_ctypes_lib.cpp: new C-ABI .so with a DIVERGENT ABI (dispatcher_run_batched carries batch_count + three batch strides). Bypasses the registry and launches the force-included kernel directly, like stream-K. - batched_gemm_utils.py: BatchedGemmKernelConfig / BatchedGemmProblem / GpuBatchedGemmRunner / setup_multiple_batched_gemm_dispatchers / expand_sweep, reusing gemm_utils codegen flags for byte-identical AMDGPU backend flags. - batched_gemm_full_benchmark.py + run_one_batched_gemm_kernel.py: 3-phase multi-GPU driver with batched operand generation. - ctypes CMakeLists: dispatcher_batched_gemm_lib target. Smoke on gfx942/MI300X: 6672 default_config kernels codegen with 0 failures; end-to-end bridge (codegen -> hipcc -> .so -> run) passes with name-parity (runner name == config name) and batched correctness vs numpy across batch counts 1/2/4/8 (max_rel ~5e-4).
ozturkosu
changed the title
✅ All Checks Passed — Ready for Review
📖 Need help? See the Policy FAQ for details on every check and how to fix failures. |
|
🎉 All checks passed! This PR is ready for review. |
… bridge Address review of the batched_gemm TileEngine->Dispatcher bridge (#9306): - Fair A/B timing (B1): thread warmup/repeat/flush_cache/rotating_count through the ctypes ABI and build the batched stream_config to match the Tile Engine batched_gemm profiler exactly (default warmup=50/repeat=100/ flush/rotating=1000, gpu timer). The .so now drives SelectedKernel::launch identically to Old-TE instead of the cold=3/repeat=10 defaults. - Non-packed batch strides (B2): resolve batch strides with <=0 -> packed and size the device allocation from the resolved batch stride so a padded (non-packed) batch_stride is honoured; the runner gains batch_pad_* params and over-allocates + scatters/gathers slabs to exercise it. - Split-K (B3): thread k_batch through the ABI/runner (default 1, Old-TE parity). The universal epilogue stores C with atomic_add when k_batch>1, so the codegen batched launch now re-zeros C before each timed iteration via a hipMemset callable (k_batch==1 path stays byte-identical to Old-TE). - Validator fix (N1): validate_kernel_config built the warp-tile lookup key as dtype_dtype_dtype (e.g. fp16_fp16_fp16), but the tables are keyed by A_B_Acc (fp16_fp16_fp32). The lookup always missed and fell back to a stale default [[32,32,16],[16,16,16]], wrongly rejecting legal warp tiles such as [16,16,32] for every dtype -- which is why the batched CI config yielded 0 configs. Key off dtype_acc; default_ci_config.json is unchanged. - N2: align the leading-stride sentinel to <=0 to match ck_tile's get_default_stride / the TE profiler exactly. Verified on gfx942/MI300X (GPU3): CI expand_sweep now yields 8 batched fp16/rcr configs (was 0); packed max_rel 4.7e-4, non-packed 4.7e-4, k_batch=2 5.7e-4.
|
Posting the round-2 review fixes here to keep the description focused on what the PR does. Here is what I addressed from review:
clang-format: verified with |
A/B perf-parity sweep vs Old-TE (MI300X / gfx942)Ran the full serialized bridge-vs-Old-TE performance parity sweep for batched GEMM on a single MI300X (gfx942), apples-to-apples:
Result (40 rows = 8 stems × 5 shapes)
Full per-(stem, shape) CSV is attached to the PR working set ( |
There was a problem hiding this comment.
Pull request overview
Adds a TileEngine → Dispatcher bridge for batched GEMM (fp16/rcr) so Python callers can codegen, build, and launch batched kernels with the correct divergent ABI (batch_count, split-K via k_batch, and per-batch strides), matching the established registry-bypass pattern used by other non-standard GEMM variants.
Changes:
- Introduces a new batched GEMM ctypes C ABI (
dispatcher_run_batched) plus a Python ctypes/runner layer to build and run per-kernel.sofiles. - Extends unified GEMM codegen with a
batchedvariant (naming, includes, launch path, andGEMM_KEY_BATCHED). - Adds CPU-only unit tests and a 3-phase multi-GPU benchmark driver/worker for the batched bridge.
Reviewed changes
Copilot reviewed 8 out of 8 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
| projects/composablekernel/tile_engine/ops/gemm/run_one_batched_gemm_kernel.py | Subprocess-isolated worker that loads a single batched .so and reports timing/verification as JSON. |
| projects/composablekernel/tile_engine/ops/gemm/batched_gemm_full_benchmark.py | 3-phase batched sweep driver (codegen/build → problem load → multi-GPU batched benchmarking). |
| projects/composablekernel/dispatcher/tests/test_batched_bridge.py | CPU-only tests for name contract, codegen JSON projection, flops scaling, and shipped config validity. |
| projects/composablekernel/dispatcher/python/ctypes_utils.py | Fixes warp-tile validation keying to use accumulator dtype (A_B_Acc). |
| projects/composablekernel/dispatcher/python/batched_gemm_utils.py | New Python bridge for batched kernel config/problem, build, ctypes wrapper, and GPU runner. |
| projects/composablekernel/dispatcher/codegen/unified_gemm_codegen.py | Adds GemmVariant.BATCHED codegen path (naming, includes, launch, epilogue specialization). |
| projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt | Adds an optional CMake-built dispatcher_batched_gemm_lib target if a generated batched header exists. |
| projects/composablekernel/dispatcher/bindings/ctypes/batched_gemm_ctypes_lib.cpp | New C ABI implementation that bypasses registry and launches the force-included batched kernel. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Fix the two unresolved Copilot review comments on the batched GEMM
TE->Dispatcher bridge, plus a parse-breaking regression:
- batched_gemm_ctypes_lib.cpp: derive the default per-batch stride from
the RESOLVED leading stride (rows*sa / cols*sa etc.) instead of the bare
M*K/K*N/M*N shape. A non-packed leading stride (padded rows/cols) with a
default batch stride previously under-sized the device buffer and could
cause out-of-bounds access.
- unified_gemm_codegen.py (split-K launch): use hipMemsetAsync on the
launch stream (sc.stream_id_) instead of the synchronous hipMemset, which
ignored stream_config and would serialise onto the default stream,
skewing timing/correctness on a non-default stream.
- unified_gemm_codegen.py (includes): restore the missing closing triple
quote on the BATCHED include block. The prior auto-fix commit merged it
into the STREAM_K block, leaving the whole module unparseable (codegen
could not run at all).
Verified on MI300X (gfx942): build + correctness for packed batch
{1,2,4,8}, split-K k_batch {2,4}, non-packed batch stride, and the exact
non-packed-leading-stride / default-batch-stride case, all max_rel < 1e-3.
Resolve conflicts in unified_gemm_codegen.py with the grouped-GEMM bridge that landed on develop. All conflicts were the batched (this branch) and grouped (develop) variants touching the same extension points; kept both: - GemmVariant enum: keep BATCHED and GROUPED. - KernelNaming variant suffix: keep both _batched and _grouped branches. - _launch_function dispatch: keep both BATCHED and GROUPED cases. - Reconstruct the two interleaved launch generators into separate complete methods (_launch_function_batched, _launch_function_grouped) -- the merge had spliced them together on their shared ave_time/scheduler/return lines. - --variants CLI choices: union of both (standard, preshuffle, multi_d, stream_k, batched, grouped).
Remove the unused _make_batched_name duplicate so the _batched suffix rule has a single source of truth (the _batched_name property already bound onto BatchedGemmKernelConfig.name, matching codegen). Nothing imported the deleted symbol. Also correct the test docstring run path to tests/test_batched_bridge.py (runnable from the dispatcher dir), matching sibling tests. No kernel logic or parity-relevant code changed.
Parity sweep — whole
|
| metric | value |
|---|---|
| median gap | −0.46% |
| mean gap | −0.66% |
| within ±15% | 2100/2100 (100%) |
| within ±5% | 1998/2100 (95.1%) |
| range | [−11.8%, +10.6%] |
gap% = (bridge − oldte)/oldte × 100 (positive = bridge faster). Every pipeline (mem −0.65%, compv3 −0.74%, compv4 −0.45%) and both epilogues are at parity, on both shapes. The only two initial >15% outliers were 4-GPU-sweep contention artifacts — they collapsed to ±0.6% on standalone re-measure.
Split-K — k_batch>1 (reported separately)
Bridge is faster here (not a regression): +16% mean on 4096³, larger on the tiny 1024×512×256 shape where Old-TE's split-K is overhead-bound. Correctness verified (max_rel ~5e-4 at k_batch=2). Old-TE additionally rejects default-epilogue split-K entirely (1178 nan) — a coverage asymmetry, not a gap.
Build time (fair 334-kernel sample, same set, 32-way)
Bridge 1.22 vs Old-TE 1.04 s/kernel (~17%). Codegen is negligible on both; the residual is the bridge's -fPIC + --shared --hip-link (required to emit a loadable .so), not the kernel compile.
Verdict: bridge at parity with Old-TE for the standard hot path across the full config.
… not default to gfx942
Fresh MI300X / gfx942 fair parity re-verify — 2026-07-28Verdict: AT PARITY. Full apples-to-apples bridge-vs-Old-TE default_config sweep on 1× MI300X (gfx942), ROCm 7.2, 5 shapes × whole matched-intersection.
gap% = (bridge − oldTE)/oldTE × 100 (+ = bridge slower). This PR (batched_gemm / #9306): 2,100 k_batch=1 rows (420 stems × 5 shapes, fp16 rcr), ALL within ±15%, median −0.46%, 95.1% within ±5%. Zero all-zeros / silent-fail rows. k_batch>1 split-K rows: bridge faster (median +27.4%) — bridge win due to better scheduling, not a regression. Fairness enforced: warmup/repeat 50/100 both sides (bridge runtime ABI); Old-TE built via CK CMake real TE Per-PR CSVs + consolidated report: https://gist.github.com/ozturkosu/1fc7e430f742b2399fce4e6e5c7cdf12 ( |
Bridge vs Old-TE Performance Parity -- gfx950/MI355X (2026-07-28)Ran A/B interleaved parity sweep of the batched-GEMM bridge vs Old-TE on MI355X (gfx950), ROCm 7.2.0, SLURM job 67616534. Result: AT PARITY
All 766 rows within +/-15%. Two initial outliers (-80.7%, -55.9%) were confirmed as contention artifacts (parallel bridge build on CPU during sweep) and collapsed to parity on standalone median-of-3 remeasure. Fairness Checklist (all verified)
Full CSV + consolidated report: https://gist.github.com/ozturkosu/75f4ad6542b27fd25ca0029dd2b42743 |
detect_devices() in batched_gemm_full_benchmark.py called `rocm-smi --showid` directly for GPU id detection. Delegate to smi_utils.detect_gpu_ids() (added in #9731), which prefers amd-smi and falls back to rocm-smi, matching gemm_full_benchmark.py. Added common/ to sys.path, dropped unused `import re`. JIRA ID : ROCM-28734
| dtype_a="fp16", | ||
| dtype_b="fp16", | ||
| dtype_c="fp16", |
There was a problem hiding this comment.
do we support other dtypes - bf16, fp8?
There was a problem hiding this comment.
Not at the moment. The batched GEMM path in Old TE only declares fp16 with rcr layout (see batched_gemm_instance_builder.py, --datatype choices=["fp16"] and --layout choices=["rcr"]), so I kept the bridge matched to exactly that set to stay at parity with the reference side.
The plumbing for bf16 and fp8 is mostly already in place though. The encode helpers and the C dtype size table handle them, so once Old TE adds those instances on the batched side, turning them on here is basically flipping the guard in is_valid() plus the codegen. For this PR I intentionally gate everything except fp16 rcr so we don't expose a dtype the reference cannot verify against.
Bridge vs Old-TE parity resultsFull data (secret gist): https://gist.github.com/ozturkosu/3a6b3c7f3daf12efba103d7599f39311 AT PARITY on both archs. gfx942 (MI300X): 2100 cases, median -0.46%, 100% within +/-15%. gfx950 (MI355X): 766 cases, median +0.09%, 100% within +/-15%. Split-K (k_batch>1) is a bridge win, not a regression. Contains: consolidated parity summary + per-PR CSV result file(s). Fair A/B methodology: warmup=50/repeat=100 both sides, interleaved per (stem,shape), Old-TE built via CK CMake with real TE |
MI300X / gfx942 A/B perf-parity — batched_gemm bridge vs Old Tile-EngineVerdict: AT PARITY. Fresh from PR tip Headline
Fairness (apples-to-apples, all applied)
Outlier diagnosis (the raw within-±15% of 66% is a measurement artifact, not codegen)The 3-shard (3-GPU) throughput sweep inflates the Old-TE side timing under contention (its per-call process spawn + 1000 rotating-buffer device allocation competes with the other shards; the resident bridge .so does not). Standalone, single-idle-GPU median-5 re-measure of 400 >15% rows: 397 collapsed to parity, only 3 persist (all
Correctness / coverage finding (not a perf gap)280 kernels — every one a 192-dimension tile (192×192×64, 256×192×64, 64×192×64, …) — failed the fp32 gate and were recorded INVALID (never timed). Old-TE's launcher also rejects these ("Wrong! Arguments not supported! Skipping gemm!"), i.e. they are invalid on both sides (shared ck_tile 192-tile / CShuffle-store limitation, cf. #9684). The bridge codegen validator, however, accepts them and the registry-bypass launch runs them to an incorrect result. Recommendation: tighten the bridge validity gate to reject 192-dimension tiles up front, matching Old-TE. Deliverables: 16,190-row CSV, summary.json/md, self-contained HTML analysis. Coverage note: of the full default_config (~5,792 bridge-codegen-valid fp16/rcr configs), 1,619 build on both sides; the remainder fail codegen/compile identically on both. |
MI355X / gfx950 (CDNA4) whole-config parity — AT PARITY ✅Complements the earlier MI300X/gfx942 run. Ran an apples-to-apples A/B perf-parity of this batched-GEMM TE→Dispatcher bridge vs Old Tile-Engine on a dedicated 8× MI355X node (ROCm 7.2, gfx950), from the PR tip Headline (k_batch=1) — AT PARITY
Per-pipeline (k1): compv3 −0.31% · compv4 −0.19% · mem −0.10% — all 100% within ±15%. Split-K (k_batch>1) — reported separately11,182 rows, median −2.45%, 99.2% within ±15%. The marginal slowdown is the correctness-required per-iter Correctnessfp32 numpy reference per stem: 1374 / 1654 bridge stems PASS (worst Scale & methodology (fairness checklist all enforced)
Verdict: batched-GEMM bridge is AT PARITY with Old-TE on MI355X/gfx950 (headline median −0.15%, 100% within ±15%), matching the MI300X result. Split-K reported separately (correctness-required overhead). |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 10 out of 10 changed files in this pull request and generated 2 comments.
Suppressed comments (3)
projects/composablekernel/dispatcher/tests/CMakeLists.txt:16
- This CTest command discovers zero tests: the module contains no
unittest.TestCase; its GPU check only runs frommain(). As a result, CTest reports success without building or running a batched GEMM. Invoke the script directly so its exit status reflects the GPU check (and its existing no-GPU skip path).
COMMAND ${Python3_EXECUTABLE} -m unittest discover -s ${CMAKE_CURRENT_SOURCE_DIR} -p test_batched_gemm_gpu_correctness.py -v
projects/composablekernel/dispatcher/python/gemm_utils.py:2112
- This is a second definition of
_cshuffle_store_ok; the same function already exists atgemm_utils.py:124-138. This later definition silently shadows the first, so future fixes to the earlier shared helper would have no effect onexpand_sweep. Remove this duplicate and keep a single implementation.
def _cshuffle_store_ok(
m_repeat: int, n_repeat: int, warp_tile_m: int, warp_tile_n: int
) -> bool:
"""Return False for the one CShuffle-store combination that is numerically
wrong (issue #9684): an ODD per-wave repeat (>1) paired with a 32-wide warp
tile in that dimension. GPU-verified on gfx942 -- e.g. tile_m=192 / wave_m=2
/ warp_tile_m=32 (MRepeat=3) returns garbage, while every other non-power-of-
two repeat (incl. MRepeat=3 with warp_tile_m=16, and even repeats like 6/12)
is correct. Only relevant for the CShuffle epilogue; the default epilogue is
exempt."""
def _dim_bad(repeat: int, warp_tile: int) -> bool:
return repeat > 1 and repeat % 2 == 1 and warp_tile == 32
return not (_dim_bad(m_repeat, warp_tile_m) or _dim_bad(n_repeat, warp_tile_n))
projects/composablekernel/dispatcher/bindings/ctypes/batched_gemm_ctypes_lib.cpp:194
k_batchis narrowed to 32-bit before the later range check. A value aboveINT32_MAXtherefore truncates first (for example,INT32_MAX + 1becomes negative), andfits_i32(kbatch)can no longer detect the original out-of-range input. Validate the resolved 64-bit value before casting so the advertised truncation guard actually covers split-K.
// split-K count: <= 0 falls back to 1 (no split), matching the TE default.
const ck_tile::index_t kbatch = static_cast<ck_tile::index_t>(k_batch > 0 ? k_batch : 1);
| const std::int64_t bsa = | ||
| batch_stride_A > 0 ? batch_stride_A : (operand_is_row_major(kLayoutA) ? M : K) * sa; | ||
| const std::int64_t bsb = | ||
| batch_stride_B > 0 ? batch_stride_B : (operand_is_row_major(kLayoutB) ? K : N) * sb; | ||
| const std::int64_t bsc = | ||
| batch_stride_C > 0 ? batch_stride_C : (operand_is_row_major(kLayoutC) ? M : N) * sc; |
| stream_config so timing matches the Tile Engine batched_gemm profiler | ||
| (B1). | ||
| """ | ||
| batch, M, N, K = problem.batch_count, problem.M, problem.N, problem.K |
Parity status — batched_gemm bridge vs Old Tile-Engine (2-arch, Aug 2026)Verdict: AT PARITY on both architectures. No performance regression; the bridge is at parity or marginally faster than Old-TE across the whole
(+ = bridge faster.) Every pipeline/epilogue stratum at parity on both archs (compv3/compv4/mem, cshuffle/default all ≈−0.2%). Split-K (k_batch>1), gfx950: median −2.45%, 99.2% within ±15% — the correctness-required per-iteration Correctness: all timed bridge kernels fp32-verified (worst max_rel < 1e-2 gfx942; 4.06e-4 gfx950). The single surviving >15% outlier (mem/256×256×64) was proven byte-identical device kernels via rocprofv3 — measurement variance, and it was the bridge running faster, not slower. Fairness (both sides): warmup=50 / repeat=100 / flush_cache / rotating_count=1000; Old-TE built via its own CMake with the real TE One correctness follow-up (not a perf issue): on gfx942 the bridge codegen accepts and runs 192-dimension tiles to wrong output, while Old-TE's launcher rejects them ("Arguments not supported"); on gfx950 both reject. Recommend tightening the bridge's validity gate to reject 192-tiles (shared #9684-class limit). 280 stems affected — excluded from the parity numbers above. Note: the lower MI300X row count reflects a contended shared node during that run (fewer shapes/kernel), not a coverage or correctness gap — the gfx950 run on a dedicated node gives the full ~50-shape coverage. Full results (CSV + per-arch summaries, secret gist): https://gist.github.com/ozturkosu/79906d0ab03ce2f665c96b2afcf4ee35
|
…idges (#10612) ISSUE ID: #8997 ## Summary This PR combines two sibling TileEngine → Dispatcher bridge ops into a single PR: - **batched GEMM** (previously #9306) - **batched_contraction** (previously #9328) Both follow the same **direct-launch, registry-bypass** pattern (as the stream-K bridge #9028), because their launch ABIs carry variable-length / batch-specific arguments the single-pointer registry backend cannot express. The two ops touch disjoint files except `dispatcher/tests/CMakeLists.txt`, where both GPU-correctness test registrations are kept. ## Motivation The dispatcher had no path for batched GEMM (same GEMM across many independent problems with per-batch strides) or batched_contraction (generalized batched tensor contraction `E[G.., M.., N..] = sum_K A[G.., M.., K..] * B[G.., N.., K..]` with multi-dim G/M/N/K index groups). Both are real Old-TE ops; this bridge lets Python callers generate, build, and launch them at parity with the legacy Tile Engine — without writing C++. ## Test Plan / Result - CPU-only unit tests for both bridges + gemm_utils: **74 passed** (`pytest dispatcher/tests/test_batched_bridge.py dispatcher/tests/test_batched_contraction_bridge.py dispatcher/tests/test_gemm_utils.py`). - Batched GEMM: end-to-end name-parity + correctness across batch counts 1/2/4/8 (`max_rel ~5e-4`), non-packed strides, split-K; full `default_config` codegen 6672 kernels / 0 failures. - Batched contraction: on-GPU verify (gfx950) across dtype × layout × shape × multi-dim × pipeline + D-tensor epilogue, all PASS. - clang-format-18 clean on both ctypes libs. ### Perf parity vs Old-TE - Batched GEMM (MI300X, fp16 rcr, batch=8): at parity / slightly ahead — median gap +4.00%, 100% within ±15%. - Batched contraction (MI350X, fp16 rcr): at parity — median gap -0.95%, 100% within ±15%. ## Scope / known limitations - Batched contraction: `rcr` only, `k_batch==1` only (split-K is a shared Old-TE kernel defect — hard-rejected, never silently-wrong), non-tile-multiple M/N/K rejected by `IsSupportedArguments`. --- Supersedes and closes #9306 (batched GEMM) and #9328 (batched_contraction). **Related PRs (TileEngine → Dispatcher GEMM bridge series):** #8997 (regular GEMM), #9000 (grouped), #9028 (stream-K), #8887 (fp8/bf8/int8), #9305 (multi-ABD), #9307 (preshuffle), #9308 (multi-D), #10439 (block-scale quant, 5 ops). --------- Co-authored-by: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com> Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Co-authored-by: Thrupti Raj Lakshmana Gowda <thruptiraj.lakshmanagowda@amd.com>
…fx1250 (MI400) The batched-GEMM (#9306) and batched-contraction (#9328) TE->Dispatcher bridges combined in merged #10612 were CDNA-only (gfx90a/gfx942/gfx950, MFMA). Both keep their OWN _SUPPORTED_ARCHES/_SUPPORTED_ARCHS tuple (they do not inherit the shared gemm_utils gate), and both ship MFMA CI configs, so neither runs on gfx1250 as merged. This folds both variants onto the shared gfx1250 foundation (base branch users/muozturk/ck/gemm-bridges-gfx1250-enable, PR #10921): - batched_gemm_utils / batched_contraction_utils: add "gfx1250" to the per-variant supported-arch tuple (arch gate; codegen is arch-agnostic). - batched_gemm: gfx1250 CI config (WMMA warp_tile 16x16x32). - batched_contraction: gfx1250 CI configs (default + bridge, WMMA 16x16x32); the merged configs pin MFMA 32x32x64 / 32x32x16 which do not run on gfx1250. - both --arch help + config help mention gfx1250. - CPU-only tests: gfx1250 in supported arches, arch resolve/detect accepts gfx1250, gfx1250 CI configs present with WMMA 16x16x32, and the contraction gfx1250 sweep expands to valid WMMA kernels. Validation: 66 CPU tests pass (incl. the new gfx1250 tests). On-GPU validation on MI400/gfx1250 is pending node availability (both Helios nodes were unreachable / GPU-less at authoring time).
…fx1250 (MI400) (#10928) JIRA ID : AICK-2104 JIRA ID : AICK-2136 ## Summary Enables the **batched GEMM (#9306)** and **batched contraction (#9328)** bridges (merged together in #10612) on **gfx1250 (MI400, RDNA4/WMMA)**. Targets `develop` directly. Python + JSON only (no C/C++). > Recreated against `develop` (replaces the earlier stacked PR #10922). The diff therefore also includes the shared gfx1250 enablement that also appears in #10921 (gemm_utils `_SUPPORTED_ARCHES`/wave tables, arch_specs gfx1250, driver dtypes) — if #10921 lands first, this will rebase down to just the batched-specific files. ## Does merged #10612 support gfx1250? — No. The batched variants keep their **own** arch tuples (`batched_gemm_utils._SUPPORTED_ARCHES`, `batched_contraction_utils._SUPPORTED_ARCHS` = gfx90a/gfx942/gfx950) and their CI configs pin CDNA MFMA tiles (contraction `default` 32x32x64, `bridge` 32x32x16) that don't run on gfx1250 WMMA. ## Changes - Add `gfx1250` to both batched utils' supported-arch tuples. - gfx1250 WMMA `16x16x32` CI configs: `batched_gemm/configs/default_ci_config_gfx1250.json`; `batched_contraction/configs/default_ci_config_gfx1250.json` + `bridge_default_ci_config_gfx1250.json`. - `--arch`/config help in both `*_full_benchmark.py`. - CPU tests (gfx1250 arch + configs) in `test_batched_bridge.py` / `test_batched_contraction_bridge.py`. ## Validation — MI400 / gfx1250 (a07-1, ROCm 10.1, gpu0) - **batched_gemm** (fp16/rcr, --verify): **32/32 Successful, 0 failed** (verify ~5e-4). - **batched_contraction**: `bridge` config **1/1 verified** (max_rel 2.6e-4); `default` config **8/12 built, 8/8 verified** (the 4 non-builds are CompV3/CompV4+interwave pipeline incompatibilities, not gfx1250 runtime bugs). - 66 batched CPU tests pass. Refs: merged #10612 (#9306 + #9328); shared gfx1250 enablement also in #10921. JIRA ID : AICK-2104 JIRA ID : AICK-2136 --------- Co-authored-by: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com>
ISSUE ID: #8997
Motivation
The TileEngine → Dispatcher bridge could generate and launch single-problem GEMM, but
had no path for batched GEMM, where the same GEMM is run across many independent
problems sharing a batch dimension with per-batch strides. Batched GEMM is a core
Old-TE op used throughout inference, so the dispatcher needed it to reach feature
parity. This PR adds a complete bridge so Python callers can generate, build, and
launch batched GEMM at parity with the legacy Tile Engine
batched_gemmop.Batched GEMM has a genuinely different ABI from single-problem GEMM (it carries
batch_count, split-K viak_batch, and three batch strides), so it follows theregistry-bypass pattern used by the stream-K bridge (#9028). The capability set
matches Old-TE exactly:
fp16,rcrlayout, plus the batch dimension with per-batchstrides.
Test Plan
python3 -m pytest dispatcher/tests/test_batched_bridge.py -vdefault_configbatched sweep and confirm 0 failures.GpuBatchedGemmRunner) withname-parity and correctness across batch counts 1/2/4/8.
k_batch > 1) against a NumPyreference.
Test Result
default_configbatched sweep: 6672 kernels, 0 failures.correctness across batch counts 1/2/4/8 at
max_rel ~5e-4.max_rel = 4.7e-4(identical to packed).k_batch = 2):max_rel = 5.7e-4.expand_sweepnow yields 8 batched fp16/rcr configs (was 0).batched_gemm_ctypes_lib.cpp.stems × 5 shapes = 40 rows, interleaved, fair 50/100/flush/rotating both sides):
at parity / slightly ahead — median gap +4.00%, mean +4.13%, 100% within
±15% (range [-1.15%, +8.20%]; positive = bridge faster). See the parity comment.
Related PRs / references (TileEngine → Dispatcher GEMM bridge series): #8997 (regular GEMM fp16/bf16 all-layout), #9000 (grouped GEMM), #9028 (stream-K), #8887 (fp8/bf8/int8). This PR is a sibling in the same bridge effort tracked across those PRs.