diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt b/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt index ce13a7d8e70c..624afaee9f38 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt +++ b/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt @@ -48,6 +48,37 @@ function(add_ctypes_library TARGET_NAME SOURCE_FILE) ) endfunction() +# Feature definitions for the two single-target grouped quant bridges below. +# Keep these aligned with dispatcher_common.arch_feature_defines(). Parent CK +# builds enable features directory-wide for the union of their GPU targets. +# Clear those definitions for these bridges before applying the selected arch, +# so host/device FP8 encoding and the backend match the standalone JIT path. +function(set_grouped_quant_arch_features TARGET_NAME GFX_ARCH) + # Compile options follow directory/target definitions on the command line. + # Keep both the undefines and replacement defines here to preserve ordering + # without changing the inherited definitions of any sibling target. + foreach(_feature IN ITEMS + CK_USE_XDL CK_USE_GFX94 CK_USE_GFX950 CK_GFX1030_SUPPORT + CK_USE_WMMA CK_TILE_USE_WMMA CK_USE_WMMA_FP8 + CK_USE_OCP_FP8 CK_TILE_USE_OCP_FP8 CK_USE_FNUZ_FP8 + CK_USE_NATIVE_MX_SUPPORT CK_GFX950_SUPPORT + CK_USE_GFX1250 CK_GFX1250_SUPPORT CK_GFX12_SUPPORT + USE_NEW_UNIFIED_FRAMEWORK) + target_compile_options(${TARGET_NAME} PRIVATE "-U${_feature}") + endforeach() + if(GFX_ARCH STREQUAL "gfx950" OR GFX_ARCH STREQUAL "gfx1250") + target_compile_options(${TARGET_NAME} PRIVATE + -DCK_USE_OCP_FP8 -DCK_TILE_USE_OCP_FP8 -DCK_USE_NATIVE_MX_SUPPORT) + endif() + if(GFX_ARCH STREQUAL "gfx950") + target_compile_options(${TARGET_NAME} PRIVATE -DCK_GFX950_SUPPORT) + elseif(GFX_ARCH STREQUAL "gfx1250") + target_compile_options(${TARGET_NAME} PRIVATE + -DCK_TILE_USE_WMMA=1 -DCK_GFX12_SUPPORT -DCK_USE_GFX1250 -DCK_GFX1250_SUPPORT + -DUSE_NEW_UNIFIED_FRAMEWORK=0) + endif() +endfunction() + # Helper function for block-scale quant ctypes targets. # # Usage: @@ -354,8 +385,18 @@ if(ROWCOLQUANT_GEMM_KERNEL_HEADERS) set(CK_TILE_ROWCOLQUANT_GFX_ARCH "${_rowcolquant_gfx_arch}" CACHE STRING "GFX arch for the CMake-built RowColQuant ctypes .so (default: first CMAKE_HIP_ARCHITECTURES or gfx942)") endif() + # Normalize both inferred and explicitly supplied cache values. + string(REGEX REPLACE ":.*$" "" CK_TILE_ROWCOLQUANT_GFX_ARCH "${CK_TILE_ROWCOLQUANT_GFX_ARCH}") message(STATUS "RowColQuant ctypes lib GFX_ARCH: ${CK_TILE_ROWCOLQUANT_GFX_ARCH}") + # Standalone dispatcher builds do not inherit the top-level CK definitions. + # Override the directory-wide ID for this single-target library. + string(REGEX REPLACE "^gfx" "0x" _rowcolquant_target_id "${CK_TILE_ROWCOLQUANT_GFX_ARCH}") + target_compile_options(dispatcher_grouped_gemm_rowcolquant_lib PRIVATE + -UCK_CMAKE_GPU_TARGET_IDS "-DCK_CMAKE_GPU_TARGET_IDS=${_rowcolquant_target_id}") + set_grouped_quant_arch_features(dispatcher_grouped_gemm_rowcolquant_lib + "${CK_TILE_ROWCOLQUANT_GFX_ARCH}") + target_compile_definitions(dispatcher_grouped_gemm_rowcolquant_lib PRIVATE CK_TILE_SINGLE_KERNEL_INCLUDE GFX_ARCH="${CK_TILE_ROWCOLQUANT_GFX_ARCH}" @@ -365,7 +406,10 @@ else() "No RowColQuant GEMM kernel header found in ${CMAKE_BINARY_DIR}/generated_kernels/. " "dispatcher_grouped_gemm_rowcolquant_lib will NOT be built. " "Run unified_grouped_gemm_rowcolquant_codegen.py --output-dir " - "${CMAKE_BINARY_DIR}/generated_kernels/ before configuring to enable this target.") + "${CMAKE_BINARY_DIR}/generated_kernels/ --gfx-arch before configuring " + "to enable this target. --gfx-arch matters: the built-in default is the gfx9 " + "MFMA tile, and on gfx1250 the ctypes source static_asserts that tile away " + "rather than let it compile into a library that returns wrong results.") # Do not define the target: the source requires force-included symbols (SelectedKernel, # ADataType, KERNEL_NAME, etc.) that only exist in the generated header. Without it # the build fails at compile time, not load time, breaking any all-target build. @@ -400,8 +444,18 @@ if(TENSORQUANT_GEMM_KERNEL_HEADERS) set(CK_TILE_TENSORQUANT_GFX_ARCH "${_tensorquant_gfx_arch}" CACHE STRING "GFX arch for the CMake-built TensorQuant ctypes .so (default: first CMAKE_HIP_ARCHITECTURES or gfx942)") endif() + # Normalize both inferred and explicitly supplied cache values. + string(REGEX REPLACE ":.*$" "" CK_TILE_TENSORQUANT_GFX_ARCH "${CK_TILE_TENSORQUANT_GFX_ARCH}") message(STATUS "TensorQuant ctypes lib GFX_ARCH: ${CK_TILE_TENSORQUANT_GFX_ARCH}") + # Standalone dispatcher builds do not inherit the top-level CK definitions. + # Override the directory-wide ID for this single-target library. + string(REGEX REPLACE "^gfx" "0x" _tensorquant_target_id "${CK_TILE_TENSORQUANT_GFX_ARCH}") + target_compile_options(dispatcher_grouped_gemm_tensorquant_lib PRIVATE + -UCK_CMAKE_GPU_TARGET_IDS "-DCK_CMAKE_GPU_TARGET_IDS=${_tensorquant_target_id}") + set_grouped_quant_arch_features(dispatcher_grouped_gemm_tensorquant_lib + "${CK_TILE_TENSORQUANT_GFX_ARCH}") + target_compile_definitions(dispatcher_grouped_gemm_tensorquant_lib PRIVATE CK_TILE_SINGLE_KERNEL_INCLUDE GFX_ARCH="${CK_TILE_TENSORQUANT_GFX_ARCH}" @@ -411,7 +465,10 @@ else() "No TensorQuant GEMM kernel header found in ${CMAKE_BINARY_DIR}/generated_kernels/. " "dispatcher_grouped_gemm_tensorquant_lib will NOT be built. " "Run unified_grouped_gemm_tensorquant_codegen.py --output-dir " - "${CMAKE_BINARY_DIR}/generated_kernels/ before configuring to enable this target.") + "${CMAKE_BINARY_DIR}/generated_kernels/ --gfx-arch before configuring " + "to enable this target. --gfx-arch matters: the built-in default is the gfx9 " + "MFMA tile, and on gfx1250 the ctypes source static_asserts that tile away " + "rather than let it compile into a library that returns wrong results.") # Do not define the target: the source requires force-included symbols (SelectedKernel, # ADataType, KERNEL_NAME, etc.) that only exist in the generated header. Without it # the build fails at compile time, not load time, breaking any all-target build. diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_rowcolquant_ctypes_lib.cpp b/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_rowcolquant_ctypes_lib.cpp index 0f0b4a6b6819..c21596c4408c 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_rowcolquant_ctypes_lib.cpp +++ b/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_rowcolquant_ctypes_lib.cpp @@ -74,7 +74,7 @@ static std::atomic g_ref_count{0}; // Architectures this bridge is known to work on. These fp8/bf8 CompV3 kernels need // native FP8, so gfx90a is deliberately absent -- it compiles but produces NaN. // Enabling a new target is a one-line addition here plus a CMake arch entry. -static constexpr const char* kSupportedArchs[] = {"gfx942", "gfx950"}; +static constexpr const char* kSupportedArchs[] = {"gfx942", "gfx950", "gfx1250"}; // True if `arch` starts with any entry of kSupportedArchs. Prefix-matching, because // hipDeviceProp_t::gcnArchName carries feature suffixes (e.g. "gfx942:sramecc+:xnack-"). @@ -88,6 +88,60 @@ static bool is_supported_arch(const std::string& arch) return false; } +// --------------------------------------------------------------------------- +// Generated-target and tile guards. Codegen records an explicit target in the +// header; a build for another processor must fail before it can run that tile. +// Headers generated without an explicit target retain the gfx9 default behavior. +static constexpr bool ct_starts_with(const char* s, const char* prefix) +{ + return *prefix == '\0' ? true : (*s == *prefix && ct_starts_with(s + 1, prefix + 1)); +} + +static constexpr bool ct_equal(const char* a, const char* b) +{ + return *a == *b && (*a == '\0' || ct_equal(a + 1, b + 1)); +} + +static_assert(SelectedKernel::GfxArch[0] == '\0' || ct_equal(SelectedKernel::GfxArch, GFX_ARCH), + "Generated kernel architecture does not match GFX_ARCH. Regenerate the " + "header with --gfx-arch matching the build target."); + +static constexpr bool kCompiledForGfx1250 = ct_starts_with(GFX_ARCH, "gfx1250"); + +// Family-wide companion for the M%4 guard below. That defect is an unpadded AQ view in +// the shared quant kernel header rather than a gfx1250-specific silicon property, so +// every gfx12xx part is subject to it. Compile-time, so on gfx9 the guard folds away +// instead of constructing a std::string per call to test a constant. +static constexpr bool kCompiledForGfx12 = ct_starts_with(GFX_ARCH, "gfx12"); + +static_assert(!kCompiledForGfx1250 || SelectedKernel::WarpTileM == 16, + "gfx1250 has no 32x32 WMMA fragment: warp_tile_m must be 16. This kernel " + "would compile and then return wrong results on the device."); + +static_assert(!kCompiledForGfx1250 || SelectedKernel::WarpTileK == 64 || + SelectedKernel::WarpTileK == 128, + "gfx1250 8-bit WMMA fragments are 16x16x64 and 16x16x128 only: " + "warp_tile_k must be 64 or 128."); + +static_assert(!kCompiledForGfx1250 || SelectedKernel::WarpTileN == 16, + "gfx1250 8-bit WMMA fragments are 16x16xK: warp_tile_n must be 16."); + +// warp_k > 1 is rejected separately from the warps-per-block cap below. The [1,2,2] +// map has a product of 4 and would pass that cap on its own, but it is measured to +// compile and then return wrong results (max_rel 1.37). It is the corrupting case, +// not merely an undependable one, so it gets its own rule. +static_assert(!kCompiledForGfx1250 || SelectedKernel::WarpPerBlock_K == 1, + "gfx1250: warp_k must be 1. A warp_k of 2 (the [1,2,2] map) compiles " + "and then returns wrong results."); + +static_assert(!kCompiledForGfx1250 || + (SelectedKernel::WarpPerBlock_M * SelectedKernel::WarpPerBlock_N * + SelectedKernel::WarpPerBlock_K) <= 4, + "gfx1250 is wave32: a block of more than four warps is not dependable " + "there. Paired with a legal tile in a 14,208-row sweep the 8-warp maps " + "aborted at launch 2,168 times and passed 192 times, with no wrong " + "answers -- unreliable rather than incorrect, so refused here."); + extern "C" { /** @@ -166,7 +220,21 @@ int dispatcher_init() { return dispatcher_initialize(); } * k_batch - split-K factor (1 = no split) * time_ms - output: kernel execution time in ms (may be NULL) * - * Returns 0 on success, negative on error. + * Shape constraints (validated on gfx1250 with the default config; violations are + * rejected with an explanatory message and no output is written): + * K % 16 == 0 - required even when the kernel was generated with pad_k=true. + * pad_k covers the K-loop tail, not the global-load vector width. + * N % TileN == 0 - required whenever the kernel was generated with pad_n=false. + * TileN is 64 in the default config. + * M % 4 == 0 - gfx12 targets only. The RowColQuant per-row A-scale tile window + * has an unpadded M tail that zeroes row M-1 of C. See the guard + * in the body for the full explanation; gfx942/gfx950 are not + * subject to this restriction. + * + * Returns 0 on success, negative on error: + * -1 argument rejected by this bridge (including the constraints above) + * -2 argument rejected by the kernel's own IsSupportedArgument() + * -3 kernel launch threw */ int dispatcher_run_gemm(const void* A, const void* B, @@ -224,6 +292,57 @@ int dispatcher_run_gemm(const void* A, return -1; } + // --------------------------------------------------------------------------- + // Documented shape constraints. + // + // These were established empirically on gfx1250 with the default config and are + // enforced here so a violation produces an explanation instead of either a bare + // "unsupported args" code or, worse, a wrong answer. + // --------------------------------------------------------------------------- + + // N must be a whole number of N-tiles unless the kernel pads N. This is derived + // from the kernel's own constants rather than hardcoded, so it stays correct if + // the tile or the padding trait changes. + if(!SelectedKernel::kPadN && (N % SelectedKernel::TileN) != 0) + { + std::cerr << "dispatcher_run_gemm: N must be a multiple of " << SelectedKernel::TileN + << " for this kernel (it was generated with pad_n=false, so N has to be a " + << "whole number of N-tiles); got N=" << N << "\n"; + return -1; + } + + // gfx12 only: the RowColQuant per-row A-scale path drops the final row of C when + // M is not a multiple of 4, writing zeros into row M-1 while leaving every other + // row correct to ~3e-4. The kernel returns success, so without this check the + // caller gets a silent wrong answer -- at M=1 the entire output is zero. + // + // Cause: MakeAQBlockWindow() builds the AQ view with + // make_naive_tensor_view(aq_ptr, {M, N}, {1, 0}, ...) + // and then, as its own comment says, creates the tile window with "no padding for + // AQ" -- unlike the A/B/C windows, which get right-pad transforms driven by + // kPadM/kPadN/kPadK. The M axis of that broadcast view therefore has an unguarded + // tail. The BQ view is the mirror image ({M, N} strided {0, 1}) and has the same + // unguarded tail on N; it is simply never exercised, because pad_n=false already + // forces N to a multiple of TileN (64) and hence of 4. + // + // The fix belongs in ck_tile's shared quant kernel header, not in this bridge: + // that header is common to gfx942/gfx950, where the behaviour has not been + // measured. Rejecting here is deliberately the narrow, reversible option -- it + // converts a silent wrong answer into a clean refusal without touching code paths + // that other architectures depend on. Remove this guard once the AQ tail is + // padded upstream. + if(kCompiledForGfx12 && (M % 4) != 0) + { + std::cerr << "dispatcher_run_gemm: M must be a multiple of 4 on " << GFX_ARCH + << " for RowColQuant; got M=" << M << ". The per-row A-scale (AQ) tile " + << "window is built without a padding transform, so row M-1 of C is " + << "written as zero while the remaining rows are correct. Refusing rather " + << "than returning a silently wrong result. Pad M up to a multiple of 4 " + << "and slice the result, or use the TensorQuant bridge, which applies " + << "scalar scales and is unaffected.\n"; + return -1; + } + // Only packed (contiguous) layouts are supported for A, B, C. // stride_AQ and stride_BQ are unused here; the kernel uses broadcast strides (0). // @@ -429,7 +548,17 @@ int dispatcher_run_gemm(const void* A, if(exec_time < 0.0f) { - std::cerr << "dispatcher_run_gemm: kernel reported unsupported args\n"; + // IsSupportedArgument() is a boolean inside ck_tile and gives no reason, so + // enumerate the constraints a caller is realistically violating. Observed on + // gfx1250 with the default config: K must be a multiple of 16 (this holds even + // with pad_k=true -- padding covers the K-loop tail, not the global-load vector + // width), and N a multiple of the N-tile when pad_n is false. + std::cerr << "dispatcher_run_gemm: kernel reported unsupported args for M=" << M + << " N=" << N << " K=" << K << " k_batch=" << k_batch + << ". Known constraints for this kernel: K must be a multiple of 16 " + << "(required even though pad_k=" << (SelectedKernel::kPadK ? "true" : "false") + << "); N must be a multiple of " << SelectedKernel::TileN + << " when pad_n is false. No output was written.\n"; cleanup(); return -2; } @@ -449,6 +578,23 @@ int dispatcher_run_gemm(const void* A, */ const char* dispatcher_get_kernel_name() { return KERNEL_NAME; } +/** + * Return the N-tile of the force-included kernel. + * + * The shape checks in dispatcher_run_gemm derive their bound from SelectedKernel::TileN + * so they survive a tile change. Exporting it lets the Python runner report the same + * number instead of hardcoding the value of today's default config. + */ +int dispatcher_get_tile_n() { return static_cast(SelectedKernel::TileN); } + +/** + * Return 1 when the force-included kernel was generated with pad_n=true. + * + * When padding is on the N % TileN constraint does not apply, so the Python + * runner needs this to phrase its diagnostics correctly. + */ +int dispatcher_get_pad_n() { return SelectedKernel::kPadN ? 1 : 0; } + // This bridge is one-.so-per-kernel by construction: the build force-includes exactly // one generated header via `hipcc -include `, giving one SelectedKernel. // Scaling to N kernels means N .so files (the pattern bquant/aquant/abquant follow), diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_tensorquant_ctypes_lib.cpp b/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_tensorquant_ctypes_lib.cpp index 321de7ef2e42..b495cb3b4749 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_tensorquant_ctypes_lib.cpp +++ b/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_tensorquant_ctypes_lib.cpp @@ -77,7 +77,7 @@ static std::atomic g_ref_count{0}; // Architectures this bridge is known to work on. These fp8/bf8 CompV3 kernels need // native FP8, so gfx90a is deliberately absent -- it compiles but produces NaN. // Enabling a new target is a one-line addition here plus a CMake arch entry. -static constexpr const char* kSupportedArchs[] = {"gfx942", "gfx950"}; +static constexpr const char* kSupportedArchs[] = {"gfx942", "gfx950", "gfx1250"}; // True if `arch` starts with any entry of kSupportedArchs. Prefix-matching, because // hipDeviceProp_t::gcnArchName carries feature suffixes (e.g. "gfx942:sramecc+:xnack-"). @@ -91,6 +91,54 @@ static bool is_supported_arch(const std::string& arch) return false; } +// --------------------------------------------------------------------------- +// Generated-target and tile guards. Codegen records an explicit target in the +// header; a build for another processor must fail before it can run that tile. +// Headers generated without an explicit target retain the gfx9 default behavior. +static constexpr bool ct_starts_with(const char* s, const char* prefix) +{ + return *prefix == '\0' ? true : (*s == *prefix && ct_starts_with(s + 1, prefix + 1)); +} + +static constexpr bool ct_equal(const char* a, const char* b) +{ + return *a == *b && (*a == '\0' || ct_equal(a + 1, b + 1)); +} + +static_assert(SelectedKernel::GfxArch[0] == '\0' || ct_equal(SelectedKernel::GfxArch, GFX_ARCH), + "Generated kernel architecture does not match GFX_ARCH. Regenerate the " + "header with --gfx-arch matching the build target."); + +static constexpr bool kCompiledForGfx1250 = ct_starts_with(GFX_ARCH, "gfx1250"); + +static_assert(!kCompiledForGfx1250 || SelectedKernel::WarpTileM == 16, + "gfx1250 has no 32x32 WMMA fragment: warp_tile_m must be 16. This kernel " + "would compile and then return wrong results on the device."); + +static_assert(!kCompiledForGfx1250 || SelectedKernel::WarpTileK == 64 || + SelectedKernel::WarpTileK == 128, + "gfx1250 8-bit WMMA fragments are 16x16x64 and 16x16x128 only: " + "warp_tile_k must be 64 or 128."); + +static_assert(!kCompiledForGfx1250 || SelectedKernel::WarpTileN == 16, + "gfx1250 8-bit WMMA fragments are 16x16xK: warp_tile_n must be 16."); + +// warp_k > 1 is rejected separately from the warps-per-block cap below. The [1,2,2] +// map has a product of 4 and would pass that cap on its own, but it is measured to +// compile and then return wrong results (max_rel 1.37). It is the corrupting case, +// not merely an undependable one, so it gets its own rule. +static_assert(!kCompiledForGfx1250 || SelectedKernel::WarpPerBlock_K == 1, + "gfx1250: warp_k must be 1. A warp_k of 2 (the [1,2,2] map) compiles " + "and then returns wrong results."); + +static_assert(!kCompiledForGfx1250 || + (SelectedKernel::WarpPerBlock_M * SelectedKernel::WarpPerBlock_N * + SelectedKernel::WarpPerBlock_K) <= 4, + "gfx1250 is wave32: a block of more than four warps is not dependable " + "there. Paired with a legal tile in a 14,208-row sweep the 8-warp maps " + "aborted at launch 2,168 times and passed 192 times, with no wrong " + "answers -- unreliable rather than incorrect, so refused here."); + extern "C" { /** @@ -165,7 +213,19 @@ int dispatcher_init() { return dispatcher_initialize(); } * k_batch - split-K factor (1 = no split) * time_ms - output: kernel execution time in ms (may be NULL) * - * Returns 0 on success, negative on error. + * Shape constraints (validated on gfx1250 with the default config; violations are + * rejected with an explanatory message and no output is written): + * K % 16 == 0 - required even when the kernel was generated with pad_k=true. + * pad_k covers the K-loop tail, not the global-load vector width. + * N % TileN == 0 - required whenever the kernel was generated with pad_n=false. + * TileN is 64 in the default config. + * There is no M constraint: TensorQuant applies scalar A/B scales rather than a + * per-row scale vector, and is correct at every M tested on gfx1250. + * + * Returns 0 on success, negative on error: + * -1 argument rejected by this bridge (including the constraints above) + * -2 argument rejected by the kernel's own IsSupportedArgument() + * -3 kernel launch threw */ int dispatcher_run_gemm(const void* A, const void* B, @@ -220,6 +280,17 @@ int dispatcher_run_gemm(const void* A, return -1; } + // N must be a whole number of N-tiles unless the kernel pads N. Derived from the + // kernel's own constants rather than hardcoded, so it stays correct if the tile or + // the padding trait changes. Without this the caller only sees a bare -2. + if(!SelectedKernel::kPadN && (N % SelectedKernel::TileN) != 0) + { + std::cerr << "dispatcher_run_gemm: N must be a multiple of " << SelectedKernel::TileN + << " for this kernel (it was generated with pad_n=false, so N has to be a " + << "whole number of N-tiles); got N=" << N << "\n"; + return -1; + } + // Only packed (contiguous) layouts are supported. // // B layout (rcr = row-major A, column-major B, row-major C): @@ -416,7 +487,17 @@ int dispatcher_run_gemm(const void* A, if(exec_time < 0.0f) { - std::cerr << "dispatcher_run_gemm: kernel reported unsupported args\n"; + // IsSupportedArgument() is a boolean inside ck_tile and gives no reason, so + // enumerate the constraints a caller is realistically violating. Observed on + // gfx1250 with the default config: K must be a multiple of 16 (this holds even + // with pad_k=true -- padding covers the K-loop tail, not the global-load vector + // width), and N a multiple of the N-tile when pad_n is false. + std::cerr << "dispatcher_run_gemm: kernel reported unsupported args for M=" << M + << " N=" << N << " K=" << K << " k_batch=" << k_batch + << ". Known constraints for this kernel: K must be a multiple of 16 " + << "(required even though pad_k=" << (SelectedKernel::kPadK ? "true" : "false") + << "); N must be a multiple of " << SelectedKernel::TileN + << " when pad_n is false. No output was written.\n"; cleanup(); return -2; } @@ -436,6 +517,23 @@ int dispatcher_run_gemm(const void* A, */ const char* dispatcher_get_kernel_name() { return KERNEL_NAME; } +/** + * Return the N-tile of the force-included kernel. + * + * The shape checks in dispatcher_run_gemm derive their bound from SelectedKernel::TileN + * so they survive a tile change. Exporting it lets the Python runner report the same + * number instead of hardcoding the value of today's default config. + */ +int dispatcher_get_tile_n() { return static_cast(SelectedKernel::TileN); } + +/** + * Return 1 when the force-included kernel was generated with pad_n=true. + * + * When padding is on the N % TileN constraint does not apply, so the Python + * runner needs this to phrase its diagnostics correctly. + */ +int dispatcher_get_pad_n() { return SelectedKernel::kPadN ? 1 : 0; } + // This bridge is one-.so-per-kernel by construction: the build force-includes exactly // one generated header via `hipcc -include `, giving one SelectedKernel. // Scaling to N kernels means N .so files (the pattern bquant/aquant/abquant follow), diff --git a/projects/composablekernel/dispatcher/codegen/codegen_common.py b/projects/composablekernel/dispatcher/codegen/codegen_common.py index ac599a86ffa8..c70dde62a156 100644 --- a/projects/composablekernel/dispatcher/codegen/codegen_common.py +++ b/projects/composablekernel/dispatcher/codegen/codegen_common.py @@ -629,6 +629,37 @@ def make_gemm_rowcolquant_kernel_name( ) +# ============================================================================ +# Arch string normalization +# ============================================================================ + + +def normalize_gfx_arch(arch: str) -> str: + """Strip feature suffixes from a gfx target string. + + ``rocm_agent_enumerator`` and ``hipDeviceProp_t::gcnArchName`` may report the + target with trailing feature flags, e.g. ``"gfx942:sramecc+:xnack-"`` or + ``"gfx1250:xnack-"``. Every arch comparison in the codegen/runtime path (and + the ``--offload-arch`` we hand to hipcc) wants the bare target, so normalize + once at the boundary instead of scattering substring tests that happen to + tolerate the suffix. + + Single source of truth *for the dispatcher tree*: everything under + ``dispatcher/`` must call this rather than open-coding ``split(":")``. + + It is deliberately not claimed to be repo-wide, because it is not. + ``tile_engine/`` cannot import it: the dependency direction is + dispatcher -> tile_engine (``dispatcher/python/gemm_utils.py`` imports + ``gemm_validation_utils``), and tile_engine is on the deprecation path, so + moving the helper there to collapse the two copies would park new shared + infrastructure in the tree that is going away. + ``tile_engine/ops/gemm/gemm_validation_utils.py`` therefore keeps its own + ``_base_gfx_arch``; the two are pinned to identical behaviour by + ``dispatcher/tests/test_codegen_common.py::TestNormalizeGfxArch``. + """ + return arch.split(":", 1)[0] + + # ============================================================================ # Arch-derived warp tile K # ============================================================================ @@ -980,6 +1011,72 @@ def make_abquant_kernel_name( "warp_tile_m": 32, "warp_tile_n": 32, "warp_tile_k": 16, } +# gfx1250 (RDNA-style WMMA, MI400) cannot use the tile above: it is sized +# for the gfx9 MFMA 32x32x16 fragment, which does not exist on WMMA hardware, so the +# kernel compiles but produces all-zero output. The 8-bit WMMA fragment is 16x16x128, +# and the FlatMM 8-bit tile below is the shape validated against it. +ROWCOL_TENSOR_QUANT_DEFAULT_TILE_GFX1250 = { + "tile_m": 16, "tile_n": 64, "tile_k": 256, + "warp_m": 1, "warp_n": 4, "warp_k": 1, + "warp_tile_m": 16, "warp_tile_n": 16, "warp_tile_k": 128, +} + + +def rowcol_tensor_quant_default_tile(gfx_arch: str = "") -> dict: + """Return the default RowColQuant/TensorQuant tile for `gfx_arch`. + + Kept here, next to the tile dicts themselves, so the rowcolquant and tensorquant + runtime helpers select the arch-specific tile through one shared code path rather + than each carrying its own copy of the gfx1250 shape. + + The gfx1250 test is EXACT, not a ``gfx12`` family test. gfx1200/gfx1201 are + also WMMA parts, but their 8-bit warp fragment is 16x16x16, not the 16x16x64 / + 16x16x128 of gfx1250, so handing them the gfx1250 tile would compile cleanly + and return garbage. Contrast the OCP-FP8 define in the rowcolquant/tensorquant + runtime helpers, which *is* correctly family-wide. + """ + if normalize_gfx_arch(gfx_arch) == "gfx1250": + return dict(ROWCOL_TENSOR_QUANT_DEFAULT_TILE_GFX1250) + return dict(ROWCOL_TENSOR_QUANT_DEFAULT_TILE) + + +# Operator-specific support: both bridges require native FP8/BF8. gfx90a +# belongs to generic GEMM support, but cannot initialize these quant bridges. +ROWCOL_TENSOR_QUANT_SUPPORTED_ARCHES = ("gfx942", "gfx950", "gfx1250") + + +def validate_rowcol_tensor_quant_gfx_arch(gfx_arch: str, *, require_explicit: bool = False) -> str: + """Normalize and check a caller-supplied gfx target; return the bare target. + + Raises ``ValueError`` for anything outside + ``ROWCOL_TENSOR_QUANT_SUPPORTED_ARCHES``. Empty is allowed and means "not + specified", which selects the gfx9 MFMA tile -- the behaviour every invocation + without the flag had before the flag existed. Custom tile configurations must + set ``require_explicit=True`` so their target check cannot be bypassed. + + This exists because ``--gfx-arch`` on the two codegen scripts is the one place a + typo is completely silent. Everywhere else a bad target eventually reaches + ``--offload-arch`` and hipcc rejects it; here the value only picks a tile, so + ``--gfx-arch gfx1205`` quietly generates the gfx9 MFMA tile and the result is a + kernel that compiles for gfx1250 and returns garbage -- which is the failure mode + this whole branch exists to close, arriving through the front door. + """ + if not gfx_arch: + if require_explicit: + raise ValueError( + "Custom tile_configs require an explicit --gfx-arch (gfx_arch in Python) " + "so the generated header can reject a mismatched build target." + ) + return "" + base = normalize_gfx_arch(gfx_arch) + if base not in ROWCOL_TENSOR_QUANT_SUPPORTED_ARCHES: + raise ValueError( + f"Unsupported GPU architecture {gfx_arch!r} (normalized to {base!r}); " + f"supported: {', '.join(ROWCOL_TENSOR_QUANT_SUPPORTED_ARCHES)}." + ) + return base + + # Default traits, shared for the same reason as the tile above. pad_m is enabled # because these kernels are used with M values that are not tile-aligned. ROWCOL_TENSOR_QUANT_DEFAULT_TRAITS = { diff --git a/projects/composablekernel/dispatcher/codegen/unified_grouped_gemm_rowcolquant_codegen.py b/projects/composablekernel/dispatcher/codegen/unified_grouped_gemm_rowcolquant_codegen.py index e12199429dc4..a9fe03b2ddaa 100644 --- a/projects/composablekernel/dispatcher/codegen/unified_grouped_gemm_rowcolquant_codegen.py +++ b/projects/composablekernel/dispatcher/codegen/unified_grouped_gemm_rowcolquant_codegen.py @@ -33,12 +33,13 @@ from codegen_common import ( ROWCOL_TENSOR_QUANT_BASE_PIPELINE_MAP, - ROWCOL_TENSOR_QUANT_DEFAULT_TILE, ROWCOL_TENSOR_QUANT_DEFAULT_TRAITS, ROWCOL_TENSOR_QUANT_EPILOGUE_MAP, ROWCOL_TENSOR_QUANT_PIPELINE_MAP, ROWCOL_TENSOR_QUANT_SUPPORTED_LAYOUTS, + validate_rowcol_tensor_quant_gfx_arch, make_rowcolquant_kernel_name, + rowcol_tensor_quant_default_tile, ) logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") @@ -144,6 +145,7 @@ class RowColQuantKernelSpec: tile: RowColQuantTileConfig block_size: int = 256 k_block_per_cu: int = 1 + gfx_arch: str = "" @property def name(self) -> str: @@ -253,6 +255,8 @@ def generate(self, spec: RowColQuantKernelSpec) -> str: using BQDataType = {ns}::BQDataType; using AccDataType = {ns}::AccDataType; + static constexpr const char* GfxArch = "{validate_rowcol_tensor_quant_gfx_arch(spec.gfx_arch)}"; + static constexpr ck_tile::index_t TileM = {t.tile_m}; static constexpr ck_tile::index_t TileN = {t.tile_n}; static constexpr ck_tile::index_t TileK = {t.tile_k}; @@ -429,14 +433,20 @@ def generate(self, spec: RowColQuantKernelSpec) -> str: # ============================================================================= -def _default_config() -> dict: +def _default_config(gfx_arch: str = "") -> dict: # Traits and tile come from codegen_common so this default and the runtime # default_{fp8,bf8}_config() in grouped_gemm_rowcolquant_utils.py cannot drift. + # + # gfx_arch selects the tile. It defaults to "" -- the gfx9 MFMA tile -- so an + # invocation with no --gfx-arch generates exactly what it generated before. + # Passing "gfx1250" yields the 16x16x128 WMMA tile instead; the gfx9 tile is + # sized for an MFMA fragment that does not exist on that part, and compiles + # there rather than erroring. return { "dtypes": ["fp8", "bf8"], "layouts": list(ROWCOLQUANT_SUPPORTED_LAYOUTS), **ROWCOL_TENSOR_QUANT_DEFAULT_TRAITS, - "tile_configs": [dict(ROWCOL_TENSOR_QUANT_DEFAULT_TILE)], + "tile_configs": [rowcol_tensor_quant_default_tile(gfx_arch)], } @@ -535,11 +545,22 @@ def generate_kernels( output_dir: Path, config: Optional[dict] = None, parallel: bool = True, + gfx_arch: str = "", ) -> List[Path]: - """Generate all RowColQuant kernel headers into output_dir. Returns list of generated .hpp paths.""" + """Generate all RowColQuant kernel headers into output_dir. Returns list of generated .hpp paths. + + `gfx_arch` selects the default tile and is embedded in every header so the + bridge can reject a mismatched build target. Custom tile_configs require an + explicit target; their tile values remain unchanged. + """ + gfx_arch = validate_rowcol_tensor_quant_gfx_arch( + gfx_arch, require_explicit=config is not None and "tile_configs" in config + ) output_dir.mkdir(parents=True, exist_ok=True) - cfg = config or _default_config() + cfg = config or _default_config(gfx_arch) specs = _build_specs(cfg) + for spec in specs: + spec.gfx_arch = gfx_arch if not specs: log.warning("No kernel specs produced from config — check dtypes and tile_configs") @@ -596,8 +617,23 @@ def main() -> int: help="Disable parallel generation") parser.add_argument("--list-names", action="store_true", help="Print kernel names that would be generated and exit") + parser.add_argument("--gfx-arch", type=str, default="", + help="GPU target the generated kernels will be compiled for, " + "e.g. gfx1250. Embedded in generated headers; selects the default tile " + "unless --config/--config-json supplies tile_configs. " + "Required with custom tile_configs; otherwise defaults to the gfx9 MFMA tile.") args = parser.parse_args() + # Validate before anything uses it. This value never reaches a compiler -- it + # only picks a tile -- so an unrecognized target is not caught downstream the way + # a bad --offload-arch would be: it silently falls through to the gfx9 MFMA tile, + # which on WMMA hardware compiles and returns garbage. Refuse instead. + try: + gfx_arch = validate_rowcol_tensor_quant_gfx_arch(args.gfx_arch) + except ValueError as e: + log.error("%s", e) + return 1 + cfg: Optional[dict] = None if args.config_json: try: @@ -609,8 +645,16 @@ def main() -> int: with open(args.config) as f: cfg = json.load(f) + try: + gfx_arch = validate_rowcol_tensor_quant_gfx_arch( + gfx_arch, require_explicit=cfg is not None and "tile_configs" in cfg + ) + except ValueError as e: + log.error("%s", e) + return 1 + if args.list_names: - specs = _build_specs(cfg or _default_config()) + specs = _build_specs(cfg or _default_config(gfx_arch)) for s in specs: print(s.name) return 0 @@ -619,6 +663,7 @@ def main() -> int: output_dir=args.output_dir, config=cfg, parallel=not args.no_parallel, + gfx_arch=gfx_arch, ) return 0 if paths else 1 diff --git a/projects/composablekernel/dispatcher/codegen/unified_grouped_gemm_tensorquant_codegen.py b/projects/composablekernel/dispatcher/codegen/unified_grouped_gemm_tensorquant_codegen.py index 09591e21e32b..bd63e47b79fb 100644 --- a/projects/composablekernel/dispatcher/codegen/unified_grouped_gemm_tensorquant_codegen.py +++ b/projects/composablekernel/dispatcher/codegen/unified_grouped_gemm_tensorquant_codegen.py @@ -33,12 +33,13 @@ from codegen_common import ( ROWCOL_TENSOR_QUANT_BASE_PIPELINE_MAP, - ROWCOL_TENSOR_QUANT_DEFAULT_TILE, ROWCOL_TENSOR_QUANT_DEFAULT_TRAITS, ROWCOL_TENSOR_QUANT_EPILOGUE_MAP, ROWCOL_TENSOR_QUANT_PIPELINE_MAP, ROWCOL_TENSOR_QUANT_SUPPORTED_LAYOUTS, + validate_rowcol_tensor_quant_gfx_arch, make_tensorquant_kernel_name, + rowcol_tensor_quant_default_tile, ) logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") @@ -144,6 +145,7 @@ class TensorQuantKernelSpec: tile: TensorQuantTileConfig block_size: int = 256 k_block_per_cu: int = 1 + gfx_arch: str = "" @property def name(self) -> str: @@ -253,6 +255,8 @@ def generate(self, spec: TensorQuantKernelSpec) -> str: using BQDataType = {ns}::BQDataType; using AccDataType = {ns}::AccDataType; + static constexpr const char* GfxArch = "{validate_rowcol_tensor_quant_gfx_arch(spec.gfx_arch)}"; + static constexpr ck_tile::index_t TileM = {t.tile_m}; static constexpr ck_tile::index_t TileN = {t.tile_n}; static constexpr ck_tile::index_t TileK = {t.tile_k}; @@ -429,14 +433,20 @@ def generate(self, spec: TensorQuantKernelSpec) -> str: # ============================================================================= -def _default_config() -> dict: +def _default_config(gfx_arch: str = "") -> dict: # Traits and tile come from codegen_common so this default and the runtime # default_{fp8,bf8}_config() in grouped_gemm_tensorquant_utils.py cannot drift. + # + # gfx_arch selects the tile. It defaults to "" -- the gfx9 MFMA tile -- so an + # invocation with no --gfx-arch generates exactly what it generated before. + # Passing "gfx1250" yields the 16x16x128 WMMA tile instead; the gfx9 tile is + # sized for an MFMA fragment that does not exist on that part, and compiles + # there rather than erroring. return { "dtypes": ["fp8", "bf8"], "layouts": list(TENSORQUANT_SUPPORTED_LAYOUTS), **ROWCOL_TENSOR_QUANT_DEFAULT_TRAITS, - "tile_configs": [dict(ROWCOL_TENSOR_QUANT_DEFAULT_TILE)], + "tile_configs": [rowcol_tensor_quant_default_tile(gfx_arch)], } @@ -535,11 +545,22 @@ def generate_kernels( output_dir: Path, config: Optional[dict] = None, parallel: bool = True, + gfx_arch: str = "", ) -> List[Path]: - """Generate all TensorQuant kernel headers into output_dir. Returns list of generated .hpp paths.""" + """Generate all TensorQuant kernel headers into output_dir. Returns list of generated .hpp paths. + + `gfx_arch` selects the default tile and is embedded in every header so the + bridge can reject a mismatched build target. Custom tile_configs require an + explicit target; their tile values remain unchanged. + """ + gfx_arch = validate_rowcol_tensor_quant_gfx_arch( + gfx_arch, require_explicit=config is not None and "tile_configs" in config + ) output_dir.mkdir(parents=True, exist_ok=True) - cfg = config or _default_config() + cfg = config or _default_config(gfx_arch) specs = _build_specs(cfg) + for spec in specs: + spec.gfx_arch = gfx_arch if not specs: log.warning("No kernel specs produced from config — check dtypes and tile_configs") @@ -596,8 +617,23 @@ def main() -> int: help="Disable parallel generation") parser.add_argument("--list-names", action="store_true", help="Print kernel names that would be generated and exit") + parser.add_argument("--gfx-arch", type=str, default="", + help="GPU target the generated kernels will be compiled for, " + "e.g. gfx1250. Embedded in generated headers; selects the default tile " + "unless --config/--config-json supplies tile_configs. " + "Required with custom tile_configs; otherwise defaults to the gfx9 MFMA tile.") args = parser.parse_args() + # Validate before anything uses it. This value never reaches a compiler -- it + # only picks a tile -- so an unrecognized target is not caught downstream the way + # a bad --offload-arch would be: it silently falls through to the gfx9 MFMA tile, + # which on WMMA hardware compiles and returns garbage. Refuse instead. + try: + gfx_arch = validate_rowcol_tensor_quant_gfx_arch(args.gfx_arch) + except ValueError as e: + log.error("%s", e) + return 1 + cfg: Optional[dict] = None if args.config_json: try: @@ -609,8 +645,16 @@ def main() -> int: with open(args.config) as f: cfg = json.load(f) + try: + gfx_arch = validate_rowcol_tensor_quant_gfx_arch( + gfx_arch, require_explicit=cfg is not None and "tile_configs" in cfg + ) + except ValueError as e: + log.error("%s", e) + return 1 + if args.list_names: - specs = _build_specs(cfg or _default_config()) + specs = _build_specs(cfg or _default_config(gfx_arch)) for s in specs: print(s.name) return 0 @@ -619,6 +663,7 @@ def main() -> int: output_dir=args.output_dir, config=cfg, parallel=not args.no_parallel, + gfx_arch=gfx_arch, ) return 0 if paths else 1 diff --git a/projects/composablekernel/dispatcher/python/gemm_utils.py b/projects/composablekernel/dispatcher/python/gemm_utils.py index 767bebc3428c..277adcb21743 100644 --- a/projects/composablekernel/dispatcher/python/gemm_utils.py +++ b/projects/composablekernel/dispatcher/python/gemm_utils.py @@ -85,6 +85,26 @@ def _ocp_arch_defines(arch: Optional[str]) -> List[str]: return ["-DCK_USE_OCP_FP8", "-DCK_TILE_USE_OCP_FP8"] +def normalize_gfx_arch(arch: str) -> str: + """Strip feature suffixes from a gfx target string. + + ``"gfx950:sramecc+:xnack-"`` -> ``"gfx950"``. Empty input passes through. + + Delegates to ``codegen_common.normalize_gfx_arch``, the dispatcher tree's + single source of truth, whenever it can be imported. It usually cannot be at + this point: nothing has put the ``codegen`` dir on ``sys.path`` yet -- only + ``ctypes_utils.get_arch_filter_data()`` does that, and it may never be called -- + so this module has to be able to answer without it, the same way it already + falls back for ``_detect_gpu_arch_via_amd_smi``. The two are pinned to identical + behaviour by tests/test_gemm_utils.py::TestArchNormalizationMatchesCodegen. + """ + try: + from codegen_common import normalize_gfx_arch as _canonical # noqa: WPS433 + except ImportError: + return arch.split(":", 1)[0] + return _canonical(arch) + + @functools.lru_cache(maxsize=1) def _get_arch() -> str: """Detect the GPU architecture from rocminfo and validate it. @@ -117,29 +137,72 @@ def _get_arch() -> str: "gfx_arch (one of " f"{', '.join(_SUPPORTED_ARCHES)})." ) - if detected not in _SUPPORTED_ARCHES: + return _validate_arch(detected) + + +def _validate_arch(arch: str) -> str: + """Normalize a gfx target, then check it against ``_SUPPORTED_ARCHES``. + + Normalization comes FIRST, and that ordering is the whole point -- but not + because autodetection produces suffixes. It does not. On a real device all + three autodetect sources report the BARE target: ``amd-smi``'s + ``TARGET_GRAPHICS_VERSION``, ``rocm_agent_enumerator``, and the agent + ``Name:`` line that ``_get_arch`` above parses all print ``gfx90a`` + (measured on a gfx90a device). ``_get_arch`` could not return a suffixed + name even if rocminfo offered one: the only suffixed string in that output + is the ISA line, ``amdgcn-amd-amdhsa--gfx90a:sramecc+:xnack-``, which fails + the ``startswith("gfx")`` check after the split. + + Suffixed names get here from CALLERS, not from detection, and they are real: + + * this repository's own CMakeLists.txt sets, on the ASAN branch, + ``CK_GPU_TARGETS "gfx908:xnack+;gfx90a:xnack+;gfx942:xnack+;gfx950:xnack+"``; + * ``hipDeviceProp_t::gcnArchName`` returns ``gfx90a:sramecc+:xnack-`` + (measured on the same device). + + Either is a plausible thing to copy into ``--gfx-arch`` or an explicit + ``arch=``, and that value reaches this function unmodified. Validating it raw + rejected it: + + gfx950:sramecc+:xnack- -> ValueError + + Bare targets were never the problem; they passed before this change and pass + now (tests/test_gemm_utils.py::test_bare_supported_arches_are_unchanged). + Normalizing first costs a supported bare name nothing and makes the suffixed + spellings above usable. + + Every caller wants the bare target regardless. It is stamped onto + ``GemmKernelConfig.gfx_arch``, handed to ``--offload-arch``, and used as the key + into the warp tables, none of which know about suffixes. + + The suffix is dropped, not carried through: ``--offload-arch=gfx90a`` rather + than ``gfx90a:xnack+``. That is not a regression -- it is the target string every + already-working path used, because a suffixed name could not get past this + function at all before. It does mean an xnack+ device is compiled for the bare + target; if a caller needs a specific xnack setting it has to pass the flag + itself, exactly as it had to before. + """ + base = normalize_gfx_arch(arch) + if base not in _SUPPORTED_ARCHES: raise ValueError( - f"Unsupported GPU architecture {detected!r}; supported: " - f"{', '.join(_SUPPORTED_ARCHES)}." + f"Unsupported GPU architecture {arch!r}" + + (f" (normalized to {base!r})" if base != arch else "") + + f"; supported: {', '.join(_SUPPORTED_ARCHES)}." ) - return detected + return base def _resolve_arch(arch: Optional[str]) -> str: """Resolve a possibly-``None`` arch to a validated, supported ``gfxNNN``. - ``None``/empty -> detect via :func:`_get_arch`. An explicit value is - validated against ``_SUPPORTED_ARCHES`` (raising ``ValueError`` if unknown) - so a typo can never silently reach the compiler. + ``None``/empty -> detect via :func:`_get_arch`. An explicit value is normalized + and then validated against ``_SUPPORTED_ARCHES`` (raising ``ValueError`` if + unknown) so a typo can never silently reach the compiler, while the suffixed + name a device actually reports resolves to its base. """ if not arch: return _get_arch() - if arch not in _SUPPORTED_ARCHES: - raise ValueError( - f"Unsupported GPU architecture {arch!r}; supported: " - f"{', '.join(_SUPPORTED_ARCHES)}." - ) - return arch + return _validate_arch(arch) def _cshuffle_store_ok( diff --git a/projects/composablekernel/dispatcher/python/grouped_gemm_rowcolquant_utils.py b/projects/composablekernel/dispatcher/python/grouped_gemm_rowcolquant_utils.py index 34a086480b06..3a08018841fb 100644 --- a/projects/composablekernel/dispatcher/python/grouped_gemm_rowcolquant_utils.py +++ b/projects/composablekernel/dispatcher/python/grouped_gemm_rowcolquant_utils.py @@ -49,8 +49,10 @@ # their name builders, and keeps the shared tile/trait defaults in one place so this # module's default_*_config() cannot drift from the codegen's _default_config(). from codegen_common import ( # noqa: E402 - ROWCOL_TENSOR_QUANT_DEFAULT_TILE, ROWCOL_TENSOR_QUANT_DEFAULT_TRAITS, + rowcol_tensor_quant_default_tile, + normalize_gfx_arch, + validate_rowcol_tensor_quant_gfx_arch, make_rowcolquant_kernel_name, ) @@ -62,6 +64,24 @@ _DEFAULT_HIPCC = "hipcc" _DEFAULT_GFX_ARCH = "gfx950" +# ABI revision of the compiled .so, folded into the artifact filename. +# +# setup_multiple_rowcolquant_dispatchers() reuses an existing .so when the name +# matches, and callers that pass a persistent output_dir (the validation +# harnesses do, precisely to get cache hits) can therefore be handed an artifact +# built by an older revision of this module. The name used to be keyed on +# (kernel name, arch) only, which does not describe the exported symbol set, so +# a .so predating the dispatcher_get_tile_n()/dispatcher_get_pad_n() exports +# would be selected and then fail at attribute-lookup time with a bare +# "undefined symbol". Bump this whenever the exported C ABI of +# bindings/ctypes/grouped_gemm_rowcolquant_ctypes_lib.cpp changes; the new name +# simply cannot collide with the stale artifact, which is then ignored rather +# than papered over with a runtime fallback. +# 1 -> original export set +# 2 -> added dispatcher_get_tile_n() / dispatcher_get_pad_n() +# 3 -> require regenerated headers with build-target validation +_SO_ABI = 3 + # ============================================================================= # RowColQuantKernelConfig — byte-exact naming with codegen @@ -194,12 +214,14 @@ class RowColQuantDispatcherLib: QK_A, QK_B, k_batch, *time_ms) char* dispatcher_get_kernel_name() int dispatcher_get_kernel_count() + int dispatcher_get_tile_n() + int dispatcher_get_pad_n() void dispatcher_cleanup() """ def __init__(self, so_path: Path): self.so_path = Path(so_path) - self._cleaned_up = False + self._cleaned_up = True if not self.so_path.exists(): raise FileNotFoundError(f"RowColQuant .so not found: {self.so_path}") self._lib = ctypes.CDLL(str(self.so_path)) @@ -207,6 +229,7 @@ def __init__(self, so_path: Path): rc = self._lib.dispatcher_initialize() if rc != 0: raise RuntimeError(f"dispatcher_initialize() returned {rc}") + self._cleaned_up = False def _setup(self): lib = self._lib @@ -241,6 +264,18 @@ def _setup(self): lib.dispatcher_get_kernel_count.restype = ctypes.c_int lib.dispatcher_get_kernel_count.argtypes = [] + # Direct callers can supply libraries outside the versioned cache. + try: + lib.dispatcher_get_tile_n.restype = ctypes.c_int + lib.dispatcher_get_tile_n.argtypes = [] + lib.dispatcher_get_pad_n.restype = ctypes.c_int + lib.dispatcher_get_pad_n.argtypes = [] + except AttributeError as exc: + raise RuntimeError( + f"Incompatible quant bridge ABI in {self.so_path}: missing TileN/PadN " + "exports (ABI 2 or newer required). Rebuild the library." + ) from exc + lib.dispatcher_cleanup.restype = None lib.dispatcher_cleanup.argtypes = [] @@ -311,6 +346,14 @@ def get_kernel_name(self) -> str: def get_kernel_count(self) -> int: return self._lib.dispatcher_get_kernel_count() + def get_tile_n(self) -> int: + """N-tile of the compiled kernel (SelectedKernel::TileN).""" + return self._lib.dispatcher_get_tile_n() + + def get_pad_n(self) -> bool: + """True when the compiled kernel was generated with pad_n=true.""" + return bool(self._lib.dispatcher_get_pad_n()) + def cleanup(self): if not self._cleaned_up: self._lib.dispatcher_cleanup() @@ -342,6 +385,16 @@ def __init__(self, so_path: Path): def kernel_name(self) -> str: return self._lib.get_kernel_name() + @property + def tile_n(self) -> int: + """N-tile of the compiled kernel, read from the .so (not hardcoded).""" + return self._lib.get_tile_n() + + @property + def pad_n(self) -> bool: + """True when the compiled kernel pads N (the N % tile_n rule then lifts).""" + return self._lib.get_pad_n() + def run(self, A, B, AQ, BQ, problem: RowColQuantGemmProblem, c_dtype=None) -> RowColQuantGemmResult: """ Run RowColQuant Grouped GEMM. @@ -352,6 +405,23 @@ def run(self, A, B, AQ, BQ, problem: RowColQuantGemmProblem, c_dtype=None) -> Ro BQ shape: (N,) dtype: float (per-col B scale) c_dtype numpy dtype for the output C buffer. Defaults to np.float16. Returns RowColQuantGemmResult with C shape (M, N). + + Shape constraints enforced by the bridge (validated on gfx1250 with the + default config). Violations raise RuntimeError; no output is written: + + K % 16 == 0 Required even though the default config sets pad_k=True -- + pad_k covers the K-loop tail, not the global-load vector + width. + N % tile_n == 0 Required whenever the kernel was generated with pad_n=False, + so N must be a whole number of N-tiles. tile_n is a property + of the compiled kernel -- read `self.tile_n` rather than + assuming the value of today's default config. + M % 4 == 0 gfx12 targets only. The per-row A-scale (AQ) tile window is + built without a padding transform, so row M-1 of C comes + back as zero while every other row is correct. The bridge + refuses these shapes rather than returning a silently wrong + result. gfx942/gfx950 are not restricted. TensorQuant uses + scalar scales and has no M constraint. """ import numpy as np @@ -412,9 +482,23 @@ def run(self, A, B, AQ, BQ, problem: RowColQuantGemmProblem, c_dtype=None) -> Ro ) if rc != 0: + # rc alone is not actionable, and the C++ explanation goes to stderr, which + # a caller capturing only the exception never sees. Restate the constraints + # here so the traceback is self-contained. tile_n/pad_n come from the .so, + # so this message stays true if the compiled tile changes. + n_rule = ( + "N is unconstrained (pad_n=True)" + if self.pad_n + else f"N % {self.tile_n} == 0 when pad_n=False" + ) raise RuntimeError( f"dispatcher_run_gemm failed with code {rc} " - f"for kernel {self.kernel_name}" + f"for kernel {self.kernel_name} at M={M} N={N} K={K}. " + f"(-1 = rejected by the bridge, -2 = rejected by the kernel, " + f"-3 = launch threw.) Shape constraints: K % 16 == 0 (required even " + f"with pad_k=True); {n_rule}; and on gfx12 targets " + f"M % 4 == 0, because the per-row AQ tile window is unpadded in M and " + f"would zero row M-1 of C. See stderr for the exact reason." ) return RowColQuantGemmResult(C=C, time_ms=time_ms, kernel_name=self.kernel_name) @@ -435,7 +519,11 @@ def _detect_gpu_arch() -> str: for line in result.stdout.splitlines(): line = line.strip() if line.startswith("gfx") and line != "gfx000": - return line + # Strip feature suffixes ("gfx1250:xnack-") here so they never reach + # --offload-arch / -DGFX_ARCH. The C++ side prefix-matches the runtime + # device name against the compile-time GFX_ARCH, so the bare target + # still matches a device that reports suffixes. + return normalize_gfx_arch(line) except Exception as e: log.warning("rocm_agent_enumerator failed (%s); defaulting to %s", e, _DEFAULT_GFX_ARCH) return _DEFAULT_GFX_ARCH @@ -466,6 +554,7 @@ def _generate_rowcolquant_kernel( str(_CODEGEN_SCRIPT), "--output-dir", str(output_dir), "--config-json", config_json, + "--gfx-arch", validate_rowcol_tensor_quant_gfx_arch(config.gfx_arch), ] try: @@ -500,16 +589,26 @@ def _compile_rowcolquant_kernel( extra_include_dirs: Optional[List[str]] = None, ) -> bool: """Compile a generated .hpp into a .so via hipcc (compile then link).""" + # Normalize once, here at the boundary, so every downstream use -- the arch + # defines below and the --offload-arch/-DGFX_ARCH we hand to hipcc -- sees the + # bare target. A caller-supplied "gfx1250:xnack-" must not reach the compiler + # flags. + gfx_arch = validate_rowcol_tensor_quant_gfx_arch(gfx_arch) + ck_include = _get_ck_include_dir() static_lib = _get_dispatcher_static_lib() obj_path = so_path.with_suffix(".o") arch_defines = arch_feature_defines(gfx_arch) + # Match top-level CK policy in both host and device compilation. + if gfx_arch == "gfx1250": + arch_defines.append("-DUSE_NEW_UNIFIED_FRAMEWORK=0") compile_cmd = [hipcc, "-c", "-fPIC", "-O3", "-std=c++17", "-DCK_TILE_SINGLE_KERNEL_INCLUDE", "-w", f"--offload-arch={gfx_arch}", + f"-DCK_CMAKE_GPU_TARGET_IDS=0x{gfx_arch[3:]}", f"-DGFX_ARCH=\"{gfx_arch}\"", *arch_defines, "-include", str(hpp_path), @@ -584,7 +683,17 @@ def setup_multiple_rowcolquant_dispatchers( if not configs: return [] - arch = gfx_arch or _detect_gpu_arch() + # Normalize the explicit branch too, not just detection: an explicitly passed + # "gfx1250:xnack-" would otherwise flow into --offload-arch, -DGFX_ARCH and the + # .so cache name. _detect_gpu_arch() already normalizes. + arch = validate_rowcol_tensor_quant_gfx_arch(gfx_arch or _detect_gpu_arch()) + for cfg in configs: + config_arch = validate_rowcol_tensor_quant_gfx_arch(cfg.gfx_arch) + if config_arch != arch: + raise ValueError( + f"Config architecture {cfg.gfx_arch!r} does not match build target {arch!r}. " + "Create configs for the selected build target." + ) base_dir = output_dir or Path(tempfile.mkdtemp(prefix="rowcolquant_dispatcher_")) base_dir.mkdir(parents=True, exist_ok=True) @@ -609,7 +718,9 @@ def _build_one(idx: int, cfg: RowColQuantKernelConfig) -> Tuple[int, Optional[Pa if hpp is None: return idx, None - so = so_dir / f"lib{cfg.name}_{arch}.so" + # Name carries the ABI revision: a pre-_SO_ABI artifact in a persistent + # output_dir can never be selected here (see _SO_ABI). + so = so_dir / f"lib{cfg.name}_{arch}_abi{_SO_ABI}.so" if so.exists(): log.info(" [cached] %s", so.name) return idx, so @@ -670,6 +781,10 @@ def _default_config(dtype: str, gfx_arch: str) -> RowColQuantKernelConfig: four pad/persistent flags) is forwarded from the shared dict; only block_size and k_block_per_cu are codegen-only and are left to the dataclass defaults. """ + # Normalize at this boundary too: default_fp8_config()/default_bf8_config() are + # public entry points, so a caller-supplied "gfx1250:xnack-" must not be stored + # on the config and observed by later consumers. + gfx_arch = validate_rowcol_tensor_quant_gfx_arch(gfx_arch) traits = ROWCOL_TENSOR_QUANT_DEFAULT_TRAITS return RowColQuantKernelConfig( dtype=dtype, @@ -677,7 +792,7 @@ def _default_config(dtype: str, gfx_arch: str) -> RowColQuantKernelConfig: pipeline=traits["pipeline"], epilogue=traits["epilogue"], scheduler=traits["scheduler"], - **ROWCOL_TENSOR_QUANT_DEFAULT_TILE, + **rowcol_tensor_quant_default_tile(gfx_arch), pad_m=traits["pad_m"], pad_n=traits["pad_n"], pad_k=traits["pad_k"], diff --git a/projects/composablekernel/dispatcher/python/grouped_gemm_tensorquant_utils.py b/projects/composablekernel/dispatcher/python/grouped_gemm_tensorquant_utils.py index ef3f02e33b42..96152fb6d82f 100644 --- a/projects/composablekernel/dispatcher/python/grouped_gemm_tensorquant_utils.py +++ b/projects/composablekernel/dispatcher/python/grouped_gemm_tensorquant_utils.py @@ -49,8 +49,10 @@ # their name builders, and keeps the shared tile/trait defaults in one place so this # module's default_*_config() cannot drift from the codegen's _default_config(). from codegen_common import ( # noqa: E402 - ROWCOL_TENSOR_QUANT_DEFAULT_TILE, ROWCOL_TENSOR_QUANT_DEFAULT_TRAITS, + rowcol_tensor_quant_default_tile, + normalize_gfx_arch, + validate_rowcol_tensor_quant_gfx_arch, make_tensorquant_kernel_name, ) @@ -62,6 +64,24 @@ _DEFAULT_HIPCC = "hipcc" _DEFAULT_GFX_ARCH = "gfx950" +# ABI revision of the compiled .so, folded into the artifact filename. +# +# setup_multiple_tensorquant_dispatchers() reuses an existing .so when the name +# matches, and callers that pass a persistent output_dir (the validation +# harnesses do, precisely to get cache hits) can therefore be handed an artifact +# built by an older revision of this module. The name used to be keyed on +# (kernel name, arch) only, which does not describe the exported symbol set, so +# a .so predating the dispatcher_get_tile_n()/dispatcher_get_pad_n() exports +# would be selected and then fail at attribute-lookup time with a bare +# "undefined symbol". Bump this whenever the exported C ABI of +# bindings/ctypes/grouped_gemm_tensorquant_ctypes_lib.cpp changes; the new name +# simply cannot collide with the stale artifact, which is then ignored rather +# than papered over with a runtime fallback. +# 1 -> original export set +# 2 -> added dispatcher_get_tile_n() / dispatcher_get_pad_n() +# 3 -> require regenerated headers with build-target validation +_SO_ABI = 3 + # ============================================================================= # TensorQuantKernelConfig — byte-exact naming with codegen @@ -184,12 +204,14 @@ class TensorQuantDispatcherLib: QK_A, QK_B, k_batch, *time_ms) char* dispatcher_get_kernel_name() int dispatcher_get_kernel_count() + int dispatcher_get_tile_n() + int dispatcher_get_pad_n() void dispatcher_cleanup() """ def __init__(self, so_path: Path): self.so_path = Path(so_path) - self._cleaned_up = False + self._cleaned_up = True if not self.so_path.exists(): raise FileNotFoundError(f"TensorQuant .so not found: {self.so_path}") self._lib = ctypes.CDLL(str(self.so_path)) @@ -197,6 +219,7 @@ def __init__(self, so_path: Path): rc = self._lib.dispatcher_initialize() if rc != 0: raise RuntimeError(f"dispatcher_initialize() returned {rc}") + self._cleaned_up = False def _setup(self): lib = self._lib @@ -231,6 +254,18 @@ def _setup(self): lib.dispatcher_get_kernel_count.restype = ctypes.c_int lib.dispatcher_get_kernel_count.argtypes = [] + # Direct callers can supply libraries outside the versioned cache. + try: + lib.dispatcher_get_tile_n.restype = ctypes.c_int + lib.dispatcher_get_tile_n.argtypes = [] + lib.dispatcher_get_pad_n.restype = ctypes.c_int + lib.dispatcher_get_pad_n.argtypes = [] + except AttributeError as exc: + raise RuntimeError( + f"Incompatible quant bridge ABI in {self.so_path}: missing TileN/PadN " + "exports (ABI 2 or newer required). Rebuild the library." + ) from exc + lib.dispatcher_cleanup.restype = None lib.dispatcher_cleanup.argtypes = [] @@ -301,6 +336,14 @@ def get_kernel_name(self) -> str: def get_kernel_count(self) -> int: return self._lib.dispatcher_get_kernel_count() + def get_tile_n(self) -> int: + """N-tile of the compiled kernel (SelectedKernel::TileN).""" + return self._lib.dispatcher_get_tile_n() + + def get_pad_n(self) -> bool: + """True when the compiled kernel was generated with pad_n=true.""" + return bool(self._lib.dispatcher_get_pad_n()) + def cleanup(self): if not self._cleaned_up: self._lib.dispatcher_cleanup() @@ -332,6 +375,16 @@ def __init__(self, so_path: Path): def kernel_name(self) -> str: return self._lib.get_kernel_name() + @property + def tile_n(self) -> int: + """N-tile of the compiled kernel, read from the .so (not hardcoded).""" + return self._lib.get_tile_n() + + @property + def pad_n(self) -> bool: + """True when the compiled kernel pads N (the N % tile_n rule then lifts).""" + return self._lib.get_pad_n() + def run(self, A, B, AQ, BQ, problem: TensorQuantGemmProblem, c_dtype=None) -> TensorQuantGemmResult: """ Run TensorQuant Grouped GEMM. @@ -342,6 +395,21 @@ def run(self, A, B, AQ, BQ, problem: TensorQuantGemmProblem, c_dtype=None) -> Te BQ shape: (1,) dtype: float (per-tensor B scale) c_dtype numpy dtype for the output C buffer. Defaults to np.float16. Returns TensorQuantGemmResult with C shape (M, N). + + Shape constraints enforced by the bridge (validated on gfx1250 with the + default config). Violations raise RuntimeError; no output is written: + + K % 16 == 0 Required even though the default config sets pad_k=True -- + pad_k covers the K-loop tail, not the global-load vector + width. + N % tile_n == 0 Required whenever the kernel was generated with pad_n=False, + so N must be a whole number of N-tiles. tile_n is a property + of the compiled kernel -- read `self.tile_n` rather than + assuming the value of today's default config. + + There is no M constraint: TensorQuant applies scalar A/B scales rather than a + per-row scale vector, so it is unaffected by the unpadded per-row AQ tail that + restricts the RowColQuant bridge to M % 4 == 0 on gfx12. """ import numpy as np @@ -400,9 +468,22 @@ def run(self, A, B, AQ, BQ, problem: TensorQuantGemmProblem, c_dtype=None) -> Te ) if rc != 0: + # rc alone is not actionable, and the C++ explanation goes to stderr, which + # a caller capturing only the exception never sees. Restate the constraints + # here so the traceback is self-contained. tile_n/pad_n come from the .so, + # so this message stays true if the compiled tile changes. + n_rule = ( + "N is unconstrained (pad_n=True)" + if self.pad_n + else f"N % {self.tile_n} == 0 when pad_n=False" + ) raise RuntimeError( f"dispatcher_run_gemm failed with code {rc} " - f"for kernel {self.kernel_name}" + f"for kernel {self.kernel_name} at M={M} N={N} K={K}. " + f"(-1 = rejected by the bridge, -2 = rejected by the kernel, " + f"-3 = launch threw.) Shape constraints: K % 16 == 0 (required even " + f"with pad_k=True); {n_rule}. " + f"See stderr for the exact reason." ) return TensorQuantGemmResult(C=C, time_ms=time_ms, kernel_name=self.kernel_name) @@ -423,7 +504,11 @@ def _detect_gpu_arch() -> str: for line in result.stdout.splitlines(): line = line.strip() if line.startswith("gfx") and line != "gfx000": - return line + # Strip feature suffixes ("gfx1250:xnack-") here so they never reach + # --offload-arch / -DGFX_ARCH. The C++ side prefix-matches the runtime + # device name against the compile-time GFX_ARCH, so the bare target + # still matches a device that reports suffixes. + return normalize_gfx_arch(line) except Exception as e: log.warning("rocm_agent_enumerator failed (%s); defaulting to %s", e, _DEFAULT_GFX_ARCH) return _DEFAULT_GFX_ARCH @@ -454,6 +539,7 @@ def _generate_tensorquant_kernel( str(_CODEGEN_SCRIPT), "--output-dir", str(output_dir), "--config-json", config_json, + "--gfx-arch", validate_rowcol_tensor_quant_gfx_arch(config.gfx_arch), ] try: @@ -488,16 +574,26 @@ def _compile_tensorquant_kernel( extra_include_dirs: Optional[List[str]] = None, ) -> bool: """Compile a generated .hpp into a .so via hipcc (compile then link).""" + # Normalize once, here at the boundary, so every downstream use -- the arch + # defines below and the --offload-arch/-DGFX_ARCH we hand to hipcc -- sees the + # bare target. A caller-supplied "gfx1250:xnack-" must not reach the compiler + # flags. + gfx_arch = validate_rowcol_tensor_quant_gfx_arch(gfx_arch) + ck_include = _get_ck_include_dir() static_lib = _get_dispatcher_static_lib() obj_path = so_path.with_suffix(".o") arch_defines = arch_feature_defines(gfx_arch) + # Match top-level CK policy in both host and device compilation. + if gfx_arch == "gfx1250": + arch_defines.append("-DUSE_NEW_UNIFIED_FRAMEWORK=0") compile_cmd = [hipcc, "-c", "-fPIC", "-O3", "-std=c++17", "-DCK_TILE_SINGLE_KERNEL_INCLUDE", "-w", f"--offload-arch={gfx_arch}", + f"-DCK_CMAKE_GPU_TARGET_IDS=0x{gfx_arch[3:]}", f"-DGFX_ARCH=\"{gfx_arch}\"", *arch_defines, "-include", str(hpp_path), @@ -572,7 +668,17 @@ def setup_multiple_tensorquant_dispatchers( if not configs: return [] - arch = gfx_arch or _detect_gpu_arch() + # Normalize the explicit branch too, not just detection: an explicitly passed + # "gfx1250:xnack-" would otherwise flow into --offload-arch, -DGFX_ARCH and the + # .so cache name. _detect_gpu_arch() already normalizes. + arch = validate_rowcol_tensor_quant_gfx_arch(gfx_arch or _detect_gpu_arch()) + for cfg in configs: + config_arch = validate_rowcol_tensor_quant_gfx_arch(cfg.gfx_arch) + if config_arch != arch: + raise ValueError( + f"Config architecture {cfg.gfx_arch!r} does not match build target {arch!r}. " + "Create configs for the selected build target." + ) base_dir = output_dir or Path(tempfile.mkdtemp(prefix="tensorquant_dispatcher_")) base_dir.mkdir(parents=True, exist_ok=True) @@ -597,7 +703,9 @@ def _build_one(idx: int, cfg: TensorQuantKernelConfig) -> Tuple[int, Optional[Pa if hpp is None: return idx, None - so = so_dir / f"lib{cfg.name}_{arch}.so" + # Name carries the ABI revision: a pre-_SO_ABI artifact in a persistent + # output_dir can never be selected here (see _SO_ABI). + so = so_dir / f"lib{cfg.name}_{arch}_abi{_SO_ABI}.so" if so.exists(): log.info(" [cached] %s", so.name) return idx, so @@ -658,6 +766,10 @@ def _default_config(dtype: str, gfx_arch: str) -> TensorQuantKernelConfig: four pad/persistent flags) is forwarded from the shared dict; only block_size and k_block_per_cu are codegen-only and are left to the dataclass defaults. """ + # Normalize at this boundary too: default_fp8_config()/default_bf8_config() are + # public entry points, so a caller-supplied "gfx1250:xnack-" must not be stored + # on the config and observed by later consumers. + gfx_arch = validate_rowcol_tensor_quant_gfx_arch(gfx_arch) traits = ROWCOL_TENSOR_QUANT_DEFAULT_TRAITS return TensorQuantKernelConfig( dtype=dtype, @@ -665,7 +777,7 @@ def _default_config(dtype: str, gfx_arch: str) -> TensorQuantKernelConfig: pipeline=traits["pipeline"], epilogue=traits["epilogue"], scheduler=traits["scheduler"], - **ROWCOL_TENSOR_QUANT_DEFAULT_TILE, + **rowcol_tensor_quant_default_tile(gfx_arch), pad_m=traits["pad_m"], pad_n=traits["pad_n"], pad_k=traits["pad_k"], diff --git a/projects/composablekernel/dispatcher/tests/CMakeLists.txt b/projects/composablekernel/dispatcher/tests/CMakeLists.txt index bbf6019e0734..662dc8444c66 100644 --- a/projects/composablekernel/dispatcher/tests/CMakeLists.txt +++ b/projects/composablekernel/dispatcher/tests/CMakeLists.txt @@ -135,6 +135,53 @@ set_tests_properties(dispatcher_test_arch_filter_constraints PROPERTIES ENVIRONMENT "PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/../python:${CMAKE_CURRENT_SOURCE_DIR}/../codegen:${CMAKE_CURRENT_SOURCE_DIR}/../scripts" ) +# codegen_common shared-infrastructure regression test. Covers normalize_gfx_arch() +# and the EXACT-gfx1250 RowColQuant/TensorQuant tile selector: a family-wide match +# there would hand gfx1200/gfx1201 a tile that compiles cleanly and returns garbage, +# which no downstream check catches. +add_test( + NAME dispatcher_test_codegen_common + COMMAND ${Python3_EXECUTABLE} -m pytest -q ${CMAKE_CURRENT_SOURCE_DIR}/test_codegen_common.py + WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} +) + +set_tests_properties(dispatcher_test_codegen_common PROPERTIES + LABELS "dispatcher;python;codegen;validation" + TIMEOUT 120 + ENVIRONMENT "PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/../python:${CMAKE_CURRENT_SOURCE_DIR}/../codegen:${CMAKE_CURRENT_SOURCE_DIR}/../scripts" +) + +# Arch-string handling in tile_engine's gemm_validation_utils, which seven GEMM +# instance builders and the dispatcher's own warp gate all share. Registered on the +# dispatcher side because the bridge/Old-TE instance-set parity that depends on it +# is a dispatcher property; tile_engine has no CTest registration of its own. +add_test( + NAME dispatcher_test_gemm_validation_arch_normalization + COMMAND ${Python3_EXECUTABLE} -m pytest -q ${CMAKE_CURRENT_SOURCE_DIR}/test_gemm_validation_arch_normalization.py + WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} +) + +set_tests_properties(dispatcher_test_gemm_validation_arch_normalization PROPERTIES + LABELS "dispatcher;python;validation" + TIMEOUT 120 + ENVIRONMENT "PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/../python:${CMAKE_CURRENT_SOURCE_DIR}/../codegen:${CMAKE_CURRENT_SOURCE_DIR}/../scripts" +) + +# Arch-gate predicate used by the RowColQuant/TensorQuant GPU tests. CPU-only: it +# exercises the string predicate, not a device. Worth its own registration because +# the failure mode it guards is a SKIP -- CI stays green while covering nothing. +add_test( + NAME dispatcher_test_gpu_test_arch_gate + COMMAND ${Python3_EXECUTABLE} -m pytest -q ${CMAKE_CURRENT_SOURCE_DIR}/test_gpu_test_arch_gate.py + WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} +) + +set_tests_properties(dispatcher_test_gpu_test_arch_gate PROPERTIES + LABELS "dispatcher;python;validation" + TIMEOUT 60 + ENVIRONMENT "PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/../python:${CMAKE_CURRENT_SOURCE_DIR}/../codegen:${CMAKE_CURRENT_SOURCE_DIR}/../scripts" +) + # dispatcher_common unit tests: path helpers, validators, and -- since this is # the only place they are covered -- the fp8-encoding arch predicates # (fp8_uses_ocp / ocp_arch_defines / arch_feature_defines). Those decide whether @@ -148,6 +195,11 @@ add_test( ) set_tests_properties(dispatcher_test_dispatcher_common PROPERTIES + LABELS "dispatcher;python;validation" + TIMEOUT 60 + ENVIRONMENT "PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/../python:${CMAKE_CURRENT_SOURCE_DIR}/../codegen:${CMAKE_CURRENT_SOURCE_DIR}/../scripts" +) + # Architecture-aware LDS staging budget regression test add_test( NAME dispatcher_test_lds_capacity_arch_aware @@ -296,7 +348,8 @@ set_tests_properties(dispatcher_test_search_space_sampling PROPERTIES # ============================================================================= foreach(_quant_test test_grouped_gemm_rowcolquant_utils - test_grouped_gemm_tensorquant_utils) + test_grouped_gemm_tensorquant_utils + test_quant_review_regressions) add_test( NAME dispatcher_${_quant_test} COMMAND ${Python3_EXECUTABLE} -m pytest ${CMAKE_CURRENT_SOURCE_DIR}/${_quant_test}.py -v diff --git a/projects/composablekernel/dispatcher/tests/test_codegen_common.py b/projects/composablekernel/dispatcher/tests/test_codegen_common.py index 6e14b201943e..ca041bfbd8e5 100644 --- a/projects/composablekernel/dispatcher/tests/test_codegen_common.py +++ b/projects/composablekernel/dispatcher/tests/test_codegen_common.py @@ -38,6 +38,12 @@ gemm_aquant_effective_epilogue, make_bquant_kernel_name, make_gemm_rowcolquant_kernel_name, + normalize_gfx_arch, + rowcol_tensor_quant_default_tile, + ROWCOL_TENSOR_QUANT_DEFAULT_TILE, + ROWCOL_TENSOR_QUANT_DEFAULT_TILE_GFX1250, + ROWCOL_TENSOR_QUANT_SUPPORTED_ARCHES, + validate_rowcol_tensor_quant_gfx_arch, make_aquant_kernel_name, # Shared quant spec-sweep plumbing (see TestQuantSpecSweepHelpers). fp8_warp_tile_k_for_arch, @@ -633,5 +639,172 @@ def test_unknown_arch_takes_the_safe_non_gfx950_branch(self): self.assertEqual(fp8_warp_tile_k_for_arch("gfx90a"), 32) +class TestNormalizeGfxArch(unittest.TestCase): + """Feature suffixes must be stripped exactly once, here.""" + + def test_bare_target_is_unchanged(self): + self.assertEqual(normalize_gfx_arch("gfx1250"), "gfx1250") + self.assertEqual(normalize_gfx_arch("gfx942"), "gfx942") + + def test_single_suffix_is_stripped(self): + self.assertEqual(normalize_gfx_arch("gfx1250:xnack-"), "gfx1250") + + def test_multiple_suffixes_are_stripped(self): + self.assertEqual(normalize_gfx_arch("gfx942:sramecc+:xnack-"), "gfx942") + + def test_unknown_or_malformed_input_is_passed_through(self): + # Not a validator: an unrecognised or empty string must come back + # unchanged (minus any suffix) so the caller's own arch test decides. + self.assertEqual(normalize_gfx_arch(""), "") + self.assertEqual(normalize_gfx_arch("notagfx:weird"), "notagfx") + self.assertEqual(normalize_gfx_arch(":xnack-"), "") + + def test_agrees_with_the_tile_engine_copy(self): + """The deliberate duplicate in tile_engine must not drift from this one. + + tile_engine cannot import codegen_common (dependency direction is + dispatcher -> tile_engine, and tile_engine is on the deprecation path), + so the rule exists twice on purpose. This test is what keeps the two + honest, and the docstrings on both sides point at it. Skipped rather than + failed if the tile_engine tree is absent, so the dispatcher tests stay + runnable on their own. + """ + te_path = ( + DISPATCHER_DIR.parent + / "tile_engine" + / "ops" + / "gemm" + / "gemm_validation_utils.py" + ) + if not te_path.exists(): + self.skipTest("tile_engine tree not present") + import importlib.util + + spec = importlib.util.spec_from_file_location( + "_te_gemm_validation_utils_for_test", te_path + ) + te = importlib.util.module_from_spec(spec) + spec.loader.exec_module(te) + + for arch in ( + "", + "gfx942", + "gfx950", + "gfx1250", + "gfx1250:xnack-", + "gfx942:sramecc+:xnack-", + "notagfx:weird", + ":xnack-", + ): + with self.subTest(arch=arch): + self.assertEqual(te._base_gfx_arch(arch), normalize_gfx_arch(arch)) + + +class TestRowColTensorQuantDefaultTile(unittest.TestCase): + """gfx1250 tile selection is EXACT, not a gfx12 family match. + + gfx1200/gfx1201 are WMMA parts too, but their 8-bit warp fragment is + 16x16x16, not gfx1250's 16x16x64/16x16x128. Handing them the gfx1250 tile + compiles cleanly and returns garbage, so a family match here is a + silent-wrong-answer bug. + """ + + def test_gfx1250_takes_the_gfx1250_tile(self): + self.assertEqual( + rowcol_tensor_quant_default_tile("gfx1250"), + dict(ROWCOL_TENSOR_QUANT_DEFAULT_TILE_GFX1250), + ) + + def test_gfx1250_with_feature_suffix_takes_the_gfx1250_tile(self): + self.assertEqual( + rowcol_tensor_quant_default_tile("gfx1250:xnack-"), + dict(ROWCOL_TENSOR_QUANT_DEFAULT_TILE_GFX1250), + ) + + def test_other_gfx12_parts_do_not_inherit_the_gfx1250_tile(self): + for arch in ("gfx1200", "gfx1201", "gfx1201:xnack-"): + with self.subTest(arch=arch): + self.assertEqual( + rowcol_tensor_quant_default_tile(arch), + dict(ROWCOL_TENSOR_QUANT_DEFAULT_TILE), + ) + + def test_gfx9_parts_take_the_mfma_tile(self): + for arch in ("gfx942", "gfx950", "gfx942:sramecc+:xnack-", ""): + with self.subTest(arch=arch): + self.assertEqual( + rowcol_tensor_quant_default_tile(arch), + dict(ROWCOL_TENSOR_QUANT_DEFAULT_TILE), + ) + + def test_returns_a_copy_not_the_shared_dict(self): + t = rowcol_tensor_quant_default_tile("gfx1250") + t["tile_m"] = -1 + self.assertNotEqual(ROWCOL_TENSOR_QUANT_DEFAULT_TILE_GFX1250["tile_m"], -1) + + +class TestValidateRowColTensorQuantGfxArch(unittest.TestCase): + """--gfx-arch has to be checked, because nothing downstream checks it. + + The flag only selects a default tile; it never reaches --offload-arch, so an + unrecognized value is not rejected by hipcc the way a bad compile target is. It + just falls through to the gfx9 MFMA tile -- which on gfx1250 compiles and returns + garbage. That is the failure this operator's gfx1250 work exists to prevent, + reachable by a single typo on the command line. + """ + + def test_supported_arches_pass_through_bare(self): + for arch in ROWCOL_TENSOR_QUANT_SUPPORTED_ARCHES: + with self.subTest(arch=arch): + self.assertEqual(validate_rowcol_tensor_quant_gfx_arch(arch), arch) + + def test_feature_suffixes_are_accepted_and_stripped(self): + # A real device reports the suffixed form; refusing it would make the flag + # unusable with the name the hardware gives. + for given, expected in [ + ("gfx1250:xnack-", "gfx1250"), + ("gfx950:sramecc+:xnack-", "gfx950"), + ("gfx942:xnack+", "gfx942"), + ]: + with self.subTest(arch=given): + self.assertEqual(validate_rowcol_tensor_quant_gfx_arch(given), expected) + + def test_empty_means_unspecified(self): + self.assertEqual(validate_rowcol_tensor_quant_gfx_arch(""), "") + + def test_a_typo_is_rejected_rather_than_silently_gfx9(self): + for arch in ("gfx1205", "gfx12500", "gfx942x", "gfx", "1250", "GFX1250"): + with self.subTest(arch=arch): + with self.assertRaises(ValueError): + validate_rowcol_tensor_quant_gfx_arch(arch) + + def test_architectures_without_a_defensible_tile_are_rejected(self): + # These are real targets and some are listed in arch_specs.json, but neither + # default tile is correct for them: gfx1200/gfx1201 are WMMA parts with a + # 16x16x16 8-bit fragment, and gfx908 has no fp8 MFMA at all. Accepting them + # would hand back a tile this operator has never validated. + for arch in ("gfx908", "gfx1100", "gfx1200", "gfx1201", "gfx1201:xnack-"): + with self.subTest(arch=arch): + with self.assertRaises(ValueError): + validate_rowcol_tensor_quant_gfx_arch(arch) + + def test_the_error_names_the_value_and_the_alternatives(self): + with self.assertRaises(ValueError) as ctx: + validate_rowcol_tensor_quant_gfx_arch("gfx1205") + message = str(ctx.exception) + self.assertIn("gfx1205", message) + for arch in ROWCOL_TENSOR_QUANT_SUPPORTED_ARCHES: + self.assertIn(arch, message) + + def test_the_set_matches_the_bridge_runtime(self): + self.assertEqual( + ROWCOL_TENSOR_QUANT_SUPPORTED_ARCHES, ("gfx942", "gfx950", "gfx1250") + ) + for arch in ("gfx90a", "gfx90a:xnack-"): + with self.assertRaises(ValueError): + validate_rowcol_tensor_quant_gfx_arch(arch) + + + if __name__ == "__main__": unittest.main() diff --git a/projects/composablekernel/dispatcher/tests/test_gemm_utils.py b/projects/composablekernel/dispatcher/tests/test_gemm_utils.py index 47856269069e..7b9b248f3268 100644 --- a/projects/composablekernel/dispatcher/tests/test_gemm_utils.py +++ b/projects/composablekernel/dispatcher/tests/test_gemm_utils.py @@ -560,3 +560,107 @@ def test_grouped_gfx1250_ci_config_uses_wmma_warp_tile(self): self.assertEqual(tc["warp_tile_m"]["values"], [16]) self.assertEqual(tc["warp_tile_n"]["values"], [16]) self.assertEqual(tc["warp_tile_k"]["values"], [32]) + + +# --- feature-suffixed arch names ------------------------------------------- +# _resolve_arch() validated its input against _SUPPORTED_ARCHES *before* +# normalizing it, so a target carrying feature flags was rejected outright: +# +# gfx950:sramecc+:xnack- -> ValueError +# +# Bare targets were unaffected -- gfx1250 is in _SUPPORTED_ARCHES and passed both +# before and after (test_bare_supported_arches_are_unchanged below). +# +# The suffixed form is what hipDeviceProp_t::gcnArchName reports, the form rocminfo prints +# for the ISA line, and the form this repository's own CMakeLists.txt uses for +# GPU_TARGETS ("gfx908:xnack+;gfx90a:xnack+;gfx942:xnack+;gfx950:xnack+"), so a +# caller copying a target from any of those could not drive the bridge with it. +# Normalization now runs first and the base name is validated. +from gemm_utils import _validate_arch, normalize_gfx_arch # noqa: E402 + + +class TestSuffixedArchNames(unittest.TestCase): + def test_supported_arch_with_feature_suffix_resolves_to_its_base(self): + for given, expected in [ + ("gfx90a:xnack+", "gfx90a"), + ("gfx942:xnack+", "gfx942"), + ("gfx942:sramecc+:xnack-", "gfx942"), + ("gfx950:sramecc+:xnack-", "gfx950"), + ("gfx1250:xnack-", "gfx1250"), + ]: + with self.subTest(arch=given): + self.assertEqual(_resolve_arch(given), expected) + + def test_bare_supported_arches_are_unchanged(self): + # The whole point of normalizing before validating is that it costs the + # existing paths nothing. + for arch in _SUPPORTED_ARCHES: + with self.subTest(arch=arch): + self.assertEqual(_resolve_arch(arch), arch) + + def test_an_unsupported_arch_still_raises_with_or_without_a_suffix(self): + # Normalizing must not turn the gate off: gfx908 is a real target, it is + # simply not one this bridge supports, and neither form may pass. + for arch in ("gfx908", "gfx908:xnack+", "gfx1200", "gfx1201:xnack-"): + with self.subTest(arch=arch): + with self.assertRaises(ValueError): + _resolve_arch(arch) + + def test_a_typo_still_raises(self): + for arch in ("gfx942x", "gfx1205", "gfx", "gfx950x:xnack-"): + with self.subTest(arch=arch): + with self.assertRaises(ValueError): + _resolve_arch(arch) + + def test_the_error_shows_both_the_input_and_the_normalized_form(self): + with self.assertRaises(ValueError) as ctx: + _validate_arch("gfx908:xnack+") + message = str(ctx.exception) + self.assertIn("gfx908:xnack+", message) + self.assertIn("gfx908", message) + + +class TestArchNormalizationMatchesCodegen(unittest.TestCase): + """gemm_utils cannot import codegen_common at module scope -- nothing has put + the codegen dir on sys.path by then -- so it carries a fallback. Pin the two + to identical behaviour, the same way this branch pins tile_engine's + _base_gfx_arch to codegen_common.normalize_gfx_arch.""" + + CASES = ( + "gfx942", + "gfx942:sramecc+:xnack-", + "gfx1250:xnack-", + "gfx950:sramecc+", + "gfx908", + "", + "not-an-arch", + ) + + def test_delegates_to_codegen_common_when_importable(self): + sys.path.insert(0, str(DISPATCHER_DIR / "codegen")) + from codegen_common import normalize_gfx_arch as canonical + + for arch in self.CASES: + with self.subTest(arch=arch): + self.assertEqual(normalize_gfx_arch(arch), canonical(arch)) + + def test_the_fallback_agrees_with_the_canonical_helper(self): + sys.path.insert(0, str(DISPATCHER_DIR / "codegen")) + from codegen_common import normalize_gfx_arch as canonical + + import builtins + + real_import = builtins.__import__ + + def no_codegen_common(name, *args, **kwargs): + if name == "codegen_common": + raise ImportError("simulated: codegen dir not on sys.path") + return real_import(name, *args, **kwargs) + + builtins.__import__ = no_codegen_common + try: + for arch in self.CASES: + with self.subTest(arch=arch): + self.assertEqual(normalize_gfx_arch(arch), canonical(arch)) + finally: + builtins.__import__ = real_import diff --git a/projects/composablekernel/dispatcher/tests/test_gemm_validation_arch_normalization.py b/projects/composablekernel/dispatcher/tests/test_gemm_validation_arch_normalization.py new file mode 100644 index 000000000000..cdc737d4d1ee --- /dev/null +++ b/projects/composablekernel/dispatcher/tests/test_gemm_validation_arch_normalization.py @@ -0,0 +1,255 @@ +#!/usr/bin/env python3 + +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""Arch-string handling in tile_engine's gemm_validation_utils. + +This module is shared by seven GEMM instance builders, and the dispatcher +bridges gate on it too (dispatcher/python/gemm_utils.py imports it directly and +mirrors its warp gate to keep bridge and Old-TE instance sets identical). A +change to how it reads an architecture string is therefore never local, and the +targets that actually reach it are not always bare: CK's own CMakeLists +configures GPU_TARGETS as "gfx908:xnack+;gfx90a:xnack+;gfx942:xnack+;gfx950:xnack+" +for ASAN builds, and that string is what gets handed to --gpu_target. + +These tests pin the two decisions taken while enabling gfx1250: + + 1. Fragment-shape tests ARE normalized. Which MFMA/WMMA fragment the silicon + has does not depend on xnack or sramecc. + 2. The WARP_SUPPORTED_COMBINATIONS lookup is NOT normalized. See + TestWarpConfigurationLookupIsUnchangedOnGfx9 for the evidence. + +Run: python3 -m pytest tests/test_gemm_validation_arch_normalization.py -v +""" + +import importlib.util +import itertools +import logging +import sys +import unittest +from pathlib import Path + +SCRIPT_DIR = Path(__file__).parent.resolve() +DISPATCHER_DIR = SCRIPT_DIR.parent +CK_ROOT = DISPATCHER_DIR.parent + +_TE_VALIDATION = CK_ROOT / "tile_engine" / "ops" / "gemm" / "gemm_validation_utils.py" + + +def _load_validation_module(): + spec = importlib.util.spec_from_file_location( + "_te_gemm_validation_utils_arch_test", _TE_VALIDATION + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +if not _TE_VALIDATION.exists(): # pragma: no cover - tile_engine not checked out + gvu = None +else: + gvu = _load_validation_module() + + +_SUFFIXES = { + "gfx90a": "gfx90a:sramecc+:xnack-", + "gfx942": "gfx942:sramecc+:xnack-", + "gfx950": "gfx950:sramecc+:xnack-", + "gfx1201": "gfx1201:xnack-", + "gfx1250": "gfx1250:xnack-", +} + + +@unittest.skipIf(gvu is None, "tile_engine tree not present") +class TestFragmentShapeChecksAreNormalized(unittest.TestCase): + """A feature suffix must not change which MFMA/WMMA fragment we believe in. + + _validate_fp8_mfma_warp_tile_k encodes ISA fragment shapes. gfx950 doubles + the fp8 K-block relative to gfx90a/gfx942 (MFMA_F32_16x16x256_F8 and + MFMA_F32_32x32x128_F8). Before normalization, "gfx950:sramecc+:xnack-" fell + through the `== "gfx950"` test into the gfx942 branch, so on an ASAN-style + configure the validator rejected gfx950's correct warp_tile_k and accepted + gfx942's wrong one. Both directions of that are pinned here. + """ + + def test_suffixed_target_matches_bare_target_for_every_arch(self): + for bare, suffixed in _SUFFIXES.items(): + for dtype, wt_m, wt_k in itertools.product( + ("fp8", "bf8", "fp16", "bf16"), (16, 32, 64), (16, 32, 64, 128, 256) + ): + with self.subTest(arch=bare, dtype=dtype, m=wt_m, k=wt_k): + self.assertEqual( + gvu._validate_fp8_mfma_warp_tile_k( + wt_m, wt_m, wt_k, dtype, bare + )[0], + gvu._validate_fp8_mfma_warp_tile_k( + wt_m, wt_m, wt_k, dtype, suffixed + )[0], + ) + + def test_suffixed_gfx950_uses_the_gfx950_k_block_not_the_gfx942_one(self): + for dtype in ("fp8", "bf8"): + for wt_m, good_k, gfx942_k in ((16, 128, 64), (32, 64, 32)): + with self.subTest(dtype=dtype, warp_tile_m=wt_m): + ok, _ = gvu._validate_fp8_mfma_warp_tile_k( + wt_m, wt_m, good_k, dtype, _SUFFIXES["gfx950"] + ) + self.assertTrue(ok, "gfx950 K-block must be accepted") + bad, _ = gvu._validate_fp8_mfma_warp_tile_k( + wt_m, wt_m, gfx942_k, dtype, _SUFFIXES["gfx950"] + ) + self.assertFalse(bad, "gfx942 K-block must not be accepted") + + def test_error_message_echoes_the_caller_supplied_target(self): + # Normalization is internal; diagnostics should show what was passed in + # so a mis-typed --gpu_target is visible in the message. + _, msg = gvu._validate_fp8_mfma_warp_tile_k( + 32, 32, 999, "fp8", _SUFFIXES["gfx950"] + ) + self.assertIn(_SUFFIXES["gfx950"], msg) + + def test_gfx1250_rejects_the_32x32_warp_tile_bare_and_suffixed(self): + # There is no 32x32 WMMA fragment on gfx1250; a 32x32 warp tile compiles + # and returns garbage, so this must be rejected however the target is + # spelled. + for target in ("gfx1250", _SUFFIXES["gfx1250"]): + for dtype in ("fp8", "bf8"): + with self.subTest(target=target, dtype=dtype): + ok, _ = gvu._validate_fp8_mfma_warp_tile_k( + 32, 32, 32, dtype, target + ) + self.assertFalse(ok) + + +@unittest.skipIf(gvu is None, "tile_engine tree not present") +class TestWarpConfigurationLookupIsUnchangedOnGfx9(unittest.TestCase): + """The WARP_SUPPORTED_COMBINATIONS lookup stays raw for gfx9. On purpose. + + Normalizing it would activate a validator that a suffixed target currently + bypasses -- which sounds like a fix until you look at what the table + contains. dispatcher/codegen/arch_specs.json, the declared single JSON + source of truth for arch data, lists seven warp maps for gfx942 and nine for + gfx950; this table lists three for each. Turning it on for suffixed names + therefore rejects configurations the repo's own spec calls valid. + + gfx1250 is the exception, and the only one: it is absent from + arch_specs.json entirely, so this table is its sole listing, and it is the + target this work exists for. + + These tests fail if someone "tidies up" the asymmetry. + """ + + _GRID = [1, 2, 4, 8] + + def _decisions(self, target): + return { + (m, n, k): gvu.validate_warp_configuration(m, n, k, target) + for m, n, k in itertools.product(self._GRID, repeat=3) + } + + def test_suffixed_gfx9_targets_stay_permissive(self): + # develop's behaviour: an unrecognized key logs and allows. Pinned so the + # gfx9 instance sets cannot move under an ASAN-style configure. + for bare in ("gfx90a", "gfx942", "gfx950"): + suffixed = _SUFFIXES[bare] + with self.subTest(arch=suffixed): + self.assertTrue( + all(self._decisions(suffixed).values()), + f"{suffixed} must remain unchecked by this table", + ) + + def test_bare_gfx9_targets_still_enforce_their_three_maps(self): + for bare in ("gfx90a", "gfx942", "gfx950"): + decisions = self._decisions(bare) + with self.subTest(arch=bare): + self.assertEqual( + {c for c, ok in decisions.items() if ok}, + {(1, 4, 1), (2, 2, 1), (4, 1, 1)}, + ) + + def test_the_table_is_a_strict_subset_of_arch_specs_json(self): + """Why the gfx9 lookup is left alone, asserted rather than asserted-in-prose. + + If the two ever agree, the argument above evaporates and the raw lookup + should be revisited -- so this test is written to fail at that point. + """ + import json + + specs_path = DISPATCHER_DIR / "codegen" / "arch_specs.json" + if not specs_path.exists(): + self.skipTest("arch_specs.json not present") + specs = json.load(open(specs_path))["architectures"] + for arch in ("gfx942", "gfx950"): + table = {tuple(c) for c in gvu.WARP_SUPPORTED_COMBINATIONS[arch]} + spec = {tuple(c) for c in specs[arch]["warp_configs"]} + with self.subTest(arch=arch): + self.assertTrue( + table < spec, + f"{arch}: gemm_validation_utils lists {sorted(table)}, " + f"arch_specs.json lists {sorted(spec)}. If these now agree, " + f"normalizing the gfx9 lookup is no longer a coverage loss.", + ) + + def test_suffixed_gfx1250_is_checked_like_bare_gfx1250(self): + self.assertEqual( + self._decisions(_SUFFIXES["gfx1250"]), self._decisions("gfx1250") + ) + + def test_suffixed_gfx1250_rejects_maps_absent_from_the_table(self): + # develop allowed all of these on "gfx1250:xnack-" because the lookup + # missed. 1x1x1 is the example the review asked about. + for combo in ((1, 1, 1), (1, 1, 2), (8, 8, 8), (4, 4, 1)): + with self.subTest(combo=combo): + self.assertFalse( + gvu.validate_warp_configuration(*combo, _SUFFIXES["gfx1250"]) + ) + + def test_gfx1201_suffix_is_not_swept_up_by_the_gfx1250_exception(self): + # The exception is exact-gfx1250, not gfx12-family. gfx1201 keeps + # develop's permissive behaviour under a suffix. + self.assertTrue(all(self._decisions(_SUFFIXES["gfx1201"]).values())) + + +@unittest.skipIf(gvu is None, "tile_engine tree not present") +class TestGroupedQuantWarpRestrictionScope(unittest.TestCase): + """Exercise operator routing, not just the shared quant validator.""" + + PREFIXES = ("gemm_rowcolquant", "grouped_gemm_rowcolquant", "grouped_gemm_tensorquant") + + @staticmethod + def config(arch, dtype, warp): + return dict( + tile_m=128, tile_n=128, tile_k=256, + warp_m=warp[0], warp_n=warp[1], warp_k=warp[2], + warp_tile_m=16, warp_tile_n=16, warp_tile_k=128, + a_datatype=dtype, b_datatype=dtype, c_datatype="fp16", + pipeline="compv3", layout="rcr", gpu_target=arch, + ) + + def test_grouped_limits_do_not_filter_plain_rowcolquant(self): + for arch, dtype, prefix, warp in itertools.product( + ("gfx1250", "gfx1250:xnack-"), ("fp8", "bf8"), self.PREFIXES, + ((4, 2, 1), (1, 2, 2), (1, 4, 1)), + ): + with self.subTest(arch=arch, dtype=dtype, prefix=prefix, warp=warp): + # Eight warps and warp_k=2 are grouped-bridge restrictions. + # Four warps with warp_k=1 remain valid for all three operators. + expected = prefix == "gemm_rowcolquant" or warp == (1, 4, 1) + self.assertEqual(gvu.is_tile_config_valid( + **self.config(arch, dtype, warp), kernel_name_prefix=prefix, + ), expected) + + def test_shared_fragment_validation_still_applies_to_all_quant_operators(self): + for arch, dtype, prefix in itertools.product( + ("gfx1250", "gfx1250:xnack-"), ("fp8", "bf8"), self.PREFIXES, + ): + with self.subTest(arch=arch, dtype=dtype, prefix=prefix): + cfg = self.config(arch, dtype, (1, 4, 1)) + cfg.update(warp_tile_m=32, warp_tile_n=32) + self.assertFalse(gvu.is_tile_config_valid(**cfg, kernel_name_prefix=prefix)) + + +if __name__ == "__main__": + logging.disable(logging.CRITICAL) + sys.exit(unittest.main()) diff --git a/projects/composablekernel/dispatcher/tests/test_gpu_test_arch_gate.py b/projects/composablekernel/dispatcher/tests/test_gpu_test_arch_gate.py new file mode 100644 index 000000000000..75baa472af81 --- /dev/null +++ b/projects/composablekernel/dispatcher/tests/test_gpu_test_arch_gate.py @@ -0,0 +1,142 @@ +#!/usr/bin/env python3 + +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""Unit tests for the arch gate used by the RowColQuant/TensorQuant GPU tests. + +These run on CPU: they exercise the pure string predicate, not a device. + +Why this file exists +-------------------- +The GPU correctness tests decide whether to run or skip by comparing the detected +architecture against a tuple of supported targets. That comparison used to be an +exact ``in`` test, which is wrong: ``rocm_agent_enumerator`` (and +``hipDeviceProp_t::gcnArchName``) may report feature suffixes, so a perfectly +supported device can enumerate as ``"gfx942:sramecc+:xnack-"``. An exact compare +then evaluates False and the test SKIPS. + +A skip is not a failure, so CI stays green while covering nothing -- the defect is +invisible by construction. That failure mode is why this predicate gets its own +regression test rather than being trusted to review. +""" + +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parent)) + +import test_rowcolquant_gpu_correctness as rcq +import test_tensorquant_gpu_correctness as tq + +# Both modules must gate identically; parametrizing over them keeps them in step. +MODULES = pytest.mark.parametrize("mod", [rcq, tq], ids=["rowcolquant", "tensorquant"]) + + +@MODULES +@pytest.mark.parametrize("arch", ["gfx942", "gfx950", "gfx1250"]) +def test_bare_supported_arch_is_accepted(mod, arch): + assert mod.arch_is_supported(arch) is True + + +@MODULES +@pytest.mark.parametrize( + "arch", + [ + "gfx942:sramecc+:xnack-", + "gfx942:sramecc-:xnack-", + "gfx950:sramecc+:xnack-", + "gfx1250:xnack-", + ], +) +def test_feature_suffixed_supported_arch_is_accepted(mod, arch): + """The regression this file exists for: suffixes must not cause a skip.""" + assert mod.arch_is_supported(arch) is True, ( + f"{arch} is a supported device but the gate rejected it; " + "the GPU test would silently skip and CI would be vacuously green" + ) + + +@MODULES +@pytest.mark.parametrize( + "arch", + [ + "gfx90a", + "gfx90a:sramecc+:xnack-", # unsupported and suffixed -> still unsupported + "gfx1030", + "gfx000", + "", + ], +) +def test_unsupported_arch_is_rejected(mod, arch): + assert mod.arch_is_supported(arch) is False + + +@MODULES +def test_no_substring_false_positives(mod): + """Prefix tolerance must not degrade into a substring match. + + ``gfx9421`` is a different (hypothetical) target and must not be accepted just + because ``gfx942`` is a prefix of it. Splitting on ``:`` and comparing the whole + base token gives us that for free; a naive ``startswith`` would not. + """ + assert mod.arch_is_supported("gfx9421") is False + assert mod.arch_is_supported("gfx12501") is False + + +@MODULES +@pytest.mark.parametrize( + "raw,expected", + [ + ("gfx942:sramecc+:xnack-", "gfx942"), + ("gfx1250", "gfx1250"), + ("gfx950:xnack-", "gfx950"), + ("", ""), + ], +) +def test_normalize_strips_feature_suffix(mod, raw, expected): + """Normalization also feeds ``--offload-arch``, so it must yield a bare target.""" + assert mod.normalize_gfx_arch(raw) == expected + + +@MODULES +def test_gate_matches_supported_tuple(mod): + """Every entry of the declared tuple must pass its own gate, bare and suffixed.""" + for arch in mod._SUPPORTED_ARCHES: + assert mod.arch_is_supported(arch) is True + assert mod.arch_is_supported(f"{arch}:sramecc+:xnack-") is True + + + + +@MODULES +@pytest.mark.parametrize("detected", ["", "gfx90a"]) +def test_standalone_reports_skip_for_unavailable_target(mod, detected, monkeypatch, tmp_path): + monkeypatch.setattr(sys, "argv", [mod.__file__, "--output-dir", str(tmp_path)]) + monkeypatch.setattr(mod, "_GFX_ARCH", detected) + assert mod.main() == 77 + + +@MODULES +def test_standalone_suffix_override_reaches_all_cases(mod, monkeypatch, tmp_path): + monkeypatch.setattr(sys, "argv", [mod.__file__, "--gfx", "gfx1250:xnack-", + "--output-dir", str(tmp_path)]) + monkeypatch.setattr(mod, "_GFX_ARCH", "") + monkeypatch.setattr(mod, "_has_hipcc", lambda: True) + monkeypatch.setattr(mod, "_ml_dtypes", object()) + encoded = [] + def inputs(M, N, K, dtype, arch): + encoded.append((M, N, K, dtype, arch)) + return () + monkeypatch.setattr(mod, "_make_inputs", inputs) + monkeypatch.setattr(mod, "_run_one", lambda *a, **kw: ("PASS", "host test stub")) + assert mod.main() == 0 + assert len(encoded) == 3 + assert all(case[-1] == "gfx1250" for case in encoded) + assert (64, 256, 128, "fp8", "gfx1250") in encoded + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"])) diff --git a/projects/composablekernel/dispatcher/tests/test_grouped_gemm_rowcolquant_utils.py b/projects/composablekernel/dispatcher/tests/test_grouped_gemm_rowcolquant_utils.py index cc19453b01c7..d4e861a1a3d8 100644 --- a/projects/composablekernel/dispatcher/tests/test_grouped_gemm_rowcolquant_utils.py +++ b/projects/composablekernel/dispatcher/tests/test_grouped_gemm_rowcolquant_utils.py @@ -16,9 +16,12 @@ import sys from pathlib import Path +import pytest + sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "python")) +import grouped_gemm_rowcolquant_utils as UTILS from grouped_gemm_rowcolquant_utils import ( RowColQuantKernelConfig, RowColQuantGemmProblem, @@ -301,22 +304,215 @@ def test_header_launch_takes_vector_of_host_args(self): class TestDefaultConfigAlignment: """Ensure default_fp8_config/default_bf8_config stay in sync with _default_config().""" - def _codegen_default_names(self): + # Both defaults are arch-dependent, so comparing them at a single architecture + # cannot see a drift that only exists on another one -- which is exactly the + # drift gfx1250 enablement introduces. Every arch the bridge supports is + # checked, plus a suffixed spelling on each side. + _ARCHES = ["gfx942", "gfx950", "gfx1250", "gfx1250:xnack-", "gfx942:sramecc+:xnack-"] + + def _codegen_default_names(self, gfx_arch=""): sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "codegen")) from unified_grouped_gemm_rowcolquant_codegen import _default_config, _build_specs - specs = _build_specs(_default_config()) + specs = _build_specs(_default_config(gfx_arch)) return {s.name for s in specs} - def test_default_fp8_config_name_in_codegen_defaults(self): - cfg = default_fp8_config() - assert cfg.name in self._codegen_default_names(), ( - f"default_fp8_config().name '{cfg.name}' is not produced by _default_config() " - f"in the codegen. The two defaults have drifted — update one to match the other." + @pytest.mark.parametrize("gfx_arch", _ARCHES) + def test_default_fp8_config_name_in_codegen_defaults(self, gfx_arch): + cfg = default_fp8_config(gfx_arch) + assert cfg.name in self._codegen_default_names(gfx_arch), ( + f"default_fp8_config('{gfx_arch}').name '{cfg.name}' is not produced by " + f"_default_config('{gfx_arch}') in the codegen. The two defaults have " + f"drifted — update one to match the other." ) - def test_default_bf8_config_name_in_codegen_defaults(self): - cfg = default_bf8_config() - assert cfg.name in self._codegen_default_names(), ( - f"default_bf8_config().name '{cfg.name}' is not produced by _default_config() " - f"in the codegen. The two defaults have drifted — update one to match the other." + @pytest.mark.parametrize("gfx_arch", _ARCHES) + def test_default_bf8_config_name_in_codegen_defaults(self, gfx_arch): + cfg = default_bf8_config(gfx_arch) + assert cfg.name in self._codegen_default_names(gfx_arch), ( + f"default_bf8_config('{gfx_arch}').name '{cfg.name}' is not produced by " + f"_default_config('{gfx_arch}') in the codegen. The two defaults have " + f"drifted — update one to match the other." + ) + + def test_gfx1250_default_actually_differs_from_the_gfx9_default(self): + """Guard against the parametrization passing for the wrong reason. + + If _default_config() ever stops honouring gfx_arch, every case above + still passes -- both sides would just fall back to the gfx9 tile + together. Assert the two architectures really do produce different + kernels, so that regression is visible. + """ + assert self._codegen_default_names("gfx1250").isdisjoint( + self._codegen_default_names("gfx942") + ) + + +# ============================================================================= +# Arch-dependent hipcc defines +# ============================================================================= + + +class TestArchDefines: + """The OCP-FP8 define is deliberately a gfx12 *family* test. + + gfx1200/gfx1201/gfx1250 all use OCP FP8 encoding, so narrowing this to an + exact gfx1250 match would be wrong. This is the opposite of the tile + selector in codegen_common.rowcol_tensor_quant_default_tile(), which must be + exact because gfx1200/gfx1201 have a different 8-bit warp fragment. Both + behaviours are pinned so neither gets "fixed" into the other. + """ + + def _compile_argv(self, monkeypatch, tmp_path, gfx_arch): + captured = [] + + class _Result: + returncode = 0 + stdout = "" + stderr = "" + + def fake_run(cmd, *a, **kw): + captured.append(list(cmd)) + return _Result() + + monkeypatch.setattr(UTILS.subprocess, "run", fake_run) + ok = UTILS._compile_rowcolquant_kernel( + hpp_path=tmp_path / "k.hpp", + so_path=tmp_path / "k.so", + gfx_arch=gfx_arch, + ) + assert ok, "compile helper should report success when hipcc succeeds" + return captured[0] + + @pytest.mark.parametrize("arch", ["gfx1250", "gfx1250:xnack-"]) + def test_supported_gfx12_targets_get_ocp_fp8(self, monkeypatch, tmp_path, arch): + argv = self._compile_argv(monkeypatch, tmp_path, arch) + assert "-DUSE_NEW_UNIFIED_FRAMEWORK=0" in argv + assert "-DCK_CMAKE_GPU_TARGET_IDS=0x1250" in argv + assert "-DCK_USE_OCP_FP8" in argv + assert "-DCK_TILE_USE_OCP_FP8" in argv + + def test_gfx950_gets_ocp_fp8_and_mx(self, monkeypatch, tmp_path): + argv = self._compile_argv(monkeypatch, tmp_path, "gfx950") + assert "-DCK_CMAKE_GPU_TARGET_IDS=0x950" in argv + assert "-DCK_USE_OCP_FP8" in argv + assert "-DCK_USE_NATIVE_MX_SUPPORT" in argv + + def test_gfx942_gets_neither(self, monkeypatch, tmp_path): + argv = self._compile_argv(monkeypatch, tmp_path, "gfx942:sramecc+:xnack-") + assert "-DCK_USE_OCP_FP8" not in argv + assert "-DCK_USE_NATIVE_MX_SUPPORT" not in argv + + +# ============================================================================= +# Compiler-flag arch normalization +# ============================================================================= + + +class TestArchNormalizationInCompileFlags: + """A suffixed target must never reach --offload-arch / -DGFX_ARCH. + + normalize_gfx_arch() used to be applied only to the *detected* arch. A caller + passing gfx_arch="gfx1250:xnack-" therefore had the raw string forwarded into + the compiler flags. Normalization now happens once at the entry boundary. + """ + + def _compile_argv(self, monkeypatch, tmp_path, gfx_arch): + captured = [] + + class _Result: + returncode = 0 + stdout = "" + stderr = "" + + def fake_run(cmd, *a, **kw): + captured.append(cmd) + return _Result() + + monkeypatch.setattr(UTILS.subprocess, "run", fake_run) + ok = UTILS._compile_rowcolquant_kernel( + hpp_path=tmp_path / "k.hpp", + so_path=tmp_path / "k.so", + gfx_arch=gfx_arch, ) + assert ok + return captured[0] + + def test_suffixed_arch_is_stripped_from_offload_arch(self, monkeypatch, tmp_path): + argv = self._compile_argv(monkeypatch, tmp_path, "gfx1250:xnack-") + assert "--offload-arch=gfx1250" in argv + assert not any("xnack" in tok for tok in argv), argv + + def test_suffixed_arch_is_stripped_from_gfx_arch_define(self, monkeypatch, tmp_path): + argv = self._compile_argv(monkeypatch, tmp_path, "gfx942:sramecc+:xnack-") + assert '-DGFX_ARCH="gfx942"' in argv + + def test_default_config_stores_bare_arch(self): + assert default_fp8_config(gfx_arch="gfx1250:xnack-").gfx_arch == "gfx1250" + assert default_bf8_config(gfx_arch="gfx942:sramecc+").gfx_arch == "gfx942" + + +# ============================================================================= +# .so cache key encodes the ABI +# ============================================================================= + + +class TestSoCacheAbiKey: + """A pre-ABI artifact in a persistent output_dir must not be reused. + + setup_multiple_rowcolquant_dispatchers() reuses a .so when the filename + matches. The name used to be lib{kernel}_{arch}.so, which says nothing about + the exported symbol set, so a .so predating the + dispatcher_get_tile_n()/dispatcher_get_pad_n() exports was selected and then + died at attribute lookup with "undefined symbol". The name now carries + _SO_ABI, so the stale artifact simply cannot be selected. + """ + + def test_abi_is_versioned(self): + assert isinstance(UTILS._SO_ABI, int) and UTILS._SO_ABI >= 2 + + @pytest.mark.parametrize("old_suffix", ["", "_abi2"]) + def test_pre_abi_artifact_is_not_reused(self, monkeypatch, tmp_path, old_suffix): + cfg = default_fp8_config(gfx_arch="gfx950") + + so_dir = tmp_path / "libs" + so_dir.mkdir(parents=True) + # Exactly the name the old code would have produced and reused. + stale = so_dir / f"lib{cfg.name}_gfx950{old_suffix}.so" + stale.write_bytes(b"stale pre-ABI artifact") + + compiled = [] + + def fake_compile(hpp_path, so_path, gfx_arch, **kw): + compiled.append(so_path) + so_path.write_bytes(b"fresh") + return True + + monkeypatch.setattr(UTILS, "_compile_rowcolquant_kernel", fake_compile) + + out = UTILS.setup_multiple_rowcolquant_dispatchers( + configs=[cfg], output_dir=tmp_path, gfx_arch="gfx950", parallel=False, + ) + + assert compiled, "stale pre-ABI .so was reused instead of rebuilt" + assert out[0] != stale + assert f"_abi{UTILS._SO_ABI}.so" in out[0].name + assert stale.read_bytes() == b"stale pre-ABI artifact", "stale file was clobbered" + + def test_second_call_hits_the_cache(self, monkeypatch, tmp_path): + cfg = default_fp8_config(gfx_arch="gfx950") + calls = [] + + def fake_compile(hpp_path, so_path, gfx_arch, **kw): + calls.append(so_path) + so_path.write_bytes(b"fresh") + return True + + monkeypatch.setattr(UTILS, "_compile_rowcolquant_kernel", fake_compile) + + first = UTILS.setup_multiple_rowcolquant_dispatchers( + configs=[cfg], output_dir=tmp_path, gfx_arch="gfx950", parallel=False) + second = UTILS.setup_multiple_rowcolquant_dispatchers( + configs=[cfg], output_dir=tmp_path, gfx_arch="gfx950", parallel=False) + + assert len(calls) == 1, "second call should have hit the cache" + assert first[0] == second[0] diff --git a/projects/composablekernel/dispatcher/tests/test_grouped_gemm_tensorquant_utils.py b/projects/composablekernel/dispatcher/tests/test_grouped_gemm_tensorquant_utils.py index 381f796af92a..359ea2a8c683 100644 --- a/projects/composablekernel/dispatcher/tests/test_grouped_gemm_tensorquant_utils.py +++ b/projects/composablekernel/dispatcher/tests/test_grouped_gemm_tensorquant_utils.py @@ -16,9 +16,12 @@ import sys from pathlib import Path +import pytest + sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "python")) +import grouped_gemm_tensorquant_utils as UTILS from grouped_gemm_tensorquant_utils import ( TensorQuantKernelConfig, TensorQuantGemmProblem, @@ -308,22 +311,215 @@ def test_header_does_not_contain_rowcolquant(self): class TestDefaultConfigAlignment: """Ensure default_fp8_config/default_bf8_config stay in sync with _default_config().""" - def _codegen_default_names(self): + # Both defaults are arch-dependent, so comparing them at a single architecture + # cannot see a drift that only exists on another one -- which is exactly the + # drift gfx1250 enablement introduces. Every arch the bridge supports is + # checked, plus a suffixed spelling on each side. + _ARCHES = ["gfx942", "gfx950", "gfx1250", "gfx1250:xnack-", "gfx942:sramecc+:xnack-"] + + def _codegen_default_names(self, gfx_arch=""): sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "codegen")) from unified_grouped_gemm_tensorquant_codegen import _default_config, _build_specs - specs = _build_specs(_default_config()) + specs = _build_specs(_default_config(gfx_arch)) return {s.name for s in specs} - def test_default_fp8_config_name_in_codegen_defaults(self): - cfg = default_fp8_config() - assert cfg.name in self._codegen_default_names(), ( - f"default_fp8_config().name '{cfg.name}' is not produced by _default_config() " - f"in the codegen. The two defaults have drifted — update one to match the other." + @pytest.mark.parametrize("gfx_arch", _ARCHES) + def test_default_fp8_config_name_in_codegen_defaults(self, gfx_arch): + cfg = default_fp8_config(gfx_arch) + assert cfg.name in self._codegen_default_names(gfx_arch), ( + f"default_fp8_config('{gfx_arch}').name '{cfg.name}' is not produced by " + f"_default_config('{gfx_arch}') in the codegen. The two defaults have " + f"drifted — update one to match the other." ) - def test_default_bf8_config_name_in_codegen_defaults(self): - cfg = default_bf8_config() - assert cfg.name in self._codegen_default_names(), ( - f"default_bf8_config().name '{cfg.name}' is not produced by _default_config() " - f"in the codegen. The two defaults have drifted — update one to match the other." + @pytest.mark.parametrize("gfx_arch", _ARCHES) + def test_default_bf8_config_name_in_codegen_defaults(self, gfx_arch): + cfg = default_bf8_config(gfx_arch) + assert cfg.name in self._codegen_default_names(gfx_arch), ( + f"default_bf8_config('{gfx_arch}').name '{cfg.name}' is not produced by " + f"_default_config('{gfx_arch}') in the codegen. The two defaults have " + f"drifted — update one to match the other." + ) + + def test_gfx1250_default_actually_differs_from_the_gfx9_default(self): + """Guard against the parametrization passing for the wrong reason. + + If _default_config() ever stops honouring gfx_arch, every case above + still passes -- both sides would just fall back to the gfx9 tile + together. Assert the two architectures really do produce different + kernels, so that regression is visible. + """ + assert self._codegen_default_names("gfx1250").isdisjoint( + self._codegen_default_names("gfx942") + ) + + +# ============================================================================= +# Arch-dependent hipcc defines +# ============================================================================= + + +class TestArchDefines: + """The OCP-FP8 define is deliberately a gfx12 *family* test. + + gfx1200/gfx1201/gfx1250 all use OCP FP8 encoding, so narrowing this to an + exact gfx1250 match would be wrong. This is the opposite of the tile + selector in codegen_common.rowcol_tensor_quant_default_tile(), which must be + exact because gfx1200/gfx1201 have a different 8-bit warp fragment. Both + behaviours are pinned so neither gets "fixed" into the other. + """ + + def _compile_argv(self, monkeypatch, tmp_path, gfx_arch): + captured = [] + + class _Result: + returncode = 0 + stdout = "" + stderr = "" + + def fake_run(cmd, *a, **kw): + captured.append(list(cmd)) + return _Result() + + monkeypatch.setattr(UTILS.subprocess, "run", fake_run) + ok = UTILS._compile_tensorquant_kernel( + hpp_path=tmp_path / "k.hpp", + so_path=tmp_path / "k.so", + gfx_arch=gfx_arch, + ) + assert ok, "compile helper should report success when hipcc succeeds" + return captured[0] + + @pytest.mark.parametrize("arch", ["gfx1250", "gfx1250:xnack-"]) + def test_supported_gfx12_targets_get_ocp_fp8(self, monkeypatch, tmp_path, arch): + argv = self._compile_argv(monkeypatch, tmp_path, arch) + assert "-DUSE_NEW_UNIFIED_FRAMEWORK=0" in argv + assert "-DCK_CMAKE_GPU_TARGET_IDS=0x1250" in argv + assert "-DCK_USE_OCP_FP8" in argv + assert "-DCK_TILE_USE_OCP_FP8" in argv + + def test_gfx950_gets_ocp_fp8_and_mx(self, monkeypatch, tmp_path): + argv = self._compile_argv(monkeypatch, tmp_path, "gfx950") + assert "-DCK_CMAKE_GPU_TARGET_IDS=0x950" in argv + assert "-DCK_USE_OCP_FP8" in argv + assert "-DCK_USE_NATIVE_MX_SUPPORT" in argv + + def test_gfx942_gets_neither(self, monkeypatch, tmp_path): + argv = self._compile_argv(monkeypatch, tmp_path, "gfx942:sramecc+:xnack-") + assert "-DCK_USE_OCP_FP8" not in argv + assert "-DCK_USE_NATIVE_MX_SUPPORT" not in argv + + +# ============================================================================= +# Compiler-flag arch normalization +# ============================================================================= + + +class TestArchNormalizationInCompileFlags: + """A suffixed target must never reach --offload-arch / -DGFX_ARCH. + + normalize_gfx_arch() used to be applied only to the *detected* arch. A caller + passing gfx_arch="gfx1250:xnack-" therefore had the raw string forwarded into + the compiler flags. Normalization now happens once at the entry boundary. + """ + + def _compile_argv(self, monkeypatch, tmp_path, gfx_arch): + captured = [] + + class _Result: + returncode = 0 + stdout = "" + stderr = "" + + def fake_run(cmd, *a, **kw): + captured.append(cmd) + return _Result() + + monkeypatch.setattr(UTILS.subprocess, "run", fake_run) + ok = UTILS._compile_tensorquant_kernel( + hpp_path=tmp_path / "k.hpp", + so_path=tmp_path / "k.so", + gfx_arch=gfx_arch, ) + assert ok + return captured[0] + + def test_suffixed_arch_is_stripped_from_offload_arch(self, monkeypatch, tmp_path): + argv = self._compile_argv(monkeypatch, tmp_path, "gfx1250:xnack-") + assert "--offload-arch=gfx1250" in argv + assert not any("xnack" in tok for tok in argv), argv + + def test_suffixed_arch_is_stripped_from_gfx_arch_define(self, monkeypatch, tmp_path): + argv = self._compile_argv(monkeypatch, tmp_path, "gfx942:sramecc+:xnack-") + assert '-DGFX_ARCH="gfx942"' in argv + + def test_default_config_stores_bare_arch(self): + assert default_fp8_config(gfx_arch="gfx1250:xnack-").gfx_arch == "gfx1250" + assert default_bf8_config(gfx_arch="gfx942:sramecc+").gfx_arch == "gfx942" + + +# ============================================================================= +# .so cache key encodes the ABI +# ============================================================================= + + +class TestSoCacheAbiKey: + """A pre-ABI artifact in a persistent output_dir must not be reused. + + setup_multiple_tensorquant_dispatchers() reuses a .so when the filename + matches. The name used to be lib{kernel}_{arch}.so, which says nothing about + the exported symbol set, so a .so predating the + dispatcher_get_tile_n()/dispatcher_get_pad_n() exports was selected and then + died at attribute lookup with "undefined symbol". The name now carries + _SO_ABI, so the stale artifact simply cannot be selected. + """ + + def test_abi_is_versioned(self): + assert isinstance(UTILS._SO_ABI, int) and UTILS._SO_ABI >= 2 + + @pytest.mark.parametrize("old_suffix", ["", "_abi2"]) + def test_pre_abi_artifact_is_not_reused(self, monkeypatch, tmp_path, old_suffix): + cfg = default_fp8_config(gfx_arch="gfx950") + + so_dir = tmp_path / "libs" + so_dir.mkdir(parents=True) + # Exactly the name the old code would have produced and reused. + stale = so_dir / f"lib{cfg.name}_gfx950{old_suffix}.so" + stale.write_bytes(b"stale pre-ABI artifact") + + compiled = [] + + def fake_compile(hpp_path, so_path, gfx_arch, **kw): + compiled.append(so_path) + so_path.write_bytes(b"fresh") + return True + + monkeypatch.setattr(UTILS, "_compile_tensorquant_kernel", fake_compile) + + out = UTILS.setup_multiple_tensorquant_dispatchers( + configs=[cfg], output_dir=tmp_path, gfx_arch="gfx950", parallel=False, + ) + + assert compiled, "stale pre-ABI .so was reused instead of rebuilt" + assert out[0] != stale + assert f"_abi{UTILS._SO_ABI}.so" in out[0].name + assert stale.read_bytes() == b"stale pre-ABI artifact", "stale file was clobbered" + + def test_second_call_hits_the_cache(self, monkeypatch, tmp_path): + cfg = default_fp8_config(gfx_arch="gfx950") + calls = [] + + def fake_compile(hpp_path, so_path, gfx_arch, **kw): + calls.append(so_path) + so_path.write_bytes(b"fresh") + return True + + monkeypatch.setattr(UTILS, "_compile_tensorquant_kernel", fake_compile) + + first = UTILS.setup_multiple_tensorquant_dispatchers( + configs=[cfg], output_dir=tmp_path, gfx_arch="gfx950", parallel=False) + second = UTILS.setup_multiple_tensorquant_dispatchers( + configs=[cfg], output_dir=tmp_path, gfx_arch="gfx950", parallel=False) + + assert len(calls) == 1, "second call should have hit the cache" + assert first[0] == second[0] diff --git a/projects/composablekernel/dispatcher/tests/test_quant_review_regressions.py b/projects/composablekernel/dispatcher/tests/test_quant_review_regressions.py new file mode 100644 index 000000000000..fda26cf79a56 --- /dev/null +++ b/projects/composablekernel/dispatcher/tests/test_quant_review_regressions.py @@ -0,0 +1,259 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""Host regressions for quant build targets, direct loading, and C++ guards.""" + +import ctypes +import importlib +import json +from pathlib import Path +import shlex +import shutil +import subprocess +import sys +from types import SimpleNamespace + +import pytest + +ROOT = Path(__file__).resolve().parents[1] +sys.path[:0] = [str(ROOT / "python"), str(ROOT / "codegen")] +OPS = ("rowcolquant", "tensorquant") +CXX = shutil.which("c++") + + +@pytest.mark.parametrize("op", OPS) +def test_custom_quant_tiles_require_explicit_arch(op, tmp_path): + codegen = importlib.import_module(f"unified_grouped_gemm_{op}_codegen") + config = codegen._default_config("gfx1250") + output = tmp_path / "headers" + with pytest.raises(ValueError, match="explicit.*gfx.arch"): + codegen.generate_kernels(output, config=config, parallel=False) + assert not output.exists() + + # The same custom tile is valid when its target is explicitly recorded. + header = codegen.generate_kernels(output, config=config, gfx_arch="gfx1250:xnack-", + parallel=False)[0].read_text() + assert 'GfxArch = "gfx1250"' in header + # Preserve legacy no-config generation without a target. + legacy = codegen.generate_kernels(tmp_path / "legacy", parallel=False)[0].read_text() + assert 'GfxArch = ""' in legacy + assert "WarpTileM = 32" in legacy + + +@pytest.mark.parametrize("op", OPS) +@pytest.mark.parametrize("config_mode", ["--config", "--config-json"]) +@pytest.mark.parametrize("list_names", [False, True]) +def test_custom_quant_cli_requires_explicit_arch(op, config_mode, list_names, tmp_path): + codegen = importlib.import_module(f"unified_grouped_gemm_{op}_codegen") + config_json = json.dumps(codegen._default_config("gfx1250")) + config_file = tmp_path / "config.json" + config_file.write_text(config_json) + output = tmp_path / "headers" + args = [sys.executable, str(ROOT / f"codegen/unified_grouped_gemm_{op}_codegen.py"), + "--output-dir", str(output), config_mode, + str(config_file) if config_mode == "--config" else config_json] + if list_names: + args.append("--list-names") + result = subprocess.run(args, capture_output=True, text=True) + assert result.returncode != 0 + assert "explicit --gfx-arch" in result.stderr + assert "Traceback" not in result.stderr + assert not output.exists() + + +@pytest.mark.parametrize("op", OPS) +@pytest.mark.parametrize("config_arch,build_arch", [("gfx1250", "gfx942"), ("gfx942", "gfx1250")]) +@pytest.mark.parametrize("explicit", [False, True]) +def test_config_build_mismatch_fails_before_codegen(op, config_arch, build_arch, explicit, monkeypatch, tmp_path): + module = importlib.import_module(f"grouped_gemm_{op}_utils") + monkeypatch.setattr(module, "_detect_gpu_arch", lambda: build_arch) + monkeypatch.setattr(module, f"_generate_{op}_kernel", lambda *a: pytest.fail("codegen must not run")) + setup = getattr(module, f"setup_multiple_{op}_dispatchers") + with pytest.raises(ValueError, match="does not match build target"): + setup([module.default_fp8_config(config_arch)], output_dir=tmp_path / "output", + gfx_arch=build_arch if explicit else None) + assert not (tmp_path / "output").exists() + + +@pytest.mark.parametrize("op", OPS) +def test_matching_suffix_config_generates_target_metadata(op, tmp_path): + module = importlib.import_module(f"grouped_gemm_{op}_utils") + config = module.default_fp8_config("gfx942:sramecc+:xnack-") + config.gfx_arch = "gfx942:xnack-" + header = getattr(module, f"_generate_{op}_kernel")(config, tmp_path) + assert header is not None + assert 'GfxArch = "gfx942"' in header.read_text() + + +@pytest.mark.parametrize("op", OPS) +@pytest.mark.parametrize("arch", ["gfx90a", "gfx90a:xnack-", "gfx1200", "gfx1201"]) +def test_unsupported_quant_targets_rejected_at_entry(op, arch, tmp_path): + module = importlib.import_module(f"grouped_gemm_{op}_utils") + with pytest.raises(ValueError, match="Unsupported GPU architecture"): + module.default_fp8_config(arch) + result = subprocess.run([sys.executable, str(ROOT / f"codegen/unified_grouped_gemm_{op}_codegen.py"), + "--gfx-arch", arch, "--output-dir", str(tmp_path)], capture_output=True, text=True) + assert result.returncode != 0 + assert "Unsupported GPU architecture" in result.stderr + assert not list(tmp_path.glob("*.hpp")) + + +@pytest.mark.parametrize("op,cls", [("rowcolquant", "RowColQuant"), ("tensorquant", "TensorQuant")]) +@pytest.mark.parametrize("missing", ["dispatcher_get_tile_n", "dispatcher_get_pad_n"]) +def test_direct_runner_reports_old_abi(op, cls, missing, monkeypatch, tmp_path): + module = importlib.import_module(f"grouped_gemm_{op}_utils") + lib = SimpleNamespace() + lifecycle_calls = [] + for name in ("dispatcher_initialize", "dispatcher_run_gemm", "dispatcher_get_kernel_name", + "dispatcher_get_kernel_count", "dispatcher_get_tile_n", "dispatcher_get_pad_n", + "dispatcher_cleanup"): + if name != missing: + setattr(lib, name, lambda *args: 0) + lib.dispatcher_initialize = lambda: lifecycle_calls.append("initialize") + lib.dispatcher_cleanup = lambda: lifecycle_calls.append("cleanup") + path = tmp_path / "old.so" + path.touch() + monkeypatch.setattr(module.ctypes, "CDLL", lambda *a: lib) + with pytest.raises(RuntimeError, match="Incompatible quant bridge ABI.*Rebuild"): + getattr(module, cls + "GpuGemmRunner")(path) + assert lifecycle_calls == [] + + +def compile_cpp(tmp_path, source, shared=False): + if not CXX: + pytest.skip("requires C++ compiler") + cpp = tmp_path / "probe.cpp" + cpp.write_text(source) + output = tmp_path / ("probe.so" if shared else "probe.o") + args = [CXX, "-std=c++17", str(cpp), "-o", str(output)] + args += ["-shared", "-fPIC"] if shared else ["-c"] + return subprocess.run(args, capture_output=True, text=True), output + + +@pytest.mark.parametrize("op", OPS) +@pytest.mark.parametrize("generated,compiled,passes", [ + ("gfx1250", "gfx942", False), ("gfx942", "gfx1250", False), + ("gfx1250", "gfx1250", True), ("gfx942", "gfx942", True), ("gfx950", "gfx950", True), +]) +def test_production_cpp_rejects_header_target_mismatch(op, generated, compiled, passes, tmp_path): + codegen = importlib.import_module(f"unified_grouped_gemm_{op}_codegen") + header = codegen.generate_kernels(tmp_path / "headers", gfx_arch=generated, parallel=False)[0].read_text() + # Compile the actual emitted constants and actual bridge guard block without + # requiring HIP. Kernel implementation and GPU execution are tested separately. + constants = header[header.index(" static constexpr const char* GfxArch"):header.index(" // Informational only:")] + constants = constants.replace("ck_tile::index_t", "int") + bridge = (ROOT / f"bindings/ctypes/grouped_gemm_{op}_ctypes_lib.cpp").read_text() + guards = bridge[bridge.index("static constexpr bool ct_starts_with"):bridge.index('extern "C" {')] + source = f'#define GFX_ARCH "{compiled}"\nstruct SelectedKernel {{\n{constants}\n}};\n{guards}' + result, _ = compile_cpp(tmp_path, source) + assert (result.returncode == 0) == passes, result.stderr + if not passes: + assert "Generated kernel architecture does not match" in result.stderr + + +@pytest.mark.parametrize("arch,pad_n,M,N,expected", [ + ("gfx1250", False, 128, 128, 0), + ("gfx1250", False, 126, 128, -1), + ("gfx1250", False, 128, 100, -1), + ("gfx942", False, 126, 128, 0), + ("gfx1250", True, 128, 100, 0), +]) +def test_rowcolquant_actual_shape_guards(arch, pad_n, M, N, expected, tmp_path): + bridge = (ROOT / "bindings/ctypes/grouped_gemm_rowcolquant_ctypes_lib.cpp").read_text() + # Execute the production host-side prefix of dispatcher_run_gemm. Stop before + # allocation/launch; valid inputs return a sentinel success in this harness. + prefix = bridge[bridge.index("int dispatcher_run_gemm("):bridge.index(" // Only packed (contiguous) layouts")] + source = f'''#include +#include +#include +#define GFX_ARCH "{arch}" +std::atomic g_ref_count{{1}}; +constexpr bool kCompiledForGfx12 = {str(arch == "gfx1250").lower()}; +struct SelectedKernel {{ static constexpr bool kPadN = {str(pad_n).lower()}; static constexpr int TileN = 64; }}; +extern "C" {{ +{prefix} +return 0; +}} +}} +''' + result, output = compile_cpp(tmp_path, source, shared=True) + assert result.returncode == 0, result.stderr + lib = ctypes.CDLL(str(output)) + run = lib.dispatcher_run_gemm + run.argtypes = [ctypes.c_void_p] * 5 + [ctypes.c_int64] * 10 + [ctypes.c_int, ctypes.POINTER(ctypes.c_float)] + run.restype = ctypes.c_int + elapsed = ctypes.c_float() + assert run(*([1] * 5), M, N, 192, 192, 192, 1, 1, N, M, N, 1, ctypes.byref(elapsed)) == expected + + +@pytest.mark.parametrize("inherited_arches", [None, "gfx942;gfx1250", "gfx950;gfx942", "gfx1250;gfx950"]) +@pytest.mark.parametrize("explicit", [True, False]) +@pytest.mark.parametrize("arch", [ + "gfx942", "gfx942:sramecc+:xnack-", + "gfx950", "gfx950:sramecc+:xnack-", + "gfx1250", "gfx1250:xnack-", +]) +def test_cmake_normalizes_explicit_and_inferred_arches(tmp_path, explicit, arch, inherited_arches): + if not shutil.which("cmake"): + pytest.skip("requires CMake") + source = tmp_path / "source" + source.mkdir() + build = tmp_path / "build" + headers = build / "generated_kernels" + headers.mkdir(parents=True) + for op in OPS: + (headers / f"grouped_gemm_{op}_test.hpp").touch() + overrides = '\n'.join(f'set(CK_TILE_{op.upper()}_GFX_ARCH "{arch}")' for op in OPS) if explicit else '' + # Execute the real parent feature-definition block, including its defaults. + parent_features = "" + if inherited_arches: + parent_cmake = (ROOT.parent / "CMakeLists.txt").read_text() + start = parent_cmake.index('if (SUPPORTED_GPU_TARGETS MATCHES "gfx9|gfx11|gfx12"') + end = parent_cmake.index('if ((SUPPORTED_GPU_TARGETS MATCHES "gfx942"', start) + parent_features = f'set(SUPPORTED_GPU_TARGETS "{inherited_arches}")\n' + parent_cmake[start:end] + configured_arches = f"{arch};{inherited_arches}" if inherited_arches else arch + (source / "CMakeLists.txt").write_text(f'''cmake_minimum_required(VERSION 3.16) +project(quant_cmake_probe LANGUAGES CXX) +add_library(hip::device INTERFACE IMPORTED) +{parent_features} +set(CMAKE_HIP_ARCHITECTURES "{configured_arches}") +{overrides} +add_subdirectory("{ROOT / 'bindings/ctypes'}" ctypes) +''') + result = subprocess.run(["cmake", "-S", str(source), "-B", str(build), "-DCMAKE_EXPORT_COMPILE_COMMANDS=ON"], capture_output=True, text=True) + assert result.returncode == 0, result.stdout + result.stderr + commands = json.loads((build / "compile_commands.json").read_text()) + from dispatcher_common import arch_feature_defines + + # Compare effective macros, preserving the generated command's -D/-U order. + # Checking only flag presence misses inherited definitions and overrides. + feature_flags = set().union(*(set(arch_feature_defines(a)) for a in ("gfx942", "gfx950", "gfx1250"))) + feature_flags.update({"-DUSE_NEW_UNIFIED_FRAMEWORK=0", "-DCK_USE_FNUZ_FP8", + "-DCK_USE_XDL", "-DCK_USE_GFX94", "-DCK_USE_GFX950", + "-DCK_USE_WMMA", "-DCK_USE_WMMA_FP8", "-DCK_GFX1030_SUPPORT"}) + feature_names = {flag[2:].split("=")[0] for flag in feature_flags} + + def effective_features(flags, compiler): + result = subprocess.run([compiler, *flags, "-dM", "-E", "-x", "c++", "-"], + input="", capture_output=True, text=True) + assert result.returncode == 0, result.stderr + macros = {} + for line in result.stdout.splitlines(): + _, name, *value = line.split(maxsplit=2) + if name in feature_names: + macros[name] = value[0] if value else "" + return macros + for op in OPS: + command = next(c["command"] for c in commands if c["file"].endswith(f"grouped_gemm_{op}_ctypes_lib.cpp")) + base = arch.split(":")[0] + assert f'GFX_ARCH=\\"{base}\\"' in command + assert f'GFX_ARCH=\\"{base}:' not in command + assert f'-DCK_CMAKE_GPU_TARGET_IDS=0x{base[3:]}' in command + expected = set(arch_feature_defines(base)) + if base == "gfx1250": + expected.add("-DUSE_NEW_UNIFIED_FRAMEWORK=0") + args = shlex.split(command) + flags = [arg for arg in args[1:] if arg.startswith(("-D", "-U"))] + actual = effective_features(flags, args[0]) + expected = effective_features(sorted(expected), args[0]) + assert actual == expected, f"{op} {arch}: CMake/JIT feature mismatch: {actual} != {expected}" diff --git a/projects/composablekernel/dispatcher/tests/test_rowcolquant_gpu_correctness.py b/projects/composablekernel/dispatcher/tests/test_rowcolquant_gpu_correctness.py index f1c4c1d9eebe..858a5316a264 100644 --- a/projects/composablekernel/dispatcher/tests/test_rowcolquant_gpu_correctness.py +++ b/projects/composablekernel/dispatcher/tests/test_rowcolquant_gpu_correctness.py @@ -6,7 +6,7 @@ """ GPU correctness tests for RowColQuant GEMM dispatcher. -Requires a gfx942 or gfx950 GPU and hipcc in PATH. Skipped automatically when neither +Requires a gfx942, gfx950 or gfx1250 GPU and hipcc in PATH. Skipped automatically when neither is available (pytest.skip) so CI without a GPU still passes. Tests: @@ -31,7 +31,11 @@ import pytest sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "python")) +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "codegen")) +# normalize_gfx_arch lives in codegen_common: single source of truth shared by +# the codegen tile selector, the runtime helpers and these tests. +from codegen_common import normalize_gfx_arch # noqa: E402 from dispatcher_common import fp8_uses_ocp as _fp8_uses_ocp from grouped_gemm_rowcolquant_utils import ( @@ -56,14 +60,29 @@ def _has_hipcc() -> bool: return shutil.which("hipcc") is not None +def arch_is_supported(arch: str, supported=None) -> bool: + """True if `arch` names a supported target, ignoring feature suffixes. + + An exact ``arch in _SUPPORTED_ARCHES`` test returns False for + ``"gfx942:sramecc+:xnack-"``, so on a machine whose enumerator reports + suffixes these GPU tests would SKIP on a fully supported device. A skip reads + as a green run, so the regression would be invisible -- the bug is the silent + skip, not the missing coverage. The C++ bridge already prefix-matches in + ``is_supported_arch()``; this keeps the Python gate consistent with it. + """ + if supported is None: + supported = _SUPPORTED_ARCHES + return normalize_gfx_arch(arch) in supported + + def _detect_gfx_arch() -> str: """Return the first usable GPU arch, or empty string if none found.""" try: r = subprocess.run(["rocm_agent_enumerator"], capture_output=True, text=True, timeout=10) for line in r.stdout.splitlines(): arch = line.strip() - if arch.startswith("gfx") and arch != "gfx000": - return arch + if arch.startswith("gfx") and normalize_gfx_arch(arch) != "gfx000": + return normalize_gfx_arch(arch) except Exception: pass return "" @@ -72,11 +91,11 @@ def _detect_gfx_arch() -> str: _GFX_ARCH = _detect_gfx_arch() # RowColQuant fp8/bf8 kernels use CK CompV3 pipelines that require native fp8 hardware. # gfx90a (MI200 series) lacks native fp8 support and produces incorrect results. -# Only gfx942 (MI300X) and gfx950 (MI350X) are validated. -_SUPPORTED_ARCHES = ("gfx942", "gfx950") +# Only gfx942 (MI300X), gfx950 (MI350X) and gfx1250 (MI400) are validated. +_SUPPORTED_ARCHES = ("gfx942", "gfx950", "gfx1250") requires_gpu = pytest.mark.skipif( - not (_has_hipcc() and _GFX_ARCH in _SUPPORTED_ARCHES), + not (_has_hipcc() and arch_is_supported(_GFX_ARCH)), reason=( f"GPU test: requires hipcc and native fp8 GPU ({', '.join(_SUPPORTED_ARCHES)}); " f"detected arch='{_GFX_ARCH}'" @@ -102,26 +121,33 @@ def _require_ml_dtypes(): ) -def _fp8_ml_dtype(dtype: str): - """Return the ml_dtypes fp8 type matching the compiled kernel's arch.""" +def _fp8_ml_dtype(dtype: str, arch: str): + """Return the ml_dtypes fp8 type matching the compiled kernel's arch. + + `arch` is the arch the kernel is being COMPILED for, which is not necessarily + the arch this host detects: main() takes --gfx. Reading the module-level + detected arch here would silently encode the host bytes in the other format, + and the kernel would decode them to NaN/Inf on device. So it is a parameter, + not a global. + """ _require_ml_dtypes() - if _fp8_uses_ocp(_GFX_ARCH): + if _fp8_uses_ocp(arch): return _ml_dtypes.float8_e4m3fn if dtype == "fp8" else _ml_dtypes.float8_e5m2 return _ml_dtypes.float8_e4m3fnuz if dtype == "fp8" else _ml_dtypes.float8_e5m2fnuz -def _encode_fp8(arr: np.ndarray, dtype: str) -> np.ndarray: +def _encode_fp8(arr: np.ndarray, dtype: str, arch: str) -> np.ndarray: """Encode float32 → fp8/bf8 bytes (uint8 view). Requires ml_dtypes. The fp8 format (OCP vs FNUZ) follows the compiled kernel's arch; see _fp8_uses_ocp. """ - ml_t = _fp8_ml_dtype(dtype) + ml_t = _fp8_ml_dtype(dtype, arch) return arr.astype(ml_t).view(np.uint8) -def _decode_fp8(arr: np.ndarray, dtype: str) -> np.ndarray: +def _decode_fp8(arr: np.ndarray, dtype: str, arch: str) -> np.ndarray: """Decode fp8/bf8 bytes (uint8 view) → float32. Requires ml_dtypes.""" - ml_t = _fp8_ml_dtype(dtype) + ml_t = _fp8_ml_dtype(dtype, arch) return arr.view(ml_t).astype(np.float32) @@ -199,16 +225,22 @@ def _run_one(label: str, config, M: int, N: int, K: int, # Input factory # --------------------------------------------------------------------------- -def _make_inputs(M, N, K, dtype="fp8", seed=42): +def _make_inputs(M, N, K, dtype="fp8", seed=42, arch=None): + """Build host inputs. `arch` must be the arch the kernel is compiled for. + + Defaults to the detected arch, which is right for the pytest cases (they + compile for _GFX_ARCH); main() passes its --gfx value explicitly. + """ + arch = _GFX_ARCH if arch is None else arch rng = np.random.default_rng(seed) A_f32 = rng.uniform(-1.0, 1.0, (M, K)).astype(np.float32) B_f32 = rng.uniform(-1.0, 1.0, (K, N)).astype(np.float32) AQ = rng.uniform(0.5, 2.0, (M,)).astype(np.float32) # per-row A scale BQ = rng.uniform(0.5, 2.0, (N,)).astype(np.float32) # per-col B scale - A_raw = _encode_fp8(A_f32, dtype) - B_raw = _encode_fp8(B_f32, dtype) - A_dec = _decode_fp8(A_raw, dtype) - B_dec = _decode_fp8(B_raw, dtype) + A_raw = _encode_fp8(A_f32, dtype, arch) + B_raw = _encode_fp8(B_f32, dtype, arch) + A_dec = _decode_fp8(A_raw, dtype, arch) + B_dec = _decode_fp8(B_raw, dtype, arch) return A_raw, A_dec, B_raw, B_dec, AQ, BQ @@ -284,16 +316,18 @@ def test_rowcolquant_timing_positive(tmp_path): # test_*_timing_positive has no entry of its own because _run_one already fails # the case when time_ms <= 0. TESTS = [ + # arch=gfx, not the detected arch: with --gfx the host encoding must follow the + # arch the kernel is compiled for (see _fp8_ml_dtype). ("C4/fp8", lambda od, gfx: _run_one( "C4/fp8", default_fp8_config(gfx_arch=gfx), 128, 128, 192, - *_make_inputs(128, 128, 192, "fp8"), Path(od), gfx_arch=gfx)), + *_make_inputs(128, 128, 192, "fp8", arch=gfx), Path(od), gfx_arch=gfx)), ("C4/bf8", lambda od, gfx: _run_one( "C4/bf8", default_bf8_config(gfx_arch=gfx), 128, 128, 192, - *_make_inputs(128, 128, 192, "bf8"), Path(od), gfx_arch=gfx)), + *_make_inputs(128, 128, 192, "bf8", arch=gfx), Path(od), gfx_arch=gfx)), # Non-square M/N/K to stress stride math. ("rect/fp8", lambda od, gfx: _run_one( "rect/fp8", default_fp8_config(gfx_arch=gfx), 64, 256, 128, - *_make_inputs(64, 256, 128, "fp8"), Path(od), gfx_arch=gfx)), + *_make_inputs(64, 256, 128, "fp8", arch=gfx), Path(od), gfx_arch=gfx)), ] @@ -319,7 +353,7 @@ def main(): # what CI drives, so it needs the same gates spelled out. SKIP_EXIT keeps a # clean skip distinguishable from a failure -- ctest maps it via # SKIP_RETURN_CODE and the Jenkins lane via its run_ok helper. - gfx = args.gfx or _GFX_ARCH + gfx = normalize_gfx_arch(args.gfx or _GFX_ARCH) if not gfx: print("SKIP: no supported GPU detected (rocm_agent_enumerator)") return SKIP_EXIT diff --git a/projects/composablekernel/dispatcher/tests/test_tensorquant_gpu_correctness.py b/projects/composablekernel/dispatcher/tests/test_tensorquant_gpu_correctness.py index 426da979811f..573a41395c5a 100644 --- a/projects/composablekernel/dispatcher/tests/test_tensorquant_gpu_correctness.py +++ b/projects/composablekernel/dispatcher/tests/test_tensorquant_gpu_correctness.py @@ -6,7 +6,7 @@ """ GPU correctness tests for TensorQuant GEMM dispatcher. -Requires a gfx942 or gfx950 GPU and hipcc in PATH. Skipped automatically when neither +Requires a gfx942, gfx950 or gfx1250 GPU and hipcc in PATH. Skipped automatically when neither is available (pytest.skip) so CI without a GPU still passes. Tests: @@ -31,7 +31,11 @@ import pytest sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "python")) +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "codegen")) +# normalize_gfx_arch lives in codegen_common: single source of truth shared by +# the codegen tile selector, the runtime helpers and these tests. +from codegen_common import normalize_gfx_arch # noqa: E402 from dispatcher_common import fp8_uses_ocp as _fp8_uses_ocp from grouped_gemm_tensorquant_utils import ( @@ -56,14 +60,29 @@ def _has_hipcc() -> bool: return shutil.which("hipcc") is not None +def arch_is_supported(arch: str, supported=None) -> bool: + """True if `arch` names a supported target, ignoring feature suffixes. + + An exact ``arch in _SUPPORTED_ARCHES`` test returns False for + ``"gfx942:sramecc+:xnack-"``, so on a machine whose enumerator reports + suffixes these GPU tests would SKIP on a fully supported device. A skip reads + as a green run, so the regression would be invisible -- the bug is the silent + skip, not the missing coverage. The C++ bridge already prefix-matches in + ``is_supported_arch()``; this keeps the Python gate consistent with it. + """ + if supported is None: + supported = _SUPPORTED_ARCHES + return normalize_gfx_arch(arch) in supported + + def _detect_gfx_arch() -> str: """Return the first usable GPU arch, or empty string if none found.""" try: r = subprocess.run(["rocm_agent_enumerator"], capture_output=True, text=True, timeout=10) for line in r.stdout.splitlines(): arch = line.strip() - if arch.startswith("gfx") and arch != "gfx000": - return arch + if arch.startswith("gfx") and normalize_gfx_arch(arch) != "gfx000": + return normalize_gfx_arch(arch) except Exception: pass return "" @@ -72,11 +91,11 @@ def _detect_gfx_arch() -> str: _GFX_ARCH = _detect_gfx_arch() # TensorQuant fp8/bf8 kernels use CK CompV3 pipelines that require native fp8 hardware. # gfx90a (MI200 series) lacks native fp8 support and produces incorrect results. -# Only gfx942 (MI300X) and gfx950 (MI350X) are validated. -_SUPPORTED_ARCHES = ("gfx942", "gfx950") +# Only gfx942 (MI300X), gfx950 (MI350X) and gfx1250 (MI400) are validated. +_SUPPORTED_ARCHES = ("gfx942", "gfx950", "gfx1250") requires_gpu = pytest.mark.skipif( - not (_has_hipcc() and _GFX_ARCH in _SUPPORTED_ARCHES), + not (_has_hipcc() and arch_is_supported(_GFX_ARCH)), reason=( f"GPU test: requires hipcc and native fp8 GPU ({', '.join(_SUPPORTED_ARCHES)}); " f"detected arch='{_GFX_ARCH}'" @@ -102,26 +121,33 @@ def _require_ml_dtypes(): ) -def _fp8_ml_dtype(dtype: str): - """Return the ml_dtypes fp8 type matching the compiled kernel's arch.""" +def _fp8_ml_dtype(dtype: str, arch: str): + """Return the ml_dtypes fp8 type matching the compiled kernel's arch. + + `arch` is the arch the kernel is being COMPILED for, which is not necessarily + the arch this host detects: main() takes --gfx. Reading the module-level + detected arch here would silently encode the host bytes in the other format, + and the kernel would decode them to NaN/Inf on device. So it is a parameter, + not a global. + """ _require_ml_dtypes() - if _fp8_uses_ocp(_GFX_ARCH): + if _fp8_uses_ocp(arch): return _ml_dtypes.float8_e4m3fn if dtype == "fp8" else _ml_dtypes.float8_e5m2 return _ml_dtypes.float8_e4m3fnuz if dtype == "fp8" else _ml_dtypes.float8_e5m2fnuz -def _encode_fp8(arr: np.ndarray, dtype: str) -> np.ndarray: +def _encode_fp8(arr: np.ndarray, dtype: str, arch: str) -> np.ndarray: """Encode float32 → fp8/bf8 bytes (uint8 view). Requires ml_dtypes. The fp8 format (OCP vs FNUZ) follows the compiled kernel's arch; see _fp8_uses_ocp. """ - ml_t = _fp8_ml_dtype(dtype) + ml_t = _fp8_ml_dtype(dtype, arch) return arr.astype(ml_t).view(np.uint8) -def _decode_fp8(arr: np.ndarray, dtype: str) -> np.ndarray: +def _decode_fp8(arr: np.ndarray, dtype: str, arch: str) -> np.ndarray: """Decode fp8/bf8 bytes (uint8 view) → float32. Requires ml_dtypes.""" - ml_t = _fp8_ml_dtype(dtype) + ml_t = _fp8_ml_dtype(dtype, arch) return arr.view(ml_t).astype(np.float32) @@ -196,17 +222,23 @@ def _run_one(label: str, config, M: int, N: int, K: int, # Input factory # --------------------------------------------------------------------------- -def _make_inputs(M, N, K, dtype="fp8", seed=42): +def _make_inputs(M, N, K, dtype="fp8", seed=42, arch=None): + """Build host inputs. `arch` must be the arch the kernel is compiled for. + + Defaults to the detected arch, which is right for the pytest cases (they + compile for _GFX_ARCH); main() passes its --gfx value explicitly. + """ + arch = _GFX_ARCH if arch is None else arch rng = np.random.default_rng(seed) A_f32 = rng.uniform(-1.0, 1.0, (M, K)).astype(np.float32) B_f32 = rng.uniform(-1.0, 1.0, (K, N)).astype(np.float32) # Single per-tensor scales in (0.5, 2.0] AQ = rng.uniform(0.5, 2.0, (1,)).astype(np.float32) BQ = rng.uniform(0.5, 2.0, (1,)).astype(np.float32) - A_raw = _encode_fp8(A_f32, dtype) - B_raw = _encode_fp8(B_f32, dtype) - A_dec = _decode_fp8(A_raw, dtype) - B_dec = _decode_fp8(B_raw, dtype) + A_raw = _encode_fp8(A_f32, dtype, arch) + B_raw = _encode_fp8(B_f32, dtype, arch) + A_dec = _decode_fp8(A_raw, dtype, arch) + B_dec = _decode_fp8(B_raw, dtype, arch) return A_raw, A_dec, B_raw, B_dec, AQ, BQ @@ -282,16 +314,18 @@ def test_tensorquant_timing_positive(tmp_path): # test_*_timing_positive has no entry of its own because _run_one already fails # the case when time_ms <= 0. TESTS = [ + # arch=gfx, not the detected arch: with --gfx the host encoding must follow the + # arch the kernel is compiled for (see _fp8_ml_dtype). ("C4/fp8", lambda od, gfx: _run_one( "C4/fp8", default_fp8_config(gfx_arch=gfx), 128, 128, 192, - *_make_inputs(128, 128, 192, "fp8"), Path(od), gfx_arch=gfx)), + *_make_inputs(128, 128, 192, "fp8", arch=gfx), Path(od), gfx_arch=gfx)), ("C4/bf8", lambda od, gfx: _run_one( "C4/bf8", default_bf8_config(gfx_arch=gfx), 128, 128, 192, - *_make_inputs(128, 128, 192, "bf8"), Path(od), gfx_arch=gfx)), + *_make_inputs(128, 128, 192, "bf8", arch=gfx), Path(od), gfx_arch=gfx)), # Non-square M/N/K to stress stride math. ("rect/fp8", lambda od, gfx: _run_one( "rect/fp8", default_fp8_config(gfx_arch=gfx), 64, 256, 128, - *_make_inputs(64, 256, 128, "fp8"), Path(od), gfx_arch=gfx)), + *_make_inputs(64, 256, 128, "fp8", arch=gfx), Path(od), gfx_arch=gfx)), ] @@ -317,7 +351,7 @@ def main(): # what CI drives, so it needs the same gates spelled out. SKIP_EXIT keeps a # clean skip distinguishable from a failure -- ctest maps it via # SKIP_RETURN_CODE and the Jenkins lane via its run_ok helper. - gfx = args.gfx or _GFX_ARCH + gfx = normalize_gfx_arch(args.gfx or _GFX_ARCH) if not gfx: print("SKIP: no supported GPU detected (rocm_agent_enumerator)") return SKIP_EXIT diff --git a/projects/composablekernel/tile_engine/ops/gemm/gemm_validation_utils.py b/projects/composablekernel/tile_engine/ops/gemm/gemm_validation_utils.py index ad262a249532..4b5a4330215e 100644 --- a/projects/composablekernel/tile_engine/ops/gemm/gemm_validation_utils.py +++ b/projects/composablekernel/tile_engine/ops/gemm/gemm_validation_utils.py @@ -4,6 +4,33 @@ import logging from typing import Tuple, List + +def _base_gfx_arch(gpu_target: str) -> str: + """Strip feature suffixes from a gfx target string. + + ``"gfx942:sramecc+:xnack-"`` -> ``"gfx942"``. Empty input is passed through + unchanged; this is a normalizer, not a validator. + + This is the ONE place in this module that knows the rule. Call it at the + boundary of a function rather than re-deriving it at each comparison site -- + the failure mode it prevents is a normalized test sitting a few lines above a + raw one, so that a suffixed target takes different branches of the same + function. + + Duplication note: ``dispatcher/codegen/codegen_common.py`` carries the + identical helper (``normalize_gfx_arch``) and cannot be imported from here. + The dependency direction is dispatcher -> tile_engine (see + ``dispatcher/python/gemm_utils.py``, which imports this module), and + tile_engine is on the deprecation path, so new shared infrastructure must not + be parked here for dispatcher to consume. The two copies are held together by + ``dispatcher/tests/test_codegen_common.py::TestNormalizeGfxArch``, which + asserts they agree on the same inputs. + """ + if not gpu_target: + return gpu_target + return gpu_target.split(":", 1)[0] + + GEMM_PIPELINES = ["mem", "compv3", "compv4"] GEMM_PRESHUFFLE_PIPELINES = ["preshufflev2"] @@ -320,7 +347,33 @@ def validate_warp_configuration( current_combination = [warp_m, warp_n, warp_k] + # Deliberately a RAW lookup. Unlike every other arch test in this module, this + # one is NOT normalized, and that asymmetry is the point. + # + # Normalizing it looks like an obvious fix -- a suffixed name misses the table, + # an empty dict comes back, and the permissive branch below turns the whole + # restriction off. But this table is not a description of the hardware. It is a + # hand-maintained subset, and it is stale: dispatcher/codegen/arch_specs.json, + # the declared single JSON source of truth feeding both the Python codegen and + # the C++ runtime filter, lists seven warp maps for gfx942 + # ([1,1,1] [1,2,1] [1,4,1] [2,1,1] [2,1,2] [2,2,1] [4,1,1]) and nine for gfx950, + # where this table lists three. Switching a suffixed gfx942 from "unchecked" to + # "checked against three of its seven maps" does not reject invalid + # configurations; it rejects valid ones, silently narrowing instance generation + # for the ASAN configure in projects/composablekernel/CMakeLists.txt, which sets + # GPU_TARGETS to "gfx908:xnack+;gfx90a:xnack+;gfx942:xnack+;gfx950:xnack+". + # + # Widening the table to match arch_specs.json would be a real fix, but it + # changes what seven instance builders generate on their primary targets and + # belongs in its own PR with its own build evidence -- not smuggled in as a + # side effect of gfx1250 enablement. + # + # The one place normalization is needed is gfx1250, the target this branch + # exists for and one that arch_specs.json does not describe at all, so this + # table is its only listing. Scope the change to exactly that. allowed_combinations = WARP_SUPPORTED_COMBINATIONS.get(gpu_name, {}) + if not allowed_combinations and _base_gfx_arch(gpu_name) == "gfx1250": + allowed_combinations = WARP_SUPPORTED_COMBINATIONS.get("gfx1250", {}) if not allowed_combinations: # If GPU not recognized, try to be permissive but log warning logging.warning(f"No warp_[m/n/k] combinations found for GPU: {gpu_name}") @@ -387,7 +440,7 @@ def validate_lds_capacity( matrix_b_size = (tile_n * tile_k) * element_size(b_datatype) total_tile_in_lds = matrix_a_size + matrix_b_size - base_gpu_target = gpu_target.split(":")[0] if gpu_target else gpu_target + base_gpu_target = _base_gfx_arch(gpu_target) hw_lds_size = LDS_SIZE_MAP.get(base_gpu_target, DEFAULT_LDS_SIZE) double_buffer = pipeline in ["preshufflev2", "compv4"] max_tile_size = hw_lds_size // 2 if double_buffer else hw_lds_size @@ -716,6 +769,21 @@ def is_tile_config_valid( return False # Additional operator-specific validation (runs after pipeline validation) + # These limits come from the two grouped quant bridges' code generation, + # not from gfx1250 hardware or the shared RowColQuant fragment rules. Keep + # them at the routing boundary so plain gemm_rowcolquant is not filtered. + if kernel_name_prefix in ("grouped_gemm_rowcolquant", "grouped_gemm_tensorquant") and ( + _base_gfx_arch(gpu_target) == "gfx1250" + ): + if warp_m * warp_n * warp_k > 4: + logging.debug( + "gfx1250 grouped quant bridges require at most four warps per block" + ) + return False + if warp_k > 1: + logging.debug("gfx1250 grouped quant bridges require warp_k=1") + return False + if ( kernel_name_prefix == "gemm_rowcolquant" or kernel_name_prefix == "grouped_gemm_rowcolquant" @@ -1397,6 +1465,14 @@ def _validate_fp8_mfma_warp_tile_k( - warp_tile_m == warp_tile_n (square MFMA requirement) - warp_tile_k matches the ISA-mandated K-block for the given warp_tile_m and gpu_target """ + # Normalize once, here at the entry. Every arch test in this function is about + # which MFMA/WMMA fragment the silicon has, and a feature suffix does not change + # that: "gfx950:sramecc+:xnack-" is a gfx950 and must take the gfx950 K-block. + # Before this, a suffixed gfx950 fell through to the gfx90a/gfx942 branch and was + # told to use warp_tile_k=32/64 instead of 64/128 -- i.e. the correct tile was + # rejected and the wrong one accepted. `gpu_target` is kept unmodified for the + # error messages, which should echo what the caller passed in. + base_gpu_target = _base_gfx_arch(gpu_target) suffix = f" ({op_label})" if op_label else "" if warp_tile_m != warp_tile_n: return False, ( @@ -1411,7 +1487,26 @@ def _validate_fp8_mfma_warp_tile_k( # gfx950 doubles the K-block: # MFMA_F32_16x16x256_F8 (warp_tile_m=16) → warp_tile_k=128 # MFMA_F32_32x32x128_F8 (warp_tile_m=32) → warp_tile_k=64 - if gpu_target == "gfx950": + if base_gpu_target == "gfx1250": + # gfx1250 is wave32 RDNA-style WMMA, not MFMA. The only 8-bit fragments + # are V_WMMA_*_16x16x64 and 16x16x128; there is no 32x32 WMMA, so a + # 32x32xK warp tile compiles and then returns garbage. This is the sole + # source of silent wrong answers on this surface: all 3,226 wrong-result + # rows of a 14,208-row sweep used 32x32x32, and no other warp tile + # produced one. Both legal K depths are GPU-validated for + # RowColQuant / TensorQuant. + if warp_tile_m != 16: + return False, ( + f"On {gpu_target} the only 8-bit WMMA warp tile is 16x16xK, " + f"got warp_tile_m={warp_tile_m}{suffix}" + ) + if warp_tile_k not in (64, 128): + return False, ( + f"For {a_datatype} on {gpu_target}, warp_tile_m=16 requires " + f"warp_tile_k in (64, 128), got warp_tile_k={warp_tile_k}{suffix}" + ) + return True, "" + if base_gpu_target == "gfx950": expected_k = 64 if warp_tile_m == 32 else 128 else: expected_k = 32 if warp_tile_m == 32 else 64