char-fork/llama.cpp
git clone https://git.t4t.associates/char-fork/llama.cpp
3d3d7c818
char/aliceai-t5-moe
1#pragma once 2 3#include "vecdotq.cuh" 4#include "mma.cuh" 5 6using namespace ggml_cuda_mma; 7 8#include "mmq.cuh" 9 10template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q4_0_q8_1_dp4a( 11 const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { 12 constexpr int warp_size = ggml_cuda_get_physical_warp_size(); 13 constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size; 14 constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback); 15 16 constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q4_0, I); 17 const int * x_qs = (const int *) x; 18 const float * x_df = (const float *) x_qs + txs.qs; 19 const int * y_qs = (const int *) y + 4; 20 const half2 * y_ds = (const half2 *) y; 21 22// #pragma unroll 23 for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QR4_0*VDR_Q4_0_Q8_1_MMQ) { 24 const int k0 = k00 + k01; 25 26#pragma unroll 27 for (int j0 = 0; j0 < J; j0 += nwarps) { 28 const int j = j0 + threadIdx.y; 29 30#pragma unroll 31 for (int i0 = 0; i0 < I; i0 += warp_size) { 32 const int i = i0 + threadIdx.x; 33 const int kyqs = QI8_1 * ((k01/2) / (QI8_1/2)) + (k01/2) % (QI8_1/2); 34 35 int u[2*VDR_Q4_0_Q8_1_MMQ]; 36 37 constexpr int max_cpy = ggml_cuda_get_max_cpy_bytes(); 38 constexpr int mcpy_int = max_cpy / sizeof(int); 39 static_assert(VDR_Q4_0_Q8_1_MMQ == 4, "bad VDR_Q4_0_Q8_1_MMQ"); 40 41 int tmp0[4], tmp1[4]; 42 43 #pragma unroll 44 for (int l0 = 0; l0 < 4 / mcpy_int; ++l0) { 45 ggml_cuda_memcpy_1<max_cpy>(tmp0 + l0 * mcpy_int, &y_qs[j*MMQ_TILE_Y_K + kyqs + l0 * mcpy_int] ); 46 ggml_cuda_memcpy_1<max_cpy>(tmp1 + l0 * mcpy_int, &y_qs[j*MMQ_TILE_Y_K + kyqs + QI4_0 + l0 * mcpy_int]); 47 } 48 49 u[0]=tmp0[0]; u[2]=tmp0[1]; u[4]=tmp0[2]; u[6]=tmp0[3]; 50 u[1]=tmp1[0]; u[3]=tmp1[1]; u[5]=tmp1[2]; u[7]=tmp1[3]; 51 52 sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q4_0_q8_1_impl<VDR_Q4_0_Q8_1_MMQ> 53 (&x_qs[i*(MMQ_TILE_NE_K + 1) + k0/QR4_0], u, 54 x_df[i*(MMQ_TILE_NE_K/QI4_0) + i/QI4_0 + k0/(QR4_0*QI4_0)], y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]); 55 } 56 } 57 } 58} 59 60template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q4_1_q8_1_dp4a( 61 const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { 62 constexpr int warp_size = ggml_cuda_get_physical_warp_size(); 63 constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size; 64 constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback); 65 66 constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q4_1, I); 67 const int * x_qs = (const int *) x; 68 const half2 * x_dm = (const half2 *) x_qs + txs.qs; 69 const int * y_qs = (const int *) y + 4; 70 const half2 * y_ds = (const half2 *) y; 71 72// #pragma unroll 73 for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QR4_1*VDR_Q4_1_Q8_1_MMQ) { 74 const int k0 = k00 + k01; 75 76#pragma unroll 77 for (int j0 = 0; j0 < J; j0 += nwarps) { 78 const int j = j0 + threadIdx.y; 79 80#pragma unroll 81 for (int i0 = 0; i0 < I; i0 += warp_size) { 82 const int i = i0 + threadIdx.x; 83 const int kyqs = QI8_1 * ((k01/2) / (QI8_1/2)) + (k01/2) % (QI8_1/2); 84 85 int u[2*VDR_Q4_1_Q8_1_MMQ]; 86 87 constexpr int max_cpy = ggml_cuda_get_max_cpy_bytes(); 88 constexpr int mcpy_int = max_cpy / sizeof(int); 89 static_assert(VDR_Q4_0_Q8_1_MMQ == 4, "bad VDR_Q4_0_Q8_1_MMQ"); 90 91 int tmp0[4], tmp1[4]; 92 93 #pragma unroll 94 for (int l0 = 0; l0 < 4 / mcpy_int; ++l0) { 95 ggml_cuda_memcpy_1<max_cpy>(tmp0 + l0 * mcpy_int, &y_qs[j*MMQ_TILE_Y_K + kyqs + l0 * mcpy_int] ); 96 ggml_cuda_memcpy_1<max_cpy>(tmp1 + l0 * mcpy_int, &y_qs[j*MMQ_TILE_Y_K + kyqs + QI4_1 + l0 * mcpy_int]); 97 } 98 99 u[0]=tmp0[0]; u[2]=tmp0[1]; u[4]=tmp0[2]; u[6]=tmp0[3]; 100 u[1]=tmp1[0]; u[3]=tmp1[1]; u[5]=tmp1[2]; u[7]=tmp1[3]; 101 102 sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q4_1_q8_1_impl<VDR_Q4_1_Q8_1_MMQ> 103 (&x_qs[i*(MMQ_TILE_NE_K + 1) + k0/QR4_1], u, 104 x_dm[i*(MMQ_TILE_NE_K/QI4_1) + i/QI4_1 + k0/(QR4_1*QI4_1)], y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]); 105 } 106 } 107 } 108} 109 110template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a( 111 const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { 112 constexpr int warp_size = ggml_cuda_get_physical_warp_size(); 113 constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size; 114 constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback); 115 116 constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q8_0, I); 117 const int * x_qs = (const int *) x; 118 const float * x_df = (const float *) x_qs + txs.qs; 119 const int * y_qs = (const int *) y + 4; 120 const float * y_df = (const float *) y; 121 122// #pragma unroll 123 for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += VDR_Q8_0_Q8_1_MMQ) { 124 const int k0 = k00 + k01; 125 126#pragma unroll 127 for (int j0 = 0; j0 < J; j0 += nwarps) { 128 const int j = j0 + threadIdx.y; 129 130#pragma unroll 131 for (int i0 = 0; i0 < I; i0 += warp_size) { 132 const int i = i0 + threadIdx.x; 133 134 sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q8_0_q8_1_impl<float, VDR_Q8_0_Q8_1_MMQ> 135 (&x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0], &y_qs[j*MMQ_TILE_Y_K + k0 % MMQ_TILE_NE_K], 136 x_df[i*(2*MMQ_TILE_NE_K/QI8_0) + i/(QI8_0/2) + k0/QI8_0], y_df[j*MMQ_TILE_Y_K + (k0/QI8_1) % (MMQ_TILE_NE_K/QI8_1)]); 137 } 138 } 139 } 140} 141 142template <ggml_type type, int J, bool fallback, mmq_q8_1_ds_layout ds_layout> 143static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma( 144 const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { 145#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) 146 constexpr data_layout input_layout = get_input_data_layout(); 147 typedef tile<16, 8, int, input_layout> tile_A; 148 typedef tile<16, 8, int, input_layout> tile_B; 149 typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C; 150 151 constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback); 152 constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback); 153 constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp. 154 155 y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K); 156 157 const int * x_qs = (const int *) x; 158 const float * x_df = (const float *) x_qs + 2*MMQ_TILE_NE_K; 159 const int * y_qs = (const int *) y + 4; 160 const float * y_df = (const float *) y; 161 const half2 * y_ds = (const half2 *) y; 162 163 const int i0 = (threadIdx.y / ntx) * rows_per_warp; 164 165 for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_0) { 166 const int k0 = k00 + k01; 167 168 tile_A A[ntx]; 169#pragma unroll 170 for (int n = 0; n < ntx; ++n) { 171 load_ldmatrix(A[n], x_qs + (i0 + n*tile_A::I)*sram_stride + k0, sram_stride); 172 } 173 174#pragma unroll 175 for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) { 176 tile_B B; 177 load_ldmatrix(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K); 178 179 float dB; 180 const int j = j0 + tile_C::get_j(0); 181 if (ds_layout == MMQ_Q8_1_DS_LAYOUT_D4) { 182 dB = y_df[j*MMQ_TILE_Y_K + k01/QI8_1]; 183 } else { 184 dB = __low2float(y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]); 185 } 186 187#pragma unroll 188 for (int n = 0; n < ntx; ++n) { 189 tile_C C; 190 mma(C, A[n], B); 191 192#pragma unroll 193 for (int l = 0; l < tile_C::ne; ++l) { 194 const int i = i0 + n*tile_A::I + tile_C::get_i(l); 195 const float dA = x_df[i*sram_stride + k0/QI8_0]; 196 sum[(j0/tile_C::J + n)*tile_C::ne + l] += C.x[l]*dA*dB; 197 } 198 } 199 } 200 } 201#else 202 typedef tile<16, 8, int> tile_A; 203 typedef tile< 8, 8, int> tile_B; 204 typedef tile<16, 8, int> tile_C; 205 206 constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback); 207 constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback); 208 constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp. 209 210 y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K); 211 212 const int * x_qs = (const int *) x; 213 const float * x_df = (const float *) x_qs + 2*MMQ_TILE_NE_K; 214 const int * y_qs = (const int *) y + 4; 215 const float * y_df = (const float *) y; 216 const half2 * y_ds = (const half2 *) y; 217 218 tile_A A[ntx][MMQ_TILE_NE_K/QI8_0]; 219 float dA[ntx][tile_C::ne/2][MMQ_TILE_NE_K/QI8_0]; 220 221 const int i0 = (threadIdx.y/ntx)*rows_per_warp; 222 223#pragma unroll 224 for (int n = 0; n < ntx; ++n) { 225#pragma unroll 226 for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_0) { 227 const int k0 = k00 + k01; 228 229 load_ldmatrix(A[n][k01/QI8_0], x_qs + (i0 + n*tile_A::I)*sram_stride + k0, sram_stride); 230 } 231 232#pragma unroll 233 for (int l = 0; l < tile_C::ne/2; ++l) { 234 const int i = i0 + n*tile_A::I + tile_C::get_i(2*l); 235 236#pragma unroll 237 for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_0) { 238 const int k0 = k00 + k01; 239 240 dA[n][l][k01/QI8_0] = x_df[i*sram_stride + k0/QI8_0]; 241 } 242 } 243 } 244 245#pragma unroll 246 for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) { 247#pragma unroll 248 for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_0) { 249 tile_B B; 250 float dB[tile_C::ne/2]; 251 252 load_generic(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K); // faster than load_ldmatrix 253 254#pragma unroll 255 for (int l = 0; l < tile_C::ne/2; ++l) { 256 const int j = j0 + tile_C::get_j(l); 257 258 if (ds_layout == MMQ_Q8_1_DS_LAYOUT_D4) { 259 dB[l] = y_df[j*MMQ_TILE_Y_K + k01/QI8_1]; 260 } else { 261 dB[l] = __low2float(y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]); 262 } 263 } 264 265#pragma unroll 266 for (int n = 0; n < ntx; ++n) { 267 tile_C C; 268 mma(C, A[n][k01/QI8_0], B); 269 270#pragma unroll 271 for (int l = 0; l < tile_C::ne; ++l) { 272 sum[(j0/tile_C::J + n)*tile_C::ne + l] += C.x[l]*dA[n][l/2][k01/QI8_0]*dB[l%2]; 273 } 274 } 275 } 276 } 277#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) 278} 279 280 281template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q8_1_q8_1_dp4a( 282 const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { 283 constexpr int warp_size = ggml_cuda_get_physical_warp_size(); 284 constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size; 285 constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback); 286 287 constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q5_1, I); 288 const int * x_qs = (const int *) x; 289 const half2 * x_dm = (const half2 *) x_qs + txs.qs; 290 const int * y_qs = (const int *) y + 4; 291 const half2 * y_ds = (const half2 *) y; 292 293// #pragma unroll 294 for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += VDR_Q8_0_Q8_1_MMQ) { 295 const int k0 = k00 + k01; 296 297#pragma unroll 298 for (int j0 = 0; j0 < J; j0 += nwarps) { 299 const int j = j0 + threadIdx.y; 300 301#pragma unroll 302 for (int i0 = 0; i0 < I; i0 += warp_size) { 303 const int i = i0 + threadIdx.x; 304 305 sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q8_1_q8_1_impl<QR5_1*VDR_Q5_1_Q8_1_MMQ> 306 (&x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0], &y_qs[j*MMQ_TILE_Y_K + k01], 307 x_dm[i*(MMQ_TILE_NE_K/QI5_1) + i/QI5_1 + k0/QI8_1], y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]); 308 } 309 } 310 } 311} 312 313template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q8_1_q8_1_mma( 314 const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { 315#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) 316 constexpr data_layout input_layout = get_input_data_layout(); 317 typedef tile<16, 8, int, input_layout> tile_A; 318 typedef tile<16, 8, int, input_layout> tile_B; 319 typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C; 320 321 constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback); 322 constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback); 323 constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp. 324 325 y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K); 326 327 const int * x_qs = (const int *) x; 328 const half2 * x_dm = (const half2 *) x_qs + 2*MMQ_TILE_NE_K; 329 const int * y_qs = (const int *) y + 4; 330 const half2 * y_dm = (const half2 *) y; 331 332 const int i0 = (threadIdx.y / ntx) * rows_per_warp; 333 334 for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1) { 335 const int k0 = k00 + k01; 336 337 tile_A A[ntx]; 338#pragma unroll 339 for (int n = 0; n < ntx; ++n) { 340 load_ldmatrix(A[n], x_qs + (i0 + n*tile_A::I)*sram_stride + k0, sram_stride); 341 } 342 343#pragma unroll 344 for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) { 345 tile_B B; 346 load_ldmatrix(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K); 347 348 const int j = j0 + tile_C::get_j(0); 349 const float2 dsB = __half22float2(y_dm[j*MMQ_TILE_Y_K + k01/QI8_1]); 350 351#pragma unroll 352 for (int n = 0; n < ntx; ++n) { 353 tile_C C; 354 mma(C, A[n], B); 355 356#pragma unroll 357 for (int l = 0; l < tile_C::ne; ++l) { 358 const int i = i0 + n*tile_A::I + tile_C::get_i(l); 359 float2 dmA = __half22float2(x_dm[i*sram_stride + k0/QI8_1]); 360 sum[(j0/tile_C::J + n)*tile_C::ne + l] += dmA.x*dsB.x*C.x[l]; 361 sum[(j0/tile_C::J + n)*tile_C::ne + l] += dmA.y*dsB.y; 362 } 363 } 364 } 365 } 366#else 367 typedef tile<16, 8, int> tile_A; 368 typedef tile< 8, 8, int> tile_B; 369 typedef tile<16, 8, int> tile_C; 370 371 constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback); 372 constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback); 373 constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp. 374 375 y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K); 376 377 const int * x_qs = (const int *) x; 378 const half2 * x_dm = (const half2 *) x_qs + 2*MMQ_TILE_NE_K; 379 const int * y_qs = (const int *) y + 4; 380 const half2 * y_dm = (const half2 *) y; 381 382 tile_A A[ntx][MMQ_TILE_NE_K/QI8_1]; 383 float2 dmA[ntx][tile_C::ne/2][MMQ_TILE_NE_K/QI8_1]; 384 385 const int i0 = (threadIdx.y/ntx)*rows_per_warp; 386 387#pragma unroll 388 for (int n = 0; n < ntx; ++n) { 389#pragma unroll 390 for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1) { 391 const int k0 = k00 + k01; 392 393 load_ldmatrix(A[n][k01/QI8_1], x_qs + (i0 + n*tile_A::I)*sram_stride + k0, sram_stride); 394 } 395 396#pragma unroll 397 for (int l = 0; l < tile_C::ne/2; ++l) { 398 const int i = i0 + n*tile_A::I + tile_C::get_i(2*l); 399 400#pragma unroll 401 for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1) { 402 const int k0 = k00 + k01; 403 404 dmA[n][l][k01/QI8_1] = __half22float2(x_dm[i*sram_stride + k0/QI8_1]); 405 } 406 } 407 } 408 409#pragma unroll 410 for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) { 411#pragma unroll 412 for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1) { 413 tile_B B; 414 float2 dsB[tile_C::ne/2]; 415 416 load_generic(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K); // faster than load_ldmatrix 417 418#pragma unroll 419 for (int l = 0; l < tile_C::ne/2; ++l) { 420 const int j = j0 + tile_C::get_j(l); 421 422 dsB[l] = __half22float2(y_dm[j*MMQ_TILE_Y_K + k01/QI8_1]); 423 } 424 425#pragma unroll 426 for (int n = 0; n < ntx; ++n) { 427 tile_C C; 428 mma(C, A[n][k01/QI8_1], B); 429 430#pragma unroll 431 for (int l = 0; l < tile_C::ne; ++l) { 432 sum[(j0/tile_C::J + n)*tile_C::ne + l] += dmA[n][l/2][k01/QI8_1].x*dsB[l%2].x*C.x[l]; 433 sum[(j0/tile_C::J + n)*tile_C::ne + l] += dmA[n][l/2][k01/QI8_1].y*dsB[l%2].y; 434 } 435 } 436 } 437 } 438#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) 439} 440 441// Used for NVFP4, Q3_K, IQ2_S, and IQ2_XS 442template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_dp4a( 443 const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { 444 constexpr int warp_size = ggml_cuda_get_physical_warp_size(); 445 constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size; 446 constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback); 447 448 constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(type, I); 449 const int * x_qs = (const int *) x; 450 const float * x_df = (const float *) x_qs + txs.qs; 451 const int * y_qs = (const int *) y + 4; 452 const float * y_df = (const float *) y; 453 454// #pragma unroll 455 for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_0) { 456 const int k0 = k00 + k01; 457 458#pragma unroll 459 for (int j0 = 0; j0 < J; j0 += nwarps) { 460 const int j = j0 + threadIdx.y; 461 462#pragma unroll 463 for (int i0 = 0; i0 < I; i0 += warp_size) { 464 const int i = i0 + threadIdx.x; 465 466 sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q8_0_16_q8_1_impl<QI8_0>( 467 &x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0], 468 &y_qs[j*MMQ_TILE_Y_K + k01], 469 &x_df[i*(2*MMQ_TILE_NE_K*2/QI8_0) + i/(QI8_0/4) + k0/(QI8_0/2)], 470 y_df[j*MMQ_TILE_Y_K + k01/QI8_1]); 471 } 472 } 473 } 474} 475 476// Used for Q3_K, IQ2_S, and IQ2_XS: 477template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_mma( 478 const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { 479#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) 480 constexpr data_layout input_layout = get_input_data_layout(); 481 typedef tile<16, 4, int, input_layout> tile_A; 482 typedef tile<16, 4, int, input_layout> tile_B; 483 typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C; 484 485 constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback); 486 constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback); 487 constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp. 488 489 y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K); 490 491 const int * x_qs = (const int *) x; 492 const float * x_df = (const float *) x_qs + MMQ_TILE_NE_K*2; 493 const int * y_qs = (const int *) y + 4; 494 const float * y_df = (const float *) y; 495 496 const int i0 = (threadIdx.y / ntx) * rows_per_warp; 497 498 for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 4) { 499 const int k0 = k00 + k01; 500 501 tile_A A[ntx]; 502#pragma unroll 503 for (int n = 0; n < ntx; ++n) { 504 load_ldmatrix(A[n], x_qs + (i0 + n*tile_A::I)*sram_stride + k0, sram_stride); 505 } 506 507#pragma unroll 508 for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) { 509 tile_B B; 510 load_ldmatrix(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K); 511 512 const int j = j0 + tile_C::get_j(0); 513 const float dB = y_df[j*MMQ_TILE_Y_K + k01/QI8_1]; 514 515#pragma unroll 516 for (int n = 0; n < ntx; ++n) { 517 tile_C C; 518 mma(C, A[n], B); 519 520#pragma unroll 521 for (int l = 0; l < tile_C::ne; ++l) { 522 const int i = i0 + n*tile_C::I + tile_C::get_i(l); 523 sum[(j0/tile_C::J + n)*tile_C::ne + l] += C.x[l] * x_df[i*sram_stride + k0/4] * dB; 524 } 525 } 526 } 527 } 528#elif defined(TURING_MMA_AVAILABLE) 529 530 typedef tile<16, 4, int> tile_A; 531 typedef tile<16, 8, int> tile_A_8; 532 typedef tile< 8, 4, int> tile_B; 533 typedef tile<16, 8, int> tile_C; 534 535 constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback); 536 constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback); 537 constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp. 538 539 y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K); 540 541 const int * x_qs = (const int *) x; 542 const float * x_df = (const float *) x_qs + MMQ_TILE_NE_K*2; 543 const int * y_qs = (const int *) y + 4; 544 const float * y_df = (const float *) y; 545 546 const int i0 = (threadIdx.y / ntx) * (ntx*tile_A::I); 547 548 tile_A A[ntx][8]; 549 float dA[ntx][tile_C::ne/2][8]; 550 551#pragma unroll 552 for (int n = 0; n < ntx; ++n) { 553#pragma unroll 554 for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 8) { 555 const int k0 = k00 + k01; 556 557 load_ldmatrix(((tile_A_8 *) A[n])[k01/8], x_qs + (i0 + n*tile_A::I)*sram_stride + k0, sram_stride); 558 } 559 560#pragma unroll 561 for (int l = 0; l < tile_C::ne/2; ++l) { 562 const int i = i0 + n*tile_C::I + tile_C::get_i(2*l); 563 564#pragma unroll 565 for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 4) { 566 const int k0 = k00 + k01; 567 568 dA[n][l][k01/4] = x_df[i*sram_stride + k0/4]; 569 } 570 } 571 } 572 573#pragma unroll 574 for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) { 575#pragma unroll 576 for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QR3_K*VDR_Q3_K_Q8_1_MMQ) { 577 tile_B B[2]; 578 float dB[tile_C::ne/2]; 579 580 // Here load_generic is faster than load_ldmatrix. 581 load_generic(B[0], y_qs + j0*MMQ_TILE_Y_K + (k01 + 0), MMQ_TILE_Y_K); 582 load_generic(B[1], y_qs + j0*MMQ_TILE_Y_K + (k01 + tile_B::J), MMQ_TILE_Y_K); 583 584#pragma unroll 585 for (int l = 0; l < tile_C::ne/2; ++l) { 586 const int j = j0 + tile_C::get_j(l); 587 588 dB[l] = y_df[j*MMQ_TILE_Y_K + k01/QI8_1]; 589 } 590 591#pragma unroll 592 for (int n = 0; n < ntx; ++n) { 593 tile_C C[2]; 594 mma(C[0], A[n][k01/4 + 0], B[0]); 595 mma(C[1], A[n][k01/4 + 1], B[1]); 596 597#pragma unroll 598 for (int l = 0; l < tile_C::ne; ++l) { 599 sum[(j0/tile_C::J + n)*tile_C::ne + l] += dB[l%2]*(C[0].x[l]*dA[n][l/2][k01/4 + 0] + C[1].x[l]*dA[n][l/2][k01/4 + 1]); 600 } 601 } 602 } 603 } 604#else 605 GGML_UNUSED_VARS(x, y, sum, k00); 606 NO_DEVICE_CODE; 607#endif // AMD_MFMA_AVAILABLE || AMD_WMMA_AVAILABLE 608} 609 610template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q2_K_q8_1_dp4a( 611 const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { 612 constexpr int warp_size = ggml_cuda_get_physical_warp_size(); 613 constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size; 614 constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback); 615 616 constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q2_K, I); 617 const int * x_qs = (const int *) x; 618 const half2 * x_dm = (const half2 *) x_qs + txs.qs; 619 const int * y_qs = (const int *) y + 4; 620 const half2 * y_ds = (const half2 *) y; 621 622 float2 y_df[J/nwarps]; 623#pragma unroll 624 for (int j0 = 0; j0 < J; j0 += nwarps) { 625 const int j = j0 + threadIdx.y; 626 627 y_df[j0/nwarps] = __half22float2(y_ds[j*MMQ_TILE_Y_K]); 628 } 629 630#pragma unroll 631 for (int k01 = 0; k01 < MMQ_TILE_NE_K/2; k01 += QR2_K*VDR_Q2_K_Q8_1_MMQ) { 632 const int k0 = k00 + k01; 633 634#pragma unroll 635 for (int j0 = 0; j0 < J; j0 += nwarps) { 636 const int j = j0 + threadIdx.y; 637 638#pragma unroll 639 for (int i0 = 0; i0 < I; i0 += warp_size) { 640 const int i = i0 + threadIdx.x; 641 642 constexpr int ns = 2; 643 sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q2_K_q8_1_impl_mmq<ns>( 644 &x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0], &y_qs[j*MMQ_TILE_Y_K + k01], 645 &x_dm[i*(MMQ_TILE_NE_K + 1) + k0/4], k01 < MMQ_TILE_NE_K/2 ? y_df[j0/nwarps].x : y_df[j0/nwarps].y, 646 &y_ds[j*MMQ_TILE_Y_K + (1 + k01/QI8_1)]); 647 } 648 } 649 } 650 651 // Some compilers fail to unroll the loop over k01 if there is a conditional statement for ns in the inner loop. 652 // As a workaround 2 separate loops are used instead. 653#pragma unroll 654 for (int k01 = MMQ_TILE_NE_K/2; k01 < MMQ_TILE_NE_K; k01 += QR2_K*VDR_Q2_K_Q8_1_MMQ) { 655 const int k0 = k00 + k01; 656 657#pragma unroll 658 for (int j0 = 0; j0 < J; j0 += nwarps) { 659 const int j = j0 + threadIdx.y; 660 661#pragma unroll 662 for (int i0 = 0; i0 < I; i0 += warp_size) { 663 const int i = i0 + threadIdx.x; 664 665 constexpr int ns = 1; 666 sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q2_K_q8_1_impl_mmq<ns>( 667 &x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0], &y_qs[j*MMQ_TILE_Y_K + k01], 668 &x_dm[i*(MMQ_TILE_NE_K + 1) + k0/4], k01 < MMQ_TILE_NE_K/2 ? y_df[j0/nwarps].x : y_df[j0/nwarps].y, 669 &y_ds[j*MMQ_TILE_Y_K + (1 + k01/QI8_1)]); 670 } 671 } 672 } 673} 674 675template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q2_K_q8_1_mma( 676 const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { 677#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) 678 constexpr data_layout input_layout = get_input_data_layout(); 679 typedef tile<16, 4, int, input_layout> tile_A; 680 typedef tile<16, 4, int, input_layout> tile_B; 681 typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C; 682 683 constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback); 684 constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback); 685 constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp. 686 687 y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K); 688 689 const int * x_qs = (const int *) x; 690 const half2 * x_dm = (const half2 *) x_qs + MMQ_TILE_NE_K*2; 691 const int * y_qs = (const int *) y + 4; 692 const half2 * y_ds = (const half2 *) y; 693 694 const int i0 = (threadIdx.y / ntx) * rows_per_warp; 695 696 for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 4) { 697 const int k0 = k00 + k01; 698 699 tile_A A[ntx]; 700#pragma unroll 701 for (int n = 0; n < ntx; ++n) { 702 load_ldmatrix(A[n], x_qs + (i0 + n*tile_A::I)*sram_stride + k0, sram_stride); 703 } 704 705#pragma unroll 706 for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) { 707 tile_B B; 708 load_ldmatrix(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K); 709 710 const int j = j0 + tile_C::get_j(0); 711 const float dB = (k01 < MMQ_TILE_NE_K/2) ? __half22float2(y_ds[j*MMQ_TILE_Y_K]).x : __half22float2(y_ds[j*MMQ_TILE_Y_K]).y; 712 const float sB = (k01 >= MMQ_TILE_NE_K * 3/4) ? 0 713 : (((k01/4)%2) ? __half22float2(y_ds[j*MMQ_TILE_Y_K + (1 + k01/QI8_1)]).y 714 : __half22float2(y_ds[j*MMQ_TILE_Y_K + (1 + k01/QI8_1)]).x); 715 716 tile_C Cm; 717 if (k01 >= MMQ_TILE_NE_K * 3/4) { 718 tile_A A1; 719#pragma unroll 720 for (int l = 0; l < tile_A::ne; ++l) { 721 A1.x[l] = 0x01010101; 722 } 723 mma(Cm, A1, B); 724 } 725 726#pragma unroll 727 for (int n = 0; n < ntx; ++n) { 728 tile_C Cd; 729 mma(Cd, A[n], B); 730 731#pragma unroll 732 for (int l = 0; l < tile_C::ne; ++l) { 733 const int i = i0 + n*tile_C::I + tile_C::get_i(l); 734 const float2 dm = __half22float2(x_dm[i*sram_stride + k0/4]); 735 float tmp = Cd.x[l]*dm.x; 736 if (k01 >= MMQ_TILE_NE_K * 3/4) { 737 tmp -= Cm.x[l]*dm.y; 738 } 739 sum[(j0/tile_C::J + n)*tile_C::ne + l] += tmp*dB; 740 sum[(j0/tile_C::J + n)*tile_C::ne + l] -= dm.y*sB; 741 } 742 } 743 } 744 } 745#elif defined(TURING_MMA_AVAILABLE) 746 747 typedef tile<16, 4, int> tile_A; 748 typedef tile<16, 8, int> tile_A_8; 749 typedef tile< 8, 4, int> tile_B; 750 typedef tile<16, 8, int> tile_C; 751 752 constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback); 753 constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback); 754 constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp. 755 756 y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K); 757 758 const int * x_qs = (const int *) x; 759 const half2 * x_dm = (const half2 *) x_qs + MMQ_TILE_NE_K*2; 760 const int * y_qs = (const int *) y + 4; 761 const half2 * y_ds = (const half2 *) y; 762 763 const int i0 = (threadIdx.y / ntx) * (ntx*tile_A::I); 764 765 tile_A A[ntx][8]; 766 float dA[ntx][tile_C::ne/2][8]; 767 float mA[ntx][tile_C::ne/2][8]; 768 769#pragma unroll 770 for (int n = 0; n < ntx; ++n) { 771#pragma unroll 772 for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1) { 773 const int k0 = k00 + k01; 774 775 load_ldmatrix(((tile_A_8 *) A[n])[k01/QI8_1], x_qs + (i0 + n*tile_A::I)*sram_stride + k0, sram_stride); 776 } 777 } 778 779#pragma unroll 780 for (int n = 0; n < ntx; ++n) { 781#pragma unroll 782 for (int l = 0; l < tile_C::ne/2; ++l) { 783 const int i = i0 + n*tile_C::I + tile_C::get_i(2*l); 784 785#pragma unroll 786 for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1/2) { 787 const int k0 = k00 + k01; 788 789 const float2 dm = __half22float2(x_dm[i*sram_stride + k0/(QI8_1/2)]); 790 791 dA[n][l][k01/(QI8_1/2)] = dm.x; 792 mA[n][l][k01/(QI8_1/2)] = dm.y; 793 } 794 } 795 } 796 797#pragma unroll 798 for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) { 799 float2 dB[tile_C::ne/2]; 800 801#pragma unroll 802 for (int l = 0; l < tile_C::ne/2; ++l) { 803 const int j = j0 + tile_C::get_j(l); 804 805 dB[l] = __half22float2(y_ds[j*MMQ_TILE_Y_K]); 806 } 807 808#pragma unroll 809 for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1) { 810 tile_B B[2]; 811 812 // Here load_generic is faster than load_ldmatrix. 813 load_generic(B[0], y_qs + j0*MMQ_TILE_Y_K + (k01 + 0), MMQ_TILE_Y_K); 814 load_generic(B[1], y_qs + j0*MMQ_TILE_Y_K + (k01 + tile_B::J), MMQ_TILE_Y_K); 815 816 tile_C Cm[2]; 817 if (k01 >= MMQ_TILE_NE_K * 3/4) { 818 tile_A A1; 819 A1.x[0] = 0x01010101; 820 A1.x[1] = 0x01010101; 821 mma(Cm[0], A1, B[0]); 822 mma(Cm[1], A1, B[1]); 823 } 824 825#pragma unroll 826 for (int n = 0; n < ntx; ++n) { 827 tile_C Cd[2]; 828 829 mma(Cd[0], A[n][k01/4 + 0], B[0]); 830 mma(Cd[1], A[n][k01/4 + 1], B[1]); 831 832#pragma unroll 833 for (int l = 0; l < tile_C::ne; ++l) { 834 float tmp = Cd[0].x[l]*dA[n][l/2][k01/4 + 0] + Cd[1].x[l]*dA[n][l/2][k01/4 + 1]; 835 if (k01 >= MMQ_TILE_NE_K * 3/4) { 836 tmp -= Cm[0].x[l]*mA[n][l/2][k01/4 + 0] + Cm[1].x[l]*mA[n][l/2][k01/4 + 1]; 837 } 838 sum[(j0/tile_C::J + n)*tile_C::ne + l] += tmp*(k01 < MMQ_TILE_NE_K/2 ? dB[l%2].x : dB[l%2].y); 839 } 840 } 841 } 842 843#pragma unroll 844 for (int k01 = 0; k01 < MMQ_TILE_NE_K * 3/4; k01 += QI8_1) { 845 float2 sB[tile_C::ne/2]; 846 847#pragma unroll 848 for (int l = 0; l < tile_C::ne/2; ++l) { 849 const int j = j0 + tile_C::get_j(l); 850 851 sB[l] = __half22float2(y_ds[j*MMQ_TILE_Y_K + (1 + k01/QI8_1)]); 852 } 853 854#pragma unroll 855 for (int n = 0; n < ntx; ++n) { 856#pragma unroll 857 for (int l = 0; l < tile_C::ne; ++l) { 858 sum[(j0/tile_C::J + n)*tile_C::ne + l] -= mA[n][l/2][k01/4 + 0]*sB[l%2].x; 859 sum[(j0/tile_C::J + n)*tile_C::ne + l] -= mA[n][l/2][k01/4 + 1]*sB[l%2].y; 860 } 861 } 862 } 863 } 864#else 865 GGML_UNUSED_VARS(x, y, sum, k00); 866 NO_DEVICE_CODE; 867#endif // AMD_MFMA_AVAILABLE || AMD_WMMA_AVAILABLE 868} 869 870template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q3_K_q8_1_dp4a( 871 const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { 872 constexpr int warp_size = ggml_cuda_get_physical_warp_size(); 873 constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size; 874 constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback); 875 876 constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q3_K, I); 877 const int * x_qs = (const int *) x; 878 const float * x_df = (const float *) x_qs + txs.qs; 879 const int * x_sc = (const int *) x_df + txs.dm; 880 const int * y_qs = (const int *) y + 4; 881 const float * y_df = (const float *) y; 882 883// #pragma unroll 884 for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QR3_K*VDR_Q3_K_Q8_1_MMQ) { 885 const int k0 = k00 + k01; 886 887#pragma unroll 888 for (int j0 = 0; j0 < J; j0 += nwarps) { 889 const int j = j0 + threadIdx.y; 890 891#pragma unroll 892 for (int i0 = 0; i0 < I; i0 += warp_size) { 893 const int i = i0 + threadIdx.x; 894 895 const int8_t * scales = ((const int8_t *) (x_sc + i*(MMQ_TILE_NE_K/8) + i/8)) + k0/4; 896 897 sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q3_K_q8_1_impl_mmq( 898 &x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0], &y_qs[j*MMQ_TILE_Y_K + k01], scales, 899 x_df[i], y_df[j*MMQ_TILE_Y_K + k01/QI8_1]); 900 } 901 } 902 } 903} 904 905template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q4_K_q8_1_dp4a( 906 const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { 907 constexpr int warp_size = ggml_cuda_get_physical_warp_size(); 908 constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size; 909 constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback); 910 911 constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q4_K, I); 912 const int * x_qs = (const int *) x; 913 const half2 * x_dm = (const half2 *) x_qs + txs.qs; 914 const int * x_sc = (const int *) x_dm + txs.dm; 915 const int * y_qs = (const int *) y + 4; 916 const half2 * y_ds = (const half2 *) y; 917 918// #pragma unroll 919 for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QR4_K*VDR_Q4_K_Q8_1_MMQ) { 920 const int k0 = k00 + k01; 921 922#pragma unroll 923 for (int j0 = 0; j0 < J; j0 += nwarps) { 924 const int j = j0 + threadIdx.y; 925 926#pragma unroll 927 for (int i0 = 0; i0 < I; i0 += warp_size) { 928 const int i = i0 + threadIdx.x; 929 930 const uint8_t * sc = (const uint8_t *) &x_sc[i * (MMQ_TILE_NE_K/8) + i/8 + k0/32] + 2*(k01/16); 931 932 sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q4_K_q8_1_impl_mmq( 933 &x_qs[i*(MMQ_TILE_NE_K + 1) + k0/2], &y_qs[j*MMQ_TILE_Y_K + k01], sc, sc+8, 934 x_dm[i], &y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]); 935 } 936 } 937 } 938} 939 940template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q5_K_q8_1_dp4a( 941 const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { 942 constexpr int warp_size = ggml_cuda_get_physical_warp_size(); 943 constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size; 944 constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback); 945 946 constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q5_K, I); 947 const int * x_qs = (const int *) x; 948 const half2 * x_dm = (const half2 *) x_qs + txs.qs; 949 const int * x_sc = (const int *) x_dm + txs.dm; 950 const int * y_qs = (const int *) y + 4; 951 const half2 * y_ds = (const half2 *) y; 952 953// #pragma unroll 954 for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QR5_K*VDR_Q5_K_Q8_1_MMQ) { 955 const int k0 = k00 + k01; 956 957#pragma unroll 958 for (int j0 = 0; j0 < J; j0 += nwarps) { 959 const int j = j0 + threadIdx.y; 960 961#pragma unroll 962 for (int i0 = 0; i0 < I; i0 += warp_size) { 963 const int i = i0 + threadIdx.x; 964 965 const uint8_t * sc = ((const uint8_t *) &x_sc[i * (MMQ_TILE_NE_K/8) + i/8 + k00/32]) + 2*(k01/16); 966 967 sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q5_K_q8_1_impl_mmq( 968 &x_qs[i*(QR5_K*MMQ_TILE_NE_K + 1) + k0], &y_qs[j*MMQ_TILE_Y_K + k01], sc, sc+8, 969 x_dm[i], &y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]); 970 } 971 } 972 } 973} 974 975template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q6_K_q8_1_dp4a( 976 const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { 977 constexpr int warp_size = ggml_cuda_get_physical_warp_size(); 978 constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size; 979 constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback); 980 981 constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q6_K, I); 982 const int * x_qs = (const int *) x; 983 const float * x_df = (const float *) x_qs + txs.qs; 984 const int * x_sc = (const int *) x_df + txs.dm; 985 const int * y_qs = (const int *) y + 4; 986 const float * y_df = (const float *) y; 987 988// #pragma unroll 989 for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QR6_K*VDR_Q6_K_Q8_1_MMQ) { 990 const int k0 = k00 + k01; 991 992#pragma unroll 993 for (int j0 = 0; j0 < J; j0 += nwarps) { 994 const int j = j0 + threadIdx.y; 995 996#pragma unroll 997 for (int i0 = 0; i0 < I; i0 += warp_size) { 998 const int i = i0 + threadIdx.x; 999 1000 const int8_t * sc = ((const int8_t *) &x_sc[i * (MMQ_TILE_NE_K/8) + i/8 + k0/16]); 1001 1002 sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q6_K_q8_1_impl_mmq( 1003 &x_qs[i*(QR6_K*MMQ_TILE_NE_K + 1) + k0], &y_qs[j*MMQ_TILE_Y_K + k01], sc, 1004 x_df[i*(MMQ_TILE_NE_K/QI6_K) + i/QI6_K], &y_df[j*MMQ_TILE_Y_K + k01/QI8_1]); 1005 } 1006 } 1007 } 1008} 1009 1010template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q6_K_q8_1_mma( 1011 const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { 1012#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) 1013 constexpr data_layout input_layout = get_input_data_layout(); 1014 typedef tile<16, 4, int, input_layout> tile_A; 1015 typedef tile<16, 4, int, input_layout> tile_B; 1016 typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C; 1017 1018 constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback); 1019 constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback); 1020 constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp. 1021 1022 y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K); 1023 1024 const int * x_qs = (const int *) x; 1025 const float * x_df = (const float *) x_qs + MMQ_TILE_NE_K*2; 1026 const int * x_sc = (const int *) x_df + MMQ_TILE_NE_K/QI6_K; 1027 const int * y_qs = (const int *) y + 4; 1028 const float * y_df = (const float *) y; 1029 1030 const int i0 = (threadIdx.y / ntx) * rows_per_warp; 1031 1032 for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 4) { 1033 const int k0 = k00 + k01; 1034 1035 tile_A A[ntx]; 1036#pragma unroll 1037 for (int n = 0; n < ntx; ++n) { 1038 load_ldmatrix(A[n], x_qs + (i0 + n*tile_A::I)*sram_stride + k0, sram_stride); 1039 } 1040 1041#pragma unroll 1042 for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) { 1043 tile_B B; 1044 load_ldmatrix(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K); 1045 1046 const int j = j0 + tile_C::get_j(0); 1047 const float dB = y_df[j*MMQ_TILE_Y_K + k01/QI8_1]; 1048 1049#pragma unroll 1050 for (int n = 0; n < ntx; ++n) { 1051 tile_C C; 1052 mma(C, A[n], B); 1053 1054#pragma unroll 1055 for (int l = 0; l < tile_C::ne; ++l) { 1056 const int i = i0 + n*tile_C::I + tile_C::get_i(l); 1057 const int8_t * sc = (const int8_t *) (x_sc + i*sram_stride + k00/16); 1058 sum[(j0/tile_C::J + n)*tile_C::ne + l] += C.x[l] * sc[k01/4] * x_df[i*sram_stride] * dB; 1059 } 1060 } 1061 } 1062 } 1063#elif defined(TURING_MMA_AVAILABLE) 1064 1065 typedef tile<16, 4, int> tile_A; 1066 typedef tile< 8, 4, int> tile_B; 1067 typedef tile<16, 8, int> tile_C; 1068 1069 constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback); 1070 constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback); 1071 constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp. 1072 1073 y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K); 1074 1075 const int * x_qs = (const int *) x; 1076 const float * x_df = (const float *) x_qs + MMQ_TILE_NE_K*2; 1077 const int * x_sc = (const int *) x_df + MMQ_TILE_NE_K/QI6_K; 1078 const int * y_qs = (const int *) y + 4; 1079 const float * y_df = (const float *) y; 1080 1081 const int i0 = (threadIdx.y / ntx) * (ntx*tile_A::I); 1082 1083 tile_A A[ntx][8]; 1084 int scA[ntx][tile_C::ne/2][8]; 1085 float dA[ntx][tile_C::ne/2]; 1086 1087#pragma unroll 1088 for (int n = 0; n < ntx; ++n) { 1089#pragma unroll 1090 for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 8) { 1091 const int k0 = k00 + k01; 1092 1093 load_ldmatrix(A[n][k01/4 + 0], x_qs + (i0 + n*tile_A::I)*sram_stride + (k0 + 0), sram_stride); 1094 load_ldmatrix(A[n][k01/4 + 1], x_qs + (i0 + n*tile_A::I)*sram_stride + (k0 + tile_A::J), sram_stride); 1095 } 1096 1097#pragma unroll 1098 for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 16) { 1099 const int k0 = k00 + k01; 1100 1101#pragma unroll 1102 for (int l = 0; l < tile_C::ne/2; ++l) { 1103 const int i = i0 + n*tile_C::I + tile_C::get_i(2*l); 1104 1105 const int sc_packed = x_sc[i*sram_stride + k0/16]; 1106 const int8_t * sc = (const int8_t *) &sc_packed; 1107 1108#pragma unroll 1109 for (int ksc = 0; ksc < sizeof(int); ++ksc) { 1110 scA[n][l][k01/4 + ksc] = sc[ksc]; 1111 } 1112 } 1113 } 1114 1115#pragma unroll 1116 for (int l = 0; l < tile_C::ne/2; ++l) { 1117 const int i = i0 + n*tile_C::I + tile_C::get_i(2*l); 1118 1119 dA[n][l] = x_df[i*sram_stride]; 1120 } 1121 } 1122 1123#pragma unroll 1124 for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) { 1125 float tmp[ntx][tile_C::ne] = {{0.0f}}; 1126 1127#pragma unroll 1128 for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 8) { 1129 tile_B B[2]; 1130 float dB[tile_C::ne/2]; 1131 1132 // Here load_generic is faster than load_ldmatrix. 1133 load_generic(B[0], y_qs + j0*MMQ_TILE_Y_K + 0 + k01, MMQ_TILE_Y_K); 1134 load_generic(B[1], y_qs + j0*MMQ_TILE_Y_K + tile_B::J + k01, MMQ_TILE_Y_K); 1135 1136#pragma unroll 1137 for (int l = 0; l < tile_C::ne/2; ++l) { 1138 const int j = j0 + tile_C::get_j(l); 1139 1140 dB[l] = y_df[j*MMQ_TILE_Y_K + k01/QI8_1]; 1141 } 1142 1143#pragma unroll 1144 for (int n = 0; n < ntx; ++n) { 1145 tile_C C[2]; 1146 mma(C[0], A[n][k01/4 + 0], B[0]); 1147 mma(C[1], A[n][k01/4 + 1], B[1]); 1148 1149#pragma unroll 1150 for (int l = 0; l < tile_C::ne; ++l) { 1151 tmp[n][l] += (C[0].x[l]*scA[n][l/2][k01/4 + 0] + C[1].x[l]*scA[n][l/2][k01/4 + 1])*dB[l%2]; 1152 } 1153 } 1154 } 1155 1156#pragma unroll 1157 for (int n = 0; n < ntx; ++n) { 1158#pragma unroll 1159 for (int l = 0; l < tile_C::ne; ++l) { 1160 sum[(j0/tile_C::J + n)*tile_C::ne + l] += tmp[n][l]*dA[n][l/2]; 1161 } 1162 } 1163 } 1164#else 1165 GGML_UNUSED_VARS(x, y, sum, k00); 1166 NO_DEVICE_CODE; 1167#endif // AMD_MFMA_AVAILABLE || AMD_WMMA_AVAILABLE 1168} 1169 1170// --------------------------------------------------------------------------------------------- 1171 1172// Shared MMA kernel for MXFP4 and NVFP4 on Blackwell. 1173// Both quantizations encode values as e2m1 (FP4) and produce one uint32 scale per 1174// m16n8k64 MMA call; only the PTX kind (scale_vec::2X ue8m0 vs scale_vec::4X ue4m3) 1175// and the per-type stride constant differ. 1176template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_fp4_fp4_mma( 1177 const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { 1178 1179 typedef tile<16, 8, int> tile_A; 1180 typedef tile<8, 8, int> tile_B; 1181 typedef tile<16, 8, float> tile_C; 1182 1183 constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback); 1184 constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback); 1185 constexpr int ntx = rows_per_warp / tile_C::I; 1186 constexpr int nfrags = MMQ_TILE_NE_K / tile_A::J; 1187 1188 y += (threadIdx.y % ntx) * (tile_C::J * MMQ_TILE_Y_K); 1189 1190 const int * x_qs = (const int *) x; 1191 const uint32_t * x_sc = (const uint32_t *) (x_qs + 2 * MMQ_TILE_NE_K); 1192 const int * y_qs = (const int *) y + 4; 1193 const uint32_t * y_sc = (const uint32_t *) y; 1194 1195 // 2 threads per quad supply the packed scale register to the block_scale MMA, 1196 // see https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-block-scaling 1197 const int tidx_A = threadIdx.x / 4 + (threadIdx.x % 2) * 8; 1198 const int tidx_B = threadIdx.x / 4; 1199 const int i0 = (threadIdx.y / ntx) * rows_per_warp; 1200 1201 tile_A A[ntx][nfrags]; 1202 uint32_t scaleA[ntx][nfrags]; 1203 1204#pragma unroll 1205 for (int n = 0; n < ntx; ++n) { 1206#pragma unroll 1207 for (int frag = 0; frag < nfrags; ++frag) { 1208 const int k0 = k00 + frag * tile_A::J; 1209 load_ldmatrix(A[n][frag], x_qs + (i0 + n * tile_A::I) * sram_stride + k0, sram_stride); 1210 scaleA[n][frag] = x_sc[(i0 + n * tile_A::I + tidx_A) * sram_stride + k0 / tile_A::J]; 1211 } 1212 } 1213 1214#pragma unroll 1215 for (int j0 = 0; j0 < J; j0 += ntx * tile_C::J) { 1216 tile_B B[nfrags]; 1217 uint32_t scaleB[nfrags]; 1218 1219#pragma unroll 1220 for (int frag = 0; frag < nfrags; ++frag) { 1221 const int k0 = frag * tile_B::J; 1222 load_generic(B[frag], y_qs + j0 * MMQ_TILE_Y_K + k0, MMQ_TILE_Y_K); 1223 scaleB[frag] = y_sc[(j0 + tidx_B) * MMQ_TILE_Y_K + frag]; 1224 } 1225 1226#pragma unroll 1227 for (int n = 0; n < ntx; ++n) { 1228#pragma unroll 1229 for (int frag = 0; frag < nfrags; ++frag) { 1230 tile_C C = {}; 1231 mma_block_scaled_fp4<type>(C, A[n][frag], B[frag], scaleA[n][frag], scaleB[frag]); 1232#pragma unroll 1233 for (int l = 0; l < tile_C::ne; ++l) { 1234 sum[(j0 / tile_C::J + n) * tile_C::ne + l] += C.x[l]; 1235 } 1236 } 1237 } 1238 } 1239} 1240