Skip to content

cuda: fuse independent quantized SET_ROWS pairs into one launch #32

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
davide221 merged 1 commit into from
Jul 3, 2026
Merged

cuda: fuse independent quantized SET_ROWS pairs into one launch #32

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 25 additions & 0 deletions ggml/src/ggml-cuda/ggml-cuda.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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];
Expand Down
165 changes: 165 additions & 0 deletions ggml/src/ggml-cuda/set-rows.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 <typename idx_t>
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 <typename idx_t, typename block_type, int qk, void (*quantize_func)(const float *, block_type *)>
static __global__ void k_set_rows_quant_dual(
const set_rows_quant_params<idx_t> pa,
const set_rows_quant_params<idx_t> 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<idx_t> 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 <typename idx_t>
static set_rows_quant_params<idx_t> 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<idx_t> 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 <typename idx_t, typename block_type, int qk, void (*quantize_func)(const float *, block_type *)>
static void set_rows_dual_launch(ggml_backend_cuda_context & ctx, ggml_tensor * a, ggml_tensor * b) {
const set_rows_quant_params<idx_t> pa = make_set_rows_quant_params<idx_t>(a, qk);
const set_rows_quant_params<idx_t> pb = make_set_rows_quant_params<idx_t>(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<idx_t, block_type, qk, quantize_func>
<<<num_blocks, CUDA_SET_ROWS_BLOCK_SIZE, 0, ctx.stream()>>>(pa, pb);
}

template <typename idx_t>
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<idx_t, block_q4_0, QK4_0, quantize_f32_q4_0_block>(ctx, a, b); break;
case GGML_TYPE_Q4_1:
set_rows_dual_launch<idx_t, block_q4_1, QK4_1, quantize_f32_q4_1_block>(ctx, a, b); break;
case GGML_TYPE_Q5_0:
set_rows_dual_launch<idx_t, block_q5_0, QK5_0, quantize_f32_q5_0_block>(ctx, a, b); break;
case GGML_TYPE_Q5_1:
set_rows_dual_launch<idx_t, block_q5_1, QK5_1, quantize_f32_q5_1_block>(ctx, a, b); break;
case GGML_TYPE_Q8_0:
set_rows_dual_launch<idx_t, block_q8_0, QK8_0, quantize_f32_q8_0_block>(ctx, a, b); break;
case GGML_TYPE_IQ4_NL:
set_rows_dual_launch<idx_t, block_iq4_nl, QK4_NL, quantize_f32_iq4_nl_block>(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) {

@cubic-dev-ai cubic-dev-ai Bot Jul 3, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2: The dual SET_ROWS fusion lacks a check for overlapping destination memory ranges. The ggml_cuda_set_rows_dual_supported guard only rejects the trivial case a->data == b->data, but GGML tensor views into the same backing buffer can have different data pointers while still mapping to overlapping regions. Because the fused kernel runs both halves concurrently without synchronization, overlapping destinations would turn the deterministic sequential overwrite semantics of separate launches into nondeterministic CUDA write races. Consider adding an explicit non-overlap assertion or range check before fusion, or document that callers must guarantee disjoint destinations.

Prompt for AI agents 
Check if this issue is valid — if so, understand the root cause and fix it. At ggml/src/ggml-cuda/set-rows.cu, line 538:

<comment>The dual SET_ROWS fusion lacks a check for overlapping destination memory ranges. The `ggml_cuda_set_rows_dual_supported` guard only rejects the trivial case `a->data == b->data`, but GGML tensor views into the same backing buffer can have different `data` pointers while still mapping to overlapping regions. Because the fused kernel runs both halves concurrently without synchronization, overlapping destinations would turn the deterministic sequential overwrite semantics of separate launches into nondeterministic CUDA write races. Consider adding an explicit non-overlap assertion or range check before fusion, or document that callers must guarantee disjoint destinations.</comment>

<file context>
@@ -405,6 +405,171 @@ static void set_rows_cuda(ggml_backend_cuda_context & ctx, const ggml_tensor * s
+    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;
+    }
</file context>

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<int64_t>(ctx, a, b);
} else {
set_rows_dual_dispatch<int32_t>(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];
Expand Down
5 changes: 5 additions & 0 deletions ggml/src/ggml-cuda/set-rows.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Loading