diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 41b3be6317d8..536dd24d3728 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -1787,7 +1787,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; @@ -1804,8 +1804,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; } @@ -2825,7 +2831,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 +2882,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 +3215,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 +3261,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) { @@ -3958,7 +3983,7 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph fusion_data.gate_scale = gate_scale; fusion_data.glu_op = ggml_get_glu_op(glu); - 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; @@ -4051,7 +4076,7 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph fusion_data.gate_scale = gate_scale; fusion_data.glu_op = ggml_get_glu_op(glu); - 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; @@ -4113,7 +4138,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; @@ -4244,7 +4269,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; @@ -4309,7 +4334,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; @@ -4372,6 +4397,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 4755d7a86d00..27101852856d 100644 --- a/ggml/src/ggml-cuda/mmvq-ptq1_0.cuh +++ b/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 { diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu index 0008fdefc4b5..f4ff7de0743f 100644 --- a/ggml/src/ggml-cuda/mmvq.cu +++ b/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); 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 4a5fb0bd7270..47a57e634c79 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -10367,6 +10367,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})); for (auto gate : {GATING_FUNC_SOFTMAX, GATING_FUNC_SIGMOID, GATING_FUNC_SOFTMAX_WEIGHT, GATING_FUNC_SQRT_SOFTPLUS}) { for (bool with_norm : {false, true}) {