cuda: gate (SwiGLU) fused PTQ1_0 mat-vec for 2-4 columns

ggml/src/ggml-cuda/ggml-cuda.cu

@@ -1804,8 +1804,8 @@ static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) {
     if (cc <= GGML_CUDA_CC_PASCAL) {
         return false;
     }
-    //we only support fusion for ncols_dst = 1
-    if (tensor->op == GGML_OP_MUL_MAT && dst->ne[1] != 1) {
+    //we only support fusion for ncols_dst = 1, except the dedicated PTQ1_0 kernel (2..4 columns)
+    if (tensor->op == GGML_OP_MUL_MAT && dst->ne[1] != 1 && !ggml_cuda_mmvq_ptq1_0_can_fuse_mc(src0, (int) dst->ne[1])) {
         return false;
     }
 
@@ -2825,7 +2825,7 @@ static int ggml_cuda_try_fwht_q8(ggml_backend_cuda_context & ctx, const ggml_cgr
         const char * b1 = b0 + ggml_nbytes(b);
         return a0 < b1 && b0 < a1;
     };
-    const bool out_aliases_in = overlaps(x, mm) || (signs && overlaps(signs, mm));
+    bool out_aliases_in = overlaps(x, mm) || (signs && overlaps(signs, mm));
 
     // every use of the transform output (directly or through a reshape view of the whole tensor)
     // must be src1 of a PTQ1_0 MUL_MAT that ggml_cuda_mul_mat routes to ggml_cuda_mul_mat_vec_q.
@@ -2876,6 +2876,14 @@ static int ggml_cuda_try_fwht_q8(ggml_backend_cuda_context & ctx, const ggml_cgr
             }
             K     = src->ne[0];
             ncols = src->ne[1];
+            // a gate/up pair that the mat-vec fuses with its SwiGLU writes that GLU output while other blocks still read the q8 rows:
+            // when the allocator put the GLU output over this buffer, the rows have to go to a pool block (see out_aliases_in)
+            for (int j2 = j + 1; j2 < std::min(j + 4, cgraph->n_nodes); ++j2) {
+                const ggml_tensor * g = cgraph->nodes[j2];
+                if (g->op == GGML_OP_GLU && (g->src[0] == t || g->src[1] == t) && overlaps(g, mm)) {
+                    out_aliases_in = true;
+                }
+            }
             // ggml_cuda_mul_mat sends a padded compute-buffer view to cuBLAS instead
             const ggml_tensor * w = t->src[0];
             if (ggml_backend_buffer_get_usage(w->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE &&
@@ -3201,6 +3209,10 @@ static bool ggml_cuda_topk_moe_fusion(const struct ggml_cgraph * cgraph, int nod
 }
 
 // returns whether the write (out) nodes overwrite the read nodes in operation
+// the evaluating context's q8 rows registry, set while a graph is evaluated; lets the fusion memory check see that an input's
+// quantized rows live in a pool block and not in the tensor's own buffer
+static thread_local const ggml_cuda_fwht_q8_context * g_fwht_q8_ctx = nullptr;
+
 static bool ggml_cuda_check_fusion_memory_ranges(const ggml_cgraph * cgraph,
                                                  const int           node_idx,
                                                  const int           node_count,
@@ -3243,6 +3255,13 @@ static bool ggml_cuda_check_fusion_memory_ranges(const ggml_cgraph * cgraph,
                 }
 
                 if (nodes_overlap(dst, src)) {
+                    // the mat-vec reads the q8 rows from a pool block, not from this tensor: overlapping it is harmless
+                    if (g_fwht_q8_ctx != nullptr) {
+                        const auto it = g_fwht_q8_ctx->entries.find(src);
+                        if (it != g_fwht_q8_ctx->entries.end() && it->second.data != src->data) {
+                            continue;
+                        }
+                    }
                     bool found = false;
 
                     for (int k = node_idx; k < j; ++k) {
@@ -4372,6 +4391,11 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph
 }
 
 static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, const bool use_cuda_graph, const bool cuda_graph_update_required, const void * graph_key) {
+    // the q8 rows registry is only valid while this graph is evaluated
+    struct fwht_q8_ctx_scope {
+        explicit fwht_q8_ctx_scope(const ggml_cuda_fwht_q8_context * c) { g_fwht_q8_ctx = c; }
+        ~fwht_q8_ctx_scope() { g_fwht_q8_ctx = nullptr; }
+    } fwht_q8_ctx_guard(&cuda_ctx->fwht_q8());
     bool graph_evaluated_or_captured = false;
 
     // flag used to determine whether it is an integrated_gpu

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

@@ -510,7 +510,7 @@ static void mul_mat_vec_ptq1_0_pt_launch(
         ggml_cuda_batch_invariant())
 
     if (has_fusion) {
-        GGML_ASSERT(ncols == 1 && "fusion only supported for ncols_dst=1");
+        GGML_ASSERT(ncols <= 4 && "fusion only supported for ncols_dst <= 4");
         if (has_gate) {
             PTQ1_0_PT_LAUNCH(true, true);
         } else {

ggml/src/ggml-cuda/mmvq.cu

@@ -1579,6 +1579,22 @@ static void mul_mat_vec_q_switch_type(
     }
 }
 
+// true when a gate (SwiGLU) fused mat-vec with 2..4 columns runs on the dedicated PTQ1_0 kernel: plain 2D, K a multiple of 128,
+// shared memory (twice the partials with a gate) within the limit. The graph check and the launcher use the same rule.
+bool ggml_cuda_mmvq_ptq1_0_can_fuse_mc(const ggml_tensor * src0, const int ncols_dst) {
+#if defined(GGML_USE_HIP)
+    GGML_UNUSED(src0); GGML_UNUSED(ncols_dst);
+    return false;
+#else
+    if (src0->type != GGML_TYPE_PTQ1_0 || ncols_dst < 2 || ncols_dst > 4 || src0->ne[2] != 1 || src0->ne[3] != 1 ||
+        src0->ne[0] % QK_PTQ1_0 != 0 || !ptq1_0_pt_enabled()) {
+        return false;
+    }
+    const size_t smem = ptq1_0_pt_smem_bytes((int) (src0->ne[0] / QK_PTQ1_0), ncols_dst, (int) src0->ne[1], true);
+    return smem <= ggml_cuda_info().devices[ggml_cuda_get_device()].smpb;
+#endif
+}
+
 void ggml_cuda_mul_mat_vec_q(
         ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst,
         const ggml_cuda_mm_fusion_args_host * fusion) {
@@ -1609,7 +1625,7 @@ void ggml_cuda_mul_mat_vec_q(
 
     if (fusion) {
         GGML_ASSERT( !ids || dst->ne[2] == 1);
-        GGML_ASSERT(  ids || dst->ne[1] == 1);
+        GGML_ASSERT(  ids || dst->ne[1] == 1 || ggml_cuda_mmvq_ptq1_0_can_fuse_mc(src0, (int) dst->ne[1]));
         // Scale fusion is only allowed for NVFP4 currently as the cost of checking this at run-time in the prologue is
         // non-negligible for some models such as gpt-oss-20b
         GGML_ASSERT((fusion->x_scale == nullptr && fusion->gate_scale == nullptr) || src0->type == GGML_TYPE_NVFP4);

ggml/src/ggml-cuda/mmvq.cuh

@@ -16,3 +16,5 @@ void ggml_cuda_op_mul_mat_vec_q(
     const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, const char * src0_dd_i, const float * src1_ddf_i,
     const char * src1_ddq_i, float * dst_dd_i, const int64_t row_low, const int64_t row_high, const int64_t src1_ncols,
     const int64_t src1_padded_row_size, cudaStream_t stream);
+
+bool ggml_cuda_mmvq_ptq1_0_can_fuse_mc(const ggml_tensor * src0, int ncols_dst);