diff --git a/ggml/src/ggml-cuda/mmvq-pq2_0.cuh b/ggml/src/ggml-cuda/mmvq-pq2_0.cuh index 3d30dce0e33c..b1c47ddb4e72 100644 --- a/ggml/src/ggml-cuda/mmvq-pq2_0.cuh +++ b/ggml/src/ggml-cuda/mmvq-pq2_0.cuh @@ -5,7 +5,7 @@ // permuted inside each 16-element group (position k*4+m holds element m*4+k), then one half2 per 32-block // with d and the raw int16 sum of q. With that, (code_word >> 2k) & 0x03030303 pairs the raw 2-bit codes 0..3 // with the right activation bytes, so there is no per-weight decode. The digit bias is one integer subtraction: -// sum(q*(c-1)) = sum(q*c) - sum(q). Each warp handles 4 rows and reuses the activation slice of a 32-block for all of them. +// sum(q*(c-1)) = sum(q*c) - sum(q). Each warp handles up to 4 rows and reuses the activation slice of a 32-block for all of them. #pragma once #include "common.cuh" @@ -18,7 +18,7 @@ // 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 +template __global__ void __launch_bounds__(PQ2_0_MC_WARPS * 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, @@ -29,7 +29,7 @@ mul_mat_vec_pq2_0_mc(const void * vx_, const void * vy_, float * dst_, 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) * PQ2_0_MC_ROWS; + const int row0 = ((blockIdx.x * blockDim.x + threadIdx.x) >> 5) * rows; if (row0 >= nrows_x) { return; } @@ -42,17 +42,17 @@ 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; - float acc[PQ2_0_MC_ROWS][ncols] = {}; - [[maybe_unused]] float accg[PQ2_0_MC_ROWS][ncols] = {}; + 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[PQ2_0_MC_ROWS], w1[PQ2_0_MC_ROWS]; - float d2[PQ2_0_MC_ROWS]; - [[maybe_unused]] uint32_t g0[PQ2_0_MC_ROWS], g1[PQ2_0_MC_ROWS]; - [[maybe_unused]] float e2[PQ2_0_MC_ROWS]; + uint32_t w0[rows], w1[rows]; + float d2[rows]; + [[maybe_unused]] uint32_t g0[rows], g1[rows]; + [[maybe_unused]] float e2[rows]; #pragma unroll - for (int r = 0; r < PQ2_0_MC_ROWS; ++r) { + 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 @@ -82,7 +82,7 @@ mul_mat_vec_pq2_0_mc(const void * vx_, const void * vy_, float * dst_, } #pragma unroll - for (int r = 0; r < PQ2_0_MC_ROWS; ++r) { + for (int r = 0; r < rows; ++r) { int t[8]; #pragma unroll for (int k = 0; k < 4; ++k) { @@ -129,7 +129,7 @@ mul_mat_vec_pq2_0_mc(const void * vx_, const void * vy_, float * dst_, ggml_cuda_pdl_lc(); #pragma unroll - for (int r = 0; r < PQ2_0_MC_ROWS; ++r) { + for (int r = 0; r < rows; ++r) { #pragma unroll for (int j = 0; j < ncols; ++j) { float v = warp_reduce_sum(acc[r][j]); @@ -163,27 +163,51 @@ mul_mat_vec_pq2_0_mc(const void * vx_, const void * vy_, float * dst_, } } -template -static void mul_mat_vec_pq2_0_mc_launch( +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 + PQ2_0_MC_ROWS - 1) / PQ2_0_MC_ROWS; + 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 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, + 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, + 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); } } +template +static void mul_mat_vec_pq2_0_mc_launch(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) { + 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; + 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); + return; + } + } + 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); +} + static void mul_mat_vec_pq2_0_mc_switch( const void * vx, const void * vy, float * dst, const int ncols_x, const int nrows_x, const int ncols_dst, const int stride_row_x, const int stride_col_y, const int stride_col_dst, cudaStream_t stream, diff --git a/ggml/src/ggml-cuda/norm.cu b/ggml/src/ggml-cuda/norm.cu index df270ad0c0c4..8e80a6e33f37 100644 --- a/ggml/src/ggml-cuda/norm.cu +++ b/ggml/src/ggml-cuda/norm.cu @@ -737,6 +737,10 @@ static __global__ void __launch_bounds__(1024, 1) add_rms_norm_fwht_f32( __shared__ float s_sum[32]; __shared__ float s_fwht[N]; + if constexpr (N == 5120) { + ggml_cuda_pdl_lc(); + } + const int64_t row = blockIdx.x; const int tid = threadIdx.x; const int lane = tid % ggml_cuda_get_physical_warp_size(); @@ -749,6 +753,10 @@ static __global__ void __launch_bounds__(1024, 1) add_rms_norm_fwht_f32( normed += row*N; } + if constexpr (N == 5120) { + ggml_cuda_pdl_sync(); + } + float x[NE]; float tmp = 0.0f; #pragma unroll @@ -827,8 +835,8 @@ static __global__ void __launch_bounds__(1024, 1) add_rms_norm_fwht_f32( // ncols: the row width; blk: the Hadamard block width (the transform matrix is blk x blk, applied to each blk-element // segment of the row; blk == ncols for a full-row transform) bool ggml_cuda_add_rms_norm_fwht_supported(const int ncols, const int blk) { - return (ncols == 1024 || ncols == 2048 || ncols == 4096) && - (blk == 512 || blk == 1024 || blk == 2048 || blk == 4096) && blk <= ncols; + return (ncols == 1024 || ncols == 2048 || ncols == 4096 || ncols == 5120) && + (blk == 512 || blk == 1024 || blk == 2048 || blk == 4096) && ncols % blk == 0; } void ggml_cuda_op_add_rms_norm_fwht_fused( @@ -854,10 +862,18 @@ void ggml_cuda_op_add_rms_norm_fwht_fused( cudaStream_t stream = ctx.stream(); float * normed_d = normed ? (float *) normed->data : nullptr; void * out = q8 ? q8 : dst->data; -#define ARF_LAUNCH(NN, BB, ADD, Q8, L) \ - add_rms_norm_fwht_f32<<<(unsigned) nrows, 1024, 0, stream>>>( \ - (const float *) a->data, ADD ? (const float *) b->data : nullptr, (const float *) weight->data, \ - (const float *) signs->data, ADD ? (float *) add->data : nullptr, normed_d, out, eps, fwht_scale) +#define ARF_ARGS(ADD) \ + (const float *) a->data, ADD ? (const float *) b->data : nullptr, (const float *) weight->data, \ + (const float *) signs->data, ADD ? (float *) add->data : nullptr, normed_d, out, eps, fwht_scale +#define ARF_LAUNCH(NN, BB, ADD, Q8, L) \ + do { \ + if constexpr (NN == 5120) { \ + const ggml_cuda_kernel_launch_params lp(dim3((unsigned) nrows), dim3(1024), 0, stream); \ + ggml_cuda_kernel_launch(add_rms_norm_fwht_f32, lp, ARF_ARGS(ADD)); \ + } else { \ + add_rms_norm_fwht_f32<<<(unsigned) nrows, 1024, 0, stream>>>(ARF_ARGS(ADD)); \ + } \ + } while (0) #define ARF_CASE_ADD(NN, BB, ADD) \ if (!q8) { \ ARF_LAUNCH(NN, BB, ADD, false, GGML_CUDA_Q8_1_AOS); \ @@ -880,12 +896,15 @@ void ggml_cuda_op_add_rms_norm_fwht_fused( case 1024: ARF_CASE_B(1024, 512) ARF_CASE_B(1024, 1024) break; case 2048: ARF_CASE_B(2048, 512) ARF_CASE_B(2048, 1024) ARF_CASE_B(2048, 2048) break; case 4096: ARF_CASE_B(4096, 512) ARF_CASE_B(4096, 1024) ARF_CASE_B(4096, 2048) ARF_CASE_B(4096, 4096) break; + case 5120: + ARF_CASE_B(5120, 512) ARF_CASE_B(5120, 1024) break; default: break; } GGML_ABORT("unsupported width %d / block %d", ncols, blk); #undef ARF_CASE_B #undef ARF_CASE_ADD #undef ARF_LAUNCH +#undef ARF_ARGS } void ggml_cuda_op_add_rms_norm_scale_fused( diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index c8f9c344ef98..db966b65815e 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -6208,6 +6208,8 @@ struct test_add_rms_norm_fwht : public test_case { bool run_whole_graph() override { return true; } + int n_eval() override { return n == 5120 ? 3 : 1; } + ggml_tensor * build_graph(ggml_context * ctx) override { const int64_t bw = blk ? blk : n; ggml_tensor * h = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, bw, bw); @@ -6302,6 +6304,8 @@ struct test_rms_norm_fwht_lookahead : public test_case { bool run_whole_graph() override { return true; } + int n_eval() override { return n == 5120 ? 3 : 1; } + ggml_tensor * build_graph(ggml_context * ctx) override { const int64_t bw = blk ? blk : n; ggml_tensor * h = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, bw, bw); @@ -6358,6 +6362,13 @@ struct test_rms_norm_fwht_lookahead : public test_case { } } ggml_backend_tensor_set(t, data.data(), 0, data.size() * sizeof(float)); + } else if (n == 5120 && strcmp(t->name, "w_alpha") == 0) { + // Keep the unrelated projection exact so it does not obscure transform errors. + std::vector data(ggml_nelements(t), 0.0f); + for (int64_t row = 0; row < t->ne[1]; ++row) { + data[row * n + row] = 1.0f; + } + ggml_backend_tensor_set(t, data.data(), 0, data.size() * sizeof(float)); } else if (strcmp(t->name, "s") == 0) { std::vector data(ggml_nelements(t)); for (size_t i = 0; i < data.size(); i++) { @@ -12300,6 +12311,15 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_rms_norm_fwht_lookahead(2048, n_tokens, GGML_TYPE_PQ2_0, blk)); } } + // Five 1024-element segments: the 27B row width is not a power of two. + for (int64_t blk : { 512, 1024 }) { + for (int64_t n_tokens : { 1, 2, 8, 9, 33 }) { + for (ggml_type type_mm : { GGML_TYPE_COUNT, GGML_TYPE_PQ2_0 }) { + test_cases.emplace_back(new test_add_rms_norm_fwht(5120, n_tokens, type_mm, blk)); + test_cases.emplace_back(new test_rms_norm_fwht_lookahead(5120, n_tokens, type_mm, blk)); + } + } + } // the same with PQ2_0 consumers: q8_1 quantized in the fused kernel (1 column: AoS, 3-8 columns: PQ2 layout) for (int64_t n : {1024, 4096}) { for (int64_t n_tokens : {1, 2, 4, 8, 9}) { @@ -12416,6 +12436,21 @@ static std::vector> make_test_cases_eval() { } } } + // Hopper two-row decode tiles: partial warps/CTAs and both sides of the row-count cutoff. + for (int64_t m : { 1, 7, 8, 9, 15, 16, 17, 17408, 17409 }) { + for (int64_t n : { 1, 2 }) { + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_PQ2_0, GGML_TYPE_F32, m, n, 1024, { 1, 1 }, { 1, 1 })); + } + } + for (bool with_bias : { false, true }) { + for (int64_t m : { 67, 17408 }) { + test_cases.emplace_back(new test_mul_mat_vec_fusion(GGML_TYPE_PQ2_0, GGML_GLU_OP_SWIGLU, 1, m, 5120, false, + 1, 1, false, with_bias, true, false, { 1, 1 })); + } + test_cases.emplace_back(new test_mul_mat_vec_fusion(GGML_TYPE_PQ2_0, GGML_GLU_OP_SWIGLU, 1, 5120, 17408, false, + 1, 1, false, with_bias, false, false, { 1, 1 })); + } + // PTQ1_0 / PQ2_0 integer-dot mat-vec: Bonsai-2 shapes, odd row counts (row tail), batches and multi-column B for (int64_t n : {1, 2, 3, 4, 5, 6, 7, 8}) { for (int64_t k : {1024, 5120, 6144, 17408}) {