diff --git a/ggml/src/ggml-cuda/mmvq-pq2_0.cuh b/ggml/src/ggml-cuda/mmvq-pq2_0.cuh index b1c47ddb4e72..850ccd834178 100644 --- a/ggml/src/ggml-cuda/mmvq-pq2_0.cuh +++ b/ggml/src/ggml-cuda/mmvq-pq2_0.cuh @@ -14,24 +14,51 @@ #define PQ2_0_MC_ROWS 4 #define PQ2_0_MC_WARPS 4 +// All weight reads of the kernel below go through this one helper. For 32-element chunk c of a row it returns the +// 8 code bytes (as two words: elements 0-15 and 16-31) and the block scale. It is the only place that knows the +// 34-byte block layout [d: half][qs: 32 bytes]; a repack of the weights only has to replace it. +static __device__ __forceinline__ void pq2_0_mc_load_chunk( + const void * GGML_CUDA_RESTRICT vx, const int64_t row, const int stride_row_x, const int c, + uint32_t & w0, uint32_t & w1, float & d) { + const block_pq2_0 * blk = (const block_pq2_0 *) vx + row*stride_row_x + (c >> 2); + const uint16_t * qp = (const uint16_t *) (blk->qs + (c & 3)*8); // blocks are only 2-byte aligned + d = __half2float(blk->d); + w0 = (uint32_t) qp[0] | ((uint32_t) qp[1] << 16); + w1 = (uint32_t) qp[2] | ((uint32_t) qp[3] << 16); +} + +// Launch geometry (all template parameters, so they can be retuned independently of the kernel body): +// rows - output rows per warp-group (activation slice of a chunk is reused for all of them) +// nwarps - warps per CTA +// unroll - chunks per lane whose loads are issued together before they are consumed (order-preserving) +// splitk - warps cooperating on one row group. 1: each group is one warp and the K reduction order of every row is +// the same as before (lane c, c+32, ... then a shuffle tree), so any (rows, nwarps) is bit-exact with +// any other splitk = 1 setting. >1: the chunk loop is strided by 32*splitk, the per-warp partial sums are +// added in warp order 0..splitk-1 through shared memory, which changes the float summation order. // Optional fusion (has_fusion, the same epilogue as mul_mat_vec_q): x_bias is added to the result; with has_gate a // second weight matrix (gate, same layout and strides) is dotted against the same activations, gate_bias is added to // it and the GLU op combines the two. Without fusion the epilogue is a plain store: a runtime bias check there costs // the 8-column verify kernel ~12% on H200. -template -__global__ void __launch_bounds__(PQ2_0_MC_WARPS * 32, 3) +template +__global__ void __launch_bounds__(nwarps * 32, 3) mul_mat_vec_pq2_0_mc(const void * vx_, const void * vy_, float * dst_, const int ncols_x, const int nrows_x, const int stride_row_x, const int stride_col_y, const int stride_col_dst, const float * x_bias, const void * vgate_, const float * gate_bias, const ggml_glu_op glu_op, const float glu_limit) { + static_assert(unroll == 1 || !has_gate, "unroll is not implemented for the fused gate variant"); + static_assert(nwarps % splitk == 0, "splitk must divide nwarps"); + static_assert(splitk == 1 || !has_gate, "split-K is not implemented for the fused gate variant"); // Plain pointer parameters, restrict-qualified locals: GGML_CUDA_RESTRICT differs between the host and device // passes, which breaks the template stub. Same pattern as mul_mat_vec_q. const void * GGML_CUDA_RESTRICT vx = vx_; const void * GGML_CUDA_RESTRICT vy = vy_; float * GGML_CUDA_RESTRICT dst = dst_; const int lane = threadIdx.x & 31; - const int row0 = ((blockIdx.x * blockDim.x + threadIdx.x) >> 5) * rows; - if (row0 >= nrows_x) { - return; + const int warp = threadIdx.x >> 5; + const int part = splitk == 1 ? 0 : warp % splitk; + const int row0 = splitk == 1 ? ((blockIdx.x * blockDim.x + threadIdx.x) >> 5) * rows + : (blockIdx.x * (nwarps / splitk) + warp / splitk) * rows; + if (splitk == 1 && row0 >= nrows_x) { + return; // with split-K every warp must reach the barrier, inactive groups only skip the stores } // launched through ggml_cuda_kernel_launch, which opts into PDL on Hopper and newer: wait for the kernel // that wrote vy (the q8_1 activation quantization) before reading it @@ -42,85 +69,90 @@ mul_mat_vec_pq2_0_mc(const void * vx_, const void * vy_, float * dst_, const size_t ds_off = (size_t) stride_col_y * QK8_1; // qs region of a column, padded length const char * y = (const char *) vy; + __shared__ float red[splitk > 1 ? nwarps*rows*ncols : 1]; float acc[rows][ncols] = {}; [[maybe_unused]] float accg[rows][ncols] = {}; [[maybe_unused]] const void * GGML_CUDA_RESTRICT vgate = vgate_; - for (int c = lane; c < nchunk; c += WARP_SIZE) { - uint32_t w0[rows], w1[rows]; - float d2[rows]; + constexpr int cstep = WARP_SIZE*splitk; + // Each lane walks chunks part*32+lane, +cstep, ... exactly as before; `unroll` only issues the loads of that many + // consecutive chunks of the lane before consuming them in the same order, so the accumulation order is unchanged. + for (int cb = part*WARP_SIZE + lane; cb < nchunk; cb += cstep*unroll) { + uint32_t w0[unroll][rows], w1[unroll][rows]; + float d2[unroll][rows]; [[maybe_unused]] uint32_t g0[rows], g1[rows]; [[maybe_unused]] float e2[rows]; + int4 a0[unroll][ncols], a1[unroll][ncols]; + float d8[unroll][ncols]; + int qsum[unroll][ncols]; #pragma unroll - for (int r = 0; r < rows; ++r) { - const int row = min(row0 + r, nrows_x - 1); // clamped rows are computed but not stored - const block_pq2_0 * blk = (const block_pq2_0 *) vx + (int64_t) row*stride_row_x + (c >> 2); - const uint16_t * qp = (const uint16_t *) (blk->qs + (c & 3)*8); // blocks are only 2-byte aligned - d2[r] = __half2float(blk->d); - w0[r] = (uint32_t) qp[0] | ((uint32_t) qp[1] << 16); - w1[r] = (uint32_t) qp[2] | ((uint32_t) qp[3] << 16); - if constexpr (has_gate) { - const block_pq2_0 * gblk = (const block_pq2_0 *) vgate + (int64_t) row*stride_row_x + (c >> 2); - const uint16_t * gp = (const uint16_t *) (gblk->qs + (c & 3)*8); - e2[r] = __half2float(gblk->d); - g0[r] = (uint32_t) gp[0] | ((uint32_t) gp[1] << 16); - g1[r] = (uint32_t) gp[2] | ((uint32_t) gp[3] << 16); + for (int u = 0; u < unroll; ++u) { + const int c = cb + u*cstep; + if (c < nchunk) { +#pragma unroll + for (int r = 0; r < rows; ++r) { + const int row = min(row0 + r, nrows_x - 1); // clamped rows are computed but not stored + pq2_0_mc_load_chunk(vx, row, stride_row_x, c, w0[u][r], w1[u][r], d2[u][r]); + if constexpr (has_gate) { + pq2_0_mc_load_chunk(vgate, row, stride_row_x, c, g0[r], g1[r], e2[r]); + } + } +#pragma unroll + for (int j = 0; j < ncols; ++j) { + const char * yc = y + j*col_bytes; + a0[u][j] = *(const int4 *) (yc + c*QK8_1); + a1[u][j] = *(const int4 *) (yc + c*QK8_1 + 16); + const half2 ds = *(const half2 *) (yc + ds_off + c*sizeof(half2)); + d8[u][j] = __low2float(ds); + qsum[u][j] = __half_as_short(__high2half(ds)); + } } } - - int4 a0[ncols], a1[ncols]; - float d8[ncols]; - int qsum[ncols]; #pragma unroll - for (int j = 0; j < ncols; ++j) { - const char * yc = y + j*col_bytes; - a0[j] = *(const int4 *) (yc + c*QK8_1); - a1[j] = *(const int4 *) (yc + c*QK8_1 + 16); - const half2 ds = *(const half2 *) (yc + ds_off + c*sizeof(half2)); - d8[j] = __low2float(ds); - qsum[j] = __half_as_short(__high2half(ds)); - } - + for (int u = 0; u < unroll; ++u) { + if (cb + u*cstep < nchunk) { #pragma unroll - for (int r = 0; r < rows; ++r) { - int t[8]; + for (int r = 0; r < rows; ++r) { + int t[8]; #pragma unroll - for (int k = 0; k < 4; ++k) { - t[k] = (int) ((w0[r] >> (2*k)) & 0x03030303u); - t[4 + k] = (int) ((w1[r] >> (2*k)) & 0x03030303u); - } + for (int k = 0; k < 4; ++k) { + t[k] = (int) ((w0[u][r] >> (2*k)) & 0x03030303u); + t[4 + k] = (int) ((w1[u][r] >> (2*k)) & 0x03030303u); + } #pragma unroll - for (int j = 0; j < ncols; ++j) { - int sumi = 0; - sumi = ggml_cuda_dp4a(t[0], a0[j].x, sumi); - sumi = ggml_cuda_dp4a(t[1], a0[j].y, sumi); - sumi = ggml_cuda_dp4a(t[2], a0[j].z, sumi); - sumi = ggml_cuda_dp4a(t[3], a0[j].w, sumi); - sumi = ggml_cuda_dp4a(t[4], a1[j].x, sumi); - sumi = ggml_cuda_dp4a(t[5], a1[j].y, sumi); - sumi = ggml_cuda_dp4a(t[6], a1[j].z, sumi); - sumi = ggml_cuda_dp4a(t[7], a1[j].w, sumi); - acc[r][j] += d2[r] * d8[j] * (float) (sumi - qsum[j]); - } - if constexpr (has_gate) { - int tg[8]; + for (int j = 0; j < ncols; ++j) { + int sumi = 0; + sumi = ggml_cuda_dp4a(t[0], a0[u][j].x, sumi); + sumi = ggml_cuda_dp4a(t[1], a0[u][j].y, sumi); + sumi = ggml_cuda_dp4a(t[2], a0[u][j].z, sumi); + sumi = ggml_cuda_dp4a(t[3], a0[u][j].w, sumi); + sumi = ggml_cuda_dp4a(t[4], a1[u][j].x, sumi); + sumi = ggml_cuda_dp4a(t[5], a1[u][j].y, sumi); + sumi = ggml_cuda_dp4a(t[6], a1[u][j].z, sumi); + sumi = ggml_cuda_dp4a(t[7], a1[u][j].w, sumi); + acc[r][j] += d2[u][r] * d8[u][j] * (float) (sumi - qsum[u][j]); + } + if constexpr (has_gate) { + int tg[8]; #pragma unroll - for (int k = 0; k < 4; ++k) { - tg[k] = (int) ((g0[r] >> (2*k)) & 0x03030303u); - tg[4 + k] = (int) ((g1[r] >> (2*k)) & 0x03030303u); - } + for (int k = 0; k < 4; ++k) { + tg[k] = (int) ((g0[r] >> (2*k)) & 0x03030303u); + tg[4 + k] = (int) ((g1[r] >> (2*k)) & 0x03030303u); + } #pragma unroll - for (int j = 0; j < ncols; ++j) { - int sumi = 0; - sumi = ggml_cuda_dp4a(tg[0], a0[j].x, sumi); - sumi = ggml_cuda_dp4a(tg[1], a0[j].y, sumi); - sumi = ggml_cuda_dp4a(tg[2], a0[j].z, sumi); - sumi = ggml_cuda_dp4a(tg[3], a0[j].w, sumi); - sumi = ggml_cuda_dp4a(tg[4], a1[j].x, sumi); - sumi = ggml_cuda_dp4a(tg[5], a1[j].y, sumi); - sumi = ggml_cuda_dp4a(tg[6], a1[j].z, sumi); - sumi = ggml_cuda_dp4a(tg[7], a1[j].w, sumi); - accg[r][j] += e2[r] * d8[j] * (float) (sumi - qsum[j]); + for (int j = 0; j < ncols; ++j) { + int sumi = 0; + sumi = ggml_cuda_dp4a(tg[0], a0[u][j].x, sumi); + sumi = ggml_cuda_dp4a(tg[1], a0[u][j].y, sumi); + sumi = ggml_cuda_dp4a(tg[2], a0[u][j].z, sumi); + sumi = ggml_cuda_dp4a(tg[3], a0[u][j].w, sumi); + sumi = ggml_cuda_dp4a(tg[4], a1[u][j].x, sumi); + sumi = ggml_cuda_dp4a(tg[5], a1[u][j].y, sumi); + sumi = ggml_cuda_dp4a(tg[6], a1[u][j].z, sumi); + sumi = ggml_cuda_dp4a(tg[7], a1[u][j].w, sumi); + accg[r][j] += e2[r] * d8[u][j] * (float) (sumi - qsum[u][j]); + } + } } } } @@ -137,8 +169,13 @@ mul_mat_vec_pq2_0_mc(const void * vx_, const void * vy_, float * dst_, if constexpr (has_gate) { gv = warp_reduce_sum(accg[r][j]); } + if constexpr (splitk > 1) { + if (lane == 0) { + red[(warp*rows + r)*ncols + j] = v; + } + } const int row = row0 + r; - if (lane == 0 && row < nrows_x) { + if (splitk == 1 && lane == 0 && row < nrows_x) { const int64_t o = (int64_t) j*stride_col_dst + row; if constexpr (has_fusion) { if (x_bias) { @@ -161,25 +198,59 @@ mul_mat_vec_pq2_0_mc(const void * vx_, const void * vy_, float * dst_, } } } + + if constexpr (splitk > 1) { + __syncthreads(); + if (part == 0 && lane == 0) { +#pragma unroll + for (int r = 0; r < rows; ++r) { + const int row = row0 + r; +#pragma unroll + for (int j = 0; j < ncols; ++j) { + if (row < nrows_x) { + float v = 0.0f; +#pragma unroll + for (int p = 0; p < splitk; ++p) { + v += red[((warp + p)*rows + r)*ncols + j]; + } + const int64_t o = (int64_t) j*stride_col_dst + row; + if constexpr (has_fusion) { + if (x_bias) { + v += x_bias[o]; + } + } + dst[o] = v; + } + } + } + } + } } -template +template static void mul_mat_vec_pq2_0_mc_launch_rows( const void * vx, const void * vy, float * dst, 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 ggml_cuda_mm_fusion_args_device & fusion) { - const int nwarps = (nrows_x + rows - 1) / rows; - const dim3 block_nums((nwarps + PQ2_0_MC_WARPS - 1) / PQ2_0_MC_WARPS, 1, 1); - const dim3 block_dims(PQ2_0_MC_WARPS * WARP_SIZE, 1, 1); + const int ngroups = (nrows_x + rows - 1) / rows; + const int groups_per_cta = nwarps / splitk; + const dim3 block_nums((ngroups + groups_per_cta - 1) / groups_per_cta, 1, 1); + const dim3 block_dims(nwarps * WARP_SIZE, 1, 1); const ggml_cuda_kernel_launch_params lp = ggml_cuda_kernel_launch_params(block_nums, block_dims, 0, stream); - if (fusion.gate) { - ggml_cuda_kernel_launch(mul_mat_vec_pq2_0_mc, lp, vx, vy, dst, ncols_x, nrows_x, stride_row_x, stride_col_y, stride_col_dst, - (const float *) fusion.x_bias, fusion.gate, (const float *) fusion.gate_bias, fusion.glu_op, fusion.glu_limit); - } else if (fusion.x_bias) { - ggml_cuda_kernel_launch(mul_mat_vec_pq2_0_mc, lp, vx, vy, dst, ncols_x, nrows_x, stride_row_x, stride_col_y, stride_col_dst, + if constexpr (splitk == 1 && unroll == 1) { + if (fusion.gate) { + ggml_cuda_kernel_launch(mul_mat_vec_pq2_0_mc, lp, vx, vy, dst, ncols_x, nrows_x, stride_row_x, stride_col_y, stride_col_dst, + (const float *) fusion.x_bias, fusion.gate, (const float *) fusion.gate_bias, fusion.glu_op, fusion.glu_limit); + return; + } + } else { + GGML_ASSERT(!fusion.gate); + } + if (fusion.x_bias) { + ggml_cuda_kernel_launch(mul_mat_vec_pq2_0_mc, lp, vx, vy, dst, ncols_x, nrows_x, stride_row_x, stride_col_y, stride_col_dst, (const float *) fusion.x_bias, nullptr, nullptr, fusion.glu_op, fusion.glu_limit); } else { - ggml_cuda_kernel_launch(mul_mat_vec_pq2_0_mc, lp, vx, vy, dst, ncols_x, nrows_x, stride_row_x, stride_col_y, stride_col_dst, + ggml_cuda_kernel_launch(mul_mat_vec_pq2_0_mc, lp, vx, vy, dst, ncols_x, nrows_x, stride_row_x, stride_col_y, stride_col_dst, nullptr, nullptr, nullptr, fusion.glu_op, fusion.glu_limit); } } @@ -198,6 +269,16 @@ static void mul_mat_vec_pq2_0_mc_launch(const void * if constexpr (ncols == 1) { // Smaller Hopper projections need more CTAs; keep four-row reuse for the vocabulary head. const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + // Opt-in split-K for the two tall ffn_down shapes on Hopper (plain mat-vec only, not the gate variant): + // 2 rows per group, 2 warps per group, 8 warps per CTA. The K sum of a row is split in two halves, so the + // float summation order (and the output bits) differ from the default path. GGML_CUDA_PQ2_SPLITK=1; default off. + static const bool splitk_env = [] { const char * e = getenv("GGML_CUDA_PQ2_SPLITK"); return e != nullptr && atoi(e) != 0; }(); + if (cc == GGML_CUDA_CC_HOPPER && splitk_env && !fusion.gate && + ((ncols_x == 12288 && nrows_x == 4096) || (ncols_x == 17408 && nrows_x == 5120))) { + mul_mat_vec_pq2_0_mc_launch_rows(vx, vy, dst, ncols_x, nrows_x, stride_row_x, stride_col_y, + stride_col_dst, stream, fusion); + return; + } if (cc == GGML_CUDA_CC_HOPPER && nrows_x <= 17408) { mul_mat_vec_pq2_0_mc_launch_rows(vx, vy, dst, ncols_x, nrows_x, stride_row_x, stride_col_y, stride_col_dst, stream, fusion);