Rework GEMM kernel tuning

Author: turboderp

exllamav3/exllamav3_ext/bindings.cpp

@@ -23,6 +23,7 @@
 #include "quant/exl3_gemm.cuh"
 #include "quant/exl3_kernel_map.cuh"
 #include "quant/util.cuh"
+#include "quant/exl3_devctx.cuh"
 
 #include "generator/strings.h"
 #include "generator/sampling_basic.cuh"
@@ -87,6 +88,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
     m.def("exl3_gemm", &exl3_gemm, "exl3_gemm");
     m.def("exl3_gemm_num_kernel_shapes", &exl3_gemm_num_kernel_shapes, "exl3_gemm_num_kernel_shapes");
     m.def("exl3_gemm_shape_compat", &exl3_gemm_shape_compat, "exl3_gemm_shape_compat");
+    m.def("g_get_cc", &g_get_cc, "g_get_cc");
+    m.def("g_get_num_sms", &g_get_num_sms, "g_get_num_sms");
     m.def("exl3_mgemm", &exl3_mgemm, "exl3_mgemm");
     m.def("hgemm", &hgemm, "hgemm");
     m.def("rope", &rope, "rope");

exllamav3/exllamav3_ext/libtorch/blocksparse_mlp.cpp

@@ -85,7 +85,8 @@ void BC_BlockSparseMLP::run_bsz1
         gate_mcg_mult,
         gate_mul1_mult,
         min_expert,
-        max_expert
+        max_expert,
+        0
     );
 
     exl3_mgemm(
@@ -102,7 +103,8 @@ void BC_BlockSparseMLP::run_bsz1
         up_mcg_mult,
         up_mul1_mult,
         min_expert,
-        max_expert
+        max_expert,
+        0
     );
 
     if (act_silu)
@@ -124,7 +126,8 @@ void BC_BlockSparseMLP::run_bsz1
         down_mcg_mult,
         down_mul1_mult,
         min_expert,
-        max_expert
+        max_expert,
+        0
     );
 
     if (shared_experts)

exllamav3/exllamav3_ext/libtorch/linear.cpp

@@ -31,12 +31,12 @@ void BC_LinearEXL3::run(const at::Tensor& x, at::Tensor& y)
 {
     if (x.numel() == x.size(-1))
     {
-        exl3_gemm(x, trellis, y, suh, xh, svh, -1, mcg_mult, mul1_mult);
+        exl3_gemm(x, trellis, y, suh, xh, svh, -1, mcg_mult, mul1_mult, 0);
     }
     else
     {
         at::Tensor xh_ = at::empty_like(x);
-        exl3_gemm(x, trellis, y, suh, xh_, svh, -1, mcg_mult, mul1_mult);
+        exl3_gemm(x, trellis, y, suh, xh_, svh, -1, mcg_mult, mul1_mult, 0);
     }
 
     if (bias) y.add_(bias.value());

exllamav3/exllamav3_ext/libtorch/mlp.cpp

@@ -31,7 +31,8 @@ void BC_GatedMLP::run_bsz1
         gu_mcg_mult,
         gu_mul1_mult,
         -1,
-        -1
+        -1,
+        0
     );
 
     at::Tensor g = gu.select(0, 0).unsqueeze(0);

exllamav3/exllamav3_ext/quant/exl3_devctx.cu

@@ -55,4 +55,14 @@ int* DevCtx::get_locks(int device)
         cudaMemset(locks[device], 0, MAX_TILES_C * sizeof(int));
     }
     return (int*) locks[device];
+}
+
+int g_get_cc(int device)
+{
+    return DevCtx::instance().get_cc(device);
+}
+
+int g_get_num_sms(int device)
+{
+    return DevCtx::instance().get_num_sms(device);
 }

exllamav3/exllamav3_ext/quant/exl3_devctx.cuh

@@ -6,12 +6,13 @@
 // Max allowable output size, in tiles. Used to allocate global lock buffer per device for sync across threadblocks
 #define MAX_TILES_C (1024 * 1024)
 
+// Treat hopper and blackwell as same arch for now
 #define MAX_DEVICES 32
 #define CC_OLD        1
 #define CC_AMPERE     2
 #define CC_ADA        3
 #define CC_HOPPER     4
-#define CC_BLACKWELL  5
+#define CC_BLACKWELL  4
 
 // Singleton to manage context for each device. Stores device attributes and a large-enough lock buffer per device
 class DevCtx
@@ -32,4 +33,7 @@ private:
     DevCtx() = default;
     DevCtx(const DevCtx&) = delete;
     DevCtx& operator=(const DevCtx&) = delete;
-};
+};
+
+int g_get_cc(int device);
+int g_get_num_sms(int device);

exllamav3/exllamav3_ext/quant/exl3_gemm.cu

@@ -12,11 +12,17 @@ namespace cg = cooperative_groups;
 #include "exl3_devctx.cuh"
 #include <set>
 
+#define NEW_TUNE_GEMM
+#define NEW_TUNE_MGEMM
+
+int exl3_gemm_tilesize_k_g[] = {EXL3_GEMM_TILESIZE_K};
+int exl3_gemm_tilesize_n_g[] = {EXL3_GEMM_TILESIZE_N};
+
 /*
 EXL3 matmul, A @ B -> C
 
 - A: row-major A tensor, shape (m, k), dtype float16, contiguous
-- B: EXL3-quantized B tensor, shape (k//16, n//16, 16*bits), dtype uint16
+- B: EXL3-quantized B tensor, shape (k//16, n//16, 16*K), dtype uint16
 - C: empty row-major C tensor, shape (m, n), dtype float16 or float32, contiguous. Does not need to be zero-initialized
 - suh: optional, packed input scales/flips, shape (k//16), dtype float16
 - A_had: required if suh given, may be reference to A, temporary storage for input transform, size and dtype as A
@@ -39,7 +45,8 @@ int exl3_gemm
     const c10::optional<at::Tensor>& svh,
     int force_shape_idx,
     uint32_t mcg_mult,
-    uint32_t mul1_mult
+    uint32_t mul1_mult,
+    int force_num_sms
 )
 {
     const at::cuda::OptionalCUDAGuard device_guard(A.device());
@@ -48,7 +55,7 @@ int exl3_gemm
     TORCH_CHECK_DIM(B, 3);
     TORCH_CHECK_SHAPES(A, -1, B, 0, 16);
     TORCH_CHECK_SHAPES(C, -1, B, 1, 16);
-//    TORCH_CHECK_SHAPES(A, 0, C, 0, 1);
+    // TORCH_CHECK_SHAPES(A, 0, C, 0, 1);
     TORCH_CHECK_DTYPE(A, kHalf);
     TORCH_CHECK_DTYPE(B, kShort);
     bool c_fp32 = C.dtype() == at::kFloat;
@@ -59,26 +66,26 @@ int exl3_gemm
     half* A_had_ptr = nullptr;
     if (suh_ptr)
     {
-//        TORCH_CHECK_SHAPES(suh.value(), 0, A, 1, 1);
+        // TORCH_CHECK_SHAPES(suh.value(), 0, A, 1, 1);
         A_had_ptr = (half*) OPTPTR(A_had);
-//        TORCH_CHECK(A_had_ptr, "Must supply A_had with suh");
-//        TORCH_CHECK_SHAPES_FULL(A_had.value(), A);
+        // TORCH_CHECK(A_had_ptr, "Must supply A_had with suh");
+        // TORCH_CHECK_SHAPES_FULL(A_had.value(), A);
     }
 
     // Get SV, optionally
     const half* svh_ptr = (const half*) OPTPTR(svh);
-//    if (svh_ptr)
-//        TORCH_CHECK_SHAPES(svh.value(), 0, B, 1, 16);
+    // if (svh_ptr)
+        // TORCH_CHECK_SHAPES(svh.value(), 0, B, 1, 16);
 
     // Device properties
     int device;
     cudaGetDevice(&device);
-    int num_sms = DevCtx::instance().get_num_sms(device);
+    int num_sms = force_num_sms ? force_num_sms : DevCtx::instance().get_num_sms(device);
     int cc = DevCtx::instance().get_cc(device);
     int* locks = DevCtx::instance().get_locks(device);
 
     // Dispatch
-    int bits = B.size(2) / 16;
+    int K = B.size(2) / 16;
     const half* A_ptr = (const half*) A.data_ptr();
     const uint16_t* B_ptr = (const uint16_t*) B.data_ptr();
     void* C_ptr = (void*) C.data_ptr();
@@ -96,21 +103,33 @@ int exl3_gemm
     if (mcg_mult) { cb = 1; mult = mcg_mult; }
     if (mul1_mult) { cb = 2; mult = mul1_mult; }
 
-    int selected_shape;
     int block_dim;
-    fp_exl3_gemm_kernel kernel = select_exl3_gemm_kernel
-    (
-        cc, size_m, size_k, size_n, bits, c_fp32,
-        force_shape_idx, &block_dim, &selected_shape,
-        &num_sms, cb
-    );
-    if (!kernel) return 0;
+    int shape_idx;
+    fp_exl3_gemm_kernel kernel;
+
+    #ifndef NEW_TUNE_GEMM
+        kernel = select_exl3_gemm_kernel
+        (
+            cc, size_m, size_k, size_n, K, c_fp32,
+            force_shape_idx, &block_dim, &shape_idx,
+            &num_sms, cb
+        );
+        if (!kernel) return 0;
+    #else
+        TResult* tr = select_exl3_gemm_mgemm_kernel_new(cc, size_m, size_k, size_n, K, c_fp32, force_shape_idx, force_num_sms, cb);
+        if (!tr) return 0;
+        num_sms = MIN(num_sms, tr->num_sms);
+        kernel = tr->kernel;
+        block_dim = tr->block_dim;
+        shape_idx = tr->shape_idx;
+    #endif
 
     // Launch
-    if (kernel_attr_set[device].find((void*)kernel) == kernel_attr_set[device].end())
+    if (kernel_attr_set[device].find((void*) kernel) == kernel_attr_set[device].end())
     {
         cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, SMEM_MAX);
-        kernel_attr_set[device].insert((void*)kernel);
+        kernel_attr_set[device].insert((void*) kernel);
+        cuda_check(cudaPeekAtLastError());
     }
     void* kernelArgs[] =
     {
@@ -128,22 +147,24 @@ int exl3_gemm
     };
     cudaLaunchCooperativeKernel
     (
-        (void*)kernel,
+        (void*) kernel,
         num_sms,
         block_dim,
         kernelArgs,
         SMEM_MAX,
         stream
     );
     cuda_check(cudaPeekAtLastError());
-    return selected_shape;
+
+    // return selected_shape;
+    return shape_idx;
 }
 
 /*
 EXL3 multi matmul, A @ B -> C
 
 - A: row-major A tensor, shape (m, k), dtype float16, contiguous
-- B: EXL3-quantized B tensor, shape (k//16, n//16, 16*bits), dtype uint16
+- B: EXL3-quantized B tensor, shape (k//16, n//16, 16*K), dtype uint16
 - C: empty row-major C tensor, shape (m, n), dtype float16 or float23, contiguous. Does not need to be zero-initialized
 - suh: optional, packed input scales/flips, shape (k//16), dtype float16
 - A_had: required if suh given, may be reference to A, temporary storage for input transform, size and dtype as A
@@ -169,7 +190,8 @@ int exl3_mgemm
     uint32_t mcg_mult,
     uint32_t mul1_mult,
     int min_index,
-    int max_index
+    int max_index,
+    int force_num_sms
 )
 {
     const at::cuda::OptionalCUDAGuard device_guard(A.device());
@@ -194,6 +216,7 @@ int exl3_mgemm
     int bsz = A.size(1);
     int bszm_in = A.size(0);
     int bszm_out = C.size(0);
+    int bszm = MAX(bszm_in, bszm_out);
 
     const long* indices_ptr = (const long*) OPTPTR(indices);
     const half* weights_ptr = (const half*) OPTPTR(weights);
@@ -219,8 +242,8 @@ int exl3_mgemm
     // Device properties
     int device;
     cudaGetDevice(&device);
-    int num_sms = DevCtx::instance().get_num_sms(device);
-    int total_sms = num_sms;
+    int total_sms = DevCtx::instance().get_num_sms(device);
+    int num_sms = force_num_sms ? force_num_sms : total_sms;
     int cc = DevCtx::instance().get_cc(device);
     int* locks = DevCtx::instance().get_locks(device);
 
@@ -239,25 +262,44 @@ int exl3_mgemm
     if (mcg_mult) { cb = 1; mult = mcg_mult; }
     if (mul1_mult) { cb = 2; mult = mul1_mult; }
 
-    int selected_shape;
+    int shape_idx;
     int block_dim;
-    fp_exl3_mgemm_kernel kernel = select_exl3_mgemm_kernel
-    (
-        cc, size_m, size_k, size_n, K, c_fp32,
-        force_shape_idx, &block_dim, &selected_shape,
-        &num_sms, cb, bszm_in, bszm_out
-    );
-    if (!kernel) return 0;
+    fp_exl3_mgemm_kernel kernel;
+    int concurrency;
+
+    #ifndef NEW_TUNE_MGEMM
+        kernel = select_exl3_mgemm_kernel
+        (
+            cc, size_m, size_k, size_n, K, c_fp32,
+            force_shape_idx, &block_dim, &shape_idx,
+            &num_sms, cb, bszm_in, bszm_out
+        );
+        if (!kernel) return 0;
+        concurrency = MIN(total_sms / num_sms, bszm_out);
+    #else
+        kernel = select_exl3_mgemm_kernel
+        (
+            cc, size_m, size_k, size_n, K, c_fp32,
+            force_shape_idx, &block_dim, &shape_idx,
+            &num_sms, cb, bszm_in, bszm_out
+        );
+        int tilesize_k = exl3_gemm_tilesize_k_g[shape_idx];
+        int tilesize_n = exl3_gemm_tilesize_n_g[shape_idx];
+        int tiles = MAX(size_k / tilesize_k * size_n / tilesize_n, 1);
+        num_sms = tiles;
+        if (num_sms * bszm > total_sms) num_sms = MAX(total_sms / bszm, 1);
+        if (num_sms <= total_sms && tiles / num_sms > 48) num_sms = MIN(total_sms, num_sms * 2);
+        concurrency = MIN(total_sms / num_sms, bszm);
+    #endif
 
     // Launch bigger grid if possible
-    int concurrency = MIN(total_sms / num_sms, bszm_out);
     dim3 block_grid(num_sms, 1, concurrency);
 
     // Launch
-    if (kernel_attr_set[device].find((void*)kernel) == kernel_attr_set[device].end())
+    if (kernel_attr_set[device].find((void*) kernel) == kernel_attr_set[device].end())
     {
         cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, SMEM_MAX);
-        kernel_attr_set[device].insert((void*)kernel);
+        kernel_attr_set[device].insert((void*) kernel);
     }
     void* kernelArgs[] =
     {
@@ -279,16 +321,15 @@ int exl3_mgemm
         (void*)& min_index,
         (void*)& max_index
     };
-
     cudaLaunchCooperativeKernel
     (
-        (void*)kernel,
+        (void*) kernel,
         block_grid,
         block_dim,
         kernelArgs,
         SMEM_MAX,
         stream
     );
     cuda_check(cudaPeekAtLastError());
-    return selected_shape;
+    return shape_idx;
 }

exllamav3/exllamav3_ext/quant/exl3_gemm.cuh

@@ -12,7 +12,8 @@ int exl3_gemm
     const c10::optional<at::Tensor>& svh,
     int force_shape_idx,
     uint32_t mcg_mult,
-    uint32_t mul1_mult
+    uint32_t mul1_mult,
+    int force_num_sms
 );
 
 int exl3_mgemm
@@ -30,5 +31,6 @@ int exl3_mgemm
     uint32_t mcg_mult,
     uint32_t mul1_mult,
     int min_index,
-    int max_index
+    int max_index,
+    int force_num_sms
 );