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 <int ncols, int nrows>
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 <int ncols>
-__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 <int ncols, int MINB, int NW>
+__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<ncols>, 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<ncols, 2, 8>, 8);
+ } else if (ncols_x >= 16384 || nrows_x >= 16384) {
+ launch(mul_mat_vec_ptq1_0_tc<ncols, 3, 4>, 4);
+ } else {
+ launch(mul_mat_vec_ptq1_0_tc<ncols, 3, 8>, 8);
+ }
}
// plain 2D PTQ1_0 x 2..8 columns without fusion, activations quantized by quantize_row_q8_1_ptq1_0_tc_cuda