diff --git a/ggml/src/ggml-cuda/mmq-ptq1_0.cuh b/ggml/src/ggml-cuda/mmq-ptq1_0.cuh
new file mode 100644
index 00000000..9100c33d
--- /dev/null
+++ b/ggml/src/ggml-cuda/mmq-ptq1_0.cuh
@@ -0,0 +1,334 @@
+#pragma once
+
+// PTQ1_0 MMQ main loop for Ampere and newer (cp.async + int8 mma), used when every warp gets a multiple of 8 columns.
+// The k loop runs one 128-value block per stage over two shared-memory stages, with one barrier per stage. During
+// stage s the warps compute from buffer s%2, decode the raw x block of stage s+1 into the other buffer between their
+// mma tiles, and copy y for stage s+1 and raw x for stage s+2 with cp.async. B fragments come from ldmatrix.
+// Every output keeps the order of ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma with MMQ_Q8_1_DS_LAYOUT_D1 (an int32 sum per
+// block, then sum = fma(C*dA, dB, sum) in k order) and the tile shape stays I x J, so the stream-k split is unchanged:
+// results are bit-identical to the generic loop, which GGML_CUDA_NO_PTQ1_0_STAGED selects.
+
+#include "mmq.cuh"
+
+#if !defined(GGML_USE_HIP)
+
+namespace ggml_cuda_mmq_ptq1_0 {
+
+using namespace ggml_cuda_mma;
+
+static constexpr int x_stride = MMQ_TILE_NE_K + 4; // ints per x row: one decoded block plus padding against bank conflicts
+static constexpr int y_stride = MMQ_TILE_Y_K; // ints per y column: four scale slots and one block of int8 values
+static constexpr int x_words_per_row = sizeof(block_ptq1_0) / sizeof(int);
+
+template <int J, bool fallback> struct layout {
+ static constexpr int nthreads = ggml_cuda_mmq_get_nthreads(GGML_TYPE_PTQ1_0, J, fallback);
+ static constexpr int nwarps = nthreads / 32;
+ static constexpr int I = ggml_cuda_mmq_get_I(GGML_TYPE_PTQ1_0, J, fallback);
+ static constexpr int warps_i = I / 32; // each warp owns 32 rows: two 16-row mma tiles
+ static constexpr int warps_j = nwarps / warps_i;
+ static constexpr int cols = J / warps_j; // columns per warp
+ static constexpr int ntiles = cols / 8; // 8-column mma tiles per warp
+ static constexpr int ntasks = I * 8 / nthreads; // x words decoded per thread and stage, 8 lanes per row
+ static constexpr int x_ints = I * x_stride + I; // decoded values, then one float scale per row
+ static constexpr int raw_ints = I * x_words_per_row; // one stage of x as stored in VRAM
+ static constexpr int ncopies = (raw_ints + nthreads - 1) / nthreads;
+ static constexpr int y_ints = J * y_stride;
+
+ static_assert(ggml_cuda_mmq_ptq1_0_staged(J, fallback), "unsupported PTQ1_0 MMQ layout");
+};
+
+static __device__ __forceinline__ void ldmatrix_x4(int (&r)[4], const int * p) {
+#ifdef TURING_MMA_AVAILABLE
+ asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];"
+ : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3])
+ : "r"(ggml_cuda_cvta_generic_to_shared((void *) p)));
+#else
+ GGML_UNUSED_VARS(r, p);
+ NO_DEVICE_CODE;
+#endif // TURING_MMA_AVAILABLE
+}
+
+static __device__ __forceinline__ void cp_async_ca_4(const unsigned int dst, const void * src) {
+#ifdef CP_ASYNC_AVAILABLE
+ asm volatile("cp.async.ca.shared.global [%0], [%1], 4;" : : "r"(dst), "l"(src));
+#else
+ GGML_UNUSED_VARS(dst, src);
+ NO_DEVICE_CODE;
+#endif // CP_ASYNC_AVAILABLE
+}
+
+// Decodes word n of this thread's share of the raw x stage (block_ptq1_0 words, x_words_per_row per row) into xs.
+// Same unpack and layout as ggml_cuda_mmq_ptq1_0_x_stage::store; lane 6 also holds the scale, lane 7 has nothing to do.
+template <int J, bool fallback>
+static __device__ __forceinline__ void decode_word(const int * __restrict__ raw, int * __restrict__ xs, const int n) {
+ using L = layout<J, fallback>;
+ const int lane = threadIdx.x % 8;
+ const int i = (n * L::nthreads + threadIdx.y * 32 + threadIdx.x) / 8;
+ const uint32_t w = raw[i * x_words_per_row + min(lane, x_words_per_row - 1)];
+ int * row = xs + i * x_stride;
+
+ // Branch-free so that decoding does not split the mma code into basic blocks: lanes 6 and 7 send the stores they
+ // do not need to the padding at the end of the row, which ldmatrix never reads.
+ uint32_t v_lo = __byte_perm(w, 0, 0x4140);
+ uint32_t v_hi = __byte_perm(w, 0, 0x4342);
+ const bool full_lane = lane < 6;
+ v_hi = full_lane ? v_hi : v_lo;
+ const int dst_base = lane < 4 ? lane : lane < 6 ? 16 + lane : MMQ_TILE_NE_K;
+ const int dst_stride = lane < 4 ? 4 : lane < 6 ? 2 : 0;
+ 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;
+ // Each byte is a trit in {0, 1, 2}: adding 0x7F never carries, and the xor gives trit - 1 as int8.
+ q[t] = (__byte_perm(w_lo, w_hi, 0x7531) + 0x7F7F7F7F) ^ 0x80808080;
+ row[dst_base + t * dst_stride] = q[t];
+ }
+ // q[t] = {qh0.t, qh1.t, qh0.t, qh1.t} on lane 6; the layout wants {qh0.t, qh1.t, qh0.t+1, qh1.t+1}.
+ const bool qh_lane = lane == 6;
+ row[qh_lane ? 30 : MMQ_TILE_NE_K + 1] = __byte_perm(q[0], q[1], 0x5410);
+ row[qh_lane ? 31 : MMQ_TILE_NE_K + 2] = __byte_perm(q[2], q[3], 0x5410);
+ float * d_dst = qh_lane ? (float *) (xs + L::I * x_stride) + i : (float *) (row + MMQ_TILE_NE_K + 3);
+ *d_dst = __half2float(__ushort_as_half((unsigned short) (w >> 16)));
+}
+
+template <int J, bool fallback, bool decode>
+static __device__ __forceinline__ void stage(const int * __restrict__ xs, const int * __restrict__ ys, float * __restrict__ sum,
+ const int * __restrict__ raw_next, int * __restrict__ xs_next) {
+ using L = layout<J, fallback>;
+ typedef tile<16, 8, int> tile_A;
+ typedef tile< 8, 8, int> tile_B;
+ typedef tile<16, 8, int> tile_C;
+
+ const int wi = threadIdx.y / L::warps_j;
+ const int wj = threadIdx.y % L::warps_j;
+ const int i0 = wi * 32;
+
+ const float * xd = (const float *) (xs + L::I * x_stride);
+ const float * yd = (const float *) ys;
+ const int * yq = ys + 4;
+
+ tile_A A[2][4];
+ float dA[2][2];
+#pragma unroll
+ for (int m = 0; m < 2; ++m) {
+#pragma unroll
+ for (int k = 0; k < 4; ++k) {
+ load_ldmatrix(A[m][k], xs + (i0 + m * 16) * x_stride + k * QI8_0, x_stride);
+ }
+#pragma unroll
+ for (int h = 0; h < 2; ++h) {
+ dA[m][h] = xd[i0 + m * 16 + tile_C::get_i(2 * h)];
+ }
+ }
+
+ // Tiles are processed in groups with the k loop outside, so that each IMMA has 2*group independent accumulators
+ // to alternate with instead of waiting on its own previous result.
+ constexpr int group = L::ntiles % 4 == 0 ? 4 : L::ntiles % 3 == 0 ? 3 : L::ntiles % 2 == 0 ? 2 : 1;
+ static_assert(L::ntiles % group == 0, "bad tile group");
+
+#pragma unroll
+ for (int t0 = 0; t0 < L::ntiles; t0 += group) {
+ // b[g][h][q] holds ints 4*q + threadIdx.x%4 (+16 for h = 1) of column j0 + threadIdx.x/4 of tile t0 + g: the
+ // B fragments of k-steps 2*h and 2*h + 1.
+ int b[group][2][4];
+#pragma unroll
+ for (int g = 0; g < group; ++g) {
+ const int j0 = wj * L::cols + (t0 + g) * 8;
+#pragma unroll
+ for (int h = 0; h < 2; ++h) {
+ ldmatrix_x4(b[g][h], yq + (j0 + threadIdx.x % 8) * y_stride + (threadIdx.x / 8) * 4 + h * 16);
+ }
+ }
+
+ tile_C C[group][2];
+#pragma unroll
+ for (int k = 0; k < 4; ++k) {
+#pragma unroll
+ for (int g = 0; g < group; ++g) {
+ tile_B B;
+ B.x[0] = b[g][k / 2][2 * (k % 2) + 0];
+ B.x[1] = b[g][k / 2][2 * (k % 2) + 1];
+#pragma unroll
+ for (int m = 0; m < 2; ++m) {
+ mma(C[g][m], A[m][k], B);
+ }
+ }
+ }
+
+#pragma unroll
+ for (int g = 0; g < group; ++g) {
+ const int t = t0 + g;
+ const int j0 = wj * L::cols + t * 8;
+ float dB[2];
+#pragma unroll
+ for (int l = 0; l < 2; ++l) {
+ dB[l] = yd[(j0 + tile_C::get_j(l)) * y_stride];
+ }
+#pragma unroll
+ for (int m = 0; m < 2; ++m) {
+#pragma unroll
+ for (int l = 0; l < tile_C::ne; ++l) {
+ // Written as the fma the generic vec_dot contracts to, so that the two stay bit-identical.
+ float & s = sum[(t * 2 + m) * tile_C::ne + l];
+ s = fmaf(C[g][m].x[l] * dA[m][l / 2], dB[l % 2], s);
+ }
+ }
+ }
+
+ if constexpr (decode) {
+#pragma unroll
+ for (int n = 0; n < L::ntasks; ++n) {
+ if ((n * (L::ntiles / group)) / L::ntasks == t0 / group) {
+ decode_word<J, fallback>(raw_next, xs_next, n);
+ }
+ }
+ }
+ }
+}
+
+template <int J, bool fallback>
+static __device__ __forceinline__ void write_back(const float * __restrict__ sum, const int * __restrict__ ids_dst,
+ float * __restrict__ dst, const int stride, const int i_max, const int j_max) {
+ using L = layout<J, fallback>;
+ typedef tile<16, 8, int> tile_C;
+
+ const int i0 = (threadIdx.y / L::warps_j) * 32;
+ const int wj = threadIdx.y % L::warps_j;
+
+#pragma unroll
+ for (int t = 0; t < L::ntiles; ++t) {
+#pragma unroll
+ for (int m = 0; m < 2; ++m) {
+#pragma unroll
+ for (int l = 0; l < tile_C::ne; ++l) {
+ const int j = wj * L::cols + t * 8 + tile_C::get_j(l);
+ if (j > j_max) {
+ continue;
+ }
+ const int i = i0 + m * 16 + tile_C::get_i(l);
+ if (fallback && i > i_max) {
+ continue;
+ }
+ dst[ids_dst[j] * stride + i] = sum[(t * 2 + m) * tile_C::ne + l];
+ }
+ }
+ }
+}
+
+} // namespace ggml_cuda_mmq_ptq1_0
+
+template <int J, bool fallback, bool fixup>
+static __device__ __forceinline__ void mul_mat_q_process_tile_ptq1_0(
+ const char * __restrict__ x, const int offset_x, const int * __restrict__ y,
+ const int * __restrict__ ids_dst, float * __restrict__ dst, float * __restrict__ tmp_fixup,
+ const int stride_row_x, const int ncols_y, const int stride_col_dst,
+ const int tile_x_max_i, const int tile_y_max_j, const int kb0_start, const int kb0_stop) {
+#if defined(CP_ASYNC_AVAILABLE) && defined(TURING_MMA_AVAILABLE)
+ using namespace ggml_cuda_mmq_ptq1_0;
+ using L = layout<J, fallback>;
+
+ extern __shared__ int data_mul_mat_q[];
+ int * ys[2];
+ int * xs[2];
+ int * raw[2];
+ ys[0] = data_mul_mat_q + J;
+ ys[1] = ys[0] + L::y_ints;
+ xs[0] = ys[1] + L::y_ints;
+ xs[1] = xs[0] + L::x_ints;
+ raw[0] = xs[1] + L::x_ints;
+ raw[1] = raw[0] + L::raw_ints;
+
+ constexpr int sz = sizeof(block_q8_1_mmq) / sizeof(int);
+ const int tid = threadIdx.y * 32 + threadIdx.x;
+
+ // Source offsets in ints of this thread's raw x words, relative to the block at kb = 0.
+ int x_src[L::ncopies];
+#pragma unroll
+ for (int c = 0; c < L::ncopies; ++c) {
+ const int l = min(c * L::nthreads + tid, L::raw_ints - 1);
+ const int i = fallback ? min(l / x_words_per_row, tile_x_max_i) : l / x_words_per_row;
+ x_src[c] = (offset_x + i * stride_row_x) * x_words_per_row + l % x_words_per_row;
+ }
+ const int * x_words = (const int *) x;
+
+ // One cp.async group per stage: the y block for the next stage and the raw x block for the one after.
+ const auto copy_y = [&](int * dst_y, const int kb) {
+ const char * src = reinterpret_cast<const char *>(y + ncols_y * kb * sz);
+ char * dst_bytes = reinterpret_cast<char *>(dst_y);
+#pragma unroll
+ for (int byte0 = 16 * tid; byte0 < L::y_ints * int(sizeof(int)); byte0 += 16 * L::nthreads) {
+ cp_async_cg_16<256>(ggml_cuda_cvta_generic_to_shared(dst_bytes + byte0), src + byte0);
+ }
+ };
+ const auto copy_x = [&](int * dst_raw, const int kb) {
+#pragma unroll
+ for (int c = 0; c < L::ncopies; ++c) {
+ if (L::raw_ints % L::nthreads == 0 || c * L::nthreads + tid < L::raw_ints) {
+ cp_async_ca_4(ggml_cuda_cvta_generic_to_shared(dst_raw + c * L::nthreads + tid),
+ x_words + x_src[c] + kb * x_words_per_row);
+ }
+ }
+ };
+
+ float sum[L::ntiles * 2 * tile<16, 8, int>::ne] = {0.0f};
+
+ const int nstages = kb0_stop - kb0_start;
+
+ copy_y(ys[0], kb0_start);
+ copy_x(raw[0], kb0_start);
+ if (nstages > 1) {
+ copy_x(raw[1], kb0_start + 1);
+ }
+ cp_async_commit_group();
+ cp_async_wait_all();
+ __syncthreads();
+#pragma unroll
+ for (int n = 0; n < L::ntasks; ++n) {
+ decode_word<J, fallback>(raw[0], xs[0], n);
+ }
+
+ // Stage s computes from buffer s%2 and decodes raw x for stage s+1 into the other one. It copies y for stage s+1
+ // and raw x for stage s+2; the barrier at the top of the next stage makes both visible.
+ const auto run_stage = [&](const int s, const int p) {
+ __syncthreads();
+ const bool next = s + 1 < nstages;
+ if (next) {
+ copy_y(ys[1 - p], kb0_start + s + 1);
+ }
+ if (s + 2 < nstages) {
+ copy_x(raw[p], kb0_start + s + 2);
+ }
+ cp_async_commit_group();
+ if (next) {
+ stage<J, fallback, true>(xs[p], ys[p], sum, raw[1 - p], xs[1 - p]);
+ } else {
+ stage<J, fallback, false>(xs[p], ys[p], sum, raw[1 - p], xs[1 - p]);
+ }
+ cp_async_wait_all();
+ };
+
+ int s = 0;
+ for (; s + 1 < nstages; s += 2) {
+ run_stage(s + 0, 0);
+ run_stage(s + 1, 1);
+ }
+ if (s < nstages) {
+ run_stage(s, 0);
+ }
+
+ if (fixup) {
+ write_back<J, fallback>(sum, ids_dst, tmp_fixup + blockIdx.x * (J * L::I), L::I, L::I, J);
+ } else {
+ write_back<J, fallback>(sum, ids_dst, dst, stride_col_dst, tile_x_max_i, tile_y_max_j);
+ }
+#else
+ GGML_UNUSED_VARS(x, offset_x, y, ids_dst, dst, tmp_fixup, stride_row_x, ncols_y, stride_col_dst,
+ tile_x_max_i, tile_y_max_j, kb0_start, kb0_stop);
+ NO_DEVICE_CODE;
+#endif // defined(CP_ASYNC_AVAILABLE) && defined(TURING_MMA_AVAILABLE)
+}
+
+#endif // !defined(GGML_USE_HIP)
diff --git a/ggml/src/ggml-cuda/mmq.cuh b/ggml/src/ggml-cuda/mmq.cuh
index 99b05d86..d6fe11f1 100644
--- a/ggml/src/ggml-cuda/mmq.cuh
+++ b/ggml/src/ggml-cuda/mmq.cuh
@@ -919,7 +919,23 @@ static constexpr __device__ ggml_cuda_mmq_write_back_t ggml_cuda_mmq_get_write_b
// ---------------------------------------------------------------------------------------------
-template <ggml_type type, int J, bool fallback, bool fixup>
+#if !defined(GGML_USE_HIP)
+// PTQ1_0 on Ampere and newer runs the staged main loop in mmq-ptq1_0.cuh when every warp gets a multiple of 8 columns.
+static constexpr __device__ bool ggml_cuda_mmq_ptq1_0_staged(const int J, const bool fallback) {
+ const int nwarps = ggml_cuda_mmq_get_nthreads(GGML_TYPE_PTQ1_0, J, fallback) / 32;
+ const int warps_i = ggml_cuda_mmq_get_I(GGML_TYPE_PTQ1_0, J, fallback) / 32;
+ return warps_i > 0 && nwarps % warps_i == 0 && nwarps / warps_i > 0 && J % (8 * (nwarps / warps_i)) == 0;
+}
+
+template <int J, bool fallback, bool fixup>
+static __device__ __forceinline__ void mul_mat_q_process_tile_ptq1_0(
+ const char * __restrict__ x, const int offset_x, const int * __restrict__ y,
+ const int * __restrict__ ids_dst, float * __restrict__ dst, float * __restrict__ tmp_fixup,
+ const int stride_row_x, const int ncols_y, const int stride_col_dst,
+ const int tile_x_max_i, const int tile_y_max_j, const int kb0_start, const int kb0_stop);
+#endif // !defined(GGML_USE_HIP)
+
+template <ggml_type type, int J, bool fallback, bool fixup, bool staged>
static __device__ __forceinline__ void mul_mat_q_process_tile(
const char * __restrict__ x, const int offset_x, const int * __restrict__ y,
const int * __restrict__ ids_dst, float * __restrict__ dst, float * __restrict__ tmp_fixup,
@@ -964,7 +980,11 @@ 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) {
+ if constexpr (staged && type == GGML_TYPE_PTQ1_0 && ggml_cuda_mmq_ptq1_0_staged(J, fallback)) {
+ mul_mat_q_process_tile_ptq1_0<J, fallback, fixup>(x, offset_x, y, ids_dst, dst, tmp_fixup, stride_row_x, ncols_y,
+ stride_col_dst, tile_x_max_i, tile_y_max_j, kb0_start, kb0_stop);
+ return;
+ } else 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.
@@ -1118,7 +1138,8 @@ static __device__ __forceinline__ void mul_mat_q_process_tile(
// The mul_mat_q kernel implements "stream-k" work partitioning as described in https://arxiv.org/abs/2301.03598
-template <ggml_type type, int J, bool fallback>
+// staged = false keeps PTQ1_0 on the older main loop (GGML_CUDA_NO_PTQ1_0_STAGED); other types ignore it.
+template <ggml_type type, int J, bool fallback, bool staged = true>
__launch_bounds__(ggml_cuda_mmq_get_nthreads(type, J, fallback), ggml_cuda_mmq_get_occupancy(type, J, fallback))
static __global__ void mul_mat_q(
const char * __restrict__ x, const int * __restrict__ y, const int32_t * __restrict__ ids_dst,
@@ -1221,7 +1242,7 @@ static __global__ void mul_mat_q(
const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*I*stride_row_x;
constexpr bool fixup = false;
- mul_mat_q_process_tile<type, J, fallback, fixup>
+ mul_mat_q_process_tile<type, J, fallback, fixup, staged>
(x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile,
stride_row_x, ncols_y, stride_col_dst,
tile_x_max_i, tile_y_max_j, 0, blocks_per_ne00.z);
@@ -1315,7 +1336,7 @@ static __global__ void mul_mat_q(
const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*I*stride_row_x;
constexpr bool fixup = false; // All but (potentially) the last iterations write their data to dst rather than the fixup buffer.
- mul_mat_q_process_tile<type, J, fallback, fixup>
+ mul_mat_q_process_tile<type, J, fallback, fixup, staged>
(x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile,
stride_row_x, ncols_y, stride_col_dst,
tile_x_max_i, tile_y_max_j, kb0_start, kb0_stop);
@@ -1399,7 +1420,7 @@ static __global__ void mul_mat_q(
const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*I*stride_row_x;
constexpr bool fixup = true; // Last index writes its data to fixup buffer to avoid data races with other blocks.
- mul_mat_q_process_tile<type, J, fallback, fixup>
+ mul_mat_q_process_tile<type, J, fallback, fixup, staged>
(x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile,
stride_row_x, ncols_y, stride_col_dst,
tile_x_max_i, tile_y_max_j, kb0_start, kb0_stop);
@@ -1561,11 +1582,13 @@ static size_t mmq_get_nbytes_shared(const ggml_cuda_mmq_config & config, const i
(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;
+ // The staged PTQ1_0 loop (mmq-ptq1_0.cuh) also stages two raw x blocks per row with cp.async.
+ const size_t nbs_x_raw = config.type == GGML_TYPE_PTQ1_0 && async_buffer_y ? 2*config.I*sizeof(block_ptq1_0) : 0;
+ return nbs_ids + nbs_x + y_buffers*nbs_y_padded + nbs_x_raw;
}
-template <ggml_type type, int J, bool fallback>
-static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) {
+template <ggml_type type, int J, bool fallback, bool staged>
+static void launch_mul_mat_q_impl(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) {
const int id = ggml_cuda_get_device();
const int cc = ggml_cuda_info().devices[id].cc;
const int nsm = ggml_cuda_info().devices[id].nsm;
@@ -1578,8 +1601,8 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a
const dim3 block_dims(warp_size, nwarps, 1);
- CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q<type, J, false>), nbytes_shared);
- CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q<type, J, true>), nbytes_shared);
+ CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q<type, J, false, staged>), nbytes_shared);
+ CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q<type, J, true, staged>), nbytes_shared);
const int nty = (args.nrows_x + config.I - 1) / config.I;
const int ntx = (args.ncols_max + config.J - 1) / config.J;
@@ -1599,7 +1622,7 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a
const uint3 sample_ratio_fd = init_fastdiv_values(sample_ratio);
if (!ggml_cuda_mmq_get_stream_k(type, J, fallback, cc)) {
- mul_mat_q<type, J, fallback><<<block_nums_xy_tiling, block_dims, nbytes_shared, stream>>>
+ mul_mat_q<type, J, fallback, staged><<<block_nums_xy_tiling, block_dims, nbytes_shared, stream>>>
(args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, args.y_scale,
blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst,
channel_ratio_fd, nchannels_y_fd, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst,
@@ -1628,7 +1651,7 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a
const dim3 block_nums_fixup(block_nums_stream_k.x, config.I/warp_size, 1);
const dim3 block_dims_fixup(block_dims.x, block_dims.y/2, block_dims.z);
- mul_mat_q<type, J, fallback><<<block_nums_stream_k, block_dims, nbytes_shared, stream>>>
+ mul_mat_q<type, J, fallback, staged><<<block_nums_stream_k, block_dims, nbytes_shared, stream>>>
(args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup.ptr, args.y_scale,
blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst,
channel_ratio_fd, nchannels_y_fd, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst,
@@ -1646,6 +1669,18 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a
ntx_fd);
}
+template <ggml_type type, int J, bool fallback>
+static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) {
+ if constexpr (type == GGML_TYPE_PTQ1_0) {
+ static const bool staged = getenv("GGML_CUDA_NO_PTQ1_0_STAGED") == nullptr;
+ if (!staged) {
+ launch_mul_mat_q_impl<type, J, fallback, false>(ctx, args, stream);
+ return;
+ }
+ }
+ launch_mul_mat_q_impl<type, J, fallback, true>(ctx, args, stream);
+}
+
template <ggml_type type, bool fallback>
void mul_mat_q_switch_J(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) {
const int id = ggml_cuda_get_device();
diff --git a/ggml/src/ggml-cuda/template-instances/generate_cu_files.py b/ggml/src/ggml-cuda/template-instances/generate_cu_files.py
index 073e4154..3958a045 100755
--- a/ggml/src/ggml-cuda/template-instances/generate_cu_files.py
+++ b/ggml/src/ggml-cuda/template-instances/generate_cu_files.py
@@ -60,6 +60,7 @@ DECL_MMQ_CASE({type});
SOURCE_MMQ_CUDA_ONLY = """// This file has been autogenerated by generate_cu_files.py, do not edit manually.
#include "../mmq.cuh"
+#include "../mmq-ptq1_0.cuh"
#if !defined(GGML_USE_HIP)
DECL_MMQ_CASE({type});
diff --git a/ggml/src/ggml-cuda/template-instances/mmq-instance-ptq1_0.cu b/ggml/src/ggml-cuda/template-instances/mmq-instance-ptq1_0.cu
index 692e9529..2011f3ac 100644
--- a/ggml/src/ggml-cuda/template-instances/mmq-instance-ptq1_0.cu
+++ b/ggml/src/ggml-cuda/template-instances/mmq-instance-ptq1_0.cu
@@ -1,6 +1,7 @@
// This file has been autogenerated by generate_cu_files.py, do not edit manually.
#include "../mmq.cuh"
+#include "../mmq-ptq1_0.cuh"
#if !defined(GGML_USE_HIP)
DECL_MMQ_CASE(GGML_TYPE_PTQ1_0);