From b416432d1896e554b0e4e50a49fc0142ee8215d9 Mon Sep 17 00:00:00 2001 From: sb32445 Date: Tue, 27 Oct 2026 15:25:00 +0100 Subject: [PATCH] cuda: gate (SwiGLU) fused PTQ1_0 mat-vec for 2-4 columns The dedicated PTQ1_0 kernel (mul_mat_vec_ptq1_0_pt) is templated for fusion with any column count, but the graph check, an assert in ggml_cuda_mul_mat_vec_q and the launcher only allowed ncols_dst == 1. So gate + up + SwiGLU ran as three kernels in every FFN block of a 2-4 token verify step (speculative decoding, small batches). Allow the fused path for the dedicated kernel (plain 2D, K multiple of 128, shared memory within the limit; the graph check calls the same helper, ggml_cuda_mmvq_ptq1_0_can_fuse_mc, which uses the per-cc rows_per_item rule of the launcher so guard and launch agree on the shared memory). With more than one column the bias tensors have to be contiguous because the kernel indexes the bias like dst. Other types keep the one-column rule. ggml_cuda_check_fusion_memory_ranges refused the fusion in 24 of 64 FFN blocks of the Bonsai 2 verify graph: the GLU output overlaps src1, the Hadamard transform output, because ggml-alloc hands it the block that src1 released. The overlap is real when the q8_1 rows written by ggml_cuda_try_fwht_q8 live in that buffer, so the check was right. Now ggml_cuda_try_fwht_q8 puts the rows into a pool block when a following GLU of a consumer pair overlaps the buffer (the mechanism it already has for out_aliases_in), and the fusion memory check ignores an overlap with an input whose q8 rows are registered in a pool block (g_fwht_q8_ctx, reset by a scope guard after the graph evaluation). All 64 FFN blocks fuse. RTX 4070, Bonsai 2 27B PTQ1_0 + MTP head (n-max 2), q4_0 K/V, greedy, 4 prompts x 256 tokens, 8 interleaved A/B pairs, base prism c96baf4ab: decode 115.98 -> 117.16 tok/s (+1.02 %, 95 % CI [+0.96, +1.08] %), outputs identical in all runs, acceptance 74.0 % in both. Measured earlier with environment switches on an older base: the fusion alone +1.01 %, keeping it over the overlap another +0.33 %. Removing the whole change from the full patch stack costs 0.65 % at depth 0 and 0.37 % at 114688 tokens. tests: gate fused PTQ1_0 mat-vec cases with 2 to 4 columns in MUL_MAT_VEC_FUSION (538 -> 546 cases, all pass on CUDA0 together with MUL_MAT 1669, SWIGLU 24, RMS_NORM 51). Rebased on prism 5628c7338: ggml_cuda_should_fuse_mul_mat refuses the fusion for GGML_GLU_OP_SWIGLU_CLAMP when there is more than one column, because the 2..4 column kernel does not implement that GLU op (test-backend-ops MUL_MAT_VEC_FUSION glu_op=6 failed with 2 and 4 columns before this guard; the graph now runs unfused). Co-Authored-By: Claude Sonnet 5.5 --- ggml/src/ggml-cuda/ggml-cuda.cu | 63 +++++++++++++++++++++++++----- ggml/src/ggml-cuda/mmvq-ptq1_0.cuh | 2 +- ggml/src/ggml-cuda/mmvq.cu | 21 +++++++++- ggml/src/ggml-cuda/mmvq.cuh | 2 + tests/test-backend-ops.cpp | 11 ++++++ 5 files changed, 88 insertions(+), 11 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 6e868e806db0..3fb292d32872 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -1845,6 +1845,11 @@ static bool ggml_cuda_should_fuse_mul_mat(const ggml_tensor * ffn_up, return false; } + // the 2..4 column PTQ1_0 kernel implements SWIGLU, GEGLU and SWIGLU_OAI, not SWIGLU_CLAMP: such a graph runs unfused + if (ffn_up->op == GGML_OP_MUL_MAT && ffn_up->ne[1] != 1 && ggml_get_glu_op(glu) == GGML_GLU_OP_SWIGLU_CLAMP) { + return false; + } + if (const bool swapped = ggml_get_op_params_i32(glu, 1); swapped) { return false; } @@ -1880,7 +1885,7 @@ static bool ggml_cuda_should_fuse_mul_mat_vec_f(const ggml_tensor * tensor) { return use_mul_mat_vec_f; } -static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) { +static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor, const ggml_tensor * bias0 = nullptr, const ggml_tensor * bias1 = nullptr) { ggml_tensor * src0 = tensor->src[0]; ggml_tensor * src1 = tensor->src[1]; const ggml_tensor * dst = tensor; @@ -1897,8 +1902,14 @@ 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; + } + + // the 2..4 column kernel indexes the bias like dst + if (tensor->op == GGML_OP_MUL_MAT && dst->ne[1] != 1 && + ((bias0 && !ggml_is_contiguous(bias0)) || (bias1 && !ggml_is_contiguous(bias1)))) { return false; } @@ -3278,7 +3289,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)); // x is only known to be contiguous with the same element count as the consumers' src1 [K, ncols]; // in the unsigned pattern it is the [n, rows] reshape @@ -3286,6 +3297,24 @@ static int ggml_cuda_try_fwht_q8(ggml_backend_cuda_context & ctx, const ggml_cgr if (!ggml_cuda_fwht_q8_find_uses(cgraph, i_mm, cc, pq2_enabled, u)) { return 0; } + // 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 j = i_mm + 1; j < cgraph->n_nodes && !out_aliases_in; ++j) { + const ggml_tensor * t = cgraph->nodes[j]; + if (t->op != GGML_OP_MUL_MAT) { + continue; + } + bool uses_x = false; + for (int a = 0; a < u.n_aliases; ++a) { + uses_x = uses_x || t->src[1] == u.aliases[a]; + } + for (int j2 = j + 1; uses_x && 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; + } + } + } const int64_t K = u.K; const int64_t ncols = u.ncols; const ggml_type wtype = u.wtype; @@ -4121,6 +4150,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, @@ -4175,6 +4208,13 @@ static bool ggml_cuda_check_fusion_memory_ranges(const ggml_cgraph * cgraph, } if (nodes_overlap(dst, src) && !is_exact_alias(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) { @@ -5859,7 +5899,7 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph fusion_data.glu_op = ggml_get_glu_op(glu); fusion_data.glu_limit = ggml_get_op_params_f32(glu, 3); - if (ggml_cuda_should_fuse_mul_mat_vec_q(up_n)) { + if (ggml_cuda_should_fuse_mul_mat_vec_q(up_n, up_bias, gate_bias)) { ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, cgraph->nodes[glu_idx], &fusion_data); fused_mul_mat_vec = true; fused_node_count = n_ops; @@ -5953,7 +5993,7 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph fusion_data.glu_op = ggml_get_glu_op(glu); fusion_data.glu_limit = ggml_get_op_params_f32(glu, 3); - if (ggml_cuda_should_fuse_mul_mat_vec_q(up_n)) { + if (ggml_cuda_should_fuse_mul_mat_vec_q(up_n, up_bias, gate_bias)) { ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, cgraph->nodes[glu_idx], &fusion_data); fused_mul_mat_vec = true; fused_node_count = n_ops; @@ -6016,7 +6056,7 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph break; } - if (ggml_cuda_should_fuse_mul_mat_vec_q(up_n)) { + if (ggml_cuda_should_fuse_mul_mat_vec_q(up_n, up_bias_tensor, gate_bias_tensor)) { ggml_cuda_mm_fusion_args_host fusion_data{}; fusion_data.gate = gate_n->src[0]; fusion_data.x_bias = up_bias_tensor; @@ -6162,7 +6202,7 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph fusion_data.x_bias = bias; fusion_data.x_scale = scale; - if (ggml_cuda_should_fuse_mul_mat_vec_q(mm_node)) { + if (ggml_cuda_should_fuse_mul_mat_vec_q(mm_node, bias)) { ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, out_node, &fusion_data); fused_mul_mat_vec = true; fused_node_count = n_ops; @@ -6227,7 +6267,7 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph break; } - if (ggml_cuda_should_fuse_mul_mat_vec_q(mm_node)) { + if (ggml_cuda_should_fuse_mul_mat_vec_q(mm_node, bias_tensor)) { ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, bias_node, &fusion_data); fused_mul_mat_vec = true; fused_node_count = 2; @@ -6326,6 +6366,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 diff --git a/ggml/src/ggml-cuda/mmvq-ptq1_0.cuh b/ggml/src/ggml-cuda/mmvq-ptq1_0.cuh index 4775f4f6be7d..2e641797af49 100644 --- a/ggml/src/ggml-cuda/mmvq-ptq1_0.cuh +++ b/ggml/src/ggml-cuda/mmvq-ptq1_0.cuh @@ -523,7 +523,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 { diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu index 4121310f0d20..e34a055aeaf9 100644 --- a/ggml/src/ggml-cuda/mmvq.cu +++ b/ggml/src/ggml-cuda/mmvq.cu @@ -1948,6 +1948,25 @@ 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 int device = ggml_cuda_get_device(); + // the launcher picks the rows per work item from the compute capability: use the same rule + const int rows_per_item = ptq1_0_pt_rows_per_item_for_cc(ncols_dst, ggml_cuda_info().devices[device].cc); + const size_t smem = ptq1_0_pt_smem_bytes((int) (src0->ne[0] / QK_PTQ1_0), ncols_dst, (int) src0->ne[1], rows_per_item, true); + return smem <= ggml_cuda_info().devices[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) { @@ -1979,7 +1998,7 @@ void ggml_cuda_mul_mat_vec_q( if (fusion) { const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; GGML_ASSERT( !ids || dst->ne[2] <= get_mmvq_mmid_max_batch(src0->type, cc)); - 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); diff --git a/ggml/src/ggml-cuda/mmvq.cuh b/ggml/src/ggml-cuda/mmvq.cuh index 5605bf7a4e60..23629d1303f0 100644 --- a/ggml/src/ggml-cuda/mmvq.cuh +++ b/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); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index db966b65815e..031a44463084 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -13850,6 +13850,17 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_mul_mat_vec_fusion(GGML_TYPE_PTQ1_0, GGML_GLU_OP_SWIGLU, 1, 64, k, false, 1, 1, false, false, true, false, {1, 1})); } + // gate fused PTQ1_0 mat-vec with 2 to 4 columns + for (int64_t m : {2, 3, 4}) { + for (int64_t k : {512, 5120}) { + test_cases.emplace_back(new test_mul_mat_vec_fusion(GGML_TYPE_PTQ1_0, GGML_GLU_OP_SWIGLU, m, 100, k, + false, 1, 1, false, false, true, false, {1, 1})); + } + } + test_cases.emplace_back(new test_mul_mat_vec_fusion(GGML_TYPE_PTQ1_0, GGML_GLU_OP_SWIGLU, 3, 17408, 5120, + false, 1, 1, false, false, true, false, {1, 1})); + test_cases.emplace_back(new test_mul_mat_vec_fusion(GGML_TYPE_PTQ1_0, GGML_GLU_OP_SWIGLU, 3, 5120, 17408, + false, 1, 1, false, false, true, false, {1, 1})); // PTQ1_0 one-column K-block-stationary mat-vec (CUDA, Hopper): at least 4096 rows with a row tail, // K = 40, 48 and 136 blocks, plain and with bias / gate / gate + bias fusion