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 <int n>
+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 <ggml_type type, int J, bool fallback> static __device__ __forceinline_
}
#if !defined(GGML_USE_HIP)
-template <ggml_type type, int J, bool fallback>
-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 <ggml_type type, int J, bool fallback> 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 <ggml_type type, int J, bool fallback>
+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<type, J, fallback> 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<std::unique_ptr<test_case>> 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<std::pair<int, int>>{{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<const char *>(y + ncols_y * (kb * qk / ne_block) * sz);
+ char * dst_bytes = reinterpret_cast<char *>(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<type, J, fallback> 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<const char *>(
@@ -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;
}