Skip to content

Commit 85ea951

Browse files
ozturkosuassistant-librarian[bot]
ozturkosu
authored and
assistant-librarian[bot]
committed
[rocm-libraries] ROCm/rocm-libraries#9329 (commit f66c060)
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
1 parent 5305cdd commit 85ea951

5 files changed

Lines changed: 1856 additions & 0 deletions

File tree

‎dispatcher/bindings/ctypes/CMakeLists.txt‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
#
1010
# Targets:
1111
# - dispatcher_gemm_lib : GEMM dispatcher library
12+
# - dispatcher_mx_gemm_lib : MX-GEMM (microscaling) dispatcher library (gfx950)
1213
# - dispatcher_conv_lib : Convolution dispatcher library (forward + bwd_data)
1314
# - dispatcher_conv_bwdw_lib : Convolution backward weight library
1415
# - gpu_helper : GPU helper executable for Python
@@ -71,6 +72,81 @@ else()
7172
target_link_libraries(dispatcher_gemm_lib PRIVATE hip::device)
7273
endif()
7374

75+
# =============================================================================
76+
# MX-GEMM ctypes Library (microscaling GEMM, gfx950/MI350 only)
77+
# =============================================================================
78+
#
79+
# Force-includes a generated mx_gemm kernel header (SelectedKernel/KERNEL_NAME/
80+
# MxGemmHostArgs/ScaleType/...). The lib also includes the Old-TE common helper
81+
# (common/utils.hpp) for the free-function is_row_major(Layout), so add the
82+
# tile_engine/ops include roots.
83+
84+
set(MX_GEMM_TE_INCLUDE_DIRS
85+
${PROJECT_SOURCE_DIR}/tile_engine/ops
86+
${PROJECT_SOURCE_DIR}/tile_engine/ops/gemm
87+
${PROJECT_SOURCE_DIR}/tile_engine/ops/gemm/mx_gemm
88+
)
89+
90+
# Resolve the configured GPU target GENERICALLY -- do NOT hardcode/default the
91+
# arch to gfx950. mx_gemm is gfx950-only (it calls the gfx950-only
92+
# preShuffleScaleBuffer_gfx950 helper), and that requirement is enforced IN THE
93+
# SOURCE: mx_gemm_ctypes_lib.cpp #errors if GFX_ARCH is unset and static_asserts
94+
# GFX_ARCH == "gfx950" (plus a runtime device guard). Here we simply pass the
95+
# real build target through as -DGFX_ARCH and only wire up the target when the
96+
# build actually targets gfx950, so a gfx942/other dispatcher build is neither
97+
# broken by the static_assert nor silently built for an unsupported arch.
98+
if(DEFINED GPU_TARGETS AND NOT GPU_TARGETS STREQUAL "")
99+
string(REPLACE ";" " " _mx_gpu_targets_space "${GPU_TARGETS}")
100+
string(REPLACE " " ";" _mx_gpu_targets_list "${_mx_gpu_targets_space}")
101+
list(GET _mx_gpu_targets_list 0 MX_GEMM_GPU_TARGET)
102+
elseif(DEFINED CK_TILE_GEMM_GPU_TARGET AND NOT CK_TILE_GEMM_GPU_TARGET STREQUAL "")
103+
set(MX_GEMM_GPU_TARGET "${CK_TILE_GEMM_GPU_TARGET}")
104+
else()
105+
set(MX_GEMM_GPU_TARGET "")
106+
endif()
107+
108+
if(NOT MX_GEMM_GPU_TARGET STREQUAL "gfx950")
109+
message(STATUS
110+
"MX-GEMM ctypes lib is gfx950-only; skipping for GPU target "
111+
"'${MX_GEMM_GPU_TARGET}' (configure with GPU_TARGETS=gfx950 to enable).")
112+
else()
113+
file(GLOB MX_GEMM_KERNEL_HEADERS "${CMAKE_BINARY_DIR}/generated_kernels/mx_gemm_*.hpp")
114+
if(MX_GEMM_KERNEL_HEADERS)
115+
list(GET MX_GEMM_KERNEL_HEADERS 0 MX_GEMM_KERNEL_HEADER)
116+
message(STATUS "Found MX-GEMM kernel for ctypes lib: ${MX_GEMM_KERNEL_HEADER}")
117+
118+
add_ctypes_library(dispatcher_mx_gemm_lib
119+
mx_gemm_ctypes_lib.cpp
120+
KERNEL_HEADER ${MX_GEMM_KERNEL_HEADER}
121+
)
122+
target_include_directories(dispatcher_mx_gemm_lib PRIVATE ${MX_GEMM_TE_INCLUDE_DIRS})
123+
# The generated header only exports SelectedKernel/KERNEL_NAME/ScaleType/...
124+
# under this define. Pass the resolved build target through as -DGFX_ARCH
125+
# (the source #errors if unset and static_asserts it is gfx950).
126+
target_compile_definitions(dispatcher_mx_gemm_lib PRIVATE
127+
CK_TILE_SINGLE_KERNEL_INCLUDE
128+
GFX_ARCH="${MX_GEMM_GPU_TARGET}"
129+
)
130+
else()
131+
message(STATUS "No MX-GEMM kernel found for ctypes lib - building without kernel")
132+
add_library(dispatcher_mx_gemm_lib SHARED mx_gemm_ctypes_lib.cpp)
133+
target_include_directories(dispatcher_mx_gemm_lib PRIVATE
134+
${PROJECT_SOURCE_DIR}/include
135+
${PROJECT_SOURCE_DIR}/dispatcher/include
136+
${MX_GEMM_TE_INCLUDE_DIRS}
137+
)
138+
target_link_libraries(dispatcher_mx_gemm_lib PRIVATE hip::device)
139+
# Pass the resolved build target through as -DGFX_ARCH (the source #errors
140+
# if unset and static_asserts it is gfx950).
141+
target_compile_definitions(dispatcher_mx_gemm_lib PRIVATE
142+
GFX_ARCH="${MX_GEMM_GPU_TARGET}")
143+
set_target_properties(dispatcher_mx_gemm_lib PROPERTIES
144+
POSITION_INDEPENDENT_CODE ON
145+
CXX_STANDARD 17
146+
)
147+
endif()
148+
endif()
149+
74150
# =============================================================================
75151
# Convolution ctypes Library (supports forward + bwd_data)
76152
# =============================================================================

0 commit comments

Comments
 (0)