cuda: PQ2_0 mat-vec kernel for 3-8 columns on Ada

ggml/src/ggml-cuda/common.cuh

@@ -1055,6 +1055,7 @@ enum ggml_cuda_q8_1_layout : int {
     GGML_CUDA_Q8_1_AOS      = 0, // plain block_q8_1 array (every type except the PTQ1_0 cases below)
     GGML_CUDA_Q8_1_SOA_ISUM = 1, // PTQ1_0, one column: warp-transposed, exact int sums (ggml_cuda_ptq1_q8_word)
     GGML_CUDA_Q8_1_PT       = 2, // PTQ1_0, 2-8 columns or MoE ids: planar-transposed (mmvq-ptq1_0.cuh)
+    GGML_CUDA_Q8_1_PQ2      = 3, // PQ2_0, 3-8 columns, plain 2D, Ada: permuted qs plus half2 (d, int sum) per block (mmvq-pq2_0.cuh)
 };
 
 // Column-count helper, not the layout decision. 2-8 columns and MoE (ids) take the planar
@@ -1285,7 +1286,23 @@ int ggml_cuda_get_device();
 // Ampere (sm_80/86, including the 3060/3090/170HX): the #218 PT kernel wins at one column
 // too (+5.9% tg128 vs SoA on a 3060). Ada and newer keep SOA_ISUM at one column (4070 win).
 // Lives here, after ggml_cuda_info(), because the body reads the current device's cc.
-static inline ggml_cuda_q8_1_layout ggml_cuda_q8_1_layout_host(ggml_type type_src0, int ncols_dst, bool has_ids) {
+// PQ2_0 multi-column mat-vec (mmvq-pq2_0.cuh): Ada only, plain 2D (no ids, no batch dims), 3-8 columns.
+// Batch invariance keeps the generic kernel.
+static inline bool ggml_cuda_pq2_multicol(int ncols_dst) {
+#if defined(GGML_USE_HIP)
+    GGML_UNUSED(ncols_dst);
+    return false;
+#else
+    const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc;
+    return !ggml_cuda_batch_invariant() && GGML_CUDA_CC_IS_NVIDIA(cc) && cc == GGML_CUDA_CC_ADA_LOVELACE &&
+        ncols_dst >= 3 && ncols_dst <= 8;
+#endif
+}
+
+static inline ggml_cuda_q8_1_layout ggml_cuda_q8_1_layout_host(ggml_type type_src0, int ncols_dst, bool has_ids, bool plain_2d = true) {
+    if (type_src0 == GGML_TYPE_PQ2_0 && plain_2d && !has_ids && ggml_cuda_pq2_multicol(ncols_dst)) {
+        return GGML_CUDA_Q8_1_PQ2;
+    }
     const ggml_cuda_q8_1_layout l = ggml_cuda_q8_1_layout_for(type_src0, ncols_dst, has_ids);
     if (l == GGML_CUDA_Q8_1_SOA_ISUM && ggml_cuda_batch_invariant()) {
         return GGML_CUDA_Q8_1_PT;

ggml/src/ggml-cuda/mmvq-pq2_0.cuh

@@ -0,0 +1,120 @@
+// PQ2_0 mat-vec for 3-8 columns (plain 2D), tuned on Ada.
+//
+// Why: the generic mmvq kernel loses most of its DRAM throughput from 3 columns on (RTX 4070, Bonsai 2 27B shapes:
+// 44% of peak at 4 columns). Here the activations use GGML_CUDA_Q8_1_PQ2: per column the qs bytes come first,
+// 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.
+#pragma once
+
+#include "common.cuh"
+
+#define PQ2_0_MC_ROWS  4
+#define PQ2_0_MC_WARPS 4
+
+template <int ncols>
+__global__ void __launch_bounds__(PQ2_0_MC_WARPS * 32, 3)
+mul_mat_vec_pq2_0_mc(const void * __restrict__ vx, const void * __restrict__ vy, float * __restrict__ 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 int lane = threadIdx.x & 31;
+    const int row0 = ((blockIdx.x * blockDim.x + threadIdx.x) >> 5) * PQ2_0_MC_ROWS;
+    if (row0 >= nrows_x) {
+        return;
+    }
+
+    const int     nchunk    = ncols_x / QK8_1;
+    const size_t  col_bytes = (size_t) stride_col_y * sizeof(block_q8_1);
+    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] = {};
+
+    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];
+#pragma unroll
+        for (int r = 0; r < PQ2_0_MC_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);
+        }
+
+        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));
+        }
+
+#pragma unroll
+        for (int r = 0; r < PQ2_0_MC_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);
+            }
+#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]);
+            }
+        }
+    }
+
+#pragma unroll
+    for (int r = 0; r < PQ2_0_MC_ROWS; ++r) {
+#pragma unroll
+        for (int j = 0; j < ncols; ++j) {
+            const float v = warp_reduce_sum<WARP_SIZE>(acc[r][j]);
+            const int row = row0 + r;
+            if (lane == 0 && row < nrows_x) {
+                dst[(int64_t) j*stride_col_dst + row] = v;
+            }
+        }
+    }
+}
+
+template <int ncols>
+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 int nwarps = (nrows_x + PQ2_0_MC_ROWS - 1) / PQ2_0_MC_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);
+    ggml_cuda_kernel_launch(mul_mat_vec_pq2_0_mc<ncols>, lp, vx, vy, dst, ncols_x, nrows_x, stride_row_x, stride_col_y, stride_col_dst);
+}
+
+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) {
+    GGML_ASSERT(ncols_x % QK_PQ2_0 == 0);
+    switch (ncols_dst) {
+        case 3: mul_mat_vec_pq2_0_mc_launch<3>(vx, vy, dst, ncols_x, nrows_x, stride_row_x, stride_col_y, stride_col_dst, stream); break;
+        case 4: mul_mat_vec_pq2_0_mc_launch<4>(vx, vy, dst, ncols_x, nrows_x, stride_row_x, stride_col_y, stride_col_dst, stream); break;
+        case 5: mul_mat_vec_pq2_0_mc_launch<5>(vx, vy, dst, ncols_x, nrows_x, stride_row_x, stride_col_y, stride_col_dst, stream); break;
+        case 6: mul_mat_vec_pq2_0_mc_launch<6>(vx, vy, dst, ncols_x, nrows_x, stride_row_x, stride_col_y, stride_col_dst, stream); break;
+        case 7: mul_mat_vec_pq2_0_mc_launch<7>(vx, vy, dst, ncols_x, nrows_x, stride_row_x, stride_col_y, stride_col_dst, stream); break;
+        case 8: mul_mat_vec_pq2_0_mc_launch<8>(vx, vy, dst, ncols_x, nrows_x, stride_row_x, stride_col_y, stride_col_dst, stream); break;
+        default: GGML_ABORT("unsupported column count for the PQ2_0 multi-column kernel");
+    }
+}

ggml/src/ggml-cuda/mmvq.cu

@@ -1,5 +1,6 @@
 #include "mmvq.cuh"
 #include "mmvq-ptq1_0.cuh"
+#include "mmvq-pq2_0.cuh"
 #include "quantize.cuh"
 #include "unary.cuh"
 #include "vecdotq.cuh"
@@ -1215,6 +1216,14 @@ static void mul_mat_vec_q_switch_ncols_dst(
     const bool y_soa = y_layout == GGML_CUDA_Q8_1_SOA_ISUM;
 
 #if !defined(GGML_USE_HIP)
+    if constexpr (type == GGML_TYPE_PQ2_0) {
+        if (y_layout == GGML_CUDA_Q8_1_PQ2) {
+            // layout_host only picks PQ2 for plain 2D calls without fusion
+            GGML_ASSERT(!ids && nchannels_dst == 1 && nsamples_dst == 1 && !fusion.gate && !fusion.x_bias && !fusion.gate_bias);
+            mul_mat_vec_pq2_0_mc_switch(vx, vy, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, stream);
+            return;
+        }
+    }
     if constexpr (type == GGML_TYPE_PTQ1_0) {
         // plain 2D PTQ1_0 mat-vec with 2-8 columns: dedicated kernel with full lane utilization,
         // see mmvq-ptq1_0.cuh. One column takes the SoA path in the generic kernel below.
@@ -1657,7 +1666,7 @@ void ggml_cuda_mul_mat_vec_q(
     }
 
     // Same (type, ncols_dst, ids) triple the kernel switch evaluates, so quantizer and kernel agree.
-    const ggml_cuda_q8_1_layout y_layout = ggml_cuda_q8_1_layout_host(src0->type, (int) (ids ? ne2 : ne11), ids != nullptr);
+    const ggml_cuda_q8_1_layout y_layout = ggml_cuda_q8_1_layout_host(src0->type, (int) (ids ? ne2 : ne11), ids != nullptr, ne12 == 1 && ne13 == 1);
 
     int64_t ne10_padded = GGML_PAD(ne10, MATRIX_ROW_PADDING);
     if (y_layout == GGML_CUDA_Q8_1_SOA_ISUM) {

ggml/src/ggml-cuda/quantize.cu

@@ -59,6 +59,9 @@ static __device__ __forceinline__ float nvfp4_native_scale_error(
 //            quantized values bit-cast into ds.y; the 1-column PTQ1_0 vec-dot folds the digit bias
 //            {0,1,2} -> {-1,0,+1} into one subtraction per 32-block instead of a SIMD byte
 //            subtract per 4 weights (bit-identical to the biased path)
+//   PQ2      PQ2_0, 3-8 columns: qs bytes of a column first (permuted inside each 16-element group so that
+//            (code_word >> 2k) & 0x03030303 lines up with the activation bytes), then one half2 (d, raw int sum
+//            of q) per 32-block (see mmvq-pq2_0.cuh)
 //   PT       planar-transposed layout consumed by the 2-8 column PTQ1_0 mat-vec path
 //            (see mmvq-ptq1_0.cuh); same quantization, same bytes per row as block_q8_1, and the
 //            same exact integer sum in ds.y as SOA_ISUM so that kernel can also use unbiased trits
@@ -124,6 +127,26 @@ static __global__ void quantize_q8_1(
         *ds = make_half2(__float2half(d), __short_as_half((short) isum));
         return;
     }
+    if constexpr (layout == GGML_CUDA_Q8_1_PQ2) {
+        const int64_t row_cont = (i3*ne2.z + i2) * ne1 + i1;
+        char * ycol = (char *) vy + row_cont * (ne0 * 9 / 8); // same row stride as block_q8_1
+        const int64_t kb = i0 / QK8_1;
+        const int     e  = i0 % QK8_1;
+        // inside each 16-element group position (e % 4)*4 + (e % 16)/4 holds element e
+        ((int8_t *) ycol)[kb*QK8_1 + (e / 16)*16 + (e % 4)*4 + (e % 16)/4] = q;
+
+        int isum = q;
+        isum = warp_reduce_sum<QK8_1>(isum);
+
+        if (iqs > 0) {
+            return;
+        }
+
+        // |isum| <= 32*127 fits int16; keep the raw bits in the half slot.
+        half2 * ds = (half2 *) (ycol + ne0) + kb;
+        *ds = make_half2(__float2half(d), __short_as_half((short) isum));
+        return;
+    }
     if constexpr (layout == GGML_CUDA_Q8_1_SOA_ISUM) {
         // Warp-transposed layout (see ggml_cuda_ptq1_q8_word); ne0 is padded to GGML_CUDA_PTQ1_K_PAD.
         int isum = q;
@@ -877,6 +900,9 @@ void quantize_row_q8_1_cuda(
             GGML_ASSERT(ne0 % GGML_CUDA_PTQ1_K_PAD == 0);
             ggml_cuda_kernel_launch(quantize_q8_1<GGML_CUDA_Q8_1_SOA_ISUM>, launch_params, x, vy, ne00, s01, s02, s03, ne0, ne1, ne2_fastdiv);
             break;
+        case GGML_CUDA_Q8_1_PQ2:
+            ggml_cuda_kernel_launch(quantize_q8_1<GGML_CUDA_Q8_1_PQ2>, launch_params, x, vy, ne00, s01, s02, s03, ne0, ne1, ne2_fastdiv);
+            break;
         default:
             ggml_cuda_kernel_launch(quantize_q8_1<GGML_CUDA_Q8_1_AOS>, launch_params, x, vy, ne00, s01, s02, s03, ne0, ne1, ne2_fastdiv);
             break;