feat(ck-tile): batched GEMM TE to dispatcher bridge#9306
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
|
Addressed the review feedback "don't default to gfx942 — use get_arch and throw an error for unsupported arch types" in this PR's own batched_gemm files only (shared files left untouched). C++ — #ifndef GFX_ARCH
#define GFX_ARCH "gfx942"
#endifwith a hard compile-time requirement: #ifndef GFX_ARCH
#error "GFX_ARCH must be defined at compile time (pass -DGFX_ARCH=<arch>); do not default to a specific GPU architecture."
#endifPython —
Note: the shared parent field Verified: |
Address Copilot review on the batched_gemm TE->Dispatcher bridge: - unified_gemm_codegen.py: the generated split-K batched launch ignored the hipMemsetAsync return. A failed C reset would let the atomic-add epilogue accumulate into a non-zero buffer and silently produce wrong results. Now checks the return and throws (mirrors the stream-K path). - batched_gemm_full_benchmark.py: --batch-size is used as the range() step over built kernels; a value of 0 raised ValueError and a negative value silently produced no batches. Reject non-positive values after parsing with a clear message.
Address the senior-review findings on the batched_gemm TE->Dispatcher bridge: - F1 (dtype/layout parity): expand_sweep now rejects any dtype != fp16 or layout != rcr. Old-TE batched_gemm_instance_builder declares --datatype choices=["fp16"] / --layout choices=["rcr"], so anything else would codegen/compile/launch a signature Old-TE never validated (claimed parity != exercised parity). Fail up front with a clear message. - F3 (observability): every HIP failure in dispatcher_run_batched now logs the failing call + hipGetErrorString before returning, and the launch catch surfaces the exception message (returns -2) instead of a bare -1. - F5 (tests): added GPU-free unit tests for _get_arch/_resolve_arch (detect, undetectable->RuntimeError, unsupported->ValueError, explicit passthrough), the ctypes ABI (19 argtypes + positional-order marshalling, status forwarding), and the fp16/rcr expand_sweep gate. 22/22 pass, no GPU. - F7 (overflow): reject any M/N/K/stride/batch value exceeding INT32_MAX before the static_casts into BatchedGemmHostArgs (ck_tile::index_t is 32-bit), so an out-of-range problem fails loudly instead of truncating.
Review follow-ups addressed (commit
|
Close the remaining senior-review findings on the batched_gemm bridge: - F2 (split-K correctness): the reduction contract is confirmed and now locked by a CPU test. For k_batch>1 the ck_tile batched kernel uses a memory_operation_enum::atomic_add epilogue, so C is zeroed pre-launch (and the memset return is checked/throws); for k_batch==1 the epilogue is `set` and C is not pre-zeroed (the Old-TE byte-identical path). New codegen test asserts both halves so a dropped memset fails here, not only on GPU. - F4 (ABI naming): documented that dispatcher_initialize is the canonical entry the Python wrapper binds and dispatcher_init is only a back-compat alias; fixed the usage-doc drift that referenced dispatcher_init(). - F6 (buffer sizing): the host numpy C dtype size must equal sizeof(CDataType) (the buffer is memcpy'd verbatim). Added a module-level _C_SIZEOF map, an assertion in run() so a bad _C_NP edit fails loudly, and a test tying the numpy itemsize to the declared kernel size. Suite: 26/26 CPU unit tests pass, no GPU required.
Remaining review findings addressed (commit
|
…9684) The dispatcher GEMM bridge dropped ALL non-power-of-2 MRepeat/NRepeat cshuffle configs. GPU verification on gfx942 shows only an ODD per-wave repeat (>1) paired with a 32-wide warp tile is numerically wrong (26 configs, e.g. tile_m=192 / wave_m=2 / warp_tile_m=32, MRepeat=3). The other 90 non-power-of-2-repeat cshuffle configs verify correct (incl. MRepeat=3 with warp_tile_m=16, and even repeats 6/12) and were needlessly excluded. Narrow the gate accordingly in gemm_utils.expand_sweep (via new _cshuffle_store_ok helper) and unified_gemm_codegen._cshuffle_repeat_ok, restoring the 90 valid configs to bridge coverage while still excluding the 26 broken ones. Add unit tests for the gate. Refs #9684
Follow-up: narrowed the CShuffle-store gate (commit bb85a0d)While auditing this PR's whole-
What changedNarrowed the gate to that exact signature so the bridge stops dropping the 90 valid configs while still excluding the 26 broken ones:
Verification
Related
|
Copilot review follow-ups on the batched_gemm bridge: - codegen: _epilogue_code() for the BATCHED variant unconditionally emitted the CShuffle epilogue, ignoring config.trait.epilogue. The batched sweep space includes both "cshuffle" and "default" (see default_config.json), so a "default"-requested config was silently built as cshuffle under an "_epi_default" name -- a config->codegen->runtime name/impl mismatch that also skipped the cshuffle-store correctness gate. It now emits DefaultGemm2DEpilogue for "default" and CShuffleEpilogue for "cshuffle", matching the standard path. Verified on GPU (gfx942, ROCm 7.2): both epilogues build as distinct kernels and match the numpy reference (max_rel=0). - ctypes doc: dispatcher_run_batched's comment said it returns only 0/-1, but it also returns -2 on a kernel-launch exception. Doc corrected.
Whole-config perf parity — batched GEMM bridge vs Old Tile-Engine (MI300X / gfx942)Extends the earlier fixed-CI (8-stem) sweep to the whole Headline (k_batch=1): AT PARITY.
gap = (bridge − old)/old × 100 (+ = bridge faster). Headline shapes: 4096³ and 1024×512×256 (1050 rows each, both 100% within ±15%). A native ROCm-7.2 single-GPU cross-check (cx71, 577 rows) agrees: median −0.451%, 97.7% within ±15%. Split-K (k_batch>1): reported separately — bridge is faster (median +27%, correctness verified max_rel ~5e-4), not a regression. Old-TE is slow on small-shape split-K and rejects default-epilogue split-K (1178 Old-TE NaN/rejects), so this is a bridge capability win rather than an apples-to-apples parity pair. Fairness: A/B interleaved median-3, warmup=50/repeat=100/flush/rotating=1000 both sides; 4-GPU sharded (rack31) with staged ROCm 7.2 runtime; same kernel set both sides; Old-TE from its own CMake; |gap|>15% outliers re-measured standalone median-5 (none survived). Build-time note (fair 334-kernel sample): bridge 1.22 vs Old-TE 1.04 s/kernel (~17%). Codegen is not the cause (both TUs #include the same 106-header umbrella and instantiate the same device kernel); the residual is attributable to the bridge's |
## Summary Adds a microscaling-GEMM (`mx_gemm`) bridge from the Old Tile-Engine into the dispatcher's ctypes path for gfx950/MI350. Supports **fp4** (`pk_fp4_t`, e2m1) and **fp8** (e4m3), **rcr** layout, `comp_async` + `cshuffle` + `intrawave`, fixed `16x16x128` warp tile, `k_batch == 1`. Per-32-K `e8m0` block scales are pre-shuffled on-host exactly as the Old-TE profiler does. ## Motivation Extend the TE→Dispatcher bridge family to the microscaling GEMM op so the dispatcher can launch byte-identical `mx_gemm` kernels with block scaling. Sibling to the other bridge PRs (see below). ## Design note - **Byte-exact kernels:** the codegen reuses Old-TE `MxGemmKernelBuilder._generate_kernel_instance` directly rather than re-implementing header assembly, guaranteeing the emitted kernel matches Old-TE. It only strips the stale `ck_tile/ops/gemm_mx.hpp` umbrella include (absent on develop; the mx pipeline is pulled in via `ck_tile/ops/gemm.hpp`). - **Registry bypass:** `mx_gemm`'s `launch` takes `ck_tile::MxGemmHostArgs` with per-32-K `e8m0` scales that the generic dispatcher backend cannot express, so the ctypes lib builds `MxGemmHostArgs` from plain C arrays and calls `SelectedKernel::launch()` directly — the same direct-launch pattern as the batched/multi-D bridges. - **Scale pre-shuffle** mirrors `mx_gemm_profiler.hpp`; pack params are derived from `SelectedKernel` tile dims at compile time (not hardcoded). - **Packing-correct byte accounting:** device buffers are sized via `HostTensor::get_element_space_size_in_bytes()`, which divides by `numeric_traits<T>::PackedSize`, so fp4 (`PackedSize==2`, two e2m1 per byte) and fp8 (`PackedSize==1`) are both correct. ## Verification - fp8 derisk: PASS (`max_rel = 0.0`, kernel-name match, output fully non-zero) on gfx950 - fp4: GPU-verified on gfx950 - `clang-format-18 -style=file` clean; Python `py_compile` clean ## Sibling PRs - #8997 — TE→dispatcher GEMM bridge (fp16/bf16, all layouts) — foundational (merged) - #8998 — fp8/bf8/int8 bridge (merged) - #9000 — grouped GEMM - #9028 — stream-K GEMM - #9305 — multi-ABD GEMM - #9306 — batched GEMM - #9307 — preshuffle GEMM - #9308 — multi-D GEMM - #9328 — batched-contraction GEMM --------- Co-authored-by: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com>
gfx950 / MI355X A/B parity + functional equivalence — AT PARITYCloses the missing arch (gfx942/MI300X already posted). Node smci355-ccs-aus-m07-13, gfx950, ROCm 7.2.0. Fresh build from PR HEAD Functional equivalence: Old-TE authoritative surface = Performance (40 measurements, 8 stems × 5 shapes, interleaved median-3):
gap = (bridge−old)/old; + = bridge slower. warmup=50/repeat=100/flush/rotating=1000 both sides. Old-TE from its own CMake ( Caveats (flagged):
Split-K note (orthogonal, k_batch>1 is outside Old-TE's k_batch=1 surface): CPU self-test |
…bridge_batched_gemm # Conflicts: # projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt
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.