diff --git a/ggml/src/ggml-cuda/mmvq-ptq1_0.cuh b/ggml/src/ggml-cuda/mmvq-ptq1_0.cuh index 29078046..95b18171 100644 --- a/ggml/src/ggml-cuda/mmvq-ptq1_0.cuh +++ b/ggml/src/ggml-cuda/mmvq-ptq1_0.cuh @@ -71,14 +71,25 @@ static __device__ __forceinline__ int ptq1_0_trits_to_weights(const int q) { return (int) (((uint32_t) q + 0x7F7F7F7Fu) ^ 0x80808080u); } -// one base-3 digit step on four bytes held as two 16-bit-lane words: -// returns the weights (-1, 0, 1) as four signed bytes, advances the remainders -static __device__ __forceinline__ int ptq1_0_trit_step(uint32_t & vlo, uint32_t & vhi) { +// one base-3 digit step on four bytes held as two 16-bit-lane words: returns the digits (0, 1, 2) as +// four bytes, advances the remainders +static __device__ __forceinline__ int ptq1_0_digit_step(uint32_t & vlo, uint32_t & vhi) { const uint32_t wlo = vlo * 3; const uint32_t whi = vhi * 3; vlo = wlo & 0x00FF00FF; vhi = whi & 0x00FF00FF; - return ptq1_0_trits_to_weights(__byte_perm(wlo, whi, 0x7531)); + return (int) __byte_perm(wlo, whi, 0x7531); +} + +// c + dot of four unsigned bytes (the digits) with four signed bytes (the activations) +static __device__ __forceinline__ int ptq1_0_dp4a_us(const int digits, const int a, const int c) { +#if defined(GGML_USE_HIP) + return ggml_cuda_dp4a(ptq1_0_trits_to_weights(digits), a, c + ggml_cuda_dp4a(0x01010101, a, 0)); +#else + int d; + asm("dp4a.u32.s32 %0, %1, %2, %3;" : "=r"(d) : "r"(digits), "r"(a), "r"(c)); + return d; +#endif } // Dot products of nrows PTQ1_0 blocks with the same block index of ncols @@ -88,6 +99,9 @@ static __device__ __forceinline__ int ptq1_0_trit_step(uint32_t & vlo, uint32_t // The integer sum of each 32-element sub-block k is folded into the fp32 // accumulator as soon as the sub-block is complete, in the order k = 0..3, // which is the expression acc = sum_k d8_k * sumi_k of the block_q8_1 kernel. +// sumi_k = sum (t - 1)*a is taken as sum t*a over the digits t minus the column's +// sum a of the sub-block, shared by the rows: the same integer, without the +// per-word mapping of the digits to -1, 0, 1. template static __device__ __forceinline__ void ptq1_0_pt_block_dot( const block_ptq1_0 * const (&bq)[nrows], @@ -101,9 +115,11 @@ static __device__ __forceinline__ void ptq1_0_pt_block_dot( } int sumi[ncols][nrows]; + int suma[ncols]; float acc[ncols][nrows]; #pragma unroll for (int j = 0; j < ncols; ++j) { + suma[j] = 0; #pragma unroll for (int i = 0; i < nrows; ++i) { sumi[j][i] = 0; @@ -111,15 +127,23 @@ static __device__ __forceinline__ void ptq1_0_pt_block_dot( } } + auto add_suma = [&](const int j, const int4 & u) { + suma[j] = ggml_cuda_dp4a(0x01010101, u.x, suma[j]); + suma[j] = ggml_cuda_dp4a(0x01010101, u.y, suma[j]); + suma[j] = ggml_cuda_dp4a(0x01010101, u.z, suma[j]); + suma[j] = ggml_cuda_dp4a(0x01010101, u.w, suma[j]); + }; + auto fold = [&](const int k) { #pragma unroll for (int j = 0; j < ncols; ++j) { const float d8 = __low2float(((const half2 *) &dsraw[j])[k]); #pragma unroll for (int i = 0; i < nrows; ++i) { - acc[j][i] = __fmaf_rn(d8, (float) sumi[j][i], acc[j][i]); // one FFMA in every instantiation + acc[j][i] = __fmaf_rn(d8, (float) (sumi[j][i] - suma[j]), acc[j][i]); // one FFMA in every instantiation sumi[j][i] = 0; } + suma[j] = 0; } }; @@ -142,15 +166,16 @@ static __device__ __forceinline__ void ptq1_0_pt_block_dot( #pragma unroll for (int j = 0; j < ncols; ++j) { u[j] = *((const int4 *) ycol[j] + t*nblk + kbx); + add_suma(j, u[j]); } #pragma unroll for (int i = 0; i < nrows; ++i) { #pragma unroll for (int g = 0; g < 4; ++g) { - const int q = ptq1_0_trit_step(vlo[i][g], vhi[i][g]); + const int q = ptq1_0_digit_step(vlo[i][g], vhi[i][g]); #pragma unroll for (int j = 0; j < ncols; ++j) { - sumi[j][i] = ggml_cuda_dp4a(q, int4_at(u[j], g), sumi[j][i]); + sumi[j][i] = ptq1_0_dp4a_us(q, int4_at(u[j], g), sumi[j][i]); } } } @@ -182,17 +207,18 @@ static __device__ __forceinline__ void ptq1_0_pt_block_dot( #pragma unroll for (int j = 0; j < ncols; ++j) { u2[j] = *((const int4 *) ycol[j] + (5 + t/2)*nblk + kbx); + add_suma(j, u2[j]); // plane 7 includes words 30, 31, used by qh below } } #pragma unroll for (int i = 0; i < nrows; ++i) { #pragma unroll for (int g = 0; g < 2; ++g) { - const int q = ptq1_0_trit_step(vlo2[i][g], vhi2[i][g]); + const int q = ptq1_0_digit_step(vlo2[i][g], vhi2[i][g]); const int w = 20 + 2*t + g; // word index within the 128-element block #pragma unroll for (int j = 0; j < ncols; ++j) { - sumi[j][i] = ggml_cuda_dp4a(q, int4_at(u2[j], w & 3), sumi[j][i]); + sumi[j][i] = ptq1_0_dp4a_us(q, int4_at(u2[j], w & 3), sumi[j][i]); } } } @@ -212,10 +238,10 @@ static __device__ __forceinline__ void ptq1_0_pt_block_dot( v = w0 & 0x00FF00FF; const uint32_t w1 = v * 3; v = w1 & 0x00FF00FF; - const int q = ptq1_0_trits_to_weights(__byte_perm(w0, w1, 0x7531)); + const int q = (int) __byte_perm(w0, w1, 0x7531); #pragma unroll for (int j = 0; j < ncols; ++j) { - sumi[j][i] = ggml_cuda_dp4a(q, int4_at(u2[j], 2 + t/2), sumi[j][i]); + sumi[j][i] = ptq1_0_dp4a_us(q, int4_at(u2[j], 2 + t/2), sumi[j][i]); } } } @@ -471,7 +497,7 @@ static __device__ __forceinline__ uint32_t ptq1_0_trits(const uint32_t h, const // acc = fma(d8_k, (float) sumi_k, acc) over k = 0..3 and the scale by d are the same // operations, and the per-row sum over K blocks keeps the dp4a kernel's order: the slot of // residue l (mod 32) adds the blocks l, l+32, ... in that order, then the slots are combined in -// the pairing of warp_reduce_sum. Warp w of the CTA owns the blocks kb = w (mod 4) and keeps +// the pairing of warp_reduce_sum. Warp w of the CTA owns the blocks kb = w (mod NW) and keeps // those residues' running sums in registers. // // A fragment of m16n8k32 (row major): lane (g = lane/4, q = lane%4) holds k = 4q..4q+3 and @@ -481,7 +507,6 @@ static __device__ __forceinline__ uint32_t ptq1_0_trits(const uint32_t h, const // need come over by shuffle. // --------------------------------------------------------------------------- -#define PTQ1_0_TC_THREADS 256 #define PTQ1_0_TC_MAX_COLS 8 // c += a*b with a the trits (u8) and b the activations (s8) @@ -549,21 +574,24 @@ static __device__ __forceinline__ void ptq1_0_tc_row_frags(const uint32_t xq, co frag[3][1] = low ? t2[4] : (q == 2 ? qh12 : qh34); // elements 112 + 4q } -template -__launch_bounds__(PTQ1_0_TC_THREADS, 2) +// NW warps per CTA split the K blocks by residue: warp w owns kb = w (mod NW), slot i of a thread +// holds residue NW*i + w. The 32 slots of a row combine in warp_reduce_sum's pairing (distance 16, +// 8, ..., 1 over the residue); distances of NW and more pair two slots of the same thread, so those +// steps run in registers and only the last log2(NW) steps go through shared memory. +template +__launch_bounds__(NW*WARP_SIZE, MINB) static __global__ void mul_mat_vec_ptq1_0_tc( const void * vx_ptr, const void * vy_ptr, float * dst_ptr, const int ncols_x, const int nrows_x, const int stride_row_x, const int stride_col_y, const int stride_col_dst) { static_assert(ncols >= 1 && ncols <= PTQ1_0_TC_MAX_COLS, "1 to 8 columns"); - constexpr int nwarps = PTQ1_0_TC_THREADS / WARP_SIZE; - static_assert(WARP_SIZE % nwarps == 0, "a warp owns the K blocks of fixed residues mod 32"); - constexpr int nslots = WARP_SIZE / nwarps; // residues mod 32 per warp + static_assert(WARP_SIZE % NW == 0, "a warp owns the K blocks of fixed residues mod 32"); + constexpr int nslots = WARP_SIZE / NW; // residues mod 32 per warp const void * GGML_CUDA_RESTRICT vx = vx_ptr; const void * GGML_CUDA_RESTRICT vy = vy_ptr; float * GGML_CUDA_RESTRICT dst = dst_ptr; - __shared__ float slots[ncols*16][WARP_SIZE + 1]; + __shared__ float part[NW][ncols*16]; const int bpr = ncols_x / QK_PTQ1_0; const int row0 = 16 * blockIdx.x; @@ -645,7 +673,7 @@ static __global__ void mul_mat_vec_ptq1_0_tc( #pragma unroll for (int i = 0; i < nslots; ++i) { float p[4]; - block(kb0 + nwarps*i + warp, p); + block(kb0 + NW*i + warp, p); #pragma unroll for (int e = 0; e < 4; ++e) { slot[i][e] += p[e]; @@ -654,7 +682,7 @@ static __global__ void mul_mat_vec_ptq1_0_tc( } #pragma unroll for (int i = 0; i < nslots; ++i) { - const int kb = bpr_full + nwarps*i + warp; + const int kb = bpr_full + NW*i + warp; if (kb < bpr) { // uniform across the warp float p[4]; block(kb, p); @@ -665,29 +693,36 @@ static __global__ void mul_mat_vec_ptq1_0_tc( } } + // the pairing steps of distance NW and more: slot i with slot i + distance/NW of the same thread #pragma unroll - for (int i = 0; i < nslots; ++i) { - if (c0 < ncols) { - slots[c0*16 + g ][nwarps*i + warp] = slot[i][0]; - slots[c0*16 + g + 8][nwarps*i + warp] = slot[i][2]; - } - if (c1 < ncols) { - slots[c1*16 + g ][nwarps*i + warp] = slot[i][1]; - slots[c1*16 + g + 8][nwarps*i + warp] = slot[i][3]; + for (int off = nslots/2; off > 0; off >>= 1) { +#pragma unroll + for (int i = 0; i < off; ++i) { +#pragma unroll + for (int e = 0; e < 4; ++e) { + slot[i][e] += slot[i + off][e]; + } } } + if (c0 < ncols) { + part[warp][c0*16 + g ] = slot[0][0]; + part[warp][c0*16 + g + 8] = slot[0][2]; + } + if (c1 < ncols) { + part[warp][c1*16 + g ] = slot[0][1]; + part[warp][c1*16 + g + 8] = slot[0][3]; + } __syncthreads(); - // one thread per (column, row): the 32 slots combined in warp_reduce_sum's pairing - // (distance 16, then 8, 4, 2, 1), which is the result that butterfly leaves in lane 0 - for (int w = threadIdx.x; w < ncols*16; w += PTQ1_0_TC_THREADS) { - float v[WARP_SIZE]; + // the steps of distance NW/2 .. 1 across the warps, one thread per (column, row) + for (int w = threadIdx.x; w < ncols*16; w += NW*WARP_SIZE) { + float v[NW]; #pragma unroll - for (int l = 0; l < WARP_SIZE; ++l) { - v[l] = slots[w][l]; + for (int l = 0; l < NW; ++l) { + v[l] = part[l][w]; } #pragma unroll - for (int off = WARP_SIZE/2; off > 0; off >>= 1) { + for (int off = NW/2; off > 0; off >>= 1) { #pragma unroll for (int l = 0; l < off; ++l) { v[l] += v[l + off]; @@ -716,9 +751,20 @@ static void mul_mat_vec_ptq1_0_tc_launch( const int ncols_x, const int nrows_x, const int stride_row_x, const int stride_col_y, const int stride_col_dst, cudaStream_t stream) { const dim3 block_nums((nrows_x + 15) / 16, 1, 1); - const dim3 block_dims(PTQ1_0_TC_THREADS, 1, 1); - const ggml_cuda_kernel_launch_params lp = ggml_cuda_kernel_launch_params(block_nums, block_dims, 0, stream); - ggml_cuda_kernel_launch(mul_mat_vec_ptq1_0_tc, lp, vx, vy, dst, ncols_x, nrows_x, stride_row_x, stride_col_y, stride_col_dst); + // every geometry gives the same bits; the choice is speed only (RTX 3060, Bonsai 2 shapes): 4 warps of + // 3 CTAs per SM for the long-K and the 17408-row and larger shapes, 8 warps otherwise, 2 CTAs at 8 + // columns where 3 would spill + auto launch = [&](auto kernel, const int nwarps) { + const ggml_cuda_kernel_launch_params lp = ggml_cuda_kernel_launch_params(block_nums, dim3(nwarps*WARP_SIZE, 1, 1), 0, stream); + ggml_cuda_kernel_launch(kernel, lp, vx, vy, dst, ncols_x, nrows_x, stride_row_x, stride_col_y, stride_col_dst); + }; + if (ncols == 8) { + launch(mul_mat_vec_ptq1_0_tc, 8); + } else if (ncols_x >= 16384 || nrows_x >= 16384) { + launch(mul_mat_vec_ptq1_0_tc, 4); + } else { + launch(mul_mat_vec_ptq1_0_tc, 8); + } } // plain 2D PTQ1_0 x 2..8 columns without fusion, activations quantized by quantize_row_q8_1_ptq1_0_tc_cuda