diff --git a/ggml/src/ggml-cuda/cp-async.cuh b/ggml/src/ggml-cuda/cp-async.cuh index 63d0c482..a9b47bd0 100644 --- a/ggml/src/ggml-cuda/cp-async.cuh +++ b/ggml/src/ggml-cuda/cp-async.cuh @@ -55,3 +55,23 @@ static __device__ __forceinline__ void cp_async_wait_all() { NO_DEVICE_CODE; #endif // CP_ASYNC_AVAILABLE } + +// Closes the group of cp.async copies issued so far by this thread. +static __device__ __forceinline__ void cp_async_commit_group() { +#ifdef CP_ASYNC_AVAILABLE + asm volatile("cp.async.commit_group;"); +#else + NO_DEVICE_CODE; +#endif // CP_ASYNC_AVAILABLE +} + +// Makes each thread wait until at most n of its committed cp.async groups are still pending. +// Like cp_async_wait_all, this does NOT synchronize threads. +template +static __device__ __forceinline__ void cp_async_wait_group() { +#ifdef CP_ASYNC_AVAILABLE + asm volatile("cp.async.wait_group %0;" : : "n"(n)); +#else + NO_DEVICE_CODE; +#endif // CP_ASYNC_AVAILABLE +} diff --git a/ggml/src/ggml-cuda/mmq-load-tiles.cuh b/ggml/src/ggml-cuda/mmq-load-tiles.cuh index 91087109..3b70f694 100644 --- a/ggml/src/ggml-cuda/mmq-load-tiles.cuh +++ b/ggml/src/ggml-cuda/mmq-load-tiles.cuh @@ -262,96 +262,125 @@ template static __device__ __forceinline_ } #if !defined(GGML_USE_HIP) -template -static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_ptq1_0(const char * __restrict__ x, - int * __restrict__ x_tile, - const int kbx0, - const int i_max, - const int stride) { - constexpr int warp_size = ggml_cuda_get_physical_warp_size(); - constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size; - constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback); - constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback); +// The PTQ1_0 x tile is loaded in two steps so the global loads for the next k-iteration can be issued before the +// current one is computed: fetch() reads the packed words and scales into registers, store() decodes them into the tile. +template struct ggml_cuda_mmq_ptq1_0_x_stage { + static constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + static constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size; + static constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback); + static constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback); + + static constexpr int blocks_per_iter = MMQ_ITER_K / QK_PTQ1_0; + static constexpr int threads_per_block = 8; + static constexpr int threads_per_row = blocks_per_iter * threads_per_block; + static constexpr int nrows = warp_size / threads_per_row; + static constexpr int scale_entries_per_block = QK_PTQ1_0 / QK8_1; + static constexpr int scale_entries_per_row = blocks_per_iter * scale_entries_per_block; + static constexpr int rows_per_warp = warp_size / scale_entries_per_row; + static constexpr int nqs = I / (nrows * nwarps); + static constexpr int nd = I / (rows_per_warp * nwarps); + + uint32_t packed[nqs]; + ggml_half d[nd]; + + static __device__ __forceinline__ int row_qs(const int i0, const int i_max) { + const int i = i0 + threadIdx.y * nrows + threadIdx.x / threads_per_row; + return fallback ? min(i, i_max) : i; + } + static __device__ __forceinline__ int row_d(const int i0, const int i_max) { + const int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / scale_entries_per_row; + return fallback ? min(i, i_max) : i; + } + + __device__ __forceinline__ void fetch(const char * __restrict__ x, const int kbx0, const int i_max, const int stride) { + const int txi = threadIdx.x % threads_per_row; + const int kbx = txi / threads_per_block; + const int lane = txi % threads_per_block; +# pragma unroll + for (int n = 0; n < nqs; ++n) { + const block_ptq1_0 * bxi = (const block_ptq1_0 *) x + kbx0 + row_qs(n * nrows * nwarps, i_max) * stride + kbx; + packed[n] = get_int_b4(bxi->qs, lane < 7 ? lane : 6); + } + const int scale_block = (threadIdx.x % scale_entries_per_row) / scale_entries_per_block; +# pragma unroll + for (int n = 0; n < nd; ++n) { + const block_ptq1_0 * bxi = + (const block_ptq1_0 *) x + kbx0 + row_d(n * nwarps * rows_per_warp, i_max) * stride + scale_block; + d[n] = bxi->d; + } + } + + __device__ __forceinline__ void store(int * __restrict__ x_tile, const int i_max) const { # if defined(TURING_MMA_AVAILABLE) - int * x_qs = (int *) x_tile; - float * x_df = (float *) (x_qs + 2 * MMQ_TILE_NE_K); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + 2 * MMQ_TILE_NE_K); # else - constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q8_0, I); - int * x_qs = (int *) x_tile; - float * x_df = (float *) (x_qs + txs.qs); + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q8_0, I); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + txs.qs); # endif - - constexpr int blocks_per_iter = MMQ_ITER_K / QK_PTQ1_0; - constexpr int threads_per_block = 8; - constexpr int threads_per_row = blocks_per_iter * threads_per_block; - constexpr int nrows = warp_size / threads_per_row; - - const int txi = threadIdx.x % threads_per_row; - const int kbx = txi / threads_per_block; - const int lane = txi % threads_per_block; + const int txi = threadIdx.x % threads_per_row; + const int kbx = txi / threads_per_block; + const int lane = txi % threads_per_block; # pragma unroll - for (int i0 = 0; i0 < I; i0 += nrows * nwarps) { - int i = i0 + threadIdx.y * nrows + threadIdx.x / threads_per_row; - if (fallback) { - i = min(i, i_max); - } - - const block_ptq1_0 * bxi = (const block_ptq1_0 *) x + kbx0 + i * stride + kbx; + for (int n = 0; n < nqs; ++n) { + const int i = row_qs(n * nrows * nwarps, i_max); # if defined(TURING_MMA_AVAILABLE) - int * row = x_qs + i * sram_stride + kbx * (QK_PTQ1_0 / 4); + int * row = x_qs + i * sram_stride + kbx * (QK_PTQ1_0 / 4); # else - int * row = x_qs + i * (2 * MMQ_TILE_NE_K + 1) + kbx * (QK_PTQ1_0 / 4); + int * row = x_qs + i * (2 * MMQ_TILE_NE_K + 1) + kbx * (QK_PTQ1_0 / 4); # endif - // Branch-free unpack: all 8 lanes run the same 5-iteration trit loop on their own 32-bit word. Only smem store offsets differ, so the warp does not diverge. - const uint32_t packed = get_int_b4(bxi->qs, lane < 7 ? lane : 6); - uint32_t v_lo = __byte_perm(packed, 0, 0x4140); // bytes 0,1 as 16-bit lanes - uint32_t v_hi = __byte_perm(packed, 0, 0x4342); // bytes 2,3 - const bool full_lane = lane < 6; - v_hi = full_lane ? v_hi : v_lo; // lane 6: both halves walk qh0/qh1 - const int dst_base = lane < 4 ? lane : 16 + lane; // lanes 4,5 -> 20,21 - const int dst_stride = lane < 4 ? 4 : 2; - int q[5]; + // Branch-free unpack: all 8 lanes run the same 5-iteration trit loop on their own 32-bit word. Only smem store offsets differ, so the warp does not diverge. + uint32_t v_lo = __byte_perm(packed[n], 0, 0x4140); // bytes 0,1 as 16-bit lanes + uint32_t v_hi = __byte_perm(packed[n], 0, 0x4342); // bytes 2,3 + const bool full_lane = lane < 6; + v_hi = full_lane ? v_hi : v_lo; // lane 6: both halves walk qh0/qh1 + const int dst_base = lane < 4 ? lane : 16 + lane; // lanes 4,5 -> 20,21 + const int dst_stride = lane < 4 ? 4 : 2; + int q[5]; # pragma unroll - for (int t = 0; t < 5; ++t) { - const uint32_t w_lo = v_lo * 3; - const uint32_t w_hi = v_hi * 3; - v_lo = w_lo & 0x00FF00FF; - v_hi = w_hi & 0x00FF00FF; - q[t] = __vsub4(__byte_perm(w_lo, w_hi, 0x7531), 0x01010101); - if (full_lane) { - row[dst_base + t * dst_stride] = q[t]; + for (int t = 0; t < 5; ++t) { + const uint32_t w_lo = v_lo * 3; + const uint32_t w_hi = v_hi * 3; + v_lo = w_lo & 0x00FF00FF; + v_hi = w_hi & 0x00FF00FF; + q[t] = __vsub4(__byte_perm(w_lo, w_hi, 0x7531), 0x01010101); + if (full_lane) { + row[dst_base + t * dst_stride] = q[t]; + } + } + if (lane == 6) { + // q[t] = {qh0.t, qh1.t, qh0.t, qh1.t}; the old layout is {qh0.t, qh1.t, qh0.t+1, qh1.t+1}. + row[30] = __byte_perm(q[0], q[1], 0x5410); + row[31] = __byte_perm(q[2], q[3], 0x5410); } } - if (lane == 6) { - // q[t] = {qh0.t, qh1.t, qh0.t, qh1.t}; the old layout is {qh0.t, qh1.t, qh0.t+1, qh1.t+1}. - row[30] = __byte_perm(q[0], q[1], 0x5410); - row[31] = __byte_perm(q[2], q[3], 0x5410); - } - } - - constexpr int scale_entries_per_block = QK_PTQ1_0 / QK8_1; - constexpr int scale_entries_per_row = blocks_per_iter * scale_entries_per_block; - constexpr int rows_per_warp = warp_size / scale_entries_per_row; - const int ksx = threadIdx.x % scale_entries_per_row; - const int scale_block = ksx / scale_entries_per_block; + const int ksx = threadIdx.x % scale_entries_per_row; # pragma unroll - for (int i0 = 0; i0 < I; i0 += nwarps * rows_per_warp) { - int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / scale_entries_per_row; - if (fallback) { - i = min(i, i_max); - } - - const block_ptq1_0 * bxi = (const block_ptq1_0 *) x + kbx0 + i * stride + scale_block; + for (int n = 0; n < nd; ++n) { + const int i = row_d(n * nwarps * rows_per_warp, i_max); # if defined(TURING_MMA_AVAILABLE) - x_df[i * sram_stride + ksx] = bxi->d; + x_df[i * sram_stride + ksx] = d[n]; # else - x_df[i * (2 * MMQ_TILE_NE_K / QI8_0) + i / (QI8_0 / 2) + ksx] = bxi->d; + x_df[i * (2 * MMQ_TILE_NE_K / QI8_0) + i / (QI8_0 / 2) + ksx] = d[n]; # endif + } } +}; + +template +static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_ptq1_0(const char * __restrict__ x, + int * __restrict__ x_tile, + const int kbx0, + const int i_max, + const int stride) { + ggml_cuda_mmq_ptq1_0_x_stage stage; + stage.fetch(x, kbx0, i_max, stride); + stage.store(x_tile, i_max); } #endif diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 462c2b71..2e610278 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -10604,6 +10604,14 @@ static std::vector> make_test_cases_perf() { } } + // Ternary Bonsai 2 27B (qwen35, PTQ1_0) projections at the prefill ubatch (MMQ path) + for (ggml_type type_a : {GGML_TYPE_PTQ1_0, GGML_TYPE_Q4_0}) { + for (auto [m, k] : std::vector>{{17408, 5120}, {5120, 17408}, {10240, 5120}, {5120, 6144}, + {6144, 5120}, {12288, 5120}, {1024, 5120}}) { + test_cases.emplace_back(new test_mul_mat(type_a, GGML_TYPE_F32, m, 1024, k, {1, 1}, {1, 1})); + } + } + // qwen3-30b-a3b for (int bs : {1, 4, 8, 32, 64, 128, 256, 512}) { for (ggml_type type_a : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_Q4_0, GGML_TYPE_Q8_0, GGML_TYPE_Q4_K, GGML_TYPE_Q6_K, GGML_TYPE_IQ2_XS}) { diff --git a/ggml/src/ggml-cuda/mmq.cuh b/ggml/src/ggml-cuda/mmq.cuh index ceb676ff..47de164c 100644 --- a/ggml/src/ggml-cuda/mmq.cuh +++ b/ggml/src/ggml-cuda/mmq.cuh @@ -930,7 +930,9 @@ static __device__ __forceinline__ void mul_mat_q_process_tile( constexpr int tile_y_stride = GGML_PAD(J*MMQ_TILE_Y_K, nwarps*warp_size); #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == GGML_CUDA_CC_DGX_SPARK constexpr bool async_buffer_y = - type == GGML_TYPE_Q1_0 || type == GGML_TYPE_Q2_0 || type == GGML_TYPE_PQ2_0; + type == GGML_TYPE_Q1_0 || type == GGML_TYPE_Q2_0 || type == GGML_TYPE_PQ2_0 || type == GGML_TYPE_PTQ1_0; +#elif defined(CP_ASYNC_AVAILABLE) && defined(TURING_MMA_AVAILABLE) + constexpr bool async_buffer_y = type == GGML_TYPE_PTQ1_0; #else constexpr bool async_buffer_y = false; #endif @@ -951,6 +953,63 @@ static __device__ __forceinline__ void mul_mat_q_process_tile( constexpr int sz = sizeof(block_q8_1_mmq) / sizeof(int); +#if !defined(GGML_USE_HIP) && defined(CP_ASYNC_AVAILABLE) && defined(TURING_MMA_AVAILABLE) + if constexpr (type == GGML_TYPE_PTQ1_0) { + // Software pipeline over the two y buffers: the y half for the next vec_dot is copied with cp.async while the + // current one is computed, and the next x tile is fetched into registers before the first vec_dot and only + // decoded into shared memory after the second, so no global load latency is exposed between k-iterations. + const int tid = threadIdx.y*warp_size + threadIdx.x; + const auto stage_y = [&](int * dst, const int kb) { + const char * src = reinterpret_cast(y + ncols_y * (kb * qk / ne_block) * sz); + char * dst_bytes = reinterpret_cast(dst); +#pragma unroll + for (int byte0 = 16*tid; byte0 < J*MMQ_TILE_Y_K*int(sizeof(int)); byte0 += 16*nwarps*warp_size) { + cp_async_cg_16<256>(ggml_cuda_cvta_generic_to_shared(dst_bytes + byte0), src + byte0); + } + cp_async_commit_group(); + }; + + ggml_cuda_mmq_ptq1_0_x_stage x_stage; + stage_y(tile_y, kb0_start); + x_stage.fetch(x, offset_x + kb0_start, tile_x_max_i, stride_row_x); + x_stage.store(tile_x, tile_x_max_i); + + for (int kb0 = kb0_start; kb0 < kb0_stop; kb0 += blocks_per_iter) { + const bool next = kb0 + blocks_per_iter < kb0_stop; + + stage_y(tile_y_next, kb0 + blocks_per_iter/2); + cp_async_wait_group<1>(); + __syncthreads(); + + if (next) { + x_stage.fetch(x, offset_x + kb0 + blocks_per_iter, tile_x_max_i, stride_row_x); + } + vec_dot(tile_x, tile_y, sum, 0); + + cp_async_wait_group<0>(); + __syncthreads(); + + if (next) { + stage_y(tile_y, kb0 + blocks_per_iter); + } + vec_dot(tile_x, tile_y_next, sum, MMQ_TILE_NE_K); + + __syncthreads(); + + if (next) { + x_stage.store(tile_x, tile_x_max_i); + } + } + + if (fixup) { + write_back(sum, ids_dst, tmp_fixup + blockIdx.x*(J*I), y_scale, I, I, J); + } else { + write_back(sum, ids_dst, dst, y_scale, stride_col_dst, tile_x_max_i, tile_y_max_j); + } + return; + } +#endif // !defined(GGML_USE_HIP) && defined(CP_ASYNC_AVAILABLE) && defined(TURING_MMA_AVAILABLE) + for (int kb0 = kb0_start; kb0 < kb0_stop; kb0 += blocks_per_iter) { if constexpr (async_buffer_y) { const char * by0 = reinterpret_cast( @@ -1488,8 +1547,9 @@ static size_t mmq_get_nbytes_shared(const ggml_cuda_mmq_config & config, const i const size_t nbs_x = ggml_cuda_mmq_get_nbytes_shared_x(config, cc); const size_t nbs_y = config.J * (sizeof(block_q8_1_mmq)); const size_t nbs_y_padded = GGML_PAD(nbs_y, config.nthreads*sizeof(int)); - const bool async_buffer_y = cc == GGML_CUDA_CC_DGX_SPARK && - (config.type == GGML_TYPE_Q1_0 || config.type == GGML_TYPE_Q2_0 || config.type == GGML_TYPE_PQ2_0); + const bool async_buffer_y = (cc == GGML_CUDA_CC_DGX_SPARK && + (config.type == GGML_TYPE_Q1_0 || config.type == GGML_TYPE_Q2_0 || config.type == GGML_TYPE_PQ2_0)) || + (config.type == GGML_TYPE_PTQ1_0 && cp_async_available(cc) && turing_mma_available(cc)); const int y_buffers = async_buffer_y ? 2 : 1; return nbs_ids + nbs_x + y_buffers*nbs_y_padded; }