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.
|
This bridge currently defaults GFX_ARCH to gfx950 and unconditionally uses preShuffleScaleBuffer_gfx950(). Can we instead make it generic and throw an error if the arch_type is not gfx950 for now? |
…default to a specific arch
Architecture handling: removed silent arch defaults (get_arch + validation)Addressed the review feedback — "don't default to gfx942. Use get_arch and throw an error for unsupported arch types." — applied to the mx_gemm bridge's own files only. What changed
Why Verification Commit: |
… preShuffleScaleBuffer_gfx950
…eShuffleScaleBuffer_gfx950
Done!! two commits:
On "make it generic": When an arch-generic |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 5 out of 5 changed files in this pull request and generated 5 comments.
Comments suppressed due to low confidence (2)
projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt:104
- mx_gemm_ctypes_lib.cpp requires GFX_ARCH to be defined (#error) and then static_asserts it equals gfx950. The CMake target never defines GFX_ARCH, so dispatcher_mx_gemm_lib will fail to compile in both the kernel and no-kernel branches.
target_compile_definitions(dispatcher_mx_gemm_lib PRIVATE
CK_TILE_SINGLE_KERNEL_INCLUDE
)
projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt:112
- The no-kernel fallback branch still compiles mx_gemm_ctypes_lib.cpp, which hard-requires GFX_ARCH (#error). Without adding a compile definition here too, the fallback target will not build when no generated header is present.
add_library(dispatcher_mx_gemm_lib SHARED mx_gemm_ctypes_lib.cpp)
target_include_directories(dispatcher_mx_gemm_lib PRIVATE
${PROJECT_SOURCE_DIR}/include
${PROJECT_SOURCE_DIR}/dispatcher/include
${MX_GEMM_TE_INCLUDE_DIRS}
)
…rch propagation in setup_multiple, grid-map contract)
…quant, parity bench defaults Addresses tough-review findings #2, #3, #5: - #2: dispatcher_run_mx_gemm now reuses persistent device allocations across calls (opt-out via CK_TILE_MX_GEMM_REUSE_BUFFERS=0), eliminating per-call hipMalloc/hipFree churn that slowed throughput sweeps and perturbed thermal steady state. A/B/scales are still uploaded and C still zeroed every call, so per-shape data is never stale; cache is freed in dispatcher_cleanup(). - #3: dequantize_fp8 replaced its interpreted per-element loop with a 256-entry LUT gather (consistent with dequantize_fp4_packed); foreign bytes map to NaN. - #5: bench warmup/repeat defaults raised to 50/100 to match the bridge parity-sweep convention for apples-to-apples bridge-vs-Old-TE timing.
Follow-up review fixes (commit
|
…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
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