feat: Backend-specific prefill chunk sizes#39
Open
Skyrion9 wants to merge 2 commits into
Open
Conversation
- The default adaptive_gpu_chunk is too conservative, For example it'd fire up 28 times for 204 tokens causing large delays in prefill with 800 ms~ prefill latency if used with FA. - By using larger chunk sizes that are adjusted for each backend, we can safely decrease this delay (378 ms with FA) to increase performance.
|
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
ChangesPrefill chunk sizing
Estimated code review effort: 2 (Simple) | ~10 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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.
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)
800-837: 🩺 Stability & Availability | 🔴 Critical | ⚡ Quick winUse runtime backend identification and fix unreachable CPU chunking.
The current chunk sizing implementation has two critical issues:
- Multi-backend segfault risk: Relying on compile-time macros (
#if defined(GGML_USE_CUDA)) ignores the actual runtime backend. In a binary compiled with both CUDA and Vulkan, selecting Vulkan at runtime will incorrectly use the CUDA chunk size (512), causing segmentation faults.- Dead-code CPU chunking: The CPU fallback block (
#else) is nested inside theif (backend_gpu_ != nullptr ...)block, making it unreachable for CPU-only runs. As a result, CPU executions process all tokens in a single giant chunk, circumventing the intended memory bandwidth optimizations.Pulling the
adaptive_chunkcalculation outside the GPU guard and checking the backend name at runtime resolves both issues. As an added benefit, wrappingn_threadswithresolve_n_threads()ensures the chunk size scales correctly whenn_threadsis passed as0(auto-detect).🔒️ Proposed fix to apply runtime checks and fix CPU sizing
- int32_t chunk_tokens = n_tokens; - if (backend_gpu_ != nullptr && n_gpu_layers_ > 0) { - int32_t semantic_prompt_tokens = 0; - for (int32_t t = 0; t < n_tokens; ++t) { - const int32_t semantic = flat_tokens[static_cast<size_t>(t) * codebook_dim]; - if (semantic >= hparams_.semantic_begin_id && - semantic <= hparams_.semantic_end_id) { - semantic_prompt_tokens++; - if (semantic_prompt_tokens > 1) { - break; - } - } - } - - if (semantic_prompt_tokens > 1 && - backend_requires_single_token_semantic_prefill(backend_gpu_)) { - chunk_tokens = 1; - } else { - // Backend-specific prefill chunk sizes for naive attention path. - // Vulkan: Limited by ggml_soft_max shader workgroup/shared memory limits - // CUDA/Metal: Highly optimized soft_max kernels handle large batches - // CPU: Memory bandwidth bound, scale with available parallelism - int32_t adaptive_gpu_chunk = 64; // Safe fallback - -#if defined(GGML_USE_CUDA) || defined(GGML_USE_METAL) - adaptive_gpu_chunk = 512; -#elif defined(GGML_USE_VULKAN) - adaptive_gpu_chunk = 64; -#else - // CPU or unknown backend: scale with thread count, capped at 256 - adaptive_gpu_chunk = std::min(256, std::max(64, n_threads * 16)); -#endif - - if (n_tokens > adaptive_gpu_chunk) { - chunk_tokens = adaptive_gpu_chunk; - } - } - } + // Backend-specific prefill chunk sizes for naive attention path. + // Vulkan: Limited by ggml_soft_max shader workgroup/shared memory limits + // CUDA/Metal: Highly optimized soft_max kernels handle large batches + // CPU: Memory bandwidth bound, scale with available parallelism + int32_t adaptive_chunk = std::min<int32_t>(256, std::max<int32_t>(64, resolve_n_threads(n_threads) * 16)); + + if (backend_gpu_ != nullptr && n_gpu_layers_ > 0) { + const char * backend_name = ggml_backend_name(backend_gpu_); + if (backend_name != nullptr) { + std::string name(backend_name); + if (name.find("CUDA") != std::string::npos || + name.find("Metal") != std::string::npos) { + adaptive_chunk = 512; + } else if (name.find("Vulkan") != std::string::npos) { + adaptive_chunk = 64; + } + } + + int32_t semantic_prompt_tokens = 0; + for (int32_t t = 0; t < n_tokens; ++t) { + const int32_t semantic = flat_tokens[static_cast<size_t>(t) * codebook_dim]; + if (semantic >= hparams_.semantic_begin_id && + semantic <= hparams_.semantic_end_id) { + semantic_prompt_tokens++; + if (semantic_prompt_tokens > 1) { + break; + } + } + } + + if (semantic_prompt_tokens > 1 && + backend_requires_single_token_semantic_prefill(backend_gpu_)) { + adaptive_chunk = 1; + } + } + + int32_t chunk_tokens = n_tokens; + if (n_tokens > adaptive_chunk) { + chunk_tokens = adaptive_chunk; + }🤖 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 800 - 837, Move adaptive chunk-size calculation out of the backend_gpu_ guard so CPU execution also applies it, and replace compile-time GGML_USE_* selection with runtime backend identification using the available backend name/type symbols. Select CUDA/Metal, Vulkan, or the CPU fallback based on the active backend, and compute the CPU size with resolve_n_threads() so an n_threads value of 0 is auto-detected. Preserve the existing single-token semantic prefill override and apply the resulting chunk_tokens limit consistently.
🤖 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 800-837: Move adaptive chunk-size calculation out of the
backend_gpu_ guard so CPU execution also applies it, and replace compile-time
GGML_USE_* selection with runtime backend identification using the available
backend name/type symbols. Select CUDA/Metal, Vulkan, or the CPU fallback based
on the active backend, and compute the CPU size with resolve_n_threads() so an
n_threads value of 0 is auto-detected. Preserve the existing single-token
semantic prefill override and apply the resulting chunk_tokens limit
consistently.
…hunking - The CPU fallback was nested inside (backend_gpu_ != nullptr) guard. We can instead establish a CPU baseline first and adapt it to backend down the line.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Before PR (FA disabled, no chunking for some reason, optimal prefill delay):
After PR (FA disabled, chunk preset of 64 for Vulkan):
Before PR (FA enabled):
After PR (FA enabled):
Don't mind the small difference in average and loop. There was some background usage so not a neat/clean comparison, the massive 500 ms~ prefill difference is what I wanted to showcase.
I'm unable to test on cuda or metal, and some higher numbers causes segmentation faults on Vulkan perhaps a driver or ggml-vulkan issue. However those should be more stable?
Needs testing by anyone with the hardware to determine if they're stable and which values might benefit performance more.
Summary by CodeRabbit
Summary by CodeRabbit