perf: replace standard attention with unified flash attention

[Skyrion9]

Replaces the standard (naive) attention implementation in SlowARModel::eval_cached with a single unified ggml_flash_attn_ext call for all token counts.

The previous attention path used a multi-step standard attention pipeline:

1. Concatenate past KV cache with current K/V
2. repeat_interleave_heads to expand KV heads for GQA
3. mul_mat to compute QK scores
4. ggml_scale + ggml_diag_mask_inf + ggml_soft_max
5. mul_mat to compute KQV output

This required 8+ separate GPU shader dispatches per layer and materialized the full $N \times N$ attention matrix in VRAM, creating an $O(N^2)$ memory bandwidth bottleneck.

Changes:

KV cache layout (init_kv_cache):

• Dimensions 1 and 2 are swapped from [head_dim, n_head_kv, max_seq_len, n_layer] to [head_dim, max_seq_len, n_head_kv, n_layer], placing the sequence dimension at ne[1] to match the view shape expected by ggml_flash_attn_ext.
• Token offset calculation updated from nb[2] to nb[1] accordingly.
• K/V write path now includes a ggml_permute + ggml_cont to transpose from the QKV projection output shape to the cache layout.

Causal mask generation (eval_cached):

• For n_tokens > 1: an F16 causal mask of shape [kv_len, n_tokens, 1, 1] is built on the CPU and uploaded to the backend before graph compute. Entry [j][i] is -INFINITY if i > n_past_ + j, 0.0f otherwise.
• For n_tokens == 1: fa_mask remains nullptr. This is not a separate code path — the same ggml_flash_attn_ext call executes regardless of token count.

Attention computation (eval_cached):

• The entire standard attention block (concat, repeat_interleave, KQ mul_mat, scale, diag_mask_inf, soft_max, KQV mul_mat) is replaced with a single ggml_flash_attn_ext call.
• Q is permuted to [head_dim, n_tokens, n_head, 1]; K/V cache views are [head_dim, kv_len, n_head_kv].
• Output is flattened to [q_size, n_tokens] for the wo projection.

As a "sideffect" #38 is no longer needed for Linux Radv (Vulkan) on RDNA2 - RX 6700 to function without segfaulting when using prefill_fast, likely due to ggml_soft_max limitations on ggml vulkan backend for this particular setup.

Impact:

Metric	Before	After
Shader dispatches per layer (attention)	~8–10	~3
Attention matrix materialized in VRAM	Yes ($O(N^2)$)	No (tiled internally)

Benchmarks:
Can be compared as the improvement over #44 as those graphs implemented the deprecated implementation of FA #40 which wasn't uniform in its approach. You can see that we're running ~20% faster across Total RTF if you compare the R1/R2 results. As before, the non-R1/R2 runs are upstream using prefill instead of prefill_fast as fast_prefill was unusable on my system before this PR.

Summary by CodeRabbit

Summary by CodeRabbit

• Performance
  • Improved attention processing for multi-token evaluations by using a flash-attention-style execution path.
  • Optimized the key/value cache layout to better match high-throughput attention access patterns.
• Compatibility
  • Maintained the same attention outputs and downstream behavior while changing the internal attention computation strategy.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: f8d0b491-f9b5-4fb2-be2d-b406d8b0c802

📥 Commits

Reviewing files that changed from the base of the PR and between a83467c and 4fb34e9.

📒 Files selected for processing (1)

• src/s2_model.cpp

🚧 Files skipped from review as they are similar to previous changes (1)

• src/s2_model.cpp

---

📝 Walkthrough

Walkthrough

SlowARModel reorders KV cache tensors and replaces cached attention computation with ggml_flash_attn_ext. Multi-token evaluation creates a causal mask, writes K/V through updated cache views, and uploads the mask before graph execution.

Changes

Flash-attention integration

Layer / File(s)	Summary
Reordered KV cache layout src/s2_model.cpp	memory_k_ and memory_v_ dimensions are reordered to match flash-attention cache indexing.
Masked cached attention execution src/s2_model.cpp	Multi-token evaluation creates a causal mask, writes K/V into updated cache views, permutes Q/K/V, invokes ggml_flash_attn_ext, and uploads the mask before graph execution.

Estimated code review effort: 4 (Complex) | ~45 minutes

Sequence Diagram(s)

sequenceDiagram
  participant eval_cached
  participant KVCache
  participant ggml_flash_attn_ext
  participant Backend
  eval_cached->>KVCache: write current K/V into cache views
  eval_cached->>ggml_flash_attn_ext: pass permuted Q/K/V and fa_mask
  ggml_flash_attn_ext-->>eval_cached: return attention output
  eval_cached->>Backend: upload fa_mask before graph execution

Loading

Possibly related PRs

• rodrigomatta/s2.cpp#40: Updates the same KV cache layout and cached attention path for flash attention.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly matches the main change: replacing standard attention with unified flash attention for performance.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

🧹 Nitpick comments (1)

src/s2_model.cpp (1)

1046-1049: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick win

Prefer reshape over ggml_cpy for the flash-attn output to avoid an extra dispatch.

attn_fa is already contiguous in [head_dim, n_head, n_tokens, 1] layout (output of ggml_flash_attn_ext), so merging the first two dims via ggml_reshape_2d/ggml_cont_2d avoids allocating a new tensor and performing a full copy. This directly furthers the PR's stated goal of reducing shader dispatches per layer.

♻️ Proposed refactor

-        // 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));
+        // Output is [head_dim, n_head, n_tokens, 1], flatten to [q_size, n_tokens] for wo
+        ggml_tensor * attn_cur = ggml_reshape_2d(ctx0, attn_fa, q_size, n_tokens);

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/s2_model.cpp` around lines 1046 - 1049, Replace the ggml_cpy call
constructing attn_cur with a reshape-based operation, using ggml_reshape_2d (or
ggml_cont_2d if required by layout) to merge attn_fa’s head_dim and n_head
dimensions into q_size while preserving n_tokens. Avoid allocating and copying
the flash-attention output, and keep the resulting [q_size, n_tokens] shape for
the subsequent wo path.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@src/s2_model.cpp`:
- Around line 1046-1049: Replace the ggml_cpy call constructing attn_cur with a
reshape-based operation, using ggml_reshape_2d (or ggml_cont_2d if required by
layout) to merge attn_fa’s head_dim and n_head dimensions into q_size while
preserving n_tokens. Avoid allocating and copying the flash-attention output,
and keep the resulting [q_size, n_tokens] shape for the subsequent wo path.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 67176eb1-7c75-4d3b-8d45-074c72b3b06b

📥 Commits

Reviewing files that changed from the base of the PR and between 2c33261 and a83467c.

📒 Files selected for processing (1)

• src/s2_model.cpp