From a83467c5480a3c049d46aaed0cde14e48f16e461 Mon Sep 17 00:00:00 2001 From: Skyrion9 Date: Wed, 22 Jul 2026 14:23:09 +0300 Subject: [PATCH 1/3] perf: replace standard attention with unified flash attention Replace the multi-step standard attention pipeline with a single ggml_flash_attn_ext call that handles all token counts uniformly. Causal masking: - For n_tokens > 1, build an F16 causal mask of shape [kv_len, n_tokens, 1, 1] on the CPU where mask[j][i] is ..-INFINITY if i > n_past_ + j and 0.0f otherwise, then upload it via ggml_backend_tensor_set before graph compute. Attention computation: - Remove the entire standard attention block: concat of past KV, repeat_interleave_heads for GQA, mul_mat for KQ scores, ..ggml_scale, ggml_diag_mask_inf, ggml_soft_max, and mul_mat for KQV. - Replaced with a single ggml_flash_attn_ext call using Q permuted to [head_dim, n_tokens, n_head, 1] and K/V cache ..views of shape [head_dim, kv_len, n_head_kv]. - Flatten the flash attention output to [q_size, n_tokens] for the wo projection via ggml_cpy. KV cache layout: - Swap dimensions 1 and 2 in init_kv_cache to match the view shape expected by ggml_flash_attn_ext. - Update token_off_k/token_off_v from nb[2] to nb[1] to reflect the new sequence-dimension stride. - Update k_slot/v_slot view shapes from [head_dim, n_head_kv, n_tokens] to [head_dim, n_tokens, n_head_kv]. - Add ggml_permute + ggml_cont on k and v before writing to the cache to transpose from the QKV projection output shape ..[head_dim, n_head_kv, n_tokens] to the cache layout [head_dim, n_tokens, n_head_kv]. - As a byproduct this fixes segfaults on Vulkan backend on Linux RADV (RDNA2) when trying to prefill_fast. Previously we had to use prefill ..with chunks instead to workaround segfaults, like due to ggml_soft_max limitations on ggml vulkan backend for this particular setup. - Significant performance gains due to FA's O(N) memory bandwidth cost as opposed to the standard O(N^2). --- src/s2_model.cpp | 92 ++++++++++++++++++++++++++++-------------------- 1 file changed, 53 insertions(+), 39 deletions(-) diff --git a/src/s2_model.cpp b/src/s2_model.cpp index 9ced623..4d3cf78 100755 --- a/src/s2_model.cpp +++ b/src/s2_model.cpp @@ -743,8 +743,9 @@ bool SlowARModel::init_kv_cache(int32_t max_seq_len) { return false; } - memory_k_ = ggml_new_tensor_4d(ctx_kv_, GGML_TYPE_F16, head_dim, n_head_kv, max_seq_len, n_layer); - memory_v_ = ggml_new_tensor_4d(ctx_kv_, GGML_TYPE_F16, head_dim, n_head_kv, max_seq_len, n_layer); + // Layout for flash attention + memory_k_ = ggml_new_tensor_4d(ctx_kv_, GGML_TYPE_F16, head_dim, max_seq_len, n_head_kv, n_layer); + memory_v_ = ggml_new_tensor_4d(ctx_kv_, GGML_TYPE_F16, head_dim, max_seq_len, n_head_kv, n_layer); ggml_backend_t kv_backend = (n_gpu_layers_ > 0 && backend_gpu_) ? backend_gpu_ : backend_cpu_; kv_buf_ = ggml_backend_alloc_ctx_tensors(ctx_kv_, kv_backend); @@ -961,6 +962,24 @@ bool SlowARModel::eval_cached(const std::vector & flat_tokens, x = ggml_mul(ctx0, x, ggml_repeat(ctx0, token_scale, x)); } + ggml_tensor * fa_mask = nullptr; + std::vector mask_data; + + if (n_tokens > 1) { + int64_t kv_len = n_past_ + n_tokens; + fa_mask = ggml_new_tensor_4d(ctx0, GGML_TYPE_F16, kv_len, n_tokens, 1, 1); + mask_data.resize(kv_len * n_tokens); + for (int j = 0; j < n_tokens; ++j) { + for (int i = 0; i < kv_len; ++i) { + if (i > n_past_ + j) { + mask_data[j * kv_len + i] = ggml_fp32_to_fp16(-INFINITY); + } else { + mask_data[j * kv_len + i] = ggml_fp32_to_fp16(0.0f); + } + } + } + } + for (int32_t il = 0; il < hparams_.block_count; ++il) { const auto & layer = weights_.layers[il]; @@ -990,53 +1009,44 @@ bool SlowARModel::eval_cached(const std::vector & flat_tokens, const size_t layer_off_k = static_cast(il) * memory_k_->nb[3]; const size_t layer_off_v = static_cast(il) * memory_v_->nb[3]; - const size_t token_off_k = static_cast(n_past_) * memory_k_->nb[2]; - const size_t token_off_v = static_cast(n_past_) * memory_v_->nb[2]; + const size_t token_off_k = static_cast(n_past_) * memory_k_->nb[1]; + const size_t token_off_v = static_cast(n_past_) * memory_v_->nb[1]; ggml_tensor * k_slot = ggml_view_3d(ctx0, memory_k_, - head_dim, n_head_kv, n_tokens, + head_dim, n_tokens, n_head_kv, memory_k_->nb[1], memory_k_->nb[2], layer_off_k + token_off_k); ggml_tensor * v_slot = ggml_view_3d(ctx0, memory_v_, - head_dim, n_head_kv, n_tokens, + head_dim, n_tokens, n_head_kv, memory_v_->nb[1], memory_v_->nb[2], layer_off_v + token_off_v); - ggml_build_forward_expand(gf, ggml_cpy(ctx0, k, k_slot)); - ggml_build_forward_expand(gf, ggml_cpy(ctx0, v, v_slot)); - - ggml_tensor * k_mem = k; - ggml_tensor * v_mem = v; - if (n_past_ > 0) { - ggml_tensor * k_past = ggml_reshape_3d(ctx0, - ggml_view_1d(ctx0, memory_k_, static_cast(n_past_) * kv_size, layer_off_k), - head_dim, n_head_kv, n_past_); - ggml_tensor * v_past = ggml_reshape_3d(ctx0, - ggml_view_1d(ctx0, memory_v_, static_cast(n_past_) * kv_size, layer_off_v), - head_dim, n_head_kv, n_past_); - if (k_past->type != k->type) k_past = ggml_cast(ctx0, k_past, k->type); - if (v_past->type != v->type) v_past = ggml_cast(ctx0, v_past, v->type); - k_mem = ggml_concat(ctx0, k_past, k, 2); - v_mem = ggml_concat(ctx0, v_past, v, 2); - } + + // Permute k and v to match the cache layout + ggml_tensor * k_perm = ggml_cont(ctx0, ggml_permute(ctx0, k, 0, 2, 1, 3)); + ggml_tensor * v_perm = ggml_cont(ctx0, ggml_permute(ctx0, v, 0, 2, 1, 3)); - if (n_head != n_head_kv && q->type != GGML_TYPE_F32) { - q = ggml_cast(ctx0, q, GGML_TYPE_F32); - } - ggml_tensor * k_rep = repeat_interleave_heads(ctx0, k_mem, n_head / n_head_kv); - ggml_tensor * v_rep = repeat_interleave_heads(ctx0, v_mem, n_head / n_head_kv); + ggml_build_forward_expand(gf, ggml_cpy(ctx0, k_perm, k_slot)); + ggml_build_forward_expand(gf, ggml_cpy(ctx0, v_perm, v_slot)); - ggml_tensor * Q = ggml_permute(ctx0, q, 0, 2, 1, 3); - ggml_tensor * K = ggml_permute(ctx0, k_rep, 0, 2, 1, 3); - ggml_tensor * KQ = mul_mat_checked(ctx0, K, Q, "mul_mat:kq"); - ggml_tensor * KQs = ggml_scale(ctx0, KQ, attn_scale); - ggml_tensor * KQm = ggml_diag_mask_inf(ctx0, KQs, n_past_); - ggml_tensor * KQf = ggml_soft_max(ctx0, KQm); + // Permute Q to [head_dim, n_tokens, n_head, 1] for flash_attn_ext + ggml_tensor * Q_fa = ggml_permute(ctx0, q, 0, 2, 1, 3); + + int64_t kv_len = n_past_ + n_tokens; + ggml_tensor * k_cache = ggml_view_3d(ctx0, memory_k_, + head_dim, kv_len, n_head_kv, + memory_k_->nb[1], memory_k_->nb[2], layer_off_k); + + ggml_tensor * v_cache = ggml_view_3d(ctx0, memory_v_, + head_dim, kv_len, n_head_kv, + memory_v_->nb[1], memory_v_->nb[2], layer_off_v); + + ggml_tensor * attn_fa = ggml_flash_attn_ext( + ctx0, Q_fa, k_cache, v_cache, fa_mask, attn_scale, 0.0f, 0.0f); + + // Output is [head_dim, n_head, n_tokens, 1], flatten to [q_size, n_tokens] for wo + ggml_tensor * attn_cur = ggml_cpy(ctx0, attn_fa, + ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, q_size, n_tokens)); - ggml_tensor * V = ggml_cont(ctx0, ggml_permute(ctx0, v_rep, 1, 2, 0, 3)); - ggml_tensor * KQV = mul_mat_checked(ctx0, V, KQf, "mul_mat:kqv"); - ggml_tensor * KQVm = ggml_permute(ctx0, KQV, 0, 2, 1, 3); - ggml_tensor * attn_cur = ggml_cpy(ctx0, KQVm, - ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, q_size, n_tokens)); ggml_tensor * attn_out = mul_mat_checked(ctx0, layer.wo, attn_cur, "mul_mat:wo"); ggml_tensor * h = ggml_add(ctx0, x, attn_out); @@ -1077,6 +1087,10 @@ bool SlowARModel::eval_cached(const std::vector & flat_tokens, for (int32_t cb = 0; cb < hparams_.num_codebooks; ++cb) { ggml_backend_tensor_set(cb_id_tensors[cb], cb_vals[cb].data(), 0, n_tokens * sizeof(int32_t)); } + + if (fa_mask) { + ggml_backend_tensor_set(fa_mask, mask_data.data(), 0, mask_data.size() * sizeof(ggml_fp16_t)); + } if (ggml_backend_sched_graph_compute(sched_, gf) != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "[eval_cached] sched compute failed\n"); From 4fb34e97b0666e47af7a5cae8f715f2cd88b66ac Mon Sep 17 00:00:00 2001 From: Skyrion9 Date: Thu, 23 Jul 2026 02:52:37 +0300 Subject: [PATCH 2/3] perf: Prefer reshape over ggml_cpy for the flash-attn output to avoid an extra dispatch. --- src/s2_model.cpp | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/s2_model.cpp b/src/s2_model.cpp index 4d3cf78..db8980a 100755 --- a/src/s2_model.cpp +++ b/src/s2_model.cpp @@ -1043,9 +1043,8 @@ bool SlowARModel::eval_cached(const std::vector & flat_tokens, ggml_tensor * attn_fa = ggml_flash_attn_ext( ctx0, Q_fa, k_cache, v_cache, fa_mask, attn_scale, 0.0f, 0.0f); - // Output is [head_dim, n_head, n_tokens, 1], flatten to [q_size, n_tokens] for wo - ggml_tensor * attn_cur = ggml_cpy(ctx0, attn_fa, - ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, q_size, n_tokens)); + // Reshape [head_dim, n_head, n_tokens, 1] → [q_size, n_tokens] (zero-copy view) + ggml_tensor * attn_cur = ggml_reshape_2d(ctx0, attn_fa, q_size, n_tokens); ggml_tensor * attn_out = mul_mat_checked(ctx0, layer.wo, attn_cur, "mul_mat:wo"); From 5aac58250507da24a4dcf57b6b3d9609391e0f2d Mon Sep 17 00:00:00 2001 From: Skyrion9 Date: Thu, 30 Jul 2026 19:25:25 +0300 Subject: [PATCH 3/3] perf: tree-reduced codebook sum and cont-free KV cache writes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Replaced the linear ggml_add chain for codebook embedding summation (depth N for N codebooks) with a pairwise tree reduction (depth ⌈log₂ N⌉), allowing the scheduler to parallelize independent adds. - Removed ggml_cont before the k/v permute into the KV cache. The cont costed an extra intermediate allocation unnecessarily. --- src/s2_model.cpp | 32 ++++++++++++++++++++++++++------ 1 file changed, 26 insertions(+), 6 deletions(-) diff --git a/src/s2_model.cpp b/src/s2_model.cpp index db8980a..2f0d518 100755 --- a/src/s2_model.cpp +++ b/src/s2_model.cpp @@ -942,14 +942,34 @@ bool SlowARModel::eval_cached(const std::vector & flat_tokens, ggml_tensor * x = ggml_get_rows(ctx0, weights_.embeddings, semantic_ids); if (x->type != GGML_TYPE_F32) x = ggml_cast(ctx0, x, GGML_TYPE_F32); - std::vector cb_id_tensors(hparams_.num_codebooks); - ggml_tensor * codebook_sum = nullptr; - for (int32_t cb = 0; cb < hparams_.num_codebooks; ++cb) { + const int32_t num_cb = hparams_.num_codebooks; + std::vector cb_id_tensors(num_cb); + std::vector cb_embs(num_cb); + + for (int32_t cb = 0; cb < num_cb; ++cb) { ggml_tensor * ids = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); cb_id_tensors[cb] = ids; ggml_tensor * emb = ggml_get_rows(ctx0, weights_.codebook_embeddings, ids); if (emb->type != GGML_TYPE_F32) emb = ggml_cast(ctx0, emb, GGML_TYPE_F32); - codebook_sum = (codebook_sum == nullptr) ? emb : ggml_add(ctx0, codebook_sum, emb); + cb_embs[cb] = emb; + } + + ggml_tensor * codebook_sum = nullptr; + if (num_cb > 0) { + std::vector level(cb_embs.begin(), cb_embs.end()); + while (level.size() > 1) { + std::vector next; + next.reserve((level.size() + 1) / 2); + for (size_t i = 0; i < level.size(); i += 2) { + if (i + 1 < level.size()) { + next.push_back(ggml_add(ctx0, level[i], level[i + 1])); + } else { + next.push_back(level[i]); + } + } + level = std::move(next); + } + codebook_sum = level[0]; } if (codebook_sum != nullptr) { @@ -1022,8 +1042,8 @@ bool SlowARModel::eval_cached(const std::vector & flat_tokens, layer_off_v + token_off_v); // Permute k and v to match the cache layout - ggml_tensor * k_perm = ggml_cont(ctx0, ggml_permute(ctx0, k, 0, 2, 1, 3)); - ggml_tensor * v_perm = ggml_cont(ctx0, ggml_permute(ctx0, v, 0, 2, 1, 3)); + ggml_tensor * k_perm = ggml_permute(ctx0, k, 0, 2, 1, 3); + ggml_tensor * v_perm = ggml_permute(ctx0, v, 0, 2, 1, 3); ggml_build_forward_expand(gf, ggml_cpy(ctx0, k_perm, k_slot)); ggml_build_forward_expand(gf, ggml_cpy(ctx0, v_perm, v_slot));