From e0e6b2c0102a73d137507435de8eaf46e591e83d Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Fri, 3 Jul 2026 12:20:42 +0200 Subject: [PATCH] cuda: fuse independent quantized SET_ROWS pairs into one launch Per-layer K/V cache appends arrive as two independent quantized SET_ROWS nodes, usually separated only by the elidable view feeding the second one. Fuse the pair into a single kernel launch: the dual kernel runs the exact k_set_rows_quant element math for each half of the grid, so outputs are bit-identical to two separate launches; only the launch and tail-latency overhead of the second kernel is saved (~0.2 ms/step at laguna w6 decode, 80 launches -> 40). Supported for matching quantized dst types (q4_0/q4_1/q5_0/q5_1/q8_0/ iq4_nl) with f32 src and i32/i64 ids; anything else keeps separate launches. Guarded by the existing GGML_CUDA_DISABLE_FUSION switch. Validated on RTX 3090 (laguna-xs2 + v24 draft, w6): output hashes byte-identical across code/math/long-context prompts, ctest 12/12. Co-Authored-By: Claude Fable 5 --- ggml/src/ggml-cuda/ggml-cuda.cu | 25 +++++ ggml/src/ggml-cuda/set-rows.cu | 165 ++++++++++++++++++++++++++++++++ ggml/src/ggml-cuda/set-rows.cuh | 5 + 3 files changed, 195 insertions(+) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index e702b209b773..b8dbe39ecf47 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -3937,6 +3937,31 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud continue; } + if (node->op == GGML_OP_SET_ROWS) { + // Per-layer K/V cache appends arrive as two independent + // quantized SET_ROWS, usually separated only by the + // elidable view that feeds the second one. Fuse the pair + // into a single launch (bit-identical per element). + int j = i + 1; + if (j < cgraph->n_nodes && + (cgraph->nodes[j]->op == GGML_OP_VIEW || + cgraph->nodes[j]->op == GGML_OP_RESHAPE || + cgraph->nodes[j]->op == GGML_OP_PERMUTE)) { + j++; + } + if (j < cgraph->n_nodes && cgraph->nodes[j]->op == GGML_OP_SET_ROWS) { + ggml_tensor * other = cgraph->nodes[j]; + const bool independent = + other->src[0] != node && other->src[1] != node && + (other->src[0]->view_src == nullptr || other->src[0]->view_src != node); + if (independent && ggml_cuda_set_rows_dual_supported(node, other)) { + ggml_cuda_op_set_rows_dual(*cuda_ctx, node, other); + i = j; + continue; + } + } + } + if (node->op == GGML_OP_ADD || node->op == GGML_OP_MUL) { int n_fuse = 0; ggml_op ops[8]; diff --git a/ggml/src/ggml-cuda/set-rows.cu b/ggml/src/ggml-cuda/set-rows.cu index db61b3995a8b..605af161535b 100644 --- a/ggml/src/ggml-cuda/set-rows.cu +++ b/ggml/src/ggml-cuda/set-rows.cu @@ -405,6 +405,171 @@ static void set_rows_cuda(ggml_backend_cuda_context & ctx, const ggml_tensor * s } + +// ---- Fused dual set_rows ------------------------------------------------- +// Two independent quantized SET_ROWS (the per-layer K and V cache appends) +// executed by one kernel launch. Each element runs the exact same math as +// k_set_rows_quant, so outputs are bit-identical to two separate launches. + +template +struct set_rows_quant_params { + const float * src0; + const idx_t * src1; + void * dst; + int64_t ne_total; + int64_t s01, s02, s03, s10, s11, s12, s1, s2, s3; + uint3 ne00, ne01, ne02, ne11_fd, ne12_fd; +}; + +template +static __global__ void k_set_rows_quant_dual( + const set_rows_quant_params pa, + const set_rows_quant_params pb) { + int64_t i = int64_t(blockDim.x) * blockIdx.x + threadIdx.x; + + const bool second = i >= pa.ne_total; + if (second) { + i -= pa.ne_total; + if (i >= pb.ne_total) { + return; + } + } + const set_rows_quant_params p = second ? pb : pa; + + const int64_t i_base = i * qk; + uint32_t tmp = (uint32_t) i_base; + uint2 div_mod; + + div_mod = fast_div_modulo(tmp, p.ne00); + const int64_t i00 = div_mod.y; + tmp = div_mod.x; + + div_mod = fast_div_modulo(tmp, p.ne01); + const int64_t i01 = div_mod.y; + tmp = div_mod.x; + + div_mod = fast_div_modulo(tmp, p.ne02); + const int64_t i02 = div_mod.y; + const int64_t i03 = div_mod.x; + + const int64_t i12 = fastmodulo((uint32_t) i03, p.ne12_fd); + const int64_t i11 = fastmodulo((uint32_t) i02, p.ne11_fd); + const int64_t i10 = i01; + + const int64_t dst_row = *(p.src1 + i10*p.s10 + i11*p.s11 + i12*p.s12); + + const float * src0_row = p.src0 + i01*p.s01 + i02*p.s02 + i03*p.s03; + block_type * dst_row_ptr = (block_type *) p.dst + (dst_row*p.s1 + i02*p.s2 + i03*p.s3) / sizeof(block_type); + + const float * src_block = src0_row + i00; + block_type * dst_block = dst_row_ptr + i00 / qk; + + quantize_func(src_block, dst_block); +} + +template +static set_rows_quant_params make_set_rows_quant_params(const ggml_tensor * dst_t, const int qk) { + const ggml_tensor * src0 = dst_t->src[0]; + const ggml_tensor * src1 = dst_t->src[1]; + const ggml_tensor * dst = dst_t; + + GGML_TENSOR_BINARY_OP_LOCALS + + set_rows_quant_params p{}; + p.src0 = (const float *) src0->data; + p.src1 = (const idx_t *) src1->data; + p.dst = dst->data; + GGML_ASSERT(ne00 % qk == 0); + p.ne_total = (ne00 * ne01 * ne02 * ne03) / qk; + p.s01 = nb01 / sizeof(float); + p.s02 = nb02 / sizeof(float); + p.s03 = nb03 / sizeof(float); + p.s10 = nb10 / sizeof(idx_t); + p.s11 = nb11 / sizeof(idx_t); + p.s12 = nb12 / sizeof(idx_t); + p.s1 = nb1; + p.s2 = nb2; + p.s3 = nb3; + p.ne00 = init_fastdiv_values((uint32_t) ne00); + p.ne01 = init_fastdiv_values((uint32_t) ne01); + p.ne02 = init_fastdiv_values((uint32_t) ne02); + p.ne11_fd = init_fastdiv_values((uint32_t) ne11); + p.ne12_fd = init_fastdiv_values((uint32_t) ne12); + return p; +} + +template +static void set_rows_dual_launch(ggml_backend_cuda_context & ctx, ggml_tensor * a, ggml_tensor * b) { + const set_rows_quant_params pa = make_set_rows_quant_params(a, qk); + const set_rows_quant_params pb = make_set_rows_quant_params(b, qk); + const int64_t total = pa.ne_total + pb.ne_total; + if (total <= 0) { + return; + } + const int num_blocks = (total + CUDA_SET_ROWS_BLOCK_SIZE - 1) / CUDA_SET_ROWS_BLOCK_SIZE; + k_set_rows_quant_dual + <<>>(pa, pb); +} + +template +static void set_rows_dual_dispatch(ggml_backend_cuda_context & ctx, ggml_tensor * a, ggml_tensor * b) { + switch (a->type) { + case GGML_TYPE_Q4_0: + set_rows_dual_launch(ctx, a, b); break; + case GGML_TYPE_Q4_1: + set_rows_dual_launch(ctx, a, b); break; + case GGML_TYPE_Q5_0: + set_rows_dual_launch(ctx, a, b); break; + case GGML_TYPE_Q5_1: + set_rows_dual_launch(ctx, a, b); break; + case GGML_TYPE_Q8_0: + set_rows_dual_launch(ctx, a, b); break; + case GGML_TYPE_IQ4_NL: + set_rows_dual_launch(ctx, a, b); break; + default: + GGML_ABORT("unsupported dual set_rows type %s", ggml_type_name(a->type)); + } +} + +bool ggml_cuda_set_rows_dual_supported(const ggml_tensor * a, const ggml_tensor * b) { + if (a->op != GGML_OP_SET_ROWS || b->op != GGML_OP_SET_ROWS) { + return false; + } + if (a->type != b->type || a->data == b->data) { + return false; + } + switch (a->type) { + case GGML_TYPE_Q4_0: case GGML_TYPE_Q4_1: case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: case GGML_TYPE_Q8_0: case GGML_TYPE_IQ4_NL: + break; + default: + return false; + } + if (a->src[0]->type != GGML_TYPE_F32 || b->src[0]->type != GGML_TYPE_F32) { + return false; + } + if (a->src[1]->type != b->src[1]->type || + (a->src[1]->type != GGML_TYPE_I32 && a->src[1]->type != GGML_TYPE_I64)) { + return false; + } + const int qk = ggml_blck_size(a->type); + if (a->src[0]->ne[0] % qk != 0 || b->src[0]->ne[0] % qk != 0) { + return false; + } + if (ggml_nelements(a->src[0]) <= 0 || ggml_nelements(b->src[0]) <= 0) { + return false; + } + return true; +} + +void ggml_cuda_op_set_rows_dual(ggml_backend_cuda_context & ctx, ggml_tensor * a, ggml_tensor * b) { + if (a->src[1]->type == GGML_TYPE_I64) { + set_rows_dual_dispatch(ctx, a, b); + } else { + set_rows_dual_dispatch(ctx, a, b); + } +} + void ggml_cuda_op_set_rows(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; const ggml_tensor * src1 = dst->src[1]; diff --git a/ggml/src/ggml-cuda/set-rows.cuh b/ggml/src/ggml-cuda/set-rows.cuh index c140c0873c8a..5da867f9ac2d 100644 --- a/ggml/src/ggml-cuda/set-rows.cuh +++ b/ggml/src/ggml-cuda/set-rows.cuh @@ -5,3 +5,8 @@ #define CUDA_SET_ROWS_BLOCK_SIZE 256 void ggml_cuda_op_set_rows(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +// Fused dual set_rows: two independent quantized SET_ROWS in one launch, +// bit-identical to running them separately. +bool ggml_cuda_set_rows_dual_supported(const ggml_tensor * a, const ggml_tensor * b); +void ggml_cuda_op_set_rows_dual(ggml_backend_cuda_context & ctx, ggml_tensor * a, ggml_tensor * b);