DFlash2 Q4_0 GEMM on gfx12: variant 26 (two warps a block, two rows a warp, unroll 1) above 640 rows, 31 (one warp a row, unroll 1) at 640 rows and below; bench variants 36-45; STRATA_DF2_GEMM4=N (N>=10) forces a variant

dflash2_gemm_bench q4 on an R9700 (gfx1201, p620, 8 columns, cold weights, bitwise vs variant 1 in every cell): gate+up 56 -> 44.5 us, down 39 -> 23.4, qkv 26.8 -> 21.4, o 23.6 -> 14.9, fc 61.5 -> 35.2, kv 15.2 -> 10.0, akp/mkp 12.1 -> 6.3, selector 11.1 -> 5.2.
Serve A/B (128K, 8 classes, 3 rounds x 2 reps): drafter 2.5-3.3 -> 2.2-2.8 ms a round, +1.0% tok/s over the old variant (n 6 per class). CLI greedy ids identical across old / new. Other cards keep 17. gfx1151 and gfx1100 compile.

src/core/dflash2_gemm_bench.cpp

@@ -49,7 +49,7 @@ int main(int argc, char** argv) {
         }
         std::vector<float> yref;
         std::vector<int> vars = {1, 2, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23};
-        if (q4) for (int v = 24; v <= 35; ++v) vars.push_back(v);
+        if (q4) for (int v = 24; v <= 45; ++v) vars.push_back(v);
         for (int var : vars) {
             strata::kernels::df2_gemm_variant(var);
             const int ncols = 8;

src/kernels/cuda/dflash2.cu

@@ -2,6 +2,7 @@
 #include "strata/kernels/dflash2.hpp"
 
 #include "strata/kernels/dp4a.hpp"
+#include "strata/kernels/gfx_arch.hpp"
 
 #include <cuda_fp16.h>
 #include <cuda_runtime.h>
@@ -978,8 +979,24 @@ void df2_rope_table_inv(const int32_t* pos, int rows, const double* inv, float*
 void df2_gemm_variant(int v) { g_gemm_variant = v; }
 void df2_attention_variant(int v) { g_attn_variant = v; }
 
+// The Q4_0 variant by matrix size (rows = all the launch's output rows).  17 (two warps a block, two rows a warp, unroll 2): dflash2_gemm_bench q4 best on l620 (sm_120,
+// 110 SMs) for every matrix, within 0-2 us of the best on a 188-SM RTX PRO 6000.  gfx12 (R9700, gfx1201, p620, 8 columns, cold weights, two bench runs): the same shape with unroll 1
+// (26) is 25-45% faster on every matrix of 1000 rows and more (gate+up 56 -> 44.5 us, down 39 -> 23.4, qkv 26.8 -> 21.4, o 23.6 -> 14.9, fc 61.5 -> 35.2, kv 15.2 -> 10.0), and one
+// warp a row with unroll 1 (31) is 2x faster on the 256 and 640 row matrices (akp / mkp 12.1 -> 6.3 us, selector 11.1 -> 5.2); all bitwise the same results.
+static int df2_q4_auto(int rows) {
+#if defined(STRATA_USE_HIP)
+    static const bool gfx12 = [] {
+        int dev = 0;
+        cudaDeviceProp p;
+        return cudaGetDevice(&dev) == cudaSuccess && cudaGetDeviceProperties(&p, dev) == cudaSuccess && strata::kernels::gfx_arch_is_gfx12(p.gcnArchName);
+    }();
+    if (gfx12) return rows <= 640 ? 31 : 26;
+#endif
+    (void) rows;
+    return 17;
+}
 #ifndef DF2_Q4_AUTO
-#define DF2_Q4_AUTO(rows) 17   // two warps a block, two rows a warp (dflash2_gemm_bench q4: best on l620 (sm_120, 110 SMs) for every matrix, within 0-2 us of the best on a 188-SM RTX PRO 6000)
+#define DF2_Q4_AUTO(rows) df2_q4_auto(rows)
 #endif
 void df2_gemm(const Df2Mat* mats, int n, const void* xq, const float* xd, int n_in, int ncols, void* stream) {
     if (n < 1 || n > 3 || n_in % 512 != 0 || ncols < 1 || ncols > 8) throw std::invalid_argument("df2_gemm: unsupported shape");
@@ -996,8 +1013,8 @@ void df2_gemm(const Df2Mat* mats, int n, const void* xq, const float* xd, int n_
     const int8_t* xqp = static_cast<const int8_t*>(xq);
     cudaStream_t st = (cudaStream_t) stream;
     if (mats[0].q4) {   // the Q4_0 pack: 4 warps a block when the rows allow, else 2; the rows-a-warp variants (bitwise the same results) by matrix size
-        static const int tune4 = [] { const char* e = std::getenv("STRATA_DF2_GEMM4"); return e ? std::atoi(e) : 0; }();   // 0: the table below; 1: 8 rows a warp as before
-        int v4 = g_gemm_variant != 0 ? g_gemm_variant : tune4 != 0 ? 1 : DF2_Q4_AUTO(rows);
+        static const int tune4 = [] { const char* e = std::getenv("STRATA_DF2_GEMM4"); return e ? std::atoi(e) : 0; }();   // 0: the table; 1: 8 rows a warp as before; 10 and up: that variant for every matrix (bench / A/B)
+        int v4 = g_gemm_variant != 0 ? g_gemm_variant : tune4 >= 10 ? tune4 : tune4 != 0 ? 1 : DF2_Q4_AUTO(rows);
 #define DF2_ROWS4_CASE(V, W, R, U) case V: gemm4_rows_kernel<W, R, U><<<(rows / 8) * (8 / R) / W, W * 32, 0, st>>>(segs, xqp, xd, n_in, ncols); break;
         switch (v4) {
             DF2_ROWS4_CASE(10, 4, 4, 2)
@@ -1026,6 +1043,16 @@ void df2_gemm(const Df2Mat* mats, int n, const void* xq, const float* xd, int n_
             DF2_ROWS4_CASE(33, 1, 1, 8)
             DF2_ROWS4_CASE(34, 2, 1, 2)
             DF2_ROWS4_CASE(35, 4, 1, 2)
+            DF2_ROWS4_CASE(36, 4, 2, 1)
+            DF2_ROWS4_CASE(37, 8, 2, 1)
+            DF2_ROWS4_CASE(38, 2, 1, 1)
+            DF2_ROWS4_CASE(39, 4, 1, 1)
+            DF2_ROWS4_CASE(40, 1, 2, 1)
+            DF2_ROWS4_CASE(41, 1, 4, 1)
+            DF2_ROWS4_CASE(42, 2, 4, 1)
+            DF2_ROWS4_CASE(43, 4, 4, 1)
+            DF2_ROWS4_CASE(44, 8, 1, 1)
+            DF2_ROWS4_CASE(45, 1, 8, 1)
             case 2: gemm4_kernel<2><<<rows / 16, 64, 0, st>>>(segs, xqp, xd, n_in, ncols); break;
             default: if (rows % 32 == 0) gemm4_kernel<4><<<rows / 32, 128, 0, st>>>(segs, xqp, xd, n_in, ncols);
                      else gemm4_kernel<2><<<rows / 16, 64, 0, st>>>(segs, xqp, xd, n_in, ncols);