char-fork/llama.cpp
git clone https://git.t4t.associates/char-fork/llama.cpp
c8edceb06
char/aliceai-t5-moe
1#pragma once 2 3#include "common.cuh" 4 5#include <climits> 6#include <cstdint> 7 8#define MMQ_DP4A_MAX_BATCH_SIZE 64 // Max. batch size to use for dp4a MMQ kernels when FP16 tensor cores are available. 9#define MMQ_ITER_K 256 10#define MMQ_ITER_K_FP4 512 11#define MMQ_NWARPS 8 12 13typedef void (*ggml_cuda_mmq_load_tiles_t)(const char * __restrict__ x, int * x_tile, const int kbx0, const int i_max, const int stride); 14typedef void (*ggml_cuda_mmq_vec_dot_t)(const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00); 15typedef void (*ggml_cuda_mmq_write_back_t)(const float * __restrict__ sum, const int32_t * __restrict__ get_rows_to_sorted, 16 float * __restrict__ dst, const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max); 17 18enum mmq_q8_1_ds_layout { 19 MMQ_Q8_1_DS_LAYOUT_D4, 20 MMQ_Q8_1_DS_LAYOUT_DS4, 21 MMQ_Q8_1_DS_LAYOUT_D2S6, 22}; 23 24static constexpr int QK8_1_MMQ = 4*QK8_1; 25static constexpr int QK_FP4_MMQ = 2*QK8_1_MMQ; 26 27struct block_q8_1_mmq { 28 // The y float data is converted to a data layout that can simply be copied to shared memory as a contiguous block. 29 // The y float data is first grouped as blocks of 128 values. 30 // These blocks are then treated as individual data values and transposed. 31 // 32 // To avoid shared memory bank conflicts each block is padded with 16 bytes. 33 // This padding is also used to store block scales/partial sums. 34 // The scales multiplied with the quantized data are equal to the unquantized values. 35 // The partial sums are obtained by summing up a subgroup of the contained values (prior to quantization) 36 // and are only needed for performance reasons. 37 // 38 // The exact data stored depends on the x data type. 39 union { 40 float d4[4]; // 1 32 bit scale per 32 values, stored as d0,d1,d2,d3 41 half2 ds4[4]; // 1 16 bit scale + 1 16 bit partial sum per 32 values, stored as d0,s0,d1,s1,d2,s2,d3,s3 42 half d2s6[8]; // 1 16 bit scale per 64 values + 1 16 bit partial sum per 16 values for the first 96 values, 43 // stored as d0,d1,s1,s2,s3,s4,s5 44 }; 45 int8_t qs[QK8_1_MMQ]; 46}; 47 48// this struct is used for fp4 data types (currently only used for Blackwell) 49// mxfp4 has block size 32, each int32 of d4 contains 2 e8m0 scales in the lower 16 bits 50// nvfp4 has block size 16, each int32 of d4 contains 4 ue4m3 scales 51struct block_fp4_mmq { 52 uint32_t d4[4]; 53 int8_t qs[QK_FP4_MMQ / 2]; 54}; 55 56static_assert(sizeof(block_q8_1_mmq) == QK8_1_MMQ + 4*sizeof(half2), "Unexpected block_q8_1_mmq size"); 57static_assert(sizeof(block_q8_1_mmq) == 4*sizeof(block_q8_1), "Unexpected block_q8_1_mmq size"); 58static_assert(sizeof(block_fp4_mmq) == sizeof(block_q8_1_mmq), "Unexpected block_fp4_mmq size"); 59 60static mmq_q8_1_ds_layout mmq_get_q8_1_ds_layout(const ggml_type type_x) { 61 switch (type_x) { 62 case GGML_TYPE_Q1_0: 63 case GGML_TYPE_Q2_0: 64 return MMQ_Q8_1_DS_LAYOUT_D4; 65 case GGML_TYPE_Q4_0: 66 case GGML_TYPE_Q4_1: 67 return MMQ_Q8_1_DS_LAYOUT_DS4; 68 case GGML_TYPE_Q5_0: 69 return MMQ_Q8_1_DS_LAYOUT_D4; 70 case GGML_TYPE_Q5_1: 71 return MMQ_Q8_1_DS_LAYOUT_DS4; 72 case GGML_TYPE_Q8_0: 73 return MMQ_Q8_1_DS_LAYOUT_D4; 74 case GGML_TYPE_MXFP4: 75 return MMQ_Q8_1_DS_LAYOUT_D4; 76 case GGML_TYPE_NVFP4: 77 return MMQ_Q8_1_DS_LAYOUT_D4; 78 case GGML_TYPE_Q2_K: 79 return MMQ_Q8_1_DS_LAYOUT_D2S6; 80 case GGML_TYPE_Q3_K: 81 return MMQ_Q8_1_DS_LAYOUT_D4; 82 case GGML_TYPE_Q4_K: 83 case GGML_TYPE_Q5_K: 84 return MMQ_Q8_1_DS_LAYOUT_DS4; 85 case GGML_TYPE_Q6_K: 86 case GGML_TYPE_IQ2_XXS: 87 case GGML_TYPE_IQ2_XS: 88 case GGML_TYPE_IQ2_S: 89 case GGML_TYPE_IQ3_XXS: 90 case GGML_TYPE_IQ3_S: 91 return MMQ_Q8_1_DS_LAYOUT_D4; 92 case GGML_TYPE_IQ1_S: 93 return MMQ_Q8_1_DS_LAYOUT_DS4; 94 case GGML_TYPE_IQ4_XS: 95 case GGML_TYPE_IQ4_NL: 96 return MMQ_Q8_1_DS_LAYOUT_D4; 97 default: 98 GGML_ABORT("fatal error"); 99 break; 100 } 101} 102 103struct tile_x_sizes { 104 int qs; 105 int dm; 106 int sc; 107}; 108 109// Decouple shared memory tile sizes from WARP_SIZE to allow for different warp sizes. 110// The K dimension of the tiles has either, 111// 1*MMQ_TILE_NE_K==32 (always for TILE_Y_K) or 2*MMQ_TILE_NE_K==64 (typically for TILE_X_K), 112// 32 bit elements for the quantized data (does not include scales). 113// In other words, the size of the quantized data in the K dimension is a multiple of MMQ_TILE_NE_K. 114// The final tile size in K direction is padded to avoid shared memory bank conflicts, 115// in terms of 32 bit elements that means K % 2 == 1 for dp4a or K % 8 == 4 for mma. 116#define MMQ_TILE_NE_K 32 117 118// block_q8_1_mmq has (128 8-bit ints == 32 32-bit ints + 4 32-bit scales) 119#define MMQ_TILE_Y_K (MMQ_TILE_NE_K + MMQ_TILE_NE_K / QI8_1) 120#define MMQ_TILE_Y_FP4_K MMQ_TILE_Y_K 121 122enum ggml_cuda_mmq_sram_layout { 123 GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 124 GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, 125 GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, 126 GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, 127 GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, 128 GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, // MXFP4 and NVFP4 on Blackwell. 129 GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, // Generic NVFP4 130}; 131 132static constexpr __host__ __device__ int ggml_cuda_mmq_get_sram_stride(ggml_cuda_mmq_sram_layout sram_layout) { 133 switch (sram_layout) { 134 case GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0: 135 return 2*MMQ_TILE_NE_K + 2*MMQ_TILE_NE_K/QI8_0 + 4; 136 case GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1: 137 return 2*MMQ_TILE_NE_K + 2*MMQ_TILE_NE_K/QI8_1 + 4; 138 case GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K: 139 return 2*MMQ_TILE_NE_K + MMQ_TILE_NE_K + 4; 140 case GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K: 141 return 2*MMQ_TILE_NE_K + MMQ_TILE_NE_K/2 + 4; 142 case GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K: 143 return 2*MMQ_TILE_NE_K + MMQ_TILE_NE_K/QI6_K + MMQ_TILE_NE_K/8 + 7; 144 case GGML_CUDA_MMQ_SRAM_LAYOUT_FP4: 145 return 2*MMQ_TILE_NE_K + 8 + 4; 146 case GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4: 147 return 2*MMQ_TILE_NE_K + MMQ_TILE_NE_K/2 + 4; 148 default: 149 return -1; 150 } 151} 152 153static_assert(ggml_cuda_mmq_get_sram_stride(GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0) % 8 == 4, "Wrong padding."); 154static_assert(ggml_cuda_mmq_get_sram_stride(GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1) % 8 == 4, "Wrong padding."); 155static_assert(ggml_cuda_mmq_get_sram_stride(GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K) % 8 == 4, "Wrong padding."); 156static_assert(ggml_cuda_mmq_get_sram_stride(GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K) % 8 == 4, "Wrong padding."); 157static_assert(ggml_cuda_mmq_get_sram_stride(GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K) % 8 == 4, "Wrong padding."); 158static_assert(ggml_cuda_mmq_get_sram_stride(GGML_CUDA_MMQ_SRAM_LAYOUT_FP4) % 8 == 4, "Wrong padding."); 159static_assert(ggml_cuda_mmq_get_sram_stride(GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4) % 8 == 4, "Wrong padding."); 160 161static_assert(ggml_cuda_mmq_get_sram_stride(GGML_CUDA_MMQ_SRAM_LAYOUT_FP4) == ggml_cuda_mmq_get_sram_stride(GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1), "Wrong tile size for MXFP4"); 162 163// Config options for the MMQ kernel. 164// Should not affect results, only speed/register pressure/shared memory use. 165struct ggml_cuda_mmq_config { 166 ggml_type type; // src0->type 167 int nthreads; // Number of threads per CUDA block. 168 int occupancy; // Targeted occupancy for the MMA kernel. 169 int I; // SRAM tile width in src0->ne[1]/dst->ne[0] direction. 170 int J; // SRAM tile width in src1->ne[1]/dst->ne[1] direction. 171 ggml_cuda_mmq_sram_layout sram_layout; // SRAM tile length in src0->ne[0]/src1->ne[0] direction (physical 32 bit elements). 172 int K_vram; // VRAM tile length in src0->ne[0]/src1->ne[0] direction (logical elements). 173 bool stream_k; // Whether or not to use stream-k decomposition. 174 bool fallback; // Whether a fallback for out-of-bounds check in src0->ne[1] direction is needed. 175 176 constexpr __host__ __device__ ggml_cuda_mmq_config( 177 ggml_type type, int nthreads, int occupancy, int I, int J, ggml_cuda_mmq_sram_layout sram_layout, int K_vram, bool stream_k, bool fallback) : 178 type(type), nthreads(nthreads), occupancy(occupancy), I(I), J(J), sram_layout(sram_layout), K_vram(K_vram), stream_k(stream_k), fallback(fallback) {} 179 180 constexpr __device__ int rows_per_warp() const { 181#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) 182 return 16; 183#else 184 return J >= 48 && J % 16 == 0 ? 32 : 16; 185#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) 186 } 187 188 // TODO transition all combinations of GPUs and quantizations to the MMA data layout. 189 __host__ int use_mma_data_layout(const int cc) const { 190 if (amd_mfma_available(cc) || amd_wmma_available(cc) || turing_mma_available(cc)) { 191 return true; 192 } 193 return false; 194 } 195 196 constexpr __device__ bool use_mma_data_layout() const { 197#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) 198 return true; 199#else 200 return false; 201#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) 202 } 203 204}; 205 206#define CASE(type_, nthreads_, occupancy_, I_, J_, sram_layout_, K_vram_, stream_k_, fallback_) \ 207 if (type == (type_) && J == (J_) && fallback == (fallback_)) { \ 208 static_assert((nthreads_) % 32 == 0 && (nthreads_) <= 512, "bad nthreads"); \ 209 static_assert( (occupancy_) <= 8, "bad occupancy"); \ 210 static_assert((I_) % 32 == 0, "bad I"); \ 211 static_assert((J_) % 8 == 0, "bad J"); \ 212 static_assert((K_vram_) % 256 == 0, "bad K_vram"); \ 213 return ggml_cuda_mmq_config((type_), (nthreads_), (occupancy_), (I_), (J_), (sram_layout_), (K_vram_), (stream_k_), (fallback_)); \ 214 } \ 215 216#include "mmq-config-pascal-older.cuh" 217#include "mmq-config-pascal-dp4a.cuh" 218#include "mmq-config-ampere.cuh" 219#include "mmq-config-blackwell.cuh" 220 221#include "mmq-config-gcn.cuh" 222#include "mmq-config-cdna.cuh" 223#include "mmq-config-rdna2.cuh" 224#include "mmq-config-rdna3.cuh" 225#include "mmq-config-rdna3-5.cuh" 226#include "mmq-config-rdna4.cuh" 227 228#undef CASE 229 230static __host__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(const ggml_type type, const int J, const bool fallback, const int cc) { 231 if (GGML_CUDA_CC_IS_AMD(cc)) { 232 if (GGML_CUDA_CC_IS_GCN(cc)) { 233 return ggml_cuda_mmq_get_config_gcn(type, J, fallback); 234 } 235 if (GGML_CUDA_CC_IS_CDNA(cc)) { 236 return ggml_cuda_mmq_get_config_cdna(type, J, fallback); 237 } 238 if (GGML_CUDA_CC_IS_RDNA4(cc)) { 239 return ggml_cuda_mmq_get_config_rdna4(type, J, fallback); 240 } 241 if (GGML_CUDA_CC_IS_RDNA3_5(cc)) { 242 return ggml_cuda_mmq_get_config_rdna3_5(type, J, fallback); 243 } 244 if (GGML_CUDA_CC_IS_RDNA3(cc)) { // covers RDNA 3.0 245 return ggml_cuda_mmq_get_config_rdna3(type, J, fallback); 246 } 247 return ggml_cuda_mmq_get_config_rdna2(type, J, fallback); 248 } 249 if (blackwell_mma_available(cc)) { 250 return ggml_cuda_mmq_get_config_blackwell(type, J, fallback); 251 } 252 if (ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) { 253 return ggml_cuda_mmq_get_config_ampere(type, J, fallback); 254 } 255 if (ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_DP4A) { 256 return ggml_cuda_mmq_get_config_pascal_dp4a(type, J, fallback); 257 } 258 return ggml_cuda_mmq_get_config_pascal_older(type, J, fallback); 259} 260 261static constexpr __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(ggml_type type, int J, bool fallback) { 262#ifdef GGML_USE_HIP 263#ifdef GCN 264 return ggml_cuda_mmq_get_config_gcn(type, J, fallback); 265#elif defined(CDNA) 266 return ggml_cuda_mmq_get_config_cdna(type, J, fallback); 267#elif defined(RDNA4) 268 return ggml_cuda_mmq_get_config_rdna4(type, J, fallback); 269#elif defined(RDNA3_5) 270 return ggml_cuda_mmq_get_config_rdna3_5(type, J, fallback); 271#elif defined(RDNA3) 272 return ggml_cuda_mmq_get_config_rdna3(type, J, fallback); 273#else 274 return ggml_cuda_mmq_get_config_rdna2(type, J, fallback); 275#endif // CDNA 276#else 277#ifdef BLACKWELL_MMA_AVAILABLE 278 return ggml_cuda_mmq_get_config_blackwell(type, J, fallback); 279#elif __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA 280 return ggml_cuda_mmq_get_config_ampere(type, J, fallback); 281#elif __CUDA_ARCH__ >= GGML_CUDA_CC_DP4A 282 return ggml_cuda_mmq_get_config_pascal_dp4a(type, J, fallback); 283#else 284 return ggml_cuda_mmq_get_config_pascal_older(type, J, fallback); 285#endif // BLACKWELL_MMA_AVAILABLE 286#endif // GGML_USE_HIP 287 GGML_UNUSED_VARS(type, J, fallback); 288} 289 290static __host__ int ggml_cuda_mmq_get_type(const ggml_type type, const int J, const bool fallback, const int cc) { 291 return ggml_cuda_mmq_get_config(type, J, fallback, cc).type; 292} 293 294static constexpr __device__ int ggml_cuda_mmq_get_type(ggml_type type, int J, bool fallback) { 295 return ggml_cuda_mmq_get_config(type, J, fallback).type; 296} 297 298static __host__ int ggml_cuda_mmq_get_nthreads(const ggml_type type, const int J, const bool fallback, const int cc) { 299 return ggml_cuda_mmq_get_config(type, J, fallback, cc).nthreads; 300} 301 302static constexpr __device__ int ggml_cuda_mmq_get_nthreads(ggml_type type, int J, bool fallback) { 303 return ggml_cuda_mmq_get_config(type, J, fallback).nthreads; 304} 305 306static __host__ int ggml_cuda_mmq_get_occupancy(const ggml_type type, const int J, const bool fallback, const int cc) { 307 return ggml_cuda_mmq_get_config(type, J, fallback, cc).occupancy; 308} 309 310static constexpr __device__ int ggml_cuda_mmq_get_occupancy(ggml_type type, int J, bool fallback) { 311 return ggml_cuda_mmq_get_config(type, J, fallback).occupancy; 312} 313 314static __host__ int ggml_cuda_mmq_get_I(const ggml_type type, const int J, const bool fallback, const int cc) { 315 return ggml_cuda_mmq_get_config(type, J, fallback, cc).I; 316} 317 318static constexpr __device__ int ggml_cuda_mmq_get_I(ggml_type type, int J, bool fallback) { 319 return ggml_cuda_mmq_get_config(type, J, fallback).I; 320} 321 322static __host__ int ggml_cuda_mmq_get_J(const ggml_type type, const int J, const bool fallback, const int cc) { 323 return ggml_cuda_mmq_get_config(type, J, fallback, cc).J; 324} 325 326static constexpr __device__ int ggml_cuda_mmq_get_J(ggml_type type, int J, bool fallback) { 327 return ggml_cuda_mmq_get_config(type, J, fallback).J; 328} 329 330static __host__ ggml_cuda_mmq_sram_layout ggml_cuda_mmq_get_sram_layout(const ggml_type type, const int J, const bool fallback, const int cc) { 331 return ggml_cuda_mmq_get_config(type, J, fallback, cc).sram_layout; 332} 333 334static constexpr __device__ ggml_cuda_mmq_sram_layout ggml_cuda_mmq_get_sram_layout(ggml_type type, int J, bool fallback) { 335 return ggml_cuda_mmq_get_config(type, J, fallback).sram_layout; 336} 337 338static __host__ int ggml_cuda_mmq_get_K_vram(const ggml_type type, const int J, const bool fallback, const int cc) { 339 return ggml_cuda_mmq_get_config(type, J, fallback, cc).K_vram; 340} 341 342static constexpr __device__ int ggml_cuda_mmq_get_K_vram(ggml_type type, int J, bool fallback) { 343 return ggml_cuda_mmq_get_config(type, J, fallback).K_vram; 344} 345 346static __host__ bool ggml_cuda_mmq_get_stream_k(const ggml_type type, const int J, const bool fallback, const int cc) { 347 return ggml_cuda_mmq_get_config(type, J, fallback, cc).stream_k; 348} 349 350static constexpr __device__ bool ggml_cuda_mmq_get_stream_k(ggml_type type, int J, bool fallback) { 351 return ggml_cuda_mmq_get_config(type, J, fallback).stream_k; 352} 353 354static __host__ int ggml_cuda_mmq_get_fallback(const ggml_type type, const int J, const bool fallback, const int cc) { 355 return ggml_cuda_mmq_get_config(type, J, fallback, cc).fallback; 356} 357 358static constexpr __device__ int ggml_cuda_mmq_get_fallback(ggml_type type, int J, bool fallback) { 359 return ggml_cuda_mmq_get_config(type, J, fallback).fallback; 360} 361 362// --------------------------------------------------------------------------------------------- 363 364static __host__ int ggml_cuda_mmq_get_sram_stride(const ggml_type type, const int J, const bool fallback, const int cc) { 365 return ggml_cuda_mmq_get_sram_stride(ggml_cuda_mmq_get_sram_layout(type, J, fallback, cc)); 366} 367 368static constexpr __device__ int ggml_cuda_mmq_get_sram_stride(ggml_type type, int J, bool fallback) { 369 return ggml_cuda_mmq_get_sram_stride(ggml_cuda_mmq_get_sram_layout(type, J, fallback)); 370} 371 372static __host__ int ggml_cuda_mmq_get_J_max(const ggml_type type, const bool fallback, const int cc, const int64_t ne11) { 373 int ret = std::min(ne11, int64_t(512)); 374 ret -= ret % 8; 375 for (;ret > 0; ret -= 8) { 376 if (ggml_cuda_mmq_get_config(type, ret, fallback, cc).type != GGML_TYPE_COUNT) { 377 return ret; 378 } 379 } 380 return ret; 381} 382 383static constexpr __device__ int ggml_cuda_mmq_get_rows_per_warp(ggml_type type, int J, bool fallback) { 384 return ggml_cuda_mmq_get_config(type, J, fallback).rows_per_warp(); 385} 386 387#define MMQ_DP4A_TXS_Q4_0 tile_x_sizes{I*MMQ_TILE_NE_K + I, I*MMQ_TILE_NE_K/QI4_0 + I/QI4_0, 0} 388#define MMQ_DP4A_TXS_Q4_1 tile_x_sizes{I*MMQ_TILE_NE_K + I, I*MMQ_TILE_NE_K/QI4_1 + I/QI4_1, 0} 389#define MMQ_DP4A_TXS_Q8_0 tile_x_sizes{I*MMQ_TILE_NE_K*2 + I, I*MMQ_TILE_NE_K*2/QI8_0 + I/(QI8_0/2), 0} 390#define MMQ_DP4A_TXS_Q8_0_16 tile_x_sizes{I*MMQ_TILE_NE_K*2 + I, I*MMQ_TILE_NE_K*4/QI8_0 + I/(QI8_0/4), 0} 391#define MMQ_DP4A_TXS_Q8_1 tile_x_sizes{I*MMQ_TILE_NE_K*2 + I, I*MMQ_TILE_NE_K*2/QI8_1 + I/(QI8_1/2), 0} 392#define MMQ_DP4A_TXS_Q2_K tile_x_sizes{I*MMQ_TILE_NE_K*2 + I, I*MMQ_TILE_NE_K + I, 0} 393#define MMQ_DP4A_TXS_Q3_K tile_x_sizes{I*MMQ_TILE_NE_K*2 + I, I, I*MMQ_TILE_NE_K/8 + I/8} 394#define MMQ_DP4A_TXS_Q4_K tile_x_sizes{I*MMQ_TILE_NE_K + I, I*MMQ_TILE_NE_K/QI4_K, I*MMQ_TILE_NE_K/8 + I/8} 395#define MMQ_DP4A_TXS_Q5_K tile_x_sizes{I*MMQ_TILE_NE_K*2 + I, I*MMQ_TILE_NE_K/QI5_K + I/QI5_K, I*MMQ_TILE_NE_K/8 + I/8} 396#define MMQ_DP4A_TXS_Q6_K tile_x_sizes{I*MMQ_TILE_NE_K*2 + I, I*MMQ_TILE_NE_K/QI6_K + I/QI6_K, I*MMQ_TILE_NE_K/8 + I/8} 397 398static constexpr __host__ __device__ tile_x_sizes mmq_get_dp4a_tile_x_sizes(ggml_type type, int I) { 399 switch (type) { 400 case GGML_TYPE_Q1_0: return MMQ_DP4A_TXS_Q8_0; 401 case GGML_TYPE_Q2_0: return MMQ_DP4A_TXS_Q8_0; 402 case GGML_TYPE_Q4_0: return MMQ_DP4A_TXS_Q4_0; 403 case GGML_TYPE_Q4_1: return MMQ_DP4A_TXS_Q4_1; 404 case GGML_TYPE_Q5_0: return MMQ_DP4A_TXS_Q8_0; 405 case GGML_TYPE_Q5_1: return MMQ_DP4A_TXS_Q8_1; 406 case GGML_TYPE_Q8_0: return MMQ_DP4A_TXS_Q8_0; 407 case GGML_TYPE_MXFP4: return MMQ_DP4A_TXS_Q8_1; 408 case GGML_TYPE_NVFP4: return MMQ_DP4A_TXS_Q8_0_16; 409 case GGML_TYPE_Q2_K: return MMQ_DP4A_TXS_Q2_K; 410 case GGML_TYPE_Q3_K: return MMQ_DP4A_TXS_Q3_K; 411 case GGML_TYPE_Q4_K: return MMQ_DP4A_TXS_Q4_K; 412 case GGML_TYPE_Q5_K: return MMQ_DP4A_TXS_Q5_K; 413 case GGML_TYPE_Q6_K: return MMQ_DP4A_TXS_Q6_K; 414 case GGML_TYPE_IQ2_XXS: return MMQ_DP4A_TXS_Q8_0; 415 case GGML_TYPE_IQ2_XS: return MMQ_DP4A_TXS_Q8_0_16; 416 case GGML_TYPE_IQ2_S: return MMQ_DP4A_TXS_Q8_0_16; 417 case GGML_TYPE_IQ3_XXS: return MMQ_DP4A_TXS_Q8_0; 418 case GGML_TYPE_IQ3_S: return MMQ_DP4A_TXS_Q8_0; 419 case GGML_TYPE_IQ1_S: return MMQ_DP4A_TXS_Q8_0; 420 case GGML_TYPE_IQ4_XS: return MMQ_DP4A_TXS_Q8_0; 421 case GGML_TYPE_IQ4_NL: return MMQ_DP4A_TXS_Q8_0; 422 default: return tile_x_sizes{0, 0, 0}; 423 } 424} 425 426// FIXME temporary until all combinations of data types and GPUs can use the MMA data layout 427static __host__ int ggml_cuda_mmq_get_nbytes_shared_x(const ggml_cuda_mmq_config & config, const int cc) { 428 if (config.use_mma_data_layout(cc)) { 429 return config.I * ggml_cuda_mmq_get_sram_stride(config.sram_layout) * 4; 430 } 431 const tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(config.type, config.I); 432 return (txs.qs + txs.dm + txs.sc) * 4; 433} 434 435// ------------------------------------------------------------ 436 437#include "mmq-load-tiles.cuh" 438#include "mmq-vec-dot.cuh" 439 440template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_write_back_dp4a( 441 const float * __restrict__ sum, const int32_t * __restrict__ ids_dst, float * __restrict__ dst, 442 const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max) { 443 constexpr int warp_size = ggml_cuda_get_physical_warp_size(); 444 constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size; 445 constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback); 446 447 const bool y_scale_used = y_scale != nullptr; 448 449#pragma unroll 450 for (int j0 = 0; j0 < J; j0 += nwarps) { 451 const int j = j0 + threadIdx.y; 452 453 if (j > j_max) { 454 return; 455 } 456 457#pragma unroll 458 for (int i0 = 0; i0 < I; i0 += warp_size) { 459 const int i = i0 + threadIdx.x; 460 461 if (fallback && i > i_max) { 462 continue; 463 } 464 465 if constexpr (type == GGML_TYPE_NVFP4) { 466 if (y_scale_used) { 467 dst[ids_dst[j]*stride + i] = y_scale[j] * sum[(j0/nwarps) * (I/warp_size) + i0/warp_size]; 468 } else { 469 dst[ids_dst[j]*stride + i] = sum[(j0/nwarps) * (I/warp_size) + i0/warp_size]; 470 } 471 } else { 472 dst[ids_dst[j]*stride + i] = sum[(j0/nwarps) * (I/warp_size) + i0/warp_size]; 473 GGML_UNUSED(y_scale_used); 474 } 475 } 476 } 477} 478 479template<ggml_type type, int J, bool fallback> 480static __device__ __forceinline__ void ggml_cuda_mmq_write_back_mma( 481 const float * __restrict__ sum, const int * __restrict__ ids_dst, float * __restrict__ dst, 482 const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max) { 483 484#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) 485 typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C; 486#else 487 typedef tile<16, 8, int> tile_C; 488#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) 489 490 constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback); 491 constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp. 492 493 const int i0 = (threadIdx.y / ntx) * (ntx*tile_C::I); 494 495 const bool y_scale_used = y_scale != nullptr; 496 497#pragma unroll 498 for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) { 499#pragma unroll 500 for (int n = 0; n < ntx; ++n) { 501#pragma unroll 502 for (int l = 0; l < tile_C::ne; ++l) { 503 const int j = j0 + (threadIdx.y % ntx) * tile_C::J + tile_C::get_j(l); 504 505 if (j > j_max) { 506 continue; 507 } 508 509 const int i = i0 + n*tile_C::I + tile_C::get_i(l); 510 511 if (fallback && i > i_max) { 512 continue; 513 } 514 515 if constexpr (type == GGML_TYPE_NVFP4) { 516 if (y_scale_used) { 517 dst[ids_dst[j]*stride + i] = y_scale[j] * sum[(j0/tile_C::J + n)*tile_C::ne + l]; 518 } else { 519 dst[ids_dst[j]*stride + i] = sum[(j0/tile_C::J + n)*tile_C::ne + l]; 520 } 521 } else { 522 dst[ids_dst[j]*stride + i] = sum[(j0/tile_C::J + n)*tile_C::ne + l]; 523 GGML_UNUSED(y_scale_used); 524 } 525 } 526 } 527 } 528} 529 530// ------------------------------------------------------------------------------------------------------------------------------------- 531 532// TODO remove this struct and use ggml_cuda_mmq_sram_layout instead. 533struct ggml_cuda_mmq_util_funcs { 534 int vdr; 535 ggml_cuda_mmq_load_tiles_t load_tiles; 536 ggml_cuda_mmq_vec_dot_t vec_dot; 537 ggml_cuda_mmq_write_back_t write_back; 538 539 constexpr __host__ __device__ ggml_cuda_mmq_util_funcs( 540 int vdr, ggml_cuda_mmq_load_tiles_t load_tiles, ggml_cuda_mmq_vec_dot_t vec_dot, ggml_cuda_mmq_write_back_t write_back) : 541 vdr(vdr), load_tiles(load_tiles), vec_dot(vec_dot), write_back(write_back) {} 542}; 543 544template <ggml_type type, int J, bool fallback> 545static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_funcs() { 546 if (!ggml_cuda_mmq_get_config(type, J, fallback).use_mma_data_layout()) { 547 switch (type) { 548 case GGML_TYPE_Q1_0: 549 return ggml_cuda_mmq_util_funcs( 550 VDR_Q1_0_Q8_1_MMQ, 551 ggml_cuda_mmq_load_tiles_q1_0<type, J, fallback>, 552 ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a<type, J, fallback>, 553 ggml_cuda_mmq_write_back_dp4a<type, J, fallback>); 554 case GGML_TYPE_Q2_0: 555 return ggml_cuda_mmq_util_funcs( 556 VDR_Q2_0_Q8_1_MMQ, 557 ggml_cuda_mmq_load_tiles_q2_0<type, J, fallback>, 558 ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a<type, J, fallback>, 559 ggml_cuda_mmq_write_back_dp4a<type, J, fallback>); 560 case GGML_TYPE_Q4_0: 561 return ggml_cuda_mmq_util_funcs( 562 VDR_Q4_0_Q8_1_MMQ, 563 ggml_cuda_mmq_load_tiles_q4_0<type, J, fallback>, 564 ggml_cuda_mmq_vec_dot_q4_0_q8_1_dp4a<type, J, fallback>, 565 ggml_cuda_mmq_write_back_dp4a<type, J, fallback>); 566 case GGML_TYPE_Q4_1: 567 return ggml_cuda_mmq_util_funcs( 568 VDR_Q4_1_Q8_1_MMQ, 569 ggml_cuda_mmq_load_tiles_q4_1<type, J, fallback>, 570 ggml_cuda_mmq_vec_dot_q4_1_q8_1_dp4a<type, J, fallback>, 571 ggml_cuda_mmq_write_back_dp4a<type, J, fallback>); 572 case GGML_TYPE_Q5_0: 573 return ggml_cuda_mmq_util_funcs( 574 VDR_Q5_0_Q8_1_MMQ, 575 ggml_cuda_mmq_load_tiles_q5_0<type, J, fallback>, 576 ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a<type, J, fallback>, 577 ggml_cuda_mmq_write_back_dp4a<type, J, fallback>); 578 case GGML_TYPE_Q5_1: 579 return ggml_cuda_mmq_util_funcs( 580 VDR_Q5_1_Q8_1_MMQ, 581 ggml_cuda_mmq_load_tiles_q5_1<type, J, fallback>, 582 ggml_cuda_mmq_vec_dot_q8_1_q8_1_dp4a<type, J, fallback>, 583 ggml_cuda_mmq_write_back_dp4a<type, J, fallback>); 584 case GGML_TYPE_Q8_0: 585 return ggml_cuda_mmq_util_funcs( 586 VDR_Q8_0_Q8_1_MMQ, 587 ggml_cuda_mmq_load_tiles_q8_0<type, J, fallback>, 588 ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a<type, J, fallback>, 589 ggml_cuda_mmq_write_back_dp4a<type, J, fallback>); 590// --------------------------------------------------------------------------------------------- 591 case GGML_TYPE_Q2_K: 592 return ggml_cuda_mmq_util_funcs( 593 VDR_Q2_K_Q8_1_MMQ, 594 ggml_cuda_mmq_load_tiles_q2_K<type, J, fallback>, 595 ggml_cuda_mmq_vec_dot_q2_K_q8_1_dp4a<type, J, fallback>, 596 ggml_cuda_mmq_write_back_dp4a<type, J, fallback>); 597 case GGML_TYPE_Q3_K: 598 return ggml_cuda_mmq_util_funcs( 599 VDR_Q3_K_Q8_1_MMQ, 600 ggml_cuda_mmq_load_tiles_q3_K<type, J, fallback>, 601 ggml_cuda_mmq_vec_dot_q3_K_q8_1_dp4a<type, J, fallback>, 602 ggml_cuda_mmq_write_back_dp4a<type, J, fallback>); 603 case GGML_TYPE_Q4_K: 604 return ggml_cuda_mmq_util_funcs( 605 VDR_Q4_K_Q8_1_MMQ, 606 ggml_cuda_mmq_load_tiles_q4_K<type, J, fallback>, 607 ggml_cuda_mmq_vec_dot_q4_K_q8_1_dp4a<type, J, fallback>, 608 ggml_cuda_mmq_write_back_dp4a<type, J, fallback>); 609 case GGML_TYPE_Q5_K: 610 return ggml_cuda_mmq_util_funcs( 611 VDR_Q5_K_Q8_1_MMQ, 612 ggml_cuda_mmq_load_tiles_q5_K<type, J, fallback>, 613 ggml_cuda_mmq_vec_dot_q5_K_q8_1_dp4a<type, J, fallback>, 614 ggml_cuda_mmq_write_back_dp4a<type, J, fallback>); 615 case GGML_TYPE_Q6_K: 616 return ggml_cuda_mmq_util_funcs( 617 VDR_Q6_K_Q8_1_MMQ, 618 ggml_cuda_mmq_load_tiles_q6_K<type, J, fallback>, 619 ggml_cuda_mmq_vec_dot_q6_K_q8_1_dp4a<type, J, fallback>, 620 ggml_cuda_mmq_write_back_dp4a<type, J, fallback>); 621// --------------------------------------------------------------------------------------------- 622 case GGML_TYPE_IQ1_S: 623 return ggml_cuda_mmq_util_funcs( 624 VDR_IQ1_S_Q8_1_MMQ, 625 ggml_cuda_mmq_load_tiles_iq1_s<type, J, fallback>, 626 ggml_cuda_mmq_vec_dot_q8_1_q8_1_dp4a<type, J, fallback>, 627 ggml_cuda_mmq_write_back_dp4a<type, J, fallback>); 628 case GGML_TYPE_IQ2_XXS: 629 return ggml_cuda_mmq_util_funcs( 630 VDR_IQ2_XXS_Q8_1_MMQ, 631 ggml_cuda_mmq_load_tiles_iq2_xxs<type, J, fallback>, 632 ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a<type, J, fallback>, 633 ggml_cuda_mmq_write_back_dp4a<type, J, fallback>); 634 case GGML_TYPE_IQ2_XS: 635 return ggml_cuda_mmq_util_funcs( 636 VDR_IQ2_XS_Q8_1_MMQ, 637 ggml_cuda_mmq_load_tiles_iq2_xs<type, J, fallback>, 638 ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_dp4a<type, J, fallback>, 639 ggml_cuda_mmq_write_back_dp4a<type, J, fallback>); 640 case GGML_TYPE_IQ2_S: 641 return ggml_cuda_mmq_util_funcs( 642 VDR_IQ2_S_Q8_1_MMQ, 643 ggml_cuda_mmq_load_tiles_iq2_s<type, J, fallback>, 644 ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_dp4a<type, J, fallback>, 645 ggml_cuda_mmq_write_back_dp4a<type, J, fallback>); 646 case GGML_TYPE_IQ3_XXS: 647 return ggml_cuda_mmq_util_funcs( 648 VDR_IQ3_XXS_Q8_1_MMQ, 649 ggml_cuda_mmq_load_tiles_iq3_xxs<type, J, fallback>, 650 ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a<type, J, fallback>, 651 ggml_cuda_mmq_write_back_dp4a<type, J, fallback>); 652 case GGML_TYPE_IQ3_S: 653 return ggml_cuda_mmq_util_funcs( 654 VDR_IQ3_S_Q8_1_MMQ, 655 ggml_cuda_mmq_load_tiles_iq3_s<type, J, fallback>, 656 ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a<type, J, fallback>, 657 ggml_cuda_mmq_write_back_dp4a<type, J, fallback>); 658 case GGML_TYPE_IQ4_XS: 659 return ggml_cuda_mmq_util_funcs( 660 VDR_IQ4_XS_Q8_1_MMQ, 661 ggml_cuda_mmq_load_tiles_iq4_xs<type, J, fallback>, 662 ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a<type, J, fallback>, 663 ggml_cuda_mmq_write_back_dp4a<type, J, fallback>); 664 case GGML_TYPE_IQ4_NL: 665 return ggml_cuda_mmq_util_funcs( 666 VDR_IQ4_NL_Q8_1_MMQ, 667 ggml_cuda_mmq_load_tiles_iq4_nl<type, J, fallback>, 668 ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a<type, J, fallback>, 669 ggml_cuda_mmq_write_back_dp4a<type, J, fallback>); 670// --------------------------------------------------------------------------------------------- 671 case GGML_TYPE_MXFP4: 672 return ggml_cuda_mmq_util_funcs( 673 VDR_MXFP4_Q8_1_MMQ, 674 ggml_cuda_mmq_load_tiles_mxfp4<type, J, fallback>, 675 ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a<type, J, fallback>, 676 ggml_cuda_mmq_write_back_dp4a<type, J, fallback>); 677 case GGML_TYPE_NVFP4: 678 return ggml_cuda_mmq_util_funcs( 679 VDR_NVFP4_Q8_1_MMQ, 680 ggml_cuda_mmq_load_tiles_nvfp4<type, J, fallback>, 681 ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_dp4a<type, J, fallback>, 682 ggml_cuda_mmq_write_back_dp4a<type, J, fallback>); 683 default: 684 return ggml_cuda_mmq_util_funcs(1, nullptr, nullptr, nullptr); 685 } 686 } 687 688// --------------------------------------------------------------------------------------------- 689 690#ifdef BLACKWELL_MMA_AVAILABLE 691 switch (type) { 692 case GGML_TYPE_MXFP4: 693 return ggml_cuda_mmq_util_funcs( 694 -1, 695 ggml_cuda_mmq_load_tiles_mxfp4_fp4<type, J, fallback>, 696 ggml_cuda_mmq_vec_dot_fp4_fp4_mma<type, J, fallback>, 697 ggml_cuda_mmq_write_back_mma<type, J, fallback>); 698 case GGML_TYPE_NVFP4: 699 return ggml_cuda_mmq_util_funcs( 700 -1, 701 ggml_cuda_mmq_load_tiles_nvfp4_nvfp4<type, J, fallback>, 702 ggml_cuda_mmq_vec_dot_fp4_fp4_mma<type, J, fallback>, 703 ggml_cuda_mmq_write_back_mma<type, J, fallback>); 704 default: 705 break; 706 } 707#endif // BLACKWELL_MMA_AVAILABLE 708 709// --------------------------------------------------------------------------------------------- 710 711 switch (type) { 712 case GGML_TYPE_Q1_0: 713 return ggml_cuda_mmq_util_funcs( 714 -1, 715 ggml_cuda_mmq_load_tiles_q1_0<type, J, fallback>, 716 ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>, 717 ggml_cuda_mmq_write_back_mma<type, J, fallback>); 718 case GGML_TYPE_Q2_0: 719 return ggml_cuda_mmq_util_funcs( 720 -1, 721 ggml_cuda_mmq_load_tiles_q2_0<type, J, fallback>, 722 ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>, 723 ggml_cuda_mmq_write_back_mma<type, J, fallback>); 724 case GGML_TYPE_Q4_0: 725 return ggml_cuda_mmq_util_funcs( 726 -1, 727 ggml_cuda_mmq_load_tiles_q4_0<type, J, fallback>, 728 ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_DS4>, 729 ggml_cuda_mmq_write_back_mma<type, J, fallback>); 730 case GGML_TYPE_Q4_1: 731 return ggml_cuda_mmq_util_funcs( 732 -1, 733 ggml_cuda_mmq_load_tiles_q4_1<type, J, fallback>, 734 ggml_cuda_mmq_vec_dot_q8_1_q8_1_mma<type, J, fallback>, 735 ggml_cuda_mmq_write_back_mma<type, J, fallback>); 736 case GGML_TYPE_Q5_0: 737 return ggml_cuda_mmq_util_funcs( 738 -1, 739 ggml_cuda_mmq_load_tiles_q5_0<type, J, fallback>, 740 ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>, 741 ggml_cuda_mmq_write_back_mma<type, J, fallback>); 742 case GGML_TYPE_Q5_1: 743 return ggml_cuda_mmq_util_funcs( 744 -1, 745 ggml_cuda_mmq_load_tiles_q5_1<type, J, fallback>, 746 ggml_cuda_mmq_vec_dot_q8_1_q8_1_mma<type, J, fallback>, 747 ggml_cuda_mmq_write_back_mma<type, J, fallback>); 748 case GGML_TYPE_Q8_0: 749 return ggml_cuda_mmq_util_funcs( 750 -1, 751 ggml_cuda_mmq_load_tiles_q8_0<type, J, fallback>, 752 ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>, 753 ggml_cuda_mmq_write_back_mma<type, J, fallback>); 754// --------------------------------------------------------------------------------------------- 755 case GGML_TYPE_Q2_K: 756 return ggml_cuda_mmq_util_funcs( 757 -1, 758 ggml_cuda_mmq_load_tiles_q2_K<type, J, fallback>, 759 ggml_cuda_mmq_vec_dot_q2_K_q8_1_mma<type, J, fallback>, 760 ggml_cuda_mmq_write_back_mma<type, J, fallback>); 761 case GGML_TYPE_Q3_K: 762 return ggml_cuda_mmq_util_funcs( 763 -1, 764 ggml_cuda_mmq_load_tiles_q3_K<type, J, fallback>, 765 ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_mma<type, J, fallback>, 766 ggml_cuda_mmq_write_back_mma<type, J, fallback>); 767 case GGML_TYPE_Q4_K: 768 return ggml_cuda_mmq_util_funcs( 769 -1, 770 ggml_cuda_mmq_load_tiles_q4_K<type, J, fallback>, 771 ggml_cuda_mmq_vec_dot_q8_1_q8_1_mma<type, J, fallback>, 772 ggml_cuda_mmq_write_back_mma<type, J, fallback>); 773 case GGML_TYPE_Q5_K: 774 return ggml_cuda_mmq_util_funcs( 775 -1, 776 ggml_cuda_mmq_load_tiles_q5_K<type, J, fallback>, 777 ggml_cuda_mmq_vec_dot_q8_1_q8_1_mma<type, J, fallback>, 778 ggml_cuda_mmq_write_back_mma<type, J, fallback>); 779 case GGML_TYPE_Q6_K: 780 return ggml_cuda_mmq_util_funcs( 781 -1, 782 ggml_cuda_mmq_load_tiles_q6_K<type, J, fallback>, 783 ggml_cuda_mmq_vec_dot_q6_K_q8_1_mma<type, J, fallback>, 784 ggml_cuda_mmq_write_back_mma<type, J, fallback>); 785// --------------------------------------------------------------------------------------------- 786 case GGML_TYPE_IQ1_S: 787 return ggml_cuda_mmq_util_funcs( 788 -1, 789 ggml_cuda_mmq_load_tiles_iq1_s<type, J, fallback>, 790 ggml_cuda_mmq_vec_dot_q8_1_q8_1_mma<type, J, fallback>, 791 ggml_cuda_mmq_write_back_mma<type, J, fallback>); 792 case GGML_TYPE_IQ2_XXS: 793 return ggml_cuda_mmq_util_funcs( 794 -1, 795 ggml_cuda_mmq_load_tiles_iq2_xxs<type, J, fallback>, 796 ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>, 797 ggml_cuda_mmq_write_back_mma<type, J, fallback>); 798 case GGML_TYPE_IQ2_XS: 799 return ggml_cuda_mmq_util_funcs( 800 -1, 801 ggml_cuda_mmq_load_tiles_iq2_xs<type, J, fallback>, 802 ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_mma<type, J, fallback>, 803 ggml_cuda_mmq_write_back_mma<type, J, fallback>); 804 case GGML_TYPE_IQ2_S: 805 return ggml_cuda_mmq_util_funcs( 806 -1, 807 ggml_cuda_mmq_load_tiles_iq2_s<type, J, fallback>, 808 ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_mma<type, J, fallback>, 809 ggml_cuda_mmq_write_back_mma<type, J, fallback>); 810 case GGML_TYPE_IQ3_XXS: 811 return ggml_cuda_mmq_util_funcs( 812 -1, 813 ggml_cuda_mmq_load_tiles_iq3_xxs<type, J, fallback>, 814 ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>, 815 ggml_cuda_mmq_write_back_mma<type, J, fallback>); 816 case GGML_TYPE_IQ3_S: 817 return ggml_cuda_mmq_util_funcs( 818 -1, 819 ggml_cuda_mmq_load_tiles_iq3_s<type, J, fallback>, 820 ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>, 821 ggml_cuda_mmq_write_back_mma<type, J, fallback>); 822 case GGML_TYPE_IQ4_XS: 823 return ggml_cuda_mmq_util_funcs( 824 -1, 825 ggml_cuda_mmq_load_tiles_iq4_xs<type, J, fallback>, 826 ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>, 827 ggml_cuda_mmq_write_back_mma<type, J, fallback>); 828 case GGML_TYPE_IQ4_NL: 829 return ggml_cuda_mmq_util_funcs( 830 -1, 831 ggml_cuda_mmq_load_tiles_iq4_nl<type, J, fallback>, 832 ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>, 833 ggml_cuda_mmq_write_back_mma<type, J, fallback>); 834// --------------------------------------------------------------------------------------------- 835 case GGML_TYPE_MXFP4: 836 return ggml_cuda_mmq_util_funcs( 837 -1, 838 ggml_cuda_mmq_load_tiles_mxfp4<type, J, fallback>, 839 ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>, 840 ggml_cuda_mmq_write_back_mma<type, J, fallback>); 841 case GGML_TYPE_NVFP4: 842 return ggml_cuda_mmq_util_funcs( 843 -1, 844 ggml_cuda_mmq_load_tiles_nvfp4<type, J, fallback>, 845 ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_mma<type, J, fallback>, 846 ggml_cuda_mmq_write_back_mma<type, J, fallback>); 847 default: 848 return ggml_cuda_mmq_util_funcs(1, nullptr, nullptr, nullptr); 849 } 850} 851 852template <ggml_type type, int J, bool fallback> 853static constexpr __device__ int ggml_cuda_mmq_get_vdr() { 854 return ggml_cuda_mmq_get_util_funcs<type, J, fallback>().vdr; 855} 856 857template <ggml_type type, int J, bool fallback> 858static constexpr __device__ ggml_cuda_mmq_load_tiles_t ggml_cuda_mmq_get_load_tiles() { 859 return ggml_cuda_mmq_get_util_funcs<type, J, fallback>().load_tiles; 860} 861 862template <ggml_type type, int J, bool fallback> 863static constexpr __device__ ggml_cuda_mmq_vec_dot_t ggml_cuda_mmq_get_vec_dot() { 864 return ggml_cuda_mmq_get_util_funcs<type, J, fallback>().vec_dot; 865} 866 867template <ggml_type type, int J, bool fallback> 868static constexpr __device__ ggml_cuda_mmq_write_back_t ggml_cuda_mmq_get_write_back() { 869 return ggml_cuda_mmq_get_util_funcs<type, J, fallback>().write_back; 870} 871 872// --------------------------------------------------------------------------------------------- 873 874template <ggml_type type, int J, bool fallback, bool fixup> 875static __device__ __forceinline__ void mul_mat_q_process_tile( 876 const char * __restrict__ x, const int offset_x, const int * __restrict__ y, 877 const int * __restrict__ ids_dst, float * __restrict__ dst, float * __restrict__ tmp_fixup, 878 const float * __restrict__ y_scale, 879 const int stride_row_x, const int ncols_y, const int stride_col_dst, 880 const int tile_x_max_i, const int tile_y_max_j, const int kb0_start, const int kb0_stop) { 881 882 constexpr int warp_size = ggml_cuda_get_physical_warp_size(); 883 constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size; 884 constexpr int qk = ggml_cuda_type_traits<type>::qk; 885 constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback); 886 constexpr ggml_cuda_mmq_load_tiles_t load_tiles = ggml_cuda_mmq_get_load_tiles<type, J, fallback>(); 887 constexpr ggml_cuda_mmq_vec_dot_t vec_dot = ggml_cuda_mmq_get_vec_dot<type, J, fallback>(); 888 constexpr ggml_cuda_mmq_write_back_t write_back = ggml_cuda_mmq_get_write_back<type, J, fallback>(); 889 890 extern __shared__ int data_mul_mat_q[]; 891 int * tile_y = data_mul_mat_q + J; 892 int * tile_x = tile_y + GGML_PAD(J*MMQ_TILE_Y_K, nwarps*warp_size); 893 894#if defined(BLACKWELL_MMA_AVAILABLE) 895 // FP4 tile stores 8 blocks 896 constexpr int ne_block = (type == GGML_TYPE_MXFP4 || type == GGML_TYPE_NVFP4) ? QK_FP4_MMQ : QK8_1_MMQ; 897#else 898 constexpr int ne_block = QK8_1_MMQ; 899#endif // defined(BLACKWELL_MMA_AVAILABLE) 900 901 constexpr int ITER_K = ggml_cuda_mmq_get_K_vram(type, J, fallback); 902 constexpr int blocks_per_iter = ITER_K / qk; 903 904 float sum[J*I / (nwarps*warp_size)] = {0.0f}; 905 906 constexpr int sz = sizeof(block_q8_1_mmq) / sizeof(int); 907 908 for (int kb0 = kb0_start; kb0 < kb0_stop; kb0 += blocks_per_iter) { 909 load_tiles(x, tile_x, offset_x + kb0, tile_x_max_i, stride_row_x); 910 { 911 const int * by0 = y + ncols_y * (kb0 * qk / ne_block) * sz; 912#pragma unroll 913 for (int l0 = 0; l0 < J * MMQ_TILE_Y_K; l0 += nwarps * warp_size) { 914 int l = l0 + threadIdx.y*warp_size + threadIdx.x; 915 916 tile_y[l] = by0[l]; 917 } 918 } 919 920 __syncthreads(); 921 922 vec_dot(tile_x, tile_y, sum, 0); 923 924 __syncthreads(); 925 926 { 927 const int * by0 = y + ncols_y * ((kb0 * qk / ne_block) * sz + sz); 928#pragma unroll 929 for (int l0 = 0; l0 < J * MMQ_TILE_Y_K; l0 += nwarps * warp_size) { 930 int l = l0 + threadIdx.y*warp_size + threadIdx.x; 931 932 tile_y[l] = by0[l]; 933 } 934 } 935 936 __syncthreads(); 937 938 vec_dot(tile_x, tile_y, sum, MMQ_TILE_NE_K); 939 940 __syncthreads(); 941 } 942 943 if (fixup) { 944 write_back(sum, ids_dst, tmp_fixup + blockIdx.x*(J*I), y_scale, I, I, J); 945 } else { 946 write_back(sum, ids_dst, dst, y_scale, stride_col_dst, tile_x_max_i, tile_y_max_j); 947 } 948} 949 950 951// The mul_mat_q kernel implements "stream-k" work partitioning as described in https://arxiv.org/abs/2301.03598 952 953template <ggml_type type, int J, bool fallback> 954__launch_bounds__(ggml_cuda_mmq_get_nthreads(type, J, fallback), ggml_cuda_mmq_get_occupancy(type, J, fallback)) 955static __global__ void mul_mat_q( 956 const char * __restrict__ x, const int * __restrict__ y, const int32_t * __restrict__ ids_dst, 957 const int32_t * __restrict__ expert_bounds, float * __restrict__ dst, float * __restrict__ tmp_fixup, 958 const float * __restrict__ y_scale, 959 const uint3 blocks_per_ne00, const int nrows_x, const int ncols_dst, const int stride_row_x, const int ncols_y, const int stride_col_dst, 960 const uint3 channel_ratio, const uint3 nchannels_y, const int stride_channel_x, const int stride_channel_y, const int stride_channel_dst, 961 const uint3 sample_ratio, const uint3 nsamples_y, const int stride_sample_x, const int stride_sample_y, const int stride_sample_dst, 962 const uint3 ntx) { 963 964 // Skip unused template specializations for faster compilation: 965 if (ggml_cuda_mmq_get_config(type, J, fallback).type == GGML_TYPE_COUNT) { 966 NO_DEVICE_CODE; 967 return; 968 } 969 970 constexpr int warp_size = ggml_cuda_get_physical_warp_size(); 971 constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size; 972 constexpr int qk = ggml_cuda_type_traits<type>::qk; 973 constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback); 974 975 const uint32_t nty = (nrows_x + I - 1) / I; // Number of tiles y 976 977 // Initialize the ids for writing back data with just the index. 978 // For regular matrix multiplications this is never changed. 979 // For MoE the correct indices are loaded from ids_dst. 980 extern __shared__ int ids_dst_shared[]; // Stored at beginning of shared memory. 981#pragma unroll 982 for (int j0 = 0; j0 < J; j0 += nwarps*warp_size) { 983 const int j = j0 + threadIdx.y*warp_size + threadIdx.x; 984 985 if (j0 + nwarps*warp_size > J && j >= J) { 986 break; 987 } 988 989 ids_dst_shared[j] = j; 990 } 991 __syncthreads(); 992 993 if constexpr (!ggml_cuda_mmq_get_stream_k(type, J, fallback)) { 994 const uint2 tmp2 = fast_div_modulo(blockIdx.z, nchannels_y); 995 const int wt = tmp2.x; 996 const int zt = tmp2.y; 997 const int jt = blockIdx.y; 998 const int it = blockIdx.x; 999 1000 // Defaults for regular matrix multiplication: 1001 int col_low = 0; 1002 int col_high = ncols_dst; 1003 int col_diff = ncols_dst; 1004 int offset_y = wt*stride_sample_y + zt*stride_channel_y; 1005 int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst; 1006 int offset_y_scale; 1007 if constexpr (type == GGML_TYPE_NVFP4) { 1008 offset_y_scale = wt*nchannels_y.z*ncols_y + zt*ncols_y; 1009 } else { 1010 GGML_UNUSED(offset_y_scale); 1011 } 1012 1013 if (ids_dst) { 1014 col_low = expert_bounds[zt + 0]; 1015 col_high = expert_bounds[zt + 1]; 1016 col_diff = col_high - col_low; 1017 1018 offset_y = 0; 1019 offset_dst = 0; 1020 if constexpr (type == GGML_TYPE_NVFP4) { 1021 offset_y_scale = 0; 1022 } 1023 1024 if (jt*J >= col_diff) { 1025 return; 1026 } 1027 1028 // __syncthreads(); // There is no previous tile that could cause a race condition. 1029#pragma unroll 1030 for (int j0 = 0; j0 < J; j0 += nwarps*warp_size) { 1031 const int j = j0 + threadIdx.y*warp_size + threadIdx.x; 1032 1033 if (j0 + nwarps*warp_size > J && j >= J) { 1034 break; 1035 } 1036 1037 ids_dst_shared[j] = ids_dst[col_low + jt*J + j]; 1038 } 1039 __syncthreads(); 1040 } 1041 1042 offset_y += (col_low + jt*J)*(sizeof(block_q8_1_mmq)/sizeof(int)); 1043 offset_dst += it*I; 1044 const float * y_scale_tile = nullptr; 1045 if constexpr (type == GGML_TYPE_NVFP4) { 1046 offset_y_scale += col_low + jt*J; 1047 y_scale_tile = y_scale ? y_scale + offset_y_scale : nullptr; 1048 } 1049 1050 const int tile_x_max_i = nrows_x - it*I - 1; 1051 const int tile_y_max_j = col_diff - jt*J - 1; 1052 1053 const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*I*stride_row_x; 1054 1055 constexpr bool fixup = false; 1056 mul_mat_q_process_tile<type, J, fallback, fixup> 1057 (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile, 1058 stride_row_x, ncols_y, stride_col_dst, 1059 tile_x_max_i, tile_y_max_j, 0, blocks_per_ne00.z); 1060 return; 1061 } 1062 1063 constexpr int ITER_K = ggml_cuda_mmq_get_K_vram(type, J, fallback); 1064 constexpr int blocks_per_iter = ITER_K / qk; 1065 1066 // kbc == k block continuous, current index in continuous ijk space. 1067 int kbc = int64_t(blockIdx.x) *(nsamples_y.z*nchannels_y.z*ntx.z*nty*blocks_per_ne00.z) / gridDim.x; 1068 int kbc_stop = int64_t(blockIdx.x + 1)*(nsamples_y.z*nchannels_y.z*ntx.z*nty*blocks_per_ne00.z) / gridDim.x; 1069 1070 kbc -= fastmodulo(kbc, blocks_per_ne00) % blocks_per_iter; 1071 kbc_stop -= fastmodulo(kbc_stop, blocks_per_ne00) % blocks_per_iter; 1072 1073 // kb0 == k index when doing the matrix multiplication for an output tile. 1074 int kb0_start = fastmodulo(kbc, blocks_per_ne00); 1075 int kb0_stop = min(blocks_per_ne00.z, uint32_t(kb0_start + kbc_stop - kbc)); 1076 while (kbc < kbc_stop && kb0_stop == int(blocks_per_ne00.z)) { 1077 int tmp = fastdiv(kbc, blocks_per_ne00); 1078 uint2 tmp2 = fast_div_modulo(tmp, ntx); 1079 const int jt = tmp2.y; 1080 tmp = tmp2.x; 1081 tmp2 = fast_div_modulo(tmp, nchannels_y); 1082 const int zt = tmp2.y; 1083 tmp = tmp2.x; 1084 tmp2 = fast_div_modulo(tmp, nsamples_y); 1085 const int wt = tmp2.y; 1086 const int it = tmp2.x; 1087 1088 // Defaults for regular matrix multiplication: 1089 int col_low = 0; 1090 int col_high = ncols_dst; 1091 int col_diff = ncols_dst; 1092 int offset_y = wt*stride_sample_y + zt*stride_channel_y; 1093 int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst; 1094 int offset_y_scale; 1095 if constexpr (type == GGML_TYPE_NVFP4) { 1096 offset_y_scale = wt*nchannels_y.z*ncols_y + zt*ncols_y; 1097 } else { 1098 GGML_UNUSED(offset_y_scale); 1099 } 1100 1101 if (ids_dst) { 1102 col_low = expert_bounds[zt + 0]; 1103 col_high = expert_bounds[zt + 1]; 1104 col_diff = col_high - col_low; 1105 1106 offset_y = 0; 1107 offset_dst = 0; 1108 if constexpr (type == GGML_TYPE_NVFP4) { 1109 offset_y_scale = 0; 1110 } 1111 1112 if (jt*J >= col_diff) { 1113 kbc += blocks_per_ne00.z; 1114 kbc -= fastmodulo(kbc, blocks_per_ne00); 1115 1116 kb0_start = 0; 1117 kb0_stop = min(blocks_per_ne00.z, uint32_t(kbc_stop - kbc)); 1118 1119 continue; 1120 } 1121 1122 __syncthreads(); 1123#pragma unroll 1124 for (int j0 = 0; j0 < J; j0 += nwarps*warp_size) { 1125 const int j = j0 + threadIdx.y*warp_size + threadIdx.x; 1126 1127 if (j0 + nwarps*warp_size > J && j >= J) { 1128 break; 1129 } 1130 1131 ids_dst_shared[j] = ids_dst[col_low + jt*J + j]; 1132 } 1133 __syncthreads(); 1134 } 1135 1136 offset_y += (col_low + jt * J) * (sizeof(block_q8_1_mmq) / sizeof(int)); 1137 offset_dst += it*I; 1138 const float * y_scale_tile = nullptr; 1139 if constexpr (type == GGML_TYPE_NVFP4) { 1140 offset_y_scale += col_low + jt * J; 1141 y_scale_tile = y_scale ? y_scale + offset_y_scale : nullptr; 1142 } 1143 1144 const int tile_x_max_i = nrows_x - it*I - 1; 1145 const int tile_y_max_j = col_diff - jt*J - 1; 1146 1147 const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*I*stride_row_x; 1148 1149 constexpr bool fixup = false; // All but (potentially) the last iterations write their data to dst rather than the fixup buffer. 1150 mul_mat_q_process_tile<type, J, fallback, fixup> 1151 (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile, 1152 stride_row_x, ncols_y, stride_col_dst, 1153 tile_x_max_i, tile_y_max_j, kb0_start, kb0_stop); 1154 1155 kbc += blocks_per_ne00.z; 1156 kbc -= fastmodulo(kbc, blocks_per_ne00); 1157 1158 kb0_start = 0; 1159 kb0_stop = min(blocks_per_ne00.z, uint32_t(kbc_stop - kbc)); 1160 } 1161 1162 if (kbc >= kbc_stop) { 1163 return; 1164 } 1165 1166 int tmp = fastdiv(kbc, blocks_per_ne00); 1167 uint2 tmp2 = fast_div_modulo(tmp, ntx); 1168 const int jt = tmp2.y; 1169 tmp = tmp2.x; 1170 tmp2 = fast_div_modulo(tmp, nchannels_y); 1171 const int zt = tmp2.y; 1172 tmp = tmp2.x; 1173 tmp2 = fast_div_modulo(tmp, nsamples_y); 1174 const int wt = tmp2.y; 1175 const int it = tmp2.x; 1176 1177 // Defaults for regular matrix multiplication: 1178 int col_low = 0; 1179 int col_high = ncols_dst; 1180 int col_diff = ncols_dst; 1181 int offset_y = wt*stride_sample_y + zt*stride_channel_y; 1182 int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst; 1183 int offset_y_scale; 1184 if constexpr (type == GGML_TYPE_NVFP4) { 1185 offset_y_scale = wt*nchannels_y.z*ncols_y + zt*ncols_y; 1186 } else { 1187 GGML_UNUSED(offset_y_scale); 1188 } 1189 1190 if (ids_dst) { 1191 col_low = expert_bounds[zt + 0]; 1192 col_high = expert_bounds[zt + 1]; 1193 col_diff = col_high - col_low; 1194 1195 offset_y = 0; 1196 offset_dst = 0; 1197 if constexpr (type == GGML_TYPE_NVFP4) { 1198 offset_y_scale = 0; 1199 } 1200 1201 if (jt*J >= col_diff) { 1202 return; 1203 } 1204 1205 // The memory layout for the fixup buffer is always contiguous, therefore reset ids: 1206 __syncthreads(); 1207#pragma unroll 1208 for (int j0 = 0; j0 < J; j0 += nwarps*warp_size) { 1209 const int j = j0 + threadIdx.y*warp_size + threadIdx.x; 1210 1211 if (j0 + nwarps*warp_size > J && j >= J) { 1212 break; 1213 } 1214 1215 ids_dst_shared[j] = j; 1216 } 1217 __syncthreads(); 1218 } 1219 1220 offset_y += (col_low + jt * J) * (sizeof(block_q8_1_mmq) / sizeof(int)); 1221 offset_dst += it*I; 1222 const float * y_scale_tile = nullptr; 1223 if constexpr (type == GGML_TYPE_NVFP4) { 1224 offset_y_scale += col_low + jt * J; 1225 y_scale_tile = y_scale ? y_scale + offset_y_scale : nullptr; 1226 } 1227 1228 const int tile_x_max_i = nrows_x - it*I - 1; 1229 const int tile_y_max_j = col_diff - jt*J - 1; 1230 1231 const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*I*stride_row_x; 1232 1233 constexpr bool fixup = true; // Last index writes its data to fixup buffer to avoid data races with other blocks. 1234 mul_mat_q_process_tile<type, J, fallback, fixup> 1235 (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile, 1236 stride_row_x, ncols_y, stride_col_dst, 1237 tile_x_max_i, tile_y_max_j, kb0_start, kb0_stop); 1238} 1239 1240template <ggml_type type, int J, bool fallback> 1241__launch_bounds__(ggml_cuda_mmq_get_nthreads(type, J, fallback)/2, 1) 1242static __global__ void mul_mat_q_stream_k_fixup( 1243 const int32_t * __restrict__ ids_dst, const int32_t * __restrict__ expert_bounds, float * __restrict__ dst, 1244 float * __restrict__ tmp_last_tile, const uint3 blocks_per_ne00, const int nrows_x, const int ncols_dst, 1245 const int stride_col_dst, const uint3 nchannels_y, const int stride_channel_dst, const uint3 nsamples_y, 1246 const int stride_sample_dst, const uint3 ntx) { 1247 constexpr int warp_size = ggml_cuda_get_physical_warp_size(); 1248 constexpr int nwarps = (ggml_cuda_mmq_get_nthreads(type, J, fallback) / 2) / warp_size; 1249 constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback); 1250 constexpr int qk = ggml_cuda_type_traits<type>::qk; 1251 constexpr int ITER_K = ggml_cuda_mmq_get_K_vram(type, J, fallback); 1252 constexpr int blocks_per_iter = ITER_K / qk; 1253 1254 float sum[J / nwarps] = {0.0f}; 1255 const int i = blockIdx.y*warp_size + threadIdx.x; 1256 1257 const int nty = (nrows_x + I - 1) / I; 1258 1259 const int bidx0 = blockIdx.x; 1260 1261 // kbc == k block continuous, current index in continuous ijk space. 1262 int kbc0 = int64_t(blockIdx.x) *(nsamples_y.z*nchannels_y.z*ntx.z*nty*blocks_per_ne00.z) / gridDim.x; 1263 int kbc0_stop = int64_t(blockIdx.x + 1)*(nsamples_y.z*nchannels_y.z*ntx.z*nty*blocks_per_ne00.z) / gridDim.x; 1264 1265 kbc0 -= fastmodulo(kbc0, blocks_per_ne00) % blocks_per_iter; 1266 kbc0_stop -= fastmodulo(kbc0_stop, blocks_per_ne00) % blocks_per_iter; 1267 1268 const bool did_not_have_any_data = kbc0 == kbc0_stop; 1269 const bool wrote_beginning_of_tile = fastmodulo(kbc0, blocks_per_ne00) == 0; 1270 const bool did_not_write_last = fastdiv(kbc0, blocks_per_ne00) == fastdiv(kbc0_stop, blocks_per_ne00) && fastmodulo(kbc0_stop, blocks_per_ne00) != 0; 1271 if (did_not_have_any_data || wrote_beginning_of_tile || did_not_write_last) { 1272 return; 1273 } 1274 1275 bool any_fixup = false; 1276 1277 // Iterate over previous blocks and sum up partial sums written to fixup buffer. 1278 // All CUDA blocks that get here must have a previous block that needs a fixup. 1279 int bidx = bidx0 - 1; 1280 int kbc_stop = kbc0; 1281 while(true) { 1282 int kbc = int64_t(bidx)*(nsamples_y.z*nchannels_y.z*ntx.z*nty*blocks_per_ne00.z) / gridDim.x; 1283 kbc -= fastmodulo(kbc, blocks_per_ne00) % blocks_per_iter; 1284 1285 if (kbc == kbc_stop) { // Did not have any data. 1286 bidx--; 1287 kbc_stop = kbc; 1288 continue; 1289 } 1290 1291 any_fixup = true; 1292 1293 1294#pragma unroll 1295 for (int j0 = 0; j0 < J; j0 += nwarps) { 1296 const int j = j0 + threadIdx.y; 1297 1298 sum[j0/nwarps] += tmp_last_tile[bidx*(J*I) + j*I + i]; 1299 } 1300 1301 // If this block started in a previous tile we are done and don't need to combine additional partial results. 1302 if (fastmodulo(kbc, blocks_per_ne00) == 0 || fastdiv(kbc, blocks_per_ne00) < fastdiv(kbc0, blocks_per_ne00)) { 1303 break; 1304 } 1305 bidx--; 1306 kbc_stop = kbc; 1307 } 1308 1309 if (!any_fixup) { 1310 return; 1311 } 1312 1313 int tmp = fastdiv(kbc0, blocks_per_ne00); 1314 uint2 tmp2 = fast_div_modulo(tmp, ntx); 1315 const int jt = tmp2.y; 1316 tmp = tmp2.x; 1317 tmp2 = fast_div_modulo(tmp, nchannels_y); 1318 const int zt = tmp2.y; 1319 tmp = tmp2.x; 1320 tmp2 = fast_div_modulo(tmp, nsamples_y); 1321 const int wt = tmp2.y; 1322 const int it = tmp2.x; 1323 1324 if (!ids_dst) { 1325 const int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst + it*I; 1326 dst += offset_dst; 1327 1328 const int i_max = nrows_x - it*I - 1; 1329 const int j_max = ncols_dst - jt*J - 1; 1330 if (fallback && i > i_max) { 1331 return; 1332 } 1333 1334#pragma unroll 1335 for (int j0 = 0; j0 < J; j0 += nwarps) { 1336 const int j = j0 + threadIdx.y; 1337 1338 if (j > j_max) { 1339 return; 1340 } 1341 1342 dst[j*stride_col_dst + i] += sum[j0/nwarps]; 1343 } 1344 return; 1345 } 1346 1347 __shared__ int ids_dst_shared[J]; 1348 const int col_low = expert_bounds[zt + 0]; 1349 const int col_high = expert_bounds[zt + 1]; 1350 const int col_diff = col_high - col_low; 1351 1352 for (int j = threadIdx.y*warp_size + threadIdx.x; j < J; j += nwarps*warp_size) { 1353 ids_dst_shared[j] = ids_dst[col_low + jt*J + j]; 1354 } 1355 __syncthreads(); 1356 1357 const int offset_dst = it*I; 1358 dst += offset_dst; 1359 1360 const int i_max = nrows_x - it*I - 1; 1361 const int j_max = col_diff - jt*J - 1; 1362 if (fallback && i > i_max) { 1363 return; 1364 } 1365 1366#pragma unroll 1367 for (int j0 = 0; j0 < J; j0 += nwarps) { 1368 const int j = j0 + threadIdx.y; 1369 1370 if (j > j_max) { 1371 return; 1372 } 1373 1374 dst[ids_dst_shared[j]*stride_col_dst + i] += sum[j0/nwarps]; 1375 } 1376} 1377 1378struct mmq_args { 1379 const char * x; ggml_type type_x; const int * y; const int32_t * ids_dst; const int32_t * expert_bounds; float * dst; 1380 const float * y_scale; 1381 int64_t ncols_x; int64_t nrows_x; int64_t ncols_dst; int64_t stride_row_x; int64_t ncols_y; int64_t nrows_dst; 1382 int64_t nchannels_x; int64_t nchannels_y; int64_t stride_channel_x; int64_t stride_channel_y; int64_t stride_channel_dst; 1383 int64_t nsamples_x; int64_t nsamples_y; int64_t stride_sample_x; int64_t stride_sample_y; int64_t stride_sample_dst; 1384 int64_t ncols_max; 1385 int64_t ncols_opt; // value to optimize the tile size against, launch grid still uses ncols_max 1386}; 1387 1388static size_t mmq_get_nbytes_shared(const ggml_cuda_mmq_config & config, const int cc) { 1389 const size_t nbs_ids = config.J*sizeof(int); 1390 const size_t nbs_x = ggml_cuda_mmq_get_nbytes_shared_x(config, cc); 1391 const size_t nbs_y = config.J * (sizeof(block_q8_1_mmq)); 1392 return nbs_ids + nbs_x + GGML_PAD(nbs_y, config.nthreads*sizeof(int)); 1393} 1394 1395template <ggml_type type, int J, bool fallback> 1396static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) { 1397 const int id = ggml_cuda_get_device(); 1398 const int cc = ggml_cuda_info().devices[id].cc; 1399 const int nsm = ggml_cuda_info().devices[id].nsm; 1400 const int warp_size = ggml_cuda_info().devices[id].warp_size; 1401 1402 const ggml_cuda_mmq_config config = ggml_cuda_mmq_get_config(type, J, fallback, cc); 1403 GGML_ASSERT(config.nthreads % warp_size == 0); 1404 const int nwarps = config.nthreads / warp_size; 1405 const int nbytes_shared = mmq_get_nbytes_shared(config, cc); 1406 1407 const dim3 block_dims(warp_size, nwarps, 1); 1408 1409 CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q<type, J, false>), nbytes_shared); 1410 CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q<type, J, true>), nbytes_shared); 1411 1412 const int nty = (args.nrows_x + config.I - 1) / config.I; 1413 const int ntx = (args.ncols_max + config.J - 1) / config.J; 1414 const int ntzw = args.nchannels_y * args.nsamples_y; 1415 const dim3 block_nums_xy_tiling(nty, ntx, ntzw); 1416 1417 GGML_ASSERT(args.nchannels_y % args.nchannels_x == 0); 1418 GGML_ASSERT(args.nsamples_y % args.nsamples_x == 0); 1419 const int channel_ratio = args.nchannels_y / args.nchannels_x; 1420 const int sample_ratio = args.nsamples_y / args.nsamples_x; 1421 1422 const uint3 blocks_per_ne00_fd = init_fastdiv_values(args.ncols_x / ggml_cuda_type_traits<type>::qk); 1423 const uint3 ntx_fd = init_fastdiv_values(ntx); 1424 const uint3 nchannels_y_fd = init_fastdiv_values(args.nchannels_y); 1425 const uint3 nsamples_y_fd = init_fastdiv_values(args.nsamples_y); 1426 const uint3 channel_ratio_fd = init_fastdiv_values(channel_ratio); 1427 const uint3 sample_ratio_fd = init_fastdiv_values(sample_ratio); 1428 1429 if (!ggml_cuda_mmq_get_stream_k(type, J, fallback, cc)) { 1430 mul_mat_q<type, J, fallback><<<block_nums_xy_tiling, block_dims, nbytes_shared, stream>>> 1431 (args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, args.y_scale, 1432 blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, 1433 channel_ratio_fd, nchannels_y_fd, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, 1434 sample_ratio_fd, nsamples_y_fd, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, 1435 ntx_fd); 1436 return; 1437 } 1438 1439 // For the stream-k kernel it is possible to run it with tiling by setting the number of CUDA blocks equal to the number of tiles. 1440 // This is worthwhile if the efficiency of tiling is high and skipping the fixup kernel is more important. 1441 const int ntiles_dst = ntx * nty * ntzw; 1442 const int tiles_nwaves = (ntiles_dst + nsm - 1) / nsm; 1443 const int tiles_efficiency_percent = 100 * ntiles_dst / (nsm*tiles_nwaves); 1444 const dim3 block_nums_stream_k(GGML_CUDA_CC_IS_NVIDIA(cc) && tiles_efficiency_percent >= 90 ? ntiles_dst : nsm, 1, 1); 1445 1446 GGML_ASSERT(ntiles_dst * blocks_per_ne00_fd.z < (1 << 30)); // Assert that variable kbc will not overflow. 1447 1448 const bool fixup_needed = ntiles_dst % block_nums_stream_k.x != 0; 1449 1450 ggml_cuda_pool & pool = ctx.pool(id); 1451 ggml_cuda_pool_alloc<float> tmp_fixup(pool); 1452 if (fixup_needed) { 1453 tmp_fixup.alloc(block_nums_stream_k.x * config.J*config.I); 1454 } 1455 1456 const dim3 block_nums_fixup(block_nums_stream_k.x, config.I/warp_size, 1); 1457 const dim3 block_dims_fixup(block_dims.x, block_dims.y/2, block_dims.z); 1458 1459 mul_mat_q<type, J, fallback><<<block_nums_stream_k, block_dims, nbytes_shared, stream>>> 1460 (args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup.ptr, args.y_scale, 1461 blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, 1462 channel_ratio_fd, nchannels_y_fd, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, 1463 sample_ratio_fd, nsamples_y_fd, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, 1464 ntx_fd); 1465 1466 if (!fixup_needed) { 1467 return; 1468 } 1469 1470 CUDA_CHECK(cudaGetLastError()); 1471 mul_mat_q_stream_k_fixup<type, J, fallback><<<block_nums_fixup, block_dims_fixup, 0, stream>>> 1472 (args.ids_dst, args.expert_bounds, args.dst, tmp_fixup.ptr, blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, 1473 args.nrows_dst, nchannels_y_fd, args.stride_channel_dst, nsamples_y_fd, args.stride_sample_dst, 1474 ntx_fd); 1475} 1476 1477template <ggml_type type, bool fallback> 1478void mul_mat_q_switch_J(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) { 1479 const int id = ggml_cuda_get_device(); 1480 const int cc = ggml_cuda_info().devices[id].cc; 1481 const size_t smpbo = ggml_cuda_info().devices[id].smpbo; 1482 1483 int J_best = 0; 1484 int ntiles_J_best = INT_MAX; 1485 1486 for (int J = 8; J <= 128 && ntiles_J_best > 1; J += 8) { 1487 const ggml_cuda_mmq_config config = ggml_cuda_mmq_get_config(type, J, fallback, cc); 1488 if (config.type == GGML_TYPE_COUNT) { 1489 continue; 1490 } 1491 1492 if (mmq_get_nbytes_shared(config, cc) > smpbo) { 1493 continue; 1494 } 1495 1496 const int ntiles_x = (args.ncols_opt + config.J - 1) / config.J; 1497 1498 if (ntiles_x < ntiles_J_best) { 1499 J_best = J; 1500 ntiles_J_best = ntiles_x; 1501 } 1502 } 1503 1504 switch (J_best) { 1505 case 8: 1506 launch_mul_mat_q<type, 8, fallback>(ctx, args, stream); 1507 break; 1508 case 16: 1509 launch_mul_mat_q<type, 16, fallback>(ctx, args, stream); 1510 break; 1511 case 24: 1512 launch_mul_mat_q<type, 24, fallback>(ctx, args, stream); 1513 break; 1514 case 32: 1515 launch_mul_mat_q<type, 32, fallback>(ctx, args, stream); 1516 break; 1517 case 40: 1518 launch_mul_mat_q<type, 40, fallback>(ctx, args, stream); 1519 break; 1520 case 48: 1521 launch_mul_mat_q<type, 48, fallback>(ctx, args, stream); 1522 break; 1523 case 56: 1524 launch_mul_mat_q<type, 56, fallback>(ctx, args, stream); 1525 break; 1526 case 64: 1527 launch_mul_mat_q<type, 64, fallback>(ctx, args, stream); 1528 break; 1529 case 72: 1530 launch_mul_mat_q<type, 72, fallback>(ctx, args, stream); 1531 break; 1532 case 80: 1533 launch_mul_mat_q<type, 80, fallback>(ctx, args, stream); 1534 break; 1535 case 88: 1536 launch_mul_mat_q<type, 88, fallback>(ctx, args, stream); 1537 break; 1538 case 96: 1539 launch_mul_mat_q<type, 96, fallback>(ctx, args, stream); 1540 break; 1541 case 104: 1542 launch_mul_mat_q<type, 104, fallback>(ctx, args, stream); 1543 break; 1544 case 112: 1545 launch_mul_mat_q<type, 112, fallback>(ctx, args, stream); 1546 break; 1547 case 120: 1548 launch_mul_mat_q<type, 120, fallback>(ctx, args, stream); 1549 break; 1550 case 128: 1551 launch_mul_mat_q<type, 128, fallback>(ctx, args, stream); 1552 break; 1553 default: 1554 fprintf(stderr, "J_best=%d\n", J_best); 1555 GGML_ABORT("fatal error"); 1556 break; 1557 } 1558} 1559 1560template <ggml_type type> 1561void mul_mat_q_case(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) { 1562 if (args.nrows_x % 128 == 0) { 1563 constexpr bool fallback = false; 1564 mul_mat_q_switch_J<type, fallback>(ctx, args, stream); 1565 } else { 1566 constexpr bool fallback = true; 1567 mul_mat_q_switch_J<type, fallback>(ctx, args, stream); 1568 } 1569} 1570 1571#define DECL_MMQ_CASE(type) \ 1572 template void mul_mat_q_case<type>(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) \ 1573 1574extern DECL_MMQ_CASE(GGML_TYPE_Q1_0); 1575extern DECL_MMQ_CASE(GGML_TYPE_Q2_0); 1576extern DECL_MMQ_CASE(GGML_TYPE_Q4_0); 1577extern DECL_MMQ_CASE(GGML_TYPE_Q4_1); 1578extern DECL_MMQ_CASE(GGML_TYPE_Q5_0); 1579extern DECL_MMQ_CASE(GGML_TYPE_Q5_1); 1580extern DECL_MMQ_CASE(GGML_TYPE_Q8_0); 1581// ----------------------------------------- 1582extern DECL_MMQ_CASE(GGML_TYPE_Q2_K); 1583extern DECL_MMQ_CASE(GGML_TYPE_Q3_K); 1584extern DECL_MMQ_CASE(GGML_TYPE_Q4_K); 1585extern DECL_MMQ_CASE(GGML_TYPE_Q5_K); 1586extern DECL_MMQ_CASE(GGML_TYPE_Q6_K); 1587// ----------------------------------------- 1588extern DECL_MMQ_CASE(GGML_TYPE_IQ1_S); 1589extern DECL_MMQ_CASE(GGML_TYPE_IQ2_XXS); 1590extern DECL_MMQ_CASE(GGML_TYPE_IQ2_XS); 1591extern DECL_MMQ_CASE(GGML_TYPE_IQ2_S); 1592extern DECL_MMQ_CASE(GGML_TYPE_IQ3_XXS); 1593extern DECL_MMQ_CASE(GGML_TYPE_IQ3_S); 1594extern DECL_MMQ_CASE(GGML_TYPE_IQ4_NL); 1595extern DECL_MMQ_CASE(GGML_TYPE_IQ4_XS); 1596// ----------------------------------------- 1597extern DECL_MMQ_CASE(GGML_TYPE_MXFP4); 1598extern DECL_MMQ_CASE(GGML_TYPE_NVFP4); 1599 1600// ------------------------------------------------------------------------------------------------------------------------- 1601 1602void ggml_cuda_mul_mat_q( 1603 ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst); 1604 1605bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t n_experts);