Skip to content

feat(ck-tile): microscaling (mx) GEMM TE to dispatcher bridge - #9329

Merged
ozturkosu merged 14 commits into
developfrom
users/muozturk/ck/bridge_mx_gemm
Jul 22, 2026
Merged

ozturkosu merged 14 commits into
developfrom
users/muozturk/ck/bridge_mx_gemm

Conversation

@ozturkosu

@ozturkosu ozturkosu commented Jul 13, 2026 •

Copy link
Copy Markdown
Contributor

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, 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

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
@therock-pr-bot

therock-pr-bot Bot commented Jul 13, 2026 •

Copy link
Copy Markdown

✅ All Checks Passed — Ready for Review

Check Status Details
🌿 Branch Name ✅ Pass —
📝 PR Title/Description ✅ Pass —
⛔ Forbidden Files ✅ Pass —
🧪 Unit Test ✅ Pass —
🔎 pre-commit ✅ Pass —
🚫 Draft PR 🔜 To Be Enabled —
🚩 Feature Flag 🔜 To Be Enabled —
📊 Code Coverage 🔜 To Be Enabled —
🤖 therock-pr-bot ✅ Pass —

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

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

@therock-pr-bot

therock-pr-bot Bot commented Jul 13, 2026 •

Copy link
Copy Markdown

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

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.py to generate one force-included, Old-TE-byte-exact kernel header per config.
  • Add mx_gemm_ctypes_lib.cpp implementing a direct-launch C ABI, including host-side scale pre-shuffle matching Old-TE behavior.
  • Add mx_gemm_utils.py providing kernel config naming, build pipeline (codegen→hipcc), and a Python runner + reference implementation.
  • Wire a new dispatcher_mx_gemm_lib CMake 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.
@ozturkosu

Copy link
Copy Markdown
Contributor Author

🔎 Automated parity review — Round 1 (mx_gemm bridge vs Old-TE)

Reviewed HEAD 3119eeee25c (clean tree). CPU unit tests: 21/21 pass.

Verdict: AT PARITY (functional/codegen)

The bridge supports exactly Old-TE's real compilable mx_gemm set: {fp4, fp8} × rcr × 16x16x128 × comp_async/cshuffle/intrawave × persistent{false,true} × k_batch=1, gfx950. Kernels are byte-identical by direct reuse of MxGemmKernelBuilder._generate_kernel_instance; ABI byte-accounting (fp4 PackedSize=2 via get_element_space_size_in_bytes), e8m0 construction (byte 127→1.0), scale pre-shuffle, and MxGemmHostArgs ordering all match mx_gemm_profiler.hpp line-for-line.

Parity gaps vs Old-TE's real (compilable) set: none.

All bridge restrictions are independently confirmed non-compilable in Old-TE:

  • bf8/e5m2: rejected by Old-TE argparse (choices=["fp4","fp8"]) + validate_gemm_mx; no bf8_* warp-tile combo. OUT-OF-SCOPE.
  • non-rcr layouts: validate_gemm_mx returns False for anything but rcr. OUT-OF-SCOPE.
  • other warp tiles / pipelines / schedulers / default epilogue: MX warp-tile map has exactly 16x16x128; TRAIT_UNSUPPORTED_COMBINATIONS collapses to comp_async/cshuffle/intrawave; default epilogue raises. OUT-OF-SCOPE.
  • split-K (k_batch>1): plumbed but never exercised/validated by Old-TE configs. OUT-OF-SCOPE for v1.

Code correctness findings (all LOW / non-blocking)

  • Change shared libs name #3 (recommend fix): mx_gemm_ctypes_lib.cpp:199-211 scale-shuffle sizing assumes M/N/scale_k divisible by the xdl pack factors (mirrors profiler, but neither validates). Add a guard if (M%m_xdl_pack || N%n_xdl_pack || scale_k%k_xdl_pack) return -1; to fail loudly instead of silently corrupting.
  • [hipDNN] Enablement #2 (recommend fix): fp8 codec (mx_gemm_utils.py:83-89,348-361) assumes OCP e4m3 unconditionally; fp8_t is gated on CK_TILE_USE_OCP_FP8. Default on gfx950 = OCP so bytes match, but assert/derive it (or document as hard precondition) so an FNUZ build can't silently disagree with the numpy reference.
  • Initial workflows #1/[hipSPARSELt] monorepo test #5 (informational): bridge default tile_k=128 vs Old-TE default_config tile_k=256 — both valid, just different default shapes; codegen accepts any valid tile.
  • Verified CORRECT: _fix_includes stripping the nonexistent gemm_mx.hpp umbrella, byte accounting, e8m0, pre-shuffle (kLast=true both A & B), numpy reference orientation, kernel-name determinism (shared --list-name).

Test coverage gaps

  • No test asserts bf8 is rejected by is_valid(); no codegen-path persistent/k_batch rejection test.

Requested changes for Round 2 (hardening — no functional gap remains)

  1. Add the M/N/scale_k divisibility guard in the ctypes lib (finding Change shared libs name #3).
  2. Make the fp8 OCP/FNUZ assumption explicit (finding [hipDNN] Enablement #2).
  3. Add a bf8-rejected unit test.
    Then GPU numeric verify (fp4+fp8) and the A/B timing parity run follow.

…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).
@ozturkosu

Copy link
Copy Markdown
Contributor Author

✅ Automated parity review — Round 2 verdict: AT PARITY

Round-1 hardening implemented and independently verified on local gfx950 (local commit 4a5b4639c2c).

Verified fixes

  • Divisibility guard (mx_gemm_ctypes_lib.cpp): returns -1 if M%m_xdl_pack || N%n_xdl_pack || scale_k%k_xdl_pack before scale-shuffle sizing; pack-factor expressions confirmed byte-identical to the sizing they protect (no off-by-one). Proven live: M=127 → prints diagnostic, rc=-1.
  • fp8 OCP/FNUZ precondition (mx_gemm_utils.py): fp8_ocp_is_default_for_arch() mirrors ck_tile/core/config.hpp exactly; assertion wired into GpuMxGemmRunner for fp8. No spurious fire on gfx950/gfx12; correctly raises on FNUZ archs (gfx942/gfx90a). Build uses no explicit -DCK_TILE_USE_OCP_FP8, so OCP is the effective gfx950 default.
  • bf8-rejected test: genuinely asserts (import path not skipped) — both is_valid() False and codegen _validate raises.
  • GPU correctness (gfx950): fp4 max_rel 0.0, fp8 max_rel 0.0 (grid-exact), and a non-unit e8m0 scale (2^2) fp8 case also 0.0 (refmax 2940) — parity holds through the scale-decode/accumulate path.
  • Tests: 22/22 pass. No regression to the byte-exact codegen/kernel-name path (codegen file untouched).

Remaining parity gaps vs Old-TE's compilable set: none.

Bridge covers exactly {fp4,fp8} × rcr × 16x16x128 × comp_async/cshuffle/intrawave × persistent × k_batch=1; bf8/non-rcr/other-tiles/split-K all independently confirmed non-compilable in Old-TE.

Verdict: full functional parity reached. Proceeding to the apples-to-apples A/B interleaved timing parity check on gfx950.

@ozturkosu

Copy link
Copy Markdown
Contributor Author

📊 A/B interleaved timing parity — gfx950/MI350 — VERDICT: AT PARITY

Apples-to-apples bridge-vs-Old-TE sweep on local gfx950, fresh build from HEAD 4a5b4639c2c, byte-identical kernels on both sides.

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%.

  • Correctness (vs numpy microscaled reference): fp8 max_rel 0.0, fp4 max_rel 0.0.
dtype M×N×K bridge_us oldte_us gap%
fp8 4096³ 138.13 144.70 −4.54
fp8 2048³ 20.05 21.00 −4.52
fp8 1024³ 11.11 11.10 +0.09
fp4 4096³ 80.00 79.20 +1.01
fp4 2048³ 14.93 15.00 −0.47
fp4 1024³ 8.77 8.80 −0.34

(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 -mllvm flag set on both bridge .so and Old-TE exe; same generated kernel header force-included into both (tile 128x128x128 / warp 2x2x1 / 16x16x128, comp_async/cshuffle/intrawave — resolves tile_k=128 vs 256 by pinning tile_k=128 identically → literally the same kernel name on both sides); stale-.so guard; interleaved A/B per (dtype,shape) median-of-5.

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).

@ozturkosu
ozturkosu marked this pull request as ready for review July 15, 2026 16:18
@ozturkosu
ozturkosu requested a review from a team as a code owner July 15, 2026 16:18

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 4 comments.

- 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.
@ozturkosu

ozturkosu commented Jul 16, 2026 •

Copy link
Copy Markdown
Contributor Author

1. CMake no-kernel fallback wouldn't compile (mx_gemm_ctypes_lib.cpp, CMakeLists) — went with option (a): the TU now compiles without a force-included kernel. KERNEL_NAME and the dispatcher_run_mx_gemm body are guarded behind CK_TILE_SINGLE_KERNEL_INCLUDE; the no-kernel build returns "" for the kernel name and -2 (unsupported) from run, so the CMake fallback branch builds a valid stub instead of erroring on SelectedKernel/MxGemmHostArgs/ScaleType.

2. _make_builder temp-dir leak (unified_mx_gemm_codegen.py) — converted _make_builder into a @contextlib.contextmanager backed by tempfile.TemporaryDirectory; kernel_name()/_generate() now use with _make_builder(cfg) as builder: so the mx_gemm_codegen_* work dir is removed after each (frequent --list-name) invocation.

3. Per-element quantize loops (mx_gemm_utils.py) — quantize_fp8/quantize_fp4_packed now share a vectorized _map_grid_to_bytes helper: round to grid resolution, then a single np.searchsorted into a sorted LUT (no Python loop). Off-grid values raise (preserving the exact-grid contract) and -0.0 still collapses to the 0.0 byte.

4. Scale-buffer divisibility → OOB (mx_gemm_ctypes_lib.cpp) — already handled in the prior hardening commit: M % m_xdl_pack, N % n_xdl_pack, scale_k % k_xdl_pack are validated before the shuffled buffers are sized. It returns -1 to match this file's existing error convention (rather than -2); happy to switch to -2 if you'd prefer the "unsupported args" code.

5. float_to_e8m0 silently mapping non-positive → 1.0 (mx_gemm_utils.py) — now validates strictly-positive, finite input and raises ValueError instead of substituting 1.0.

6. run() silently copying a non-contiguous output C (mx_gemm_utils.py) — MxGemmDispatcherLib.run no longer ascontiguousarray-copies C; it validates that C is a contiguous, writeable ndarray and raises otherwise, so an output view can't silently be dropped. (A/B remain inputs and are still normalized.)

Added CPU-only unit tests in test_mx_gemm_bridge.py for the new behaviors (e8m0 non-positive/non-finite raise, fp8/fp4 off-grid raise, -0.0 collapse, shape preservation) — full suite 28 passed.

ozturkosu and others added 3 commits July 16, 2026 23:44
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.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 10 comments.

…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.
@ozturkosu

ozturkosu commented Jul 22, 2026 •

Copy link
Copy Markdown
Contributor Author

@yraparti your note on defaulting GFX_ARCH to gfx950 and unconditionally using preShuffleScaleBuffer_gfx950(): fixed so the arch is resolved generically and the code throws if it is not gfx950 details in the inline replies. Summary: source #errors if GFX_ARCH unset + static_assert(GFX_ARCH=="gfx950") + runtime device guard; CMake now passes the resolved build target (-DGFX_ARCH=${MX_GEMM_GPU_TARGET} from GPU_TARGETS), not a literal gfx950; Python _get_arch() detects and raises if != gfx950. Commit 1a6765e.

@ozturkosu

Copy link
Copy Markdown
Contributor Author

GPU verification @ latest commit 1a6765ea14 (gfx950 / MI350X, ROCm 7.1)

Built and ran the bridge end-to-end on a gfx950 node (asrock-1w300-e2-3), one fp8 and one fp4 kernel, comparing GPU output against the numpy microscaled reference:

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::DeviceMem alloc/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 source static_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).

@ozturkosu

Copy link
Copy Markdown
Contributor Author

Comprehensive whole-default_config parity — mx_gemm bridge vs Old-TE (MI350 / gfx950)

Ran the full mx_gemm default_config kernel space (not the prior fixed-config subset) on
MI350X / gfx950, ROCm 7.1.0. The config space was enumerated from the op's own
mx_gemm_instance_builder (tile × trait): 144 valid stems per dtype (72 valid tile configs
× persistent {false,true}). Both sides were built from the same generated device-kernel header
with the same real TE -mllvm flag set (byte-identical device kernel), so any gap reflects
launch/host path only.

Headline: AT PARITY across the whole config space.

dtype comparisons median gap% within ±5% within ±15%
fp8 510 +0.45 95.7% 100.0%
fp4 720 +0.16 100.0% 100.0%
all 1230 +0.23 98.2% 100.0%

gap% = (bridge_tflops − oldte_tflops)/oldte_tflops × 100 (positive = bridge faster).

Fairness checklist (all enforced):

  • warmup=50 / repeat=100 both sides; -timer=false -flush_cache=false -rotating_count=1 both.
  • Old-TE built via its real TE -mllvm flag set (force-included the SAME generated header into
    mx_gemm_benchmark_single.cpp) — byte-identical device kernel to the bridge .so.
  • Timing normalized on full-precision tflops (not rounded latency ms).
  • A/B interleaved per (stem, shape), alternating order, median-of-3.
  • Stale-.so/exe guard (artifact mtime ≥ its generated header).
  • Zero |gap|>15% rows to remeasure — every kernel is within ±15%.
  • Correctness (numpy microscaled reference, max_abs) PASS for every stem.

Flag-parity note: on ROCm 7.1.0 the -amdgpu-coerce-illegal-types=1 probe passes (it failed on
7.2 in earlier runs), so it is applied — but SYMMETRICALLY to both sides, so parity is preserved.
The rest of the flag set (-amdgpu-early-inline-all, -amdgpu-function-calls=false,
--lsr-drop-solution=1, -enable-post-misched=0, -fno-offload-uniform-block, --offload-compress)
is byte-identical to build_oldte.sh.

Coverage: fp8 = 102/144 stems buildable on both sides; fp4 = 144/144. The 42 fp8 build failures
are large-tile configs (e.g. 256×256×256, 192×128×256) whose LDS use (163904 B) exceeds the gfx950
limit (163840 B). The device kernel is byte-identical on both sides, so these fail on Old-TE too —
reported as coverage-only, NOT a bridge regression.

@ozturkosu
ozturkosu merged commit f66c060 into develop Jul 22, 2026
26 checks passed
@ozturkosu
ozturkosu deleted the users/muozturk/ck/bridge_mx_gemm branch July 22, 2026 19:17
assistant-librarian Bot pushed a commit to ROCm/composable_kernel that referenced this pull request Jul 22, 2026
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
ThruptiRajLakshmanaGowda added a commit that referenced this pull request Jul 24, 2026
… 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>
ThruptiRajLakshmanaGowda added a commit that referenced this pull request Jul 27, 2026
…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>
shumway pushed a commit to ROCm/composable_kernel that referenced this pull request Aug 18, 2026
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>
ThruptiRajLakshmanaGowda added a commit that referenced this pull request Aug 19, 2026
…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>
ozturkosu added a commit that referenced this pull request Oct 21, 2026
…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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants