diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index d1f12fcc..83cd41fb 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -2818,6 +2818,92 @@ static bool ggml_cuda_should_fuse_rms_norm_mul_rope(const ggml_tensor * rms_norm return true; } +// match concat(conv state, transposed x) -> ssm_conv -> silu, followed by the cpys that write views of the concat into +// the conv-state cache (build_conv_state with a conv kernel emits them in this order). One kernel then reads the conv +// input from the concat's two sources and does the writes, so the concat is never materialized. Returns the nodes +// to skip, 0 if the pattern does not match. +static int ggml_cuda_try_ssm_conv_concat_fusion( + const ggml_cgraph * cgraph, int node_idx, ggml_cuda_ssm_conv_state_writes & writes) { + static const bool disabled = getenv("GGML_CUDA_SSM_CONV_CONCAT_FUSION") != nullptr && + atoi(getenv("GGML_CUDA_SSM_CONV_CONCAT_FUSION")) == 0; + const ggml_tensor * concat = cgraph->nodes[node_idx]; + if (disabled || concat->op != GGML_OP_CONCAT || ggml_get_op_params_i32(concat, 0) != 0 || + node_idx + 2 >= cgraph->n_nodes) { + return 0; + } + const ggml_tensor * conv = cgraph->nodes[node_idx + 1]; + const ggml_tensor * silu = cgraph->nodes[node_idx + 2]; + if (conv->op != GGML_OP_SSM_CONV || conv->src[0] != concat || silu->op != GGML_OP_UNARY || + ggml_get_unary_op(silu) != GGML_UNARY_OP_SILU || silu->src[0] != conv || + !ggml_node_has_n_uses(cgraph, node_idx + 1, 1) || (concat->flags & GGML_TENSOR_FLAG_OUTPUT)) { + return 0; + } + + const ggml_tensor * state = concat->src[0]; + const ggml_tensor * x = concat->src[1]; + const ggml_tensor * weight = conv->src[1]; + const int64_t d_conv = weight->ne[0]; + if (concat->type != GGML_TYPE_F32 || x->type != GGML_TYPE_F32 || state->type != GGML_TYPE_F32 || + weight->type != GGML_TYPE_F32 || silu->type != GGML_TYPE_F32 || d_conv != 4 || + !ggml_is_contiguous(state) || state->ne[0] != d_conv - 1 || x->nb[1] != sizeof(float) || + x->ne[1] != state->ne[1] || x->ne[2] != state->ne[2] || x->ne[3] != 1 || state->ne[3] != 1 || + weight->nb[0] != sizeof(float) || silu->nb[0] != sizeof(float) || silu->ne[0] % 128 != 0) { + return 0; + } + + // the allocator may already have handed the memory of the concat's sources to the conv output (they are dead + // once the concat ran); the kernel then writes a scratch buffer that is copied to the output afterwards + auto overlaps = [](const ggml_tensor * a, const ggml_tensor * b) { + const char * a0 = (const char *) a->data; + const char * b0 = (const char *) b->data; + return a0 < b0 + ggml_nbytes(b) && b0 < a0 + ggml_nbytes(a); + }; + writes.staged = overlaps(silu, state) || overlaps(silu, x) || overlaps(silu, weight); + if (writes.staged && !ggml_is_contiguous(silu)) { + return 0; + } + + // every other use of the concat must be one of the cpy sources below + writes.n = 0; + int last = node_idx + 2; + for (int j = node_idx + 3; j < cgraph->n_nodes; ++j) { + const ggml_tensor * n = cgraph->nodes[j]; + if (n->op == GGML_OP_VIEW && n->view_src == concat) { + if (ggml_node_get_use_count(cgraph, j) != 1) { + return 0; + } + continue; + } + if (ggml_cuda_is_view_or_noop(n)) { + continue; + } + if (n->op != GGML_OP_CPY || n->src[0]->op != GGML_OP_VIEW || n->src[0]->view_src != concat) { + break; + } + const ggml_tensor * src = n->src[0]; + const ggml_tensor * dst = n->src[1]; + const size_t col = src->view_offs / sizeof(float); + if (writes.n == GGML_CUDA_SSM_CONV_MAX_STATE_WRITES || (n->flags & GGML_TENSOR_FLAG_OUTPUT) || + src->view_offs % sizeof(float) != 0 || col + d_conv - 1 > (size_t) concat->ne[0] || + src->ne[0] != d_conv - 1 || src->ne[1] != concat->ne[1] || src->ne[2] != concat->ne[2] || + src->nb[1] != concat->nb[1] || src->nb[2] != concat->nb[2] || dst->type != GGML_TYPE_F32 || + dst->data == nullptr || dst->nb[0] != sizeof(float) || dst->ne[0] != (d_conv - 1) * concat->ne[1] || + dst->ne[1] != concat->ne[2] || dst->ne[2] != 1 || dst->ne[3] != 1) { + return 0; + } + // the kernel writes the snapshots before it reads its conv window, so they must not land on an input + if (overlaps(dst, state) || overlaps(dst, x) || overlaps(dst, weight)) { + return 0; + } + writes.w[writes.n++] = { (char *) dst->data, dst->nb[1], (int) col }; + last = j; + } + if (ggml_node_get_use_count(cgraph, node_idx) != 1 + writes.n) { + return 0; + } + return last - node_idx; +} + // match gated_delta_net + the strided cpy that scatters its state snapshots into the cache // (slot i -> rollback group i, slot 0 newest), so the kernel can write them and skip the cpy. static int ggml_cuda_try_gdn_cache_fusion( @@ -3541,6 +3627,16 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph } } + // concat -> ssm_conv -> silu -> conv-state cpys + if (node->op == GGML_OP_CONCAT) { + ggml_cuda_ssm_conv_state_writes writes; + const int nodes_to_skip = ggml_cuda_try_ssm_conv_concat_fusion(cgraph, i, writes); + if (nodes_to_skip > 0) { + ggml_cuda_op_ssm_conv_concat(*cuda_ctx, node, cgraph->nodes[i + 1], cgraph->nodes[i + 2], writes); + return nodes_to_skip; + } + } + // gated_delta_net -> cpy: scatter recurrent-state snapshots into the cache if (node->op == GGML_OP_GATED_DELTA_NET) { ggml_cuda_gated_delta_net_fused_cache fused_state_cpy; diff --git a/ggml/src/ggml-cuda/ssm-conv.cu b/ggml/src/ggml-cuda/ssm-conv.cu index 1463169c..df0b1d59 100644 --- a/ggml/src/ggml-cuda/ssm-conv.cu +++ b/ggml/src/ggml-cuda/ssm-conv.cu @@ -204,3 +204,95 @@ void ggml_cuda_op_ssm_conv(ggml_backend_cuda_context & ctx, ggml_tensor * dst, g out->nb[2], nc, nr, n_t, n_s, stream); } } + +// one thread per (channel, seq), split_n_t tokens per block along z. Column p of conv_input is state column p for +// p < d_conv - 1, else token p - (d_conv - 1) of x. The sum, bias and silu repeat ssm_conv_f32 operation for operation. +template +static __global__ void ssm_conv_concat_silu_f32(const float * __restrict__ state, const char * __restrict__ x, + const float * __restrict__ w, const float * __restrict__ bias, + float * __restrict__ y, const ggml_cuda_ssm_conv_state_writes writes, + const int nr, const int64_t n_t, const size_t x_nb0, const size_t x_nb2, + const int w_stride, const size_t y_nb1, const size_t y_nb2) { + ggml_cuda_pdl_lc(); + const int c = blockIdx.y * blockDim.x + threadIdx.x; + const int s = blockIdx.x; + const int64_t t0 = blockIdx.z * split_n_t; + + const float * st = state + ((int64_t) s * nr + c) * (d_conv - 1); + const char * xs = x + s * x_nb2 + c * sizeof(float); + auto column = [&](int64_t p) { + return p < (int64_t) d_conv - 1 ? st[p] : *(const float *) (xs + (p - (int64_t) (d_conv - 1)) * x_nb0); + }; + + ggml_cuda_pdl_sync(); + + float wv[d_conv]; +#pragma unroll + for (size_t j = 0; j < d_conv; j++) { + wv[j] = w[c * w_stride + j]; + } + const float b = bias != nullptr ? bias[c] : 0.0f; + + if (blockIdx.z == 0) { + for (int k = 0; k < writes.n; k++) { + float * dst = (float *) (writes.w[k].dst + s * writes.w[k].nb1) + c * (d_conv - 1); +#pragma unroll + for (size_t j = 0; j < d_conv - 1; j++) { + dst[j] = column(writes.w[k].col + j); + } + } + } + + float win[d_conv]; +#pragma unroll + for (size_t j = 0; j < d_conv - 1; j++) { + win[j] = column(t0 + j); + } + + float * ys = (float *) ((char *) y + s * y_nb2) + c; + const int64_t t1 = min(t0 + split_n_t, n_t); + for (int64_t i = t0; i < t1; i++) { + win[d_conv - 1] = column(i + d_conv - 1); + + float sumf = 0.0f; +#pragma unroll + for (size_t j = 0; j < d_conv; j++) { + sumf += win[j] * wv[j]; + } + sumf += b; + ys[i * (y_nb1 / sizeof(float))] = ggml_cuda_op_silu_single(sumf); + +#pragma unroll + for (size_t j = 0; j + 1 < d_conv; j++) { + win[j] = win[j + 1]; + } + } +} + +void ggml_cuda_op_ssm_conv_concat(ggml_backend_cuda_context & ctx, const ggml_tensor * concat, const ggml_tensor * ssm_conv, + ggml_tensor * silu, const ggml_cuda_ssm_conv_state_writes & writes) { + const ggml_tensor * state = concat->src[0]; // [d_conv - 1, channels, n_s], contiguous + const ggml_tensor * x = concat->src[1]; // [n_t, channels, n_s], channels contiguous + const ggml_tensor * weight = ssm_conv->src[1]; + + const int64_t nr = silu->ne[0]; + const int64_t n_t = silu->ne[1]; + const int64_t n_s = silu->ne[2]; + + constexpr int threads = 128; + constexpr int64_t split_n_t = 32; + GGML_ASSERT(weight->ne[0] == 4 && nr % threads == 0); + + ggml_cuda_pool_alloc staged(ctx.pool()); + float * y = writes.staged ? staged.alloc(ggml_nelements(silu)) : (float *) silu->data; + + const dim3 blocks(n_s, nr / threads, (n_t + split_n_t - 1) / split_n_t); + const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(blocks, threads, 0, ctx.stream()); + ggml_cuda_kernel_launch(ssm_conv_concat_silu_f32<4, split_n_t>, launch_params, + (const float *) state->data, (const char *) x->data, (const float *) weight->data, (const float *) nullptr, + y, writes, (int) nr, n_t, x->nb[0], x->nb[2], (int) (weight->nb[1] / sizeof(float)), + silu->nb[1], silu->nb[2]); + if (writes.staged) { + CUDA_CHECK(cudaMemcpyAsync(silu->data, y, ggml_nbytes(silu), cudaMemcpyDeviceToDevice, ctx.stream())); + } +} diff --git a/ggml/src/ggml-cuda/ssm-conv.cuh b/ggml/src/ggml-cuda/ssm-conv.cuh index 8514ca84..8ca468b2 100644 --- a/ggml/src/ggml-cuda/ssm-conv.cuh +++ b/ggml/src/ggml-cuda/ssm-conv.cuh @@ -1,3 +1,24 @@ #include "common.cuh" void ggml_cuda_op_ssm_conv(ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_tensor * bias_add_node = nullptr, ggml_tensor * silu_dst = nullptr); + +// a cpy of conv_input columns [col, col + d_conv - 1) into the conv-state cache: seq s goes to dst + s * nb1 bytes, +// laid out as [d_conv - 1, channels] like the cpy it replaces +struct ggml_cuda_ssm_conv_state_write { + char * dst; + size_t nb1; + int col; +}; + +#define GGML_CUDA_SSM_CONV_MAX_STATE_WRITES 8 + +struct ggml_cuda_ssm_conv_state_writes { + ggml_cuda_ssm_conv_state_write w[GGML_CUDA_SSM_CONV_MAX_STATE_WRITES]; + int n; + bool staged; // the output overlaps an input: compute into a scratch buffer, then copy +}; + +// silu(ssm_conv(concat(state, x))) read straight from the concat's sources, plus the conv-state writes, in one +// kernel; same arithmetic as ggml_cuda_op_ssm_conv (see ggml_cuda_try_ssm_conv_concat_fusion) +void ggml_cuda_op_ssm_conv_concat(ggml_backend_cuda_context & ctx, const ggml_tensor * concat, const ggml_tensor * ssm_conv, + ggml_tensor * silu, const ggml_cuda_ssm_conv_state_writes & writes); diff --git a/src/models/delta-net-base.cpp b/src/models/delta-net-base.cpp index c4ccbd51..c81bf629 100644 --- a/src/models/delta-net-base.cpp +++ b/src/models/delta-net-base.cpp @@ -457,7 +457,8 @@ ggml_tensor * llm_build_delta_net_base::build_conv_state( ggml_tensor * qkv_mixed, int64_t conv_kernel_size, int64_t conv_channels, - int il) { + int il, + ggml_tensor * conv_kernel) { const auto * mctx_cur = inp->mctx; const auto kv_head = mctx_cur->get_head(); @@ -477,6 +478,17 @@ ggml_tensor * llm_build_delta_net_base::build_conv_state( ggml_tensor * conv_input = ggml_concat(ctx0, conv_states, qkv_mixed, 0); cb(conv_input, "conv_input", il); + // the conv goes into the graph before the conv-state writes: concat, conv, silu and the writes are then + // adjacent, which lets the CUDA backend run them as one kernel without materializing conv_input + ggml_tensor * conv_out = nullptr; + if (conv_kernel) { + conv_out = ggml_ssm_conv(ctx0, conv_input, conv_kernel); + cb(conv_out, "conv_output_raw", il); + conv_out = ggml_silu(ctx0, conv_out); + cb(conv_out, "conv_output_silu", il); + ggml_build_forward_expand(gf, conv_out); + } + const int64_t row_count = (conv_kernel_size - 1) * conv_channels; const size_t row_size = ggml_row_size(conv_states_all->type, row_count); @@ -526,7 +538,7 @@ ggml_tensor * llm_build_delta_net_base::build_conv_state( } } - return conv_input; + return conv_kernel ? conv_out : conv_input; } ggml_tensor * llm_build_delta_net_base::build_recurrent_attn( diff --git a/src/models/models.h b/src/models/models.h index c5d630ba..ce254ac4 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -67,14 +67,16 @@ struct llm_build_delta_net_base : public llm_graph_context { int il); // read conv state from cache, concat with qkv_mixed, write back (single slot or per-token) - // qkv_mixed: (qkv_dim, n_seq_tokens, n_seqs); returns conv_input: (kernel_size + n_seq_tokens - 1, channels, n_seqs) + // qkv_mixed: (qkv_dim, n_seq_tokens, n_seqs); returns conv_input: (kernel_size + n_seq_tokens - 1, channels, n_seqs), + // or silu(ssm_conv(conv_input, conv_kernel)): (channels, n_seq_tokens, n_seqs) when conv_kernel is given ggml_tensor * build_conv_state( llm_graph_input_rs * inp, ggml_tensor * conv_states_all, ggml_tensor * qkv_mixed, int64_t conv_kernel_size, int64_t conv_channels, - int il); + int il, + ggml_tensor * conv_kernel = nullptr); // run delta-net attention and write the new recurrent state(s) back to ssm_states_all // s: (head_v_dim, head_v_dim, num_v_heads, n_seqs); returns output: (head_v_dim, num_v_heads, n_seq_tokens, n_seqs) diff --git a/src/models/qwen35.cpp b/src/models/qwen35.cpp index 6291ef28..e473b73a 100644 --- a/src/models/qwen35.cpp +++ b/src/models/qwen35.cpp @@ -471,7 +471,7 @@ ggml_tensor * llama_model_qwen35::graph::build_layer_attn_linear( const int64_t conv_kernel_size = conv_kernel->ne[0]; const int64_t conv_channels = d_inner + 2 * hparams.ssm_n_group * hparams.ssm_d_state; - ggml_tensor * conv_input = build_conv_state(inp, conv_states_all, qkv_mixed, conv_kernel_size, conv_channels, il); + ggml_tensor * conv_qkv_mix = build_conv_state(inp, conv_states_all, qkv_mixed, conv_kernel_size, conv_channels, il, conv_kernel); // ring path: read per-seq live state directly from the cache inside the // fused GDN op (rows mode) instead of a gather per layer. @@ -496,14 +496,6 @@ ggml_tensor * llama_model_qwen35::graph::build_layer_attn_linear( cb(state, "state_predelta", il); } - ggml_tensor * conv_output_proper = ggml_ssm_conv(ctx0, conv_input, conv_kernel); - cb(conv_output_proper, "conv_output_raw", il); - - ggml_tensor * conv_output_silu = ggml_silu(ctx0, conv_output_proper); - cb(conv_output_silu, "conv_output_silu", il); - - ggml_tensor * conv_qkv_mix = conv_output_silu; - // Calculate the total conv dimension int64_t qkv_dim = head_k_dim * num_k_heads * 2 + head_v_dim * num_v_heads; int64_t nb1_qkv = ggml_row_size(conv_qkv_mix->type, qkv_dim);