From: opencoti
Subject: [PATCH 0094] tinyblas small-GEMM second config (bug-2104, #614)

opencoti-hook: tinyblas-small-gemm — skinny f32 GEMMs (Gemma MoE router
ffn_gate_inp [2816x128], qwen35 DeltaNet chunk GEMMs) hit a single 128x64-tile
config = 8 blocks on a 142-SM RTX 6000 (5% occupancy, 274us). Add a second
<32,32,64,...> config selected when the big config yields <24 blocks (90.7us;
prefill +25-35% on MoE models).

diff --git a/llamafile/tinyblas.cu b/llamafile/tinyblas.cu
index 758f68a..76ba099 100644
--- a/llamafile/tinyblas.cu
+++ b/llamafile/tinyblas.cu
@@ -539,6 +539,17 @@ tinyblasStatus_t tinyblasGE_launch(tinyblasHandle_t handle, tinyblasOperation_t
                                    int ldc) {
     if (can_use_matvec(aT, bT, m, n, k, alpha, beta))
         return matvec_launch<WORD>(handle, m, k, A, lda, B, C);
+    // opencoti-hook: tinyblas-small-gemm (#614/bug-2104) — the single 128x64-tile config
+    // launches only ceil(m/64)*ceil(n/128) blocks after the row-major swap below; a skinny
+    // GEMM like the Gemma MoE router ([2816x128] x [2816x512] per layer per ubatch) gets 8
+    // blocks on a 142-SM GPU (~6% occupancy, ~274 us vs cuBLAS ~8 us) and was 35% of A4B
+    // prefill GPU time. When the big config cannot fill the GPU, use a 32x32-tile config
+    // (same deterministic FFMA warp2d kernel, 4x1 thread tiles, 2 warps) so the block
+    // count grows 8x. Threshold 24 keeps every launch that already saturates on the big
+    // tiles untouched.
+    if ((long long)CEIL_DIV(m, 64) * CEIL_DIV(n, 128) < 24)
+        return tinyblasGE_launcher<32, 32, 64, 8, 32, 16, 1, 4, 4, 64>(
+            handle, bT, aT, n, m, k, alpha, B, ldb, A, lda, beta, C, ldc);
     constexpr int TT = 256;
     constexpr int BM = 128;
     constexpr int BN = 64;