feat(ck-tile): grouped rowcolquant/tensorquant GEMM bridges on gfx1250 [Experimental- Dont Review]

projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt

@@ -351,6 +351,11 @@ if(ROWCOLQUANT_GEMM_KERNEL_HEADERS)
         else()
             set(_rowcolquant_gfx_arch "gfx942")
         endif()
+        # Strip feature suffixes ("gfx942:sramecc+:xnack-"). CMAKE_HIP_ARCHITECTURES
+        # carries them for ASAN configures, and GFX_ARCH is compared against
+        # kSupportedArchs and the compile-time tile guards in the .cpp, which want the
+        # bare target. Mirrors normalize_gfx_arch() on the Python side.
+        string(REGEX REPLACE ":.*$" "" _rowcolquant_gfx_arch "${_rowcolquant_gfx_arch}")
         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()
@@ -365,7 +370,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 <target> 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.
@@ -397,6 +405,11 @@ if(TENSORQUANT_GEMM_KERNEL_HEADERS)
         else()
             set(_tensorquant_gfx_arch "gfx942")
         endif()
+        # Strip feature suffixes ("gfx942:sramecc+:xnack-"). CMAKE_HIP_ARCHITECTURES
+        # carries them for ASAN configures, and GFX_ARCH is compared against
+        # kSupportedArchs and the compile-time tile guards in the .cpp, which want the
+        # bare target. Mirrors normalize_gfx_arch() on the Python side.
+        string(REGEX REPLACE ":.*$" "" _tensorquant_gfx_arch "${_tensorquant_gfx_arch}")
         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()
@@ -411,7 +424,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 <target> 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.

projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_rowcolquant_ctypes_lib.cpp

@@ -74,7 +74,7 @@ static std::atomic<int> 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,87 @@ static bool is_supported_arch(const std::string& arch)
     return false;
 }
 
+// ---------------------------------------------------------------------------
+// Compile-time gfx1250 tile guard.
+//
+// The runtime helper in grouped_gemm_*_utils.py picks an arch-correct tile, but the
+// CMake/codegen path does not: unified_grouped_gemm_*_codegen.py::_default_config()
+// takes no arch and hard-codes the gfx9 tile. Listing gfx1250 in kSupportedArchs above
+// therefore means a CMake-built gfx1250 library would initialize cleanly and then run a
+// tile that does not exist on the device -- which compiles and returns garbage rather
+// than failing.
+//
+// Rather than trust the build to pick correctly, refuse to produce such a library at
+// all. gfx1250 is wave32 with RDNA-style WMMA: the only 8-bit fragments are 16x16x64
+// and 16x16x128 (there is no 32x32 WMMA), and a block of more than four warps has no
+// launchable kernel entry. These are the same rules the Python validator enforces.
+//
+// The test is exact-gfx1250, not gfx12-family: gfx1200/gfx1201 have a different 8-bit
+// fragment (16x16x16) and are not in kSupportedArchs anyway.
+// --- BEGIN gfx1250 compile-time tile guard ---
+// Extracted verbatim and compiled against a mock kernel struct by
+// dispatcher/tests/test_gfx1250_compile_guard.py. Keep the markers: the test
+// reads what is between them, so it exercises the shipped guard rather than a
+// copy that can drift. Nothing outside these markers may be needed to compile
+// the block except GFX_ARCH and SelectedKernel.
+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 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.
+//
+// Honest statement of what the wider predicate buys today: nothing. gfx1200 and gfx1201
+// are the only other gfx12xx parts, neither is in kSupportedArchs, so
+// dispatcher_initialize() returns -1 on them long before dispatcher_run_gemm is
+// reachable -- this constant is equal to kCompiledForGfx1250 on every architecture that
+// can execute the guard. It is kept rather than replaced by the exact predicate for one
+// reason: enabling another architecture is advertised two blocks up as "a one-line
+// addition here plus a CMake arch entry", and with the exact predicate that one line
+// would silently switch the M%4 guard off for the new gfx12xx target, reintroducing
+// exactly the silent wrong answer it exists to prevent. The static_assert below pins the
+// relationship so this note cannot quietly go stale.
+static constexpr bool kCompiledForGfx12 = ct_starts_with(GFX_ARCH, "gfx12");
+
+static_assert(!kCompiledForGfx1250 || kCompiledForGfx12,
+              "kCompiledForGfx12 must be implied by kCompiledForGfx1250. If this fires, "
+              "the two arch predicates have diverged and the M%4 guard no longer covers "
+              "the target the tile guards above are protecting.");
+
+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.");
+// --- END gfx1250 compile-time tile guard ---
+
 extern "C" {
 
 /**
@@ -166,7 +247,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 +319,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 +575,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 +605,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<int>(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 <kernel.hpp>`, giving one SelectedKernel.
 // Scaling to N kernels means N .so files (the pattern bquant/aquant/abquant follow),