feat: Flash Attention and modernized KV Cache layout for Slow-AR#40
feat: Flash Attention and modernized KV Cache layout for Slow-AR#40Skyrion9 wants to merge 3 commits into
Conversation
- 17% speedup on Q8 (per-frame ms) - Allows for longer text input before OOM. For example, 593 prefill + 1000/1024 tokens crashes on non-FA version, while FA consumes far less VRAM allowing for successfuly completion (10 GB RX 6700 Vulkan, 1.5 GB background usage, Archlinux)
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthrough
ChangesAttention cache execution
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant eval_cached
participant memory_k_v
participant ggml_flash_attn_ext
eval_cached->>memory_k_v: compute offsets and write permuted k/v
alt n_tokens == 1
eval_cached->>ggml_flash_attn_ext: compute single-token attention
else n_tokens != 1
eval_cached->>memory_k_v: build prefill/concat attention inputs
memory_k_v-->>eval_cached: provide attention inputs
end
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (2)
src/s2_model.cpp (2)
1029-1033: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueRedundant permute in the single-token flash path.
Since this branch only runs when
n_tokens == 1,attn_merged = ggml_permute(ctx0, attn_fa, 0, 2, 1, 3)swaps a size-1 axis with the head axis, which cannot change element ordering. It only adds a non-contiguous view thatggml_cpythen has to stride through.attn_facan likely be copied directly.♻️ Proposed simplification
- ggml_tensor * attn_merged = ggml_permute(ctx0, attn_fa, 0, 2, 1, 3); - - // ggml_cpy safely handles reading from the permuted view into a fresh 2D tensor. - attn_cur = ggml_cpy(ctx0, attn_merged, - ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, q_size, n_tokens)); + // attn_fa is [head_dim, n_head, n_tokens=1, 1]; with n_tokens==1 no reorder is needed. + attn_cur = ggml_cpy(ctx0, attn_fa, + ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 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 1029 - 1033, In the single-token flash-attention path around attn_merged, remove the redundant ggml_permute call and copy directly from attn_fa into the existing fresh 2D F32 tensor via ggml_cpy. Preserve the current output shape and attn_cur assignment.
481-497: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueInconsistent indentation violates the 4-space rule.
The
full_model_offload/get_weight_backendblock (Lines 481-497) and the layer-loop opening (Line 964) mix 0/4/8/12/16-space indentation for logically equivalent nesting levels. Logic itself is unchanged (same routing rules per the AI summary), this is purely a formatting regression introduced by the reformat.🧹 Proposed cleanup
- const bool full_model_offload = - backend_gpu_ != nullptr && - backend_type != BackendType::CUDA && - n_gpu_layers_ == hparams_.block_count; - - auto get_weight_backend = [&](const std::string & name) -> ggml_backend_t { - - if (full_model_offload) { - return backend_gpu_; - } - - if (name == "embeddings.weight" || - name == "codebook_embeddings.weight" || - name == "fast_embeddings.weight" || - name == "norm.weight") { - return backend_cpu_; - } + const bool full_model_offload = + backend_gpu_ != nullptr && + backend_type != BackendType::CUDA && + n_gpu_layers_ == hparams_.block_count; + + auto get_weight_backend = [&](const std::string & name) -> ggml_backend_t { + if (full_model_offload) { + return backend_gpu_; + } + + if (name == "embeddings.weight" || + name == "codebook_embeddings.weight" || + name == "fast_embeddings.weight" || + name == "norm.weight") { + return backend_cpu_; + }As per coding guidelines, "Use 4-space indentation in C++ code" for
**/*.{cpp,h,hpp}.Also applies to: 964-964
🤖 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 481 - 497, Normalize indentation to 4 spaces throughout the full_model_offload declaration and get_weight_backend lambda, including its conditional blocks and returns, without changing routing logic. Apply the same 4-space indentation rule to the layer-loop opening around the referenced loop symbol.Source: Coding guidelines
🤖 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 1029-1033: In the single-token flash-attention path around
attn_merged, remove the redundant ggml_permute call and copy directly from
attn_fa into the existing fresh 2D F32 tensor via ggml_cpy. Preserve the current
output shape and attn_cur assignment.
- Around line 481-497: Normalize indentation to 4 spaces throughout the
full_model_offload declaration and get_weight_backend lambda, including its
conditional blocks and returns, without changing routing logic. Apply the same
4-space indentation rule to the layer-loop opening around the referenced loop
symbol.
…depath. - Copying attn_fa directly into the [q_size, 1] destination is mathematically equivalent and removes a useless node from the compute graph.
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/s2_model.cpp (1)
1064-1064: 🎯 Functional Correctness | 🔴 Critical | ⚡ Quick winFix the
Vpermutation in both attention paths.ggml_mul_mat_checked(V, KQf, ...)needsV->ne[0]to matchKQf->ne[0], butggml_permute(ctx0, v_rep, 1, 2, 0, 3)makesVstart withn_headinstead ofn_tokens. Use2, 0, 1, 3here and atsrc/s2_model.cpp:1231.🤖 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` at line 1064, Update the V tensor permutation in both attention paths, at the visible V construction and its corresponding path near the second occurrence, changing the permutation order to 2, 0, 1, 3 so V->ne[0] matches KQf->ne[0] (n_tokens). Preserve the surrounding ggml_cont and matrix multiplication logic.
🤖 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.
Outside diff comments:
In `@src/s2_model.cpp`:
- Line 1064: Update the V tensor permutation in both attention paths, at the
visible V construction and its corresponding path near the second occurrence,
changing the permutation order to 2, 0, 1, 3 so V->ne[0] matches KQf->ne[0]
(n_tokens). Preserve the surrounding ggml_cont and matrix multiplication logic.
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/s2_model.cpp (1)
1009-1027: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winMake the KV-cache write a dependency of flash attention
src/s2_model.cpp:1009-1027
ggml_build_forward_expand()follows data-flow, not aliasing, soggml_flash_attn_ext()can be scheduled before theggml_cpy()intomemory_k_/memory_v_completes. Chain the copy outputs into the flash-attn path, or add an explicit dependency, and cover bothn_past_ == 0andn_past_ > 0.🤖 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 1009 - 1027, Make the KV-cache writes created by the ggml_cpy calls in the n_tokens == 1 flash-attention path explicit dependencies of ggml_flash_attn_ext, rather than relying on memory aliasing. Ensure both memory_k_ and memory_v_ copies complete before attention for n_past_ == 0 and n_past_ > 0, while preserving the existing cache views and attention behavior.
🧹 Nitpick comments (1)
src/s2_model.cpp (1)
964-964: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winFix the loop indentation.
Line 964 is indented by 8 spaces while the surrounding function body uses 4-space indentation. Align the
forstatement with the other top-level statements ineval_cached.🤖 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` at line 964, Correct the indentation of the for loop in eval_cached so it uses 4 spaces and aligns with other top-level statements in the function; do not change the loop logic.Source: Coding guidelines
🤖 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.
Outside diff comments:
In `@src/s2_model.cpp`:
- Around line 1009-1027: Make the KV-cache writes created by the ggml_cpy calls
in the n_tokens == 1 flash-attention path explicit dependencies of
ggml_flash_attn_ext, rather than relying on memory aliasing. Ensure both
memory_k_ and memory_v_ copies complete before attention for n_past_ == 0 and
n_past_ > 0, while preserving the existing cache views and attention behavior.
---
Nitpick comments:
In `@src/s2_model.cpp`:
- Line 964: Correct the indentation of the for loop in eval_cached so it uses 4
spaces and aligns with other top-level statements in the function; do not change
the loop logic.
AI seems to be confusing pytorch permute with ggml permute. When I tried its suggestion anyways it failed. |
1 participant
The only downside is that prefill takes significantly longer unless you're also using #39 with which it's negligible.
Here are some numbers:
WITH #39 (Notice smaller prefill gap):
FA run #39 :
no FA run #39 :
WITHOUT #39 (bigger prefill gap which #39 remedies & longer text showcases growing gap with FA ahead) :
FA run:
no FA run:
As you can see, the gap increases between them the longer your text is. FA doesn't only allow longer text with smaller memory footprint it also speeds things up. If you TTS many short sentences consecutively you will definitely want #39 .
First graph text was
Second graph text was
https://hazyresearch.stanford.edu/blog/2023-01-12-flashattention-long-sequences
Summary by CodeRabbit
Release Notes