GEMM kernel optimizations

turboderp · turboderp · commit c4867edd0d31 · 2025-05-08T22:44:26.000+02:00
diff --git a/exllamav3/exllamav3_ext/quant/codebook.cuh b/exllamav3/exllamav3_ext/quant/codebook.cuh
@@ -6,8 +6,11 @@ __device__ inline half decode_3inst(uint32_t x)
 {
     x *= 89226354u;
     x += 64248484u;
-    x &= 0b10001111111111111000111111111111u;
-    x ^= 0b00111011011000000011101101100000u;
+    x = (x & 0x8FFF8FFFu) ^ 0x3B603B60u;
+    // x &= 0b10001111111111111000111111111111u;
+    // x ^= 0b00111011011000000011101101100000u;
+    // Compiler doesn't automatically generate LOP3
+    asm volatile ("lop3.b32 %0, %0, 0x8fff8fff, 0x3b603b60, 0x6a;" : "+r"(x));
     half2_uint32 xu(x);
     return __hadd(__low2half(xu.as_half2), __high2half(xu.as_half2));
 }
@@ -18,10 +21,13 @@ __device__ inline half2 decode_3inst_2(uint32_t x0, uint32_t x1)
     x1 *= 89226354u;
     x0 += 64248484u;
     x1 += 64248484u;
-    x0 &= 0b10001111111111111000111111111111u;
-    x1 &= 0b10001111111111111000111111111111u;
-    x0 ^= 0b00111011011000000011101101100000u;
-    x1 ^= 0b00111011011000000011101101100000u;
+    // x0 &= 0b10001111111111111000111111111111u;
+    // x1 &= 0b10001111111111111000111111111111u;
+    // x0 ^= 0b00111011011000000011101101100000u;
+    // x1 ^= 0b00111011011000000011101101100000u;
+    // Compiler doesn't automatically generate LOP3
+    asm volatile ("lop3.b32 %0, %0, 0x8fff8fff, 0x3b603b60, 0x6a;" : "+r"(x0));
+    asm volatile ("lop3.b32 %0, %0, 0x8fff8fff, 0x3b603b60, 0x6a;" : "+r"(x1));
     half2_uint32 xu0(x0);
     half2_uint32 xu1(x1);
     half2 d0 = __halves2half2(__low2half(xu0.as_half2), __low2half(xu1.as_half2));
diff --git a/exllamav3/exllamav3_ext/quant/comp_units/exl3_comp_unit_1.cu b/exllamav3/exllamav3_ext/quant/comp_units/exl3_comp_unit_1.cu
@@ -7,22 +7,15 @@ namespace cg = cooperative_groups;
 #include "../../util.cuh"
 #include "../../ptx.cuh"
 #include "../exl3_gemm_kernel.cuh"
+#include "../exl3_gemv_kernel.cuh"
 #include "exl3_comp_unit_1.cuh"
 
 fp_exl3_gemm_kernel tfp_exl3_gemm_kernel_fp32_b1[] = {
-    nullptr,
-    exl3_gemm_kernel<1, true, EXL3_GEMM_SHAPE_1>,
-    exl3_gemm_kernel<1, true, EXL3_GEMM_SHAPE_2>,
-    exl3_gemm_kernel<1, true, EXL3_GEMM_SHAPE_3>,
-    exl3_gemm_kernel<1, true, EXL3_GEMM_SHAPE_4>
+    EXL3_GEMM_KERNEL_INSTANCES(1, true)
 };
 
 fp_exl3_gemm_kernel tfp_exl3_gemm_kernel_fp16_b1[] = {
-    nullptr,
-    exl3_gemm_kernel<1, false, EXL3_GEMM_SHAPE_1>,
-    exl3_gemm_kernel<1, false, EXL3_GEMM_SHAPE_2>,
-    exl3_gemm_kernel<1, false, EXL3_GEMM_SHAPE_3>,
-    exl3_gemm_kernel<1, false, EXL3_GEMM_SHAPE_4>
+    EXL3_GEMM_KERNEL_INSTANCES(1, false)
 };
 
 
diff --git a/exllamav3/exllamav3_ext/quant/comp_units/exl3_comp_unit_2.cu b/exllamav3/exllamav3_ext/quant/comp_units/exl3_comp_unit_2.cu
@@ -7,22 +7,15 @@ namespace cg = cooperative_groups;
 #include "../../util.cuh"
 #include "../../ptx.cuh"
 #include "../exl3_gemm_kernel.cuh"
+#include "../exl3_gemv_kernel.cuh"
 #include "exl3_comp_unit_2.cuh"
 
 fp_exl3_gemm_kernel tfp_exl3_gemm_kernel_fp32_b2[] = {
-    nullptr,
-    exl3_gemm_kernel<2, true, EXL3_GEMM_SHAPE_1>,
-    exl3_gemm_kernel<2, true, EXL3_GEMM_SHAPE_2>,
-    exl3_gemm_kernel<2, true, EXL3_GEMM_SHAPE_3>,
-    exl3_gemm_kernel<2, true, EXL3_GEMM_SHAPE_4>
+    EXL3_GEMM_KERNEL_INSTANCES(2, true)
 };
 
 fp_exl3_gemm_kernel tfp_exl3_gemm_kernel_fp16_b2[] = {
-    nullptr,
-    exl3_gemm_kernel<2, false, EXL3_GEMM_SHAPE_1>,
-    exl3_gemm_kernel<2, false, EXL3_GEMM_SHAPE_2>,
-    exl3_gemm_kernel<2, false, EXL3_GEMM_SHAPE_3>,
-    exl3_gemm_kernel<2, false, EXL3_GEMM_SHAPE_4>
+    EXL3_GEMM_KERNEL_INSTANCES(2, false)
 };
 
 
diff --git a/exllamav3/exllamav3_ext/quant/comp_units/exl3_comp_unit_3.cu b/exllamav3/exllamav3_ext/quant/comp_units/exl3_comp_unit_3.cu
@@ -7,22 +7,15 @@ namespace cg = cooperative_groups;
 #include "../../util.cuh"
 #include "../../ptx.cuh"
 #include "../exl3_gemm_kernel.cuh"
+#include "../exl3_gemv_kernel.cuh"
 #include "exl3_comp_unit_3.cuh"
 
 fp_exl3_gemm_kernel tfp_exl3_gemm_kernel_fp32_b3[] = {
-    nullptr,
-    exl3_gemm_kernel<3, true, EXL3_GEMM_SHAPE_1>,
-    exl3_gemm_kernel<3, true, EXL3_GEMM_SHAPE_2>,
-    exl3_gemm_kernel<3, true, EXL3_GEMM_SHAPE_3>,
-    exl3_gemm_kernel<3, true, EXL3_GEMM_SHAPE_4>
+    EXL3_GEMM_KERNEL_INSTANCES(3, true)
 };
 
 fp_exl3_gemm_kernel tfp_exl3_gemm_kernel_fp16_b3[] = {
-    nullptr,
-    exl3_gemm_kernel<3, false, EXL3_GEMM_SHAPE_1>,
-    exl3_gemm_kernel<3, false, EXL3_GEMM_SHAPE_2>,
-    exl3_gemm_kernel<3, false, EXL3_GEMM_SHAPE_3>,
-    exl3_gemm_kernel<3, false, EXL3_GEMM_SHAPE_4>
+    EXL3_GEMM_KERNEL_INSTANCES(3, false)
 };
 
 
diff --git a/exllamav3/exllamav3_ext/quant/comp_units/exl3_comp_unit_4.cu b/exllamav3/exllamav3_ext/quant/comp_units/exl3_comp_unit_4.cu
@@ -7,22 +7,15 @@ namespace cg = cooperative_groups;
 #include "../../util.cuh"
 #include "../../ptx.cuh"
 #include "../exl3_gemm_kernel.cuh"
+#include "../exl3_gemv_kernel.cuh"
 #include "exl3_comp_unit_4.cuh"
 
 fp_exl3_gemm_kernel tfp_exl3_gemm_kernel_fp32_b4[] = {
-    nullptr,
-    exl3_gemm_kernel<4, true, EXL3_GEMM_SHAPE_1>,
-    exl3_gemm_kernel<4, true, EXL3_GEMM_SHAPE_2>,
-    exl3_gemm_kernel<4, true, EXL3_GEMM_SHAPE_3>,
-    exl3_gemm_kernel<4, true, EXL3_GEMM_SHAPE_4>
+    EXL3_GEMM_KERNEL_INSTANCES(4, true)
 };
 
 fp_exl3_gemm_kernel tfp_exl3_gemm_kernel_fp16_b4[] = {
-    nullptr,
-    exl3_gemm_kernel<4, false, EXL3_GEMM_SHAPE_1>,
-    exl3_gemm_kernel<4, false, EXL3_GEMM_SHAPE_2>,
-    exl3_gemm_kernel<4, false, EXL3_GEMM_SHAPE_3>,
-    exl3_gemm_kernel<4, false, EXL3_GEMM_SHAPE_4>
+    EXL3_GEMM_KERNEL_INSTANCES(4, false)
 };
 
 
diff --git a/exllamav3/exllamav3_ext/quant/comp_units/exl3_comp_unit_5.cu b/exllamav3/exllamav3_ext/quant/comp_units/exl3_comp_unit_5.cu
@@ -7,22 +7,15 @@ namespace cg = cooperative_groups;
 #include "../../util.cuh"
 #include "../../ptx.cuh"
 #include "../exl3_gemm_kernel.cuh"
+#include "../exl3_gemv_kernel.cuh"
 #include "exl3_comp_unit_5.cuh"
 
 fp_exl3_gemm_kernel tfp_exl3_gemm_kernel_fp32_b5[] = {
-    nullptr,
-    exl3_gemm_kernel<5, true, EXL3_GEMM_SHAPE_1>,
-    exl3_gemm_kernel<5, true, EXL3_GEMM_SHAPE_2>,
-    exl3_gemm_kernel<5, true, EXL3_GEMM_SHAPE_3>,
-    exl3_gemm_kernel<5, true, EXL3_GEMM_SHAPE_4>
+    EXL3_GEMM_KERNEL_INSTANCES(5, true)
 };
 
 fp_exl3_gemm_kernel tfp_exl3_gemm_kernel_fp16_b5[] = {
-    nullptr,
-    exl3_gemm_kernel<5, false, EXL3_GEMM_SHAPE_1>,
-    exl3_gemm_kernel<5, false, EXL3_GEMM_SHAPE_2>,
-    exl3_gemm_kernel<5, false, EXL3_GEMM_SHAPE_3>,
-    exl3_gemm_kernel<5, false, EXL3_GEMM_SHAPE_4>
+    EXL3_GEMM_KERNEL_INSTANCES(5, false)
 };
 
 
diff --git a/exllamav3/exllamav3_ext/quant/comp_units/exl3_comp_unit_6.cu b/exllamav3/exllamav3_ext/quant/comp_units/exl3_comp_unit_6.cu
@@ -7,20 +7,13 @@ namespace cg = cooperative_groups;
 #include "../../util.cuh"
 #include "../../ptx.cuh"
 #include "../exl3_gemm_kernel.cuh"
+#include "../exl3_gemv_kernel.cuh"
 #include "exl3_comp_unit_6.cuh"
 
 fp_exl3_gemm_kernel tfp_exl3_gemm_kernel_fp32_b6[] = {
-    nullptr,
-    exl3_gemm_kernel<6, true, EXL3_GEMM_SHAPE_1>,
-    exl3_gemm_kernel<6, true, EXL3_GEMM_SHAPE_2>,
-    exl3_gemm_kernel<6, true, EXL3_GEMM_SHAPE_3>,
-    exl3_gemm_kernel<6, true, EXL3_GEMM_SHAPE_4>
+    EXL3_GEMM_KERNEL_INSTANCES(6, true)
 };
 
 fp_exl3_gemm_kernel tfp_exl3_gemm_kernel_fp16_b6[] = {
-    nullptr,
-    exl3_gemm_kernel<6, false, EXL3_GEMM_SHAPE_1>,
-    exl3_gemm_kernel<6, false, EXL3_GEMM_SHAPE_2>,
-    exl3_gemm_kernel<6, false, EXL3_GEMM_SHAPE_3>,
-    exl3_gemm_kernel<6, false, EXL3_GEMM_SHAPE_4>
+    EXL3_GEMM_KERNEL_INSTANCES(6, false)
 };
diff --git a/exllamav3/exllamav3_ext/quant/comp_units/exl3_comp_unit_7.cu b/exllamav3/exllamav3_ext/quant/comp_units/exl3_comp_unit_7.cu
@@ -7,20 +7,13 @@ namespace cg = cooperative_groups;
 #include "../../util.cuh"
 #include "../../ptx.cuh"
 #include "../exl3_gemm_kernel.cuh"
+#include "../exl3_gemv_kernel.cuh"
 #include "exl3_comp_unit_7.cuh"
 
 fp_exl3_gemm_kernel tfp_exl3_gemm_kernel_fp32_b7[] = {
-    nullptr,
-    exl3_gemm_kernel<7, true, EXL3_GEMM_SHAPE_1>,
-    exl3_gemm_kernel<7, true, EXL3_GEMM_SHAPE_2>,
-    exl3_gemm_kernel<7, true, EXL3_GEMM_SHAPE_3>,
-    exl3_gemm_kernel<7, true, EXL3_GEMM_SHAPE_4>
+    EXL3_GEMM_KERNEL_INSTANCES(7, true)
 };
 
 fp_exl3_gemm_kernel tfp_exl3_gemm_kernel_fp16_b7[] = {
-    nullptr,
-    exl3_gemm_kernel<7, false, EXL3_GEMM_SHAPE_1>,
-    exl3_gemm_kernel<7, false, EXL3_GEMM_SHAPE_2>,
-    exl3_gemm_kernel<7, false, EXL3_GEMM_SHAPE_3>,
-    exl3_gemm_kernel<7, false, EXL3_GEMM_SHAPE_4>
+    EXL3_GEMM_KERNEL_INSTANCES(7, false)
 };
diff --git a/exllamav3/exllamav3_ext/quant/comp_units/exl3_comp_unit_8.cu b/exllamav3/exllamav3_ext/quant/comp_units/exl3_comp_unit_8.cu
@@ -7,20 +7,13 @@ namespace cg = cooperative_groups;
 #include "../../util.cuh"
 #include "../../ptx.cuh"
 #include "../exl3_gemm_kernel.cuh"
+#include "../exl3_gemv_kernel.cuh"
 #include "exl3_comp_unit_8.cuh"
 
 fp_exl3_gemm_kernel tfp_exl3_gemm_kernel_fp32_b8[] = {
-    nullptr,
-    exl3_gemm_kernel<8, true, EXL3_GEMM_SHAPE_1>,
-    exl3_gemm_kernel<8, true, EXL3_GEMM_SHAPE_2>,
-    exl3_gemm_kernel<8, true, EXL3_GEMM_SHAPE_3>,
-    exl3_gemm_kernel<8, true, EXL3_GEMM_SHAPE_4>
+    EXL3_GEMM_KERNEL_INSTANCES(8, true)
 };
 
 fp_exl3_gemm_kernel tfp_exl3_gemm_kernel_fp16_b8[] = {
-    nullptr,
-    exl3_gemm_kernel<8, false, EXL3_GEMM_SHAPE_1>,
-    exl3_gemm_kernel<8, false, EXL3_GEMM_SHAPE_2>,
-    exl3_gemm_kernel<8, false, EXL3_GEMM_SHAPE_3>,
-    exl3_gemm_kernel<8, false, EXL3_GEMM_SHAPE_4>
+    EXL3_GEMM_KERNEL_INSTANCES(8, false)
 };
diff --git a/exllamav3/exllamav3_ext/quant/exl3_dq.cuh b/exllamav3/exllamav3_ext/quant/exl3_dq.cuh
@@ -4,12 +4,12 @@
 
 __device__ __forceinline__ uint32_t fshift(uint32_t b, uint32_t a, int shift)
 {
-    // uint64_t merged = ((uint64_t)b << 32) | (uint64_t) a;
-    // return (uint32_t)(merged >> shift);
+     uint64_t merged = ((uint64_t)a << 32) | (uint64_t) b;
+     return (uint32_t)(merged >> shift);
 
-    // Conditional funnel shift is somehow faster
-    if (shift < 32) return __funnelshift_r(b, a, shift);
-    return a >> (shift - 32);
+    // Conditional funnel shift is somehow no longer faster
+    // if (shift < 32) return __funnelshift_r(b, a, shift);
+    // return a >> (shift - 32);
 }
 
 template <int bits>
diff --git a/exllamav3/exllamav3_ext/quant/exl3_gemm_inner.cuh b/exllamav3/exllamav3_ext/quant/exl3_gemm_inner.cuh
@@ -30,7 +30,6 @@ void exl3_gemm_kernel_inner
     const int sh_a_stage_size = TILESIZE_M * TILESIZE_K;                         // in halfs
     const int sh_b_stage_size = TILEBLOCKS_K * TILEBLOCKS_N * 256 / 16 * bits;   // in uint16s
     const int sh_c_size = 4 * EXL3_GEMM_BASE_THREADS;                            // in floats
-    // TODO: Maybe flush global->shared pipeline before reduction step so sh_c can share memory with sh_a and sh_b
 
     // Sanity checks
     static_assert(EXL3_GEMM_BASE_THREADS == 256);
@@ -213,7 +212,6 @@ void exl3_gemm_kernel_inner
                 for (int i = 0; i < load_a_iters; ++i)
                 {
                     // TODO: Rearrange into ldmatrix friendly layout while loading?
-                    // @p seems to crash on Blackwell but does not perform better on Ampere and Ada anyway
                     // cp_async_pred(sh + EXL3_GEMM_BASE_THREADS * i + t, gl + load_a_gl[i], pred_a_gl[i]);
                     if (pred_a_gl[i]) cp_async(sh + EXL3_GEMM_BASE_THREADS * i + t, gl + load_a_gl[i]);
                 }
@@ -226,7 +224,6 @@ void exl3_gemm_kernel_inner
                 #pragma unroll
                 for (int i = 0; i < load_b_iters; ++i)
                 {
-                    // @p seems to crash on Blackwell but does not perform better on Ampere and Ada anyway
                     // cp_async_pred(sh + EXL3_GEMM_BASE_THREADS * i + t, gl + load_b_gl[i], pred_b_gl[i]);
                     if (pred_b_gl[i]) cp_async(sh + EXL3_GEMM_BASE_THREADS * i + t, gl + load_b_gl[i]);
                 }
@@ -590,4 +587,42 @@ void exl3_gemm_kernel_inner
             if (!slice2_iters) break;
         }
     }
+
+    if constexpr (FRAG_STAGES == 4)
+    {
+        while (true)
+        {
+            async_load_gl();
+            wait_stage();
+            load_frags(1);
+            matmul(0);
+            if (slice2_k == tiles.k - 1 || slice2_iters == 1) { reduce(); slice2_k0 = slice2_k + 1; }
+            advance2();
+            if (!slice2_iters) break;
+
+            async_load_gl();
+            wait_stage();
+            load_frags(2);
+            matmul(1);
+            if (slice2_k == tiles.k - 1 || slice2_iters == 1) { reduce(); slice2_k0 = slice2_k + 1; }
+            advance2();
+            if (!slice2_iters) break;
+
+            async_load_gl();
+            wait_stage();
+            load_frags(3);
+            matmul(2);
+            if (slice2_k == tiles.k - 1 || slice2_iters == 1) { reduce(); slice2_k0 = slice2_k + 1; }
+            advance2();
+            if (!slice2_iters) break;
+
+            async_load_gl();
+            wait_stage();
+            load_frags(0);
+            matmul(3);
+            if (slice2_k == tiles.k - 1 || slice2_iters == 1) { reduce(); slice2_k0 = slice2_k + 1; }
+            advance2();
+            if (!slice2_iters) break;
+        }
+    }
 }
diff --git a/exllamav3/exllamav3_ext/quant/exl3_kernel_map.cu b/exllamav3/exllamav3_ext/quant/exl3_kernel_map.cu
@@ -33,15 +33,15 @@ int select_gemm_shape(int cc, int size_m, int size_k, int size_n, int bits)
                 if (size_n <= 2048) return 2;
                 return 3;
             }
-            if (size_n <= 4096) return size_k > 8192 ? 3 : 2;
-            if (mod_512) return 4;
+            if (size_n < 4096) return size_k > 8192 ? 3 : 2;
+            if (mod_512 && (size_n * size_k) > (4096 * 4096) && bits <= 6) return 4;
             if (mod_256) return 3;
             return 2;
 
         case CC_ADA:
-            if (bits <= 4)
+            if (bits <= 3)
             {
-                if (size_n <= 4096 && size_k <= 12288) return 2;
+                if (size_n < 4096 && size_k <= 12288) return 2;
                 return 3;
             }
             if (size_n <= 16384) return 2;
@@ -51,28 +51,28 @@ int select_gemm_shape(int cc, int size_m, int size_k, int size_n, int bits)
 
         case CC_HOPPER:
         case CC_BLACKWELL:
-            if (bits == 8)
+            if (bits >= 7)
             {
                 if (size_n <= 8192) return size_k > 32768 ? 3 : 2;
                 if (mod_512 && size_n > 32768) return 4;
                 return 2;
             }
-            if (size_n <= 2048 && size_k > 4096 && bits >= 5) return 1;
             if (size_n <= 4096) return (size_k && bits >= 3) > 8192 ? 3 : 2;
-            if (mod_512 && size_n > 16384 && bits >= 6) return 4;
-            if (mod_256 && bits >= 3) return 3;
+            if (mod_512 && size_n > 16384) return 4;
+            if (mod_256) return 3;
             return 2;
     }
     return 0;
 }
 
 int exl3_gemm_num_kernel_shapes()
 {
-    return 4;
+    return EXL3_GEMM_NUM_SHAPES;
 }
 
 int exl3_gemm_tilesize_k[] = {EXL3_GEMM_TILESIZE_K};
 int exl3_gemm_tilesize_n[] = {EXL3_GEMM_TILESIZE_N};
+int exl3_gemm_blockdim[] = {EXL3_GEMM_BLOCKDIM};
 
 bool exl3_gemm_shape_compat(int shape_idx, int size_m, int size_k, int size_n, int bits)
 {
@@ -98,7 +98,7 @@ fp_exl3_gemm_kernel select_exl3_gemm_kernel
     int shape_idx = force_shape_idx <= 0 ? select_gemm_shape(cc, size_m, size_k, size_n, bits) : force_shape_idx;
     TORCH_CHECK(shape_idx > 0, "exl3_gemm: no compatible kernel");
     if (out_shape_idx) *out_shape_idx = shape_idx;
-    if (out_block_dim) *out_block_dim = EXL3_GEMM_BASE_THREADS * exl3_gemm_tilesize_k[shape_idx] / 16;
+    if (out_block_dim) *out_block_dim = exl3_gemm_blockdim[shape_idx];
 
     // Avoid empty blocks
     if (num_sms)
@@ -122,6 +122,7 @@ fp_exl3_gemm_kernel select_exl3_gemm_kernel
             case 6: return tfp_exl3_gemm_kernel_fp32_b6[shape_idx];
             case 7: return tfp_exl3_gemm_kernel_fp32_b7[shape_idx];
             case 8: return tfp_exl3_gemm_kernel_fp32_b8[shape_idx];
+            default: TORCH_CHECK(false, "No kernel for GEMM shape");
         }
     }
     else
@@ -136,6 +137,7 @@ fp_exl3_gemm_kernel select_exl3_gemm_kernel
             case 6: return tfp_exl3_gemm_kernel_fp16_b6[shape_idx];
             case 7: return tfp_exl3_gemm_kernel_fp16_b7[shape_idx];
             case 8: return tfp_exl3_gemm_kernel_fp16_b8[shape_idx];
+            default: TORCH_CHECK(false, "No kernel for GEMM shape");
         }
     }
 }
diff --git a/exllamav3/exllamav3_ext/quant/exl3_kernel_map.cuh b/exllamav3/exllamav3_ext/quant/exl3_kernel_map.cuh
diff --git a/exllamav3/exllamav3_ext/quant/hadamard_inner.cuh b/exllamav3/exllamav3_ext/quant/hadamard_inner.cuh
diff --git a/science/qgemm_benchmark.py b/science/qgemm_benchmark.py
