Repository navigation
feat(ck-tile): microscaling (mx) GEMM TE to dispatcher bridge - #9329
Conversation
Summary: Adds a microscaling-GEMM (mx_gemm) bridge from the Old Tile-Engine into the dispatcher's ctypes path for gfx950/MI350. fp4 (pk_fp4_t) and fp8 (e4m3), rcr layout, comp_async + cshuffle + intrawave, fixed 16x16x128 warp tile, k_batch==1. Motivation: Extend the TE->Dispatcher bridge family to the microscaling GEMM op so the dispatcher can launch byte-identical mx_gemm kernels with per-32-K e8m0 block scales. Design note: Codegen reuses Old-TE MxGemmKernelBuilder._generate_kernel_instance directly (no re-implemented header assembly) to guarantee byte-exact kernels. The ctypes lib bypasses the generic registry and calls SelectedKernel::launch() with a hand-built MxGemmHostArgs, mirroring mx_gemm_profiler.hpp for the gfx950 scale pre-shuffle. Byte accounting uses HostTensor get_element_space_size_in_bytes() so fp4 (PackedSize==2) and fp8 are both packing-correct. Changes: - dispatcher/codegen/unified_mx_gemm_codegen.py: per-config .hpp generator - dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp: direct-launch C API - dispatcher/python/mx_gemm_utils.py: config/ctypes/runner + build pipeline - dispatcher/bindings/ctypes/CMakeLists.txt: dispatcher_mx_gemm_lib target
✅ 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. |
There was a problem hiding this comment.
Pull request overview
This PR adds a TileEngine → Dispatcher bridge for microscaling GEMM (mx_gemm) on gfx950/MI350, spanning codegen (per-kernel header emission), a ctypes C ABI library (direct-launch / registry bypass with scale preshuffle), and a Python utility layer for build + run + NumPy reference.
Changes:
- Add
unified_mx_gemm_codegen.pyto generate one force-included, Old-TE-byte-exact kernel header per config. - Add
mx_gemm_ctypes_lib.cppimplementing a direct-launch C ABI, including host-side scale pre-shuffle matching Old-TE behavior. - Add
mx_gemm_utils.pyproviding kernel config naming, build pipeline (codegen→hipcc), and a Python runner + reference implementation. - Wire a new
dispatcher_mx_gemm_libCMake target.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
| projects/composablekernel/dispatcher/python/mx_gemm_utils.py | Python build + ctypes wrapper + input generation/reference helpers for mx_gemm |
| projects/composablekernel/dispatcher/codegen/unified_mx_gemm_codegen.py | Codegen script that reuses Old-TE builder to emit per-config mx_gemm headers |
| projects/composablekernel/dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp | C ABI library that builds MxGemmHostArgs, pre-shuffles scales, and launches SelectedKernel |
| projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt | Adds dispatcher_mx_gemm_lib target and generated-kernel discovery logic |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Adds dispatcher/tests/test_mx_gemm_bridge.py covering the config name format, codegen-JSON projection, dtype/layout/warp-tile validity gate, e8m0 scale codec, fp8/fp4 quantization round-trips, and the numpy microscaled reference. Satisfies the repo unit-test policy for the mx_gemm bridge source files. No GPU required.
🔎 Automated parity review — Round 1 (mx_gemm bridge vs Old-TE)Reviewed HEAD Verdict: AT PARITY (functional/codegen)The bridge supports exactly Old-TE's real compilable mx_gemm set: Parity gaps vs Old-TE's real (compilable) set: none.All bridge restrictions are independently confirmed non-compilable in Old-TE:
Code correctness findings (all LOW / non-blocking)
Test coverage gaps
Requested changes for Round 2 (hardening — no functional gap remains)
|
…ion, bf8-rejected test Round-2 review hardening for the mx_gemm TE->Dispatcher bridge (no functional gap remained; these are LOW-severity robustness fixes): - mx_gemm_ctypes_lib.cpp: guard M/N/scale_k divisibility by the xdl pack factors before sizing the shuffled scale buffers; return -1 (matching the existing error convention) instead of silently truncating the buffers and feeding the kernel a corrupt scale layout. - mx_gemm_utils.py: make the fp8 OCP e4m3 assumption explicit. Add fp8_ocp_is_default_for_arch()/assert_fp8_ocp_supported() mirroring config.hpp (OCP is the gfx950/gfx12 device default; the lib compiles without an explicit -DCK_TILE_USE_OCP_FP8), document the codec precondition, and have GpuMxGemmRunner assert it for fp8 so an FNUZ arch cannot silently diverge from the numpy reference. - test_mx_gemm_bridge.py: add test_bf8_rejected asserting is_valid() and the codegen _validate() both refuse dtype="bf8" (Old-TE does not compile bf8).
✅ Automated parity review — Round 2 verdict: AT PARITYRound-1 hardening implemented and independently verified on local gfx950 (local commit Verified fixes
Remaining parity gaps vs Old-TE's compilable set: none.Bridge covers exactly Verdict: full functional parity reached. Proceeding to the apples-to-apples A/B interleaved timing parity check on gfx950. |
📊 A/B interleaved timing parity — gfx950/MI350 — VERDICT: AT PARITYApples-to-apples bridge-vs-Old-TE sweep on local gfx950, fresh build from HEAD Result: median gap −0.40%, mean −1.28%, range −4.54%..+1.01%, 100% of rows within ±5%, zero |gap|>15% outliers. Per-dtype median: fp8 −2.32%, fp4 −0.09%.
(12 rows total across fp8/fp4 × {512³,1024³,2048³,4096³, 1024×512×256, 4096×4096×1024}.) Fairness (all enforced): warmup=50/repeat=100 both sides; identical measurement mode (CPU timer, no cache flush, rotating_count=1); identical TE Combined with the Round-2 functional verdict, PR #9329 is at full functional + timing parity with Old-TE across its entire compilable set (fp4/fp8 × rcr × 16x16x128 × comp_async/cshuffle/intrawave × k_batch=1, gfx950). |
- mx_gemm_ctypes_lib.cpp: make the TU compile in the CMake "no-kernel" fallback
branch. Guard the kernel-dependent code (KERNEL_NAME, dispatcher_run_mx_gemm
body) behind CK_TILE_SINGLE_KERNEL_INCLUDE; without a force-included kernel the
lib now returns "" for the name and -2 (unsupported) from run, instead of
failing to compile on SelectedKernel/MxGemmHostArgs/ScaleType.
- unified_mx_gemm_codegen.py: stop leaking mx_gemm_codegen_* temp dirs. Turn
_make_builder into a context manager backed by TemporaryDirectory so the work
dir is cleaned up after each name/code generation (kernel_name/generate are
shelled frequently via --list-name).
- mx_gemm_utils.py:
* float_to_e8m0: validate strictly-positive, finite input and raise instead of
silently mapping non-positive scales to 1.0 (wrong scale byte).
* quantize_fp8/quantize_fp4_packed: replace per-element Python loops with a
vectorized sorted-grid LUT + searchsorted (off-grid values now raise). Removes
the O(M*K) Python overhead that made the bridge look slow in sweeps.
* MxGemmDispatcherLib.run: C is an in-place output buffer; validate it is a
contiguous, writeable ndarray and raise rather than silently ascontiguousarray-
copying a non-contiguous view (which would drop the caller's results).
- test_mx_gemm_bridge.py: cover the new behaviors (e8m0 non-positive/non-finite
raise, fp8/fp4 off-grid raise, -0.0 collapse, shape preservation).
Divisibility guard (scale-buffer OOB) was already handled in the prior hardening
commit; it returns -1 to match this file's existing error convention.
|
1. CMake no-kernel fallback wouldn't compile ( 2. 3. Per-element quantize loops ( 4. Scale-buffer divisibility → OOB ( 5. 6. Added CPU-only unit tests in |
Two optional polish items from the round-1 reviewer (both LOW severity): 1. Divisibility guard returns -2 instead of -1 The xdl-pack divisibility guard in dispatcher_run_mx_gemm() previously returned -1 (bad-args convention). -2 is semantically more accurate for "shape unsupported by the selected warp tile", consistent with the IsSupportedArguments-style rejects used elsewhere in the ctypes layer. Update the comment and the function docstring to match. 2. fp8 OCP/FNUZ guard added at MxGemmDispatcherLib.run() level assert_fp8_ocp_supported() was only called in GpuMxGemmRunner.__init__. A direct caller of MxGemmDispatcherLib.run() with fp8 on an FNUZ arch would bypass it and silently diverge from the numpy reference. Add optional dtype/arch params to MxGemmDispatcherLib.__init__(); if dtype=="fp8", call assert_fp8_ocp_supported at run() entry. Backward compatible (both params default to None; existing call sites unchanged). All 28 CPU tests pass; no test asserted on the specific -1 return code.
…ernel The bridge runtime kernel compile path _compile_kernel built the byte-identical mx_gemm device kernel with only "-O3 -std=c++17", while Old-TE compiles the same kernel with the CK global AMDGPU codegen flags (C++20 + -mllvm tuning + -fno-offload-uniform-block + --offload-compress). The missing flags changed inlining/register allocation and occupancy, making bridge-built fp8 kernels run ~30% slower than Old-TE in production (700 TF vs 1076 TF at 4096^3 on a byte-identical 128x128 fp8 kernel). Add the exact flag set from projects/composablekernel/CMakeLists.txt (add_compile_options: -std=c++20, -mllvm -amdgpu-early-inline-all=true, -mllvm -amdgpu-function-calls=false, -mllvm --lsr-drop-solution=1, -mllvm -enable-post-misched=0, -mllvm -amdgpu-coerce-illegal-types=1, -fno-offload-uniform-block) plus mx_gemm/CMakeLists.txt's --offload-compress, matching the pattern already used by the gemm bridge's _tile_engine_codegen_flags. This brings fp8 to parity (measured +4.8% at 4096^3 on the same kernel).
…ain safety
The previous commit added -mllvm -amdgpu-coerce-illegal-types=1 unconditionally,
but that LLVM option is rejected by clang on ROCm 7.2 ("Unknown command line
argument"), hard-failing the kernel compile. CK's CMake only adds that flag when
check_cxx_compiler_flag passes; the shipping gemm bridge mirrors this via a probe.
Restructure _compile_kernel to use the same pattern as
gemm_utils._tile_engine_codegen_flags: an unconditional core set
(-amdgpu-early-inline-all, -amdgpu-function-calls, --lsr-drop-solution,
-enable-post-misched, -fno-offload-uniform-block, --offload-compress) plus a
probe-gated -amdgpu-coerce-illegal-types=1 added only when hipcc accepts it.
Keeps the bridge matched to Old-TE on toolchains that support the flag and
buildable on those that do not. Verified: builds on ROCm 7.2/gfx950 and brings
fp8 128x128 to parity with Old-TE at 4096^3.
…sts/cache) - CMake: define GFX_ARCH="gfx950" on both dispatcher_mx_gemm_lib branches. mx_gemm_ctypes_lib.cpp #errors when GFX_ARCH is unset, so the CMake build path (kernel and no-kernel fallback) failed to compile; gfx950 is the only arch this op supports (static_assert on preShuffleScaleBuffer_gfx950). - Python arch validation: restrict _get_arch()/setup_multiple_mx_gemm_dispatchers to gfx950 only, matching the C++ static_assert. A gfx942/gfx90a caller no longer passes detection only to fail later with a less actionable error. - Tests: pin gpu_target="gfx950" at the three to_codegen_config() call sites so the CPU-only suite never shells out to rocminfo on non-ROCm runners. - Thread-safety: remove the g_buf_cache device-buffer reuse path entirely and go back to per-call DeviceMem alloc/free (as the batched/multi_abd bridges do). The cache was a shared global with no lock; ctypes.CDLL releases the GIL so concurrent callers could race across upload->launch->D2H. Removing it (vs adding a mutex) also drops an unrequested, untested optimization. CPU tests: 27 passed, 1 skipped (codegen import unavailable off-node).
…ed gfx950)
Per reviewer feedback ("no hardcoding to gfx950 arch"; "make it generic and
throw if the arch is not gfx950"): the mx_gemm ctypes CMake no longer passes a
literal GFX_ARCH="gfx950". It now resolves the configured GPU target from
GPU_TARGETS (falling back to CK_TILE_GEMM_GPU_TARGET) and passes it through as
-DGFX_ARCH="${MX_GEMM_GPU_TARGET}". The gfx950-only requirement stays enforced
in the source (mx_gemm_ctypes_lib.cpp #errors if GFX_ARCH is unset and
static_asserts GFX_ARCH == "gfx950", plus a runtime device guard) and in Python
(_get_arch detects generically and raises if != gfx950). The target is only
wired up when the build targets gfx950, so a gfx942/other dispatcher build is
neither broken by the static_assert nor silently built for an unsupported arch.
|
@yraparti your note on defaulting GFX_ARCH to gfx950 and unconditionally using |
GPU verification @ latest commit
|
| dtype | time | max_rel vs reference | result |
|---|---|---|---|
| fp8 | 0.0063 ms | 0.000 | PASS |
| fp4 | 0.0057 ms | 0.000 | PASS |
Kernel: mx_gemm_{fp8,fp4}_rcr_comp_async_cshuffle_intrawave_..._128x128x128_2x2x1_16x16x128, shape M=N=K=256. max_rel=0 is exact (inputs are drawn from a grid exact in both fp8 e4m3 and fp4 e2m1), so this is a clean functional pass.
This confirms the latest commit — after the review fixes — still builds and runs correctly:
- Buffer-cache removal: the reverted per-call
ck_tile::DeviceMemalloc/free path produces correct output (no all-zeros, non-negative timing). - Generic arch resolution: the build resolved gfx950 via
_get_arch()(rocminfo) and compiled with the passed-DGFX_ARCH; the sourcestatic_assert(GFX_ARCH=="gfx950")compiled clean (no hardcoded default). - fp8 and fp4 both correct end-to-end.
Scope note: this is a functional/correctness smoke test of the latest commit, not the full default_config A/B perf-parity sweep (that gfx950 gate remains outstanding and will run on the MI355X 8-GPU allocation).
Comprehensive whole-default_config parity — mx_gemm bridge vs Old-TE (MI350 / gfx950)Ran the full mx_gemm Headline: AT PARITY across the whole config space.
gap% = (bridge_tflops − oldte_tflops)/oldte_tflops × 100 (positive = bridge faster). Fairness checklist (all enforced):
Flag-parity note: on ROCm 7.1.0 the Coverage: fp8 = 102/144 stems buildable on both sides; fp4 = 144/144. The 42 fp8 build failures |
feat(ck-tile): microscaling (mx) GEMM TE to dispatcher bridge (#9329) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## 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
… and grouped_gemm_tensorquant Implements the ctypes dispatcher bridge for two new quant GEMM operators, following the same pattern as the mxgemm bridge (PR #9329) and bquant. Each bridge consists of four layers: - C++ ctypes library: extern "C" API calling SelectedKernel::launch( vector<QuantGroupedGemmHostArgs>, stream_config, kargs_ptr) directly, bypassing the dispatcher registry (incompatible with grouped quant args) - Python codegen: generates per-kernel .hpp headers with the correct QuantType (RowColQuant / TensorQuant) and GemmSpatiallyLocalTilePartitioner - Python utils: KernelConfig (byte-exact name), DispatcherLib (ctypes wrapper), GpuRunner (high-level numpy interface), setup_multiple_*_dispatchers - CPU-only tests: 55 tests covering name generation, config round-trips, header codegen correctness; all passing RowColQuant: per-row A scales [M,1], per-col B scales [1,N], QK_A=M, QK_B=N. TensorQuant: single scalar scale per tensor, QK_A=QK_B=1. CMakeLists.txt updated with dispatcher_grouped_gemm_rowcolquant_lib and dispatcher_grouped_gemm_tensorquant_lib targets. Co-Authored-By: Claude <noreply@anthropic.com>
…emm operators Add registry-bypass, single-kernel-per-.so TE→Dispatcher bridges for all five block_scale_gemm operators, following the pattern established by the merged grouped_gemm_bquant bridge (PR #9166) and mxgemm bridge (PR #9329). Each operator gets four files: - codegen/unified_gemm_{op}_codegen.py — per-kernel .hpp generator - bindings/ctypes/gemm_{op}_ctypes_lib.cpp — C ABI calling SelectedKernel::launch(QuantGemmHostArgs) - python/gemm_{op}_utils.py — KernelConfig / DispatcherLib / GpuRunner / build helpers - tests/test_gemm_{op}_utils.py — 164 CPU-only unit tests total, all passing Operators and QuantType mappings: gemm_bquant → QuantType::BQuantGrouped (gemm_bquant_* naming, block_scale_gemm) gemm_aquant → QuantType::AQuantGrouped (AQ=[M,QK_A], BQ=nullptr) gemm_abquant → QuantType::ABQuantGrouped (AQ row-major, BQ col-major [QK_B,QN_B]) gemm_rowcolquant → QuantType::RowColQuant (AQ=[M,1], BQ=[1,N], no preshuffle) gemm_tensorquant → QuantType::TensorQuant (AQ and BQ single scalars) Key implementation notes: - All ctypes libs exclude dispatcher/include to avoid QuantGemmHostArgs conflicts - rowcolquant TileGemmQuantTraits needs extra UsePersistentKernel+16 args - tensorquant uses BaseWeightPreshufflePipelineAGmemBGmemCRegV2 (not BaseGemmPipelineAgBgCrCompV3) - gemm_bquant uses "default" epilogue (DefaultGemm2DEpilogue), distinct from grouped_gemm_bquant Also adds: - codegen_common.py: make_gemm_bquant/aquant/abquant/rowcolquant/tensorquant_kernel_name helpers - CMakeLists.txt: 5 new double-guarded CMake targets (dispatcher_gemm_{op}_lib) Co-Authored-By: Claude <noreply@anthropic.com>
feat(ck-tile): microscaling (mx) GEMM TE to dispatcher bridge ## 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>
…ops (#10132) ISSUE ID: #8997 ## Summary Adds **on-device GPU correctness tests** for the merged bridge operators that had **no in-tree GPU test**: grouped (#9000), multi-ABD (#9305), multi-D (#9308), mx (#9329). Closes part of the GPU-test coverage gap (11 of 15 bridge ops previously shipped only CPU-only unit tests). ## What each test does Builds the op's dispatcher `.so`, runs it **on-device** via the op's `Gpu*Runner`, and compares output to an independent **numpy fp32 reference** with a scale-relative tolerance (`max|err| / max|ref|`). Real math — reference is `A@B` per group / `A@B + D0 + D1` / the runner's built-in multi-ABD verify. No trivial always-pass asserts. They **SKIP cleanly** (not fail) when no GPU / hipcc / static-lib is present, so they add **zero cost on a CPU-only runner** and only execute on a GPU runner. ## Verification (MI300X / gfx942) - grouped — PASS (max_rel 7.6e-3 bf16 / 4.8e-4 fp16) - multi-ABD — PASS (6.4e-4) - multi-D — PASS (1.8e-3) - mx — **gfx950-only**, so it arch-gates and SKIPs on gfx942; verify-pending on a gfx950 runner. Registered in `dispatcher/tests/CMakeLists.txt` under label `dispatcher;python;bridge;gpu`. ## Scope Merged ops only (their runners are on develop). The 7 open bridge ops (batched/preshuffle/batched-contraction + 5 block-scale) get their GPU tests in their own PR branches. ## Test plan - [x] grouped / multi-ABD / multi-D PASS on gfx942 - [x] all SKIP cleanly with no GPU (CPU gate stays green) - [ ] mx verified on a gfx950 runner - [ ] CI runs the new `dispatcher_*_gpu_correctness` entries on a GPU runner --------- Co-authored-by: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com> Co-authored-by: Thrupti Raj Lakshmana Gowda <thruptiraj.lakshmanagowda@amd.com>
…atcher bridge (#13968) ## Summary Adds `mx_gemm_full_benchmark.py`, a JSON-driven sweep driver for microscaling (MX) GEMM through the TE to dispatcher bridge. It is the MX counterpart of `gemm_full_benchmark.py`. The Dispatcher stays the single source of truth for codegen, build and runtime. The Tile Engine keeps only the config search space (the same `mx_gemm/configs/*.json` the Old-TE CMake build reads) and the benchmark loop. Until now MX GEMM had the bridge Python API (#9329, #10916) but no sweep driver, so sweeping MXFP4/MXFP8 from a config JSON meant the Old-TE CMake flow. ## Why MX needs its own driver - The kernel space comes from `MxGemmKernelBuilder` (MX pipelines per arch, no K padding or persistence on gfx1250), not from `gemm_utils.expand_sweep`. - Inputs are packed fp4 or fp8 bytes with one e8m0 scale per 32 K elements, so the worker generates scale tensors and checks against the microscaled numpy reference instead of `A @ B`. - Kernels reject unsupported shapes at run time (rc -2, e.g. K not a multiple of tile_k). The worker reports these as skipped, not failed. ## Design note `expand_mx_sweep()` enumerates with the Old-TE `MxGemmKernelBuilder._get_sampled_kernel_list()`. That is the list the CMake `--list_kernels` step builds, including `--max-instances` sampling. A bridge sweep therefore runs exactly the kernels of an Old-TE build of the same JSON: all 472 kernels of `default_config.json` and `default_config_gfx1250.json` (fp4 and fp8) have byte-identical names. To avoid a second copy of the multi-GPU Phase 3 loop, `gemm_full_benchmark.py` now exposes `unique_built_kernels()`, `benchmark_on_devices()` and `print_summary()`. Its own `main()` uses them, and the MX driver imports them. `_run_batch_on_device()` gained an `extra` payload dict and a `skipped` result state; the regular GEMM worker never emits `skipped`, so its behaviour is unchanged. ## Changes **New** - `tile_engine/ops/gemm/mx_gemm_full_benchmark.py`: driver. Phase 1 expands the JSON and compiles the .so files in parallel, Phase 2 loads problems, Phase 3 runs on all GPUs. Supports `--dtype fp4|fp8`, `--verify`, `--max-instances/--seed`, `--output-dir` (build cache) and `--compile-only`. - `tile_engine/ops/gemm/run_one_mx_gemm_kernel.py`: disposable per-GPU worker. Uses random e8m0 scales (2^-3 to 2^2) and the microscaled reference check. - `mx_gemm/configs/bridge_ci_config_gfx950.json`, `bridge_ci_config_gfx1250.json`: small sweeps covering every MX pipeline of the arch (gfx950: 16 fp4 / 9 fp8 kernels; gfx1250: 13 fp4 / 11 fp8). These are the driver defaults. - `mx_gemm/configs/example_problems.json`: 5 tile-aligned M/N/K problems. **Modified** - `dispatcher/python/mx_gemm_utils.py`: adds `expand_mx_sweep()`. - `tile_engine/ops/gemm/gemm_full_benchmark.py`: shared Phase 3 helpers (no behaviour change). - `dispatcher/tests/test_mx_gemm_bridge.py`: Old-TE name parity, per-arch pipeline coverage, sampling, problem/tile fit, and worker verdicts (exact, wrong, unsupported). - `dispatcher/tests/test_gemm_driver_verify.py`: includes the new driver in the exit-code check; updates the 4-tuple mock. - `dispatcher/docs/mx_gemm.md`: usage section. Related bridge PRs: #8997 (regular GEMM), #9000 (grouped), #9028 (stream-K), #8887 (fp8/bf8/int8), #9329 (MX bridge), #10916 (MX gfx1250 pipelines). ## Capability | Axis | Supported | |---|---| | dtype | MXFP4 (e2m1, packed), MXFP8 (OCP e4m3); e8m0 scales per 32 K; fp16 output | | layout | rcr | | arch | gfx950 (comp_async, comp_async_eight_waves, weight_preshuffle), gfx1250 (+ comp_tdm, comp_tdm_v2) | ## Validation (gfx950, 8 GPUs) | Run | Kernels | Measurements | Result | |---|---|---|---| | CPU tests (`test_mx_gemm_bridge.py`, `test_gemm_driver_verify.py`, `test_gemm_driver_fp32.py`) | - | - | 207 passed | | `bridge_ci_config_gfx950.json` fp4 `--verify` | 16 | 80 | 80 verified, max_rel 0 | | `bridge_ci_config_gfx950.json` fp8 `--verify` | 9 | 45 | 45 verified | | Old-TE `default_ci_config.json` fp4 / fp8 `--verify` | 2 / 2 | 10 / 10 | all verified | | Regular `gemm_full_benchmark.py` fp16 `--verify --max-kernels 4` (refactor check) | 4 | 20 | all verified | | Old-TE `default_config.json` fp4 `--verify` | 194 (all built) | 970 | 928 verified, 4 mismatch, 42 errors (see below) | gfx1250 has CPU coverage only (expansion, names, pipeline coverage); it has not been run on hardware yet. ## Open items (why this is a draft) The full fp4 sweep found two groups of failures that I'm still investigating: 1. **rc=-1 for 14 kernels** (comp_async 256x64 tiles and weight_preshuffle 192x512) on the 2048^3, 4096^3 and 4096x2048x8192 problems. Run alone in a fresh worker, the same kernel/problem passes and verifies, so this looks like interaction between kernels batched in one worker process, not a kernel defect. 2. **Wrong results for comp_async_eight_waves 128x128x128 4x2x1** (max_rel 0.2 to 0.33) at K = 4096 and 8192; K <= 2048 passes. Next step is to check whether the Old-TE build of the same kernel gives the same mismatch (kernel issue) or not (bridge host-side issue). gfx1250 hardware validation is also pending. --------- 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: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
JIRA ID : AICK-2088
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, fixed16x16x128warp tile,k_batch == 1. Per-32-Ke8m0block 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_gemmkernels with block scaling. Sibling to the other bridge PRs (see below).Design note
MxGemmKernelBuilder._generate_kernel_instancedirectly rather than re-implementing header assembly, guaranteeing the emitted kernel matches Old-TE. It only strips the staleck_tile/ops/gemm_mx.hppumbrella include (absent on develop; the mx pipeline is pulled in viack_tile/ops/gemm.hpp).mx_gemm'slaunchtakesck_tile::MxGemmHostArgswith per-32-Ke8m0scales that the generic dispatcher backend cannot express, so the ctypes lib buildsMxGemmHostArgsfrom plain C arrays and callsSelectedKernel::launch()directly — the same direct-launch pattern as the batched/multi-D bridges.mx_gemm_profiler.hpp; pack params are derived fromSelectedKerneltile dims at compile time (not hardcoded).HostTensor::get_element_space_size_in_bytes(), which divides bynumeric_traits<T>::PackedSize, so fp4 (PackedSize==2, two e2m1 per byte) and fp8 (PackedSize==1) are both correct.Verification
max_rel = 0.0, kernel-name match, output fully non-zero) on gfx950clang-format-18 -style=fileclean; Pythonpy_compilecleanSibling PRs