perf: Conditional gallocr caching for Fast-AR decode on pure CPU runs (%4~ faster) - #43
perf: Conditional gallocr caching for Fast-AR decode on pure CPU runs (%4~ faster)#43Skyrion9 wants to merge 3 commits into
Conversation
… %4~ faster CPU inference. - Implemented ggml_gallocr_t with proper lifecycle management. - We check if we're running purely on CPU and use the cached codepath, or fallback to old behavior when any GPU backend is involved. - We effectively eliminate redundant allocation cycles in Fast-AR decoding loop for pure CPU configs.
|
Warning Review limit reached
Next review available in: 59 minutes You've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (3)
📝 WalkthroughWalkthrough
ChangesSlowARModel execution
Estimated code review effort: 3 (Moderate) | ~15–30 minutes Suggested reviewers: 🚥 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.
🧹 Nitpick comments (1)
src/s2_model.cpp (1)
1238-1289: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDeduplicate the tensor-set calls shared by both branches.
The
ggml_backend_tensor_setcalls forhidden0,positions, andprefix_ids(Lines 1254-1258) are byte-for-byte identical to the ones in the GPU-offload branch (Lines 1275-1279). Only the allocation and compute calls actually differ between branches; the input-setting step can be pulled out into a single shared block after allocation succeeds in either branch, avoiding the two copies drifting apart later.♻️ Proposed refactor
if (is_cpu_only) { if (!fast_gallocr_) { fast_gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_cpu_)); if (!fast_gallocr_) { std::fprintf(stderr, "[fast_decode] gallocr creation failed\n"); ggml_free(ctx0); return false; } } if (!ggml_gallocr_alloc_graph(fast_gallocr_, gf)) { std::fprintf(stderr, "[fast_decode] gallocr alloc failed\n"); ggml_free(ctx0); return false; } - - ggml_backend_tensor_set(hidden0, hidden_in.data(), 0, hidden_in.size() * sizeof(float)); - ggml_backend_tensor_set(positions, pos_vals.data(), 0, pos_vals.size() * sizeof(int32_t)); - if (prefix_ids) { - ggml_backend_tensor_set(prefix_ids, prefix_tokens.data(), 0, prefix_tokens.size() * sizeof(int32_t)); - } - if (ggml_backend_graph_compute(backend_cpu_, gf) != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "[fast_decode] cpu compute failed\n"); ggml_free(ctx0); return false; } } else { ggml_backend_sched_reset(fast_sched_); if (!ggml_backend_sched_alloc_graph(fast_sched_, gf)) { std::fprintf(stderr, "[fast_decode] sched alloc failed\n"); ggml_backend_sched_reset(fast_sched_); ggml_free(ctx0); return false; } - - ggml_backend_tensor_set(hidden0, hidden_in.data(), 0, hidden_in.size() * sizeof(float)); - ggml_backend_tensor_set(positions, pos_vals.data(), 0, pos_vals.size() * sizeof(int32_t)); - if (prefix_ids) { - ggml_backend_tensor_set(prefix_ids, prefix_tokens.data(), 0, prefix_tokens.size() * sizeof(int32_t)); - } - + + } + + ggml_backend_tensor_set(hidden0, hidden_in.data(), 0, hidden_in.size() * sizeof(float)); + ggml_backend_tensor_set(positions, pos_vals.data(), 0, pos_vals.size() * sizeof(int32_t)); + if (prefix_ids) { + ggml_backend_tensor_set(prefix_ids, prefix_tokens.data(), 0, prefix_tokens.size() * sizeof(int32_t)); + } + + if (is_cpu_only) { + if (ggml_backend_graph_compute(backend_cpu_, gf) != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "[fast_decode] cpu compute failed\n"); + ggml_free(ctx0); + return false; + } + } else { if (ggml_backend_sched_graph_compute(fast_sched_, gf) != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "[fast_decode] sched compute failed\n"); ggml_backend_sched_reset(fast_sched_); ggml_free(ctx0); return false; } ggml_backend_sched_reset(fast_sched_); }🤖 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 1238 - 1289, Deduplicate the input tensor setup in the fast decode flow: remove the identical ggml_backend_tensor_set calls for hidden0, positions, and optional prefix_ids from both is_cpu_only branches, and place one shared block after the branch-specific graph allocation succeeds and before computation. Preserve the existing allocation, compute, cleanup, and failure behavior for both CPU and scheduled backends.
🤖 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 1238-1289: Deduplicate the input tensor setup in the fast decode
flow: remove the identical ggml_backend_tensor_set calls for hidden0, positions,
and optional prefix_ids from both is_cpu_only branches, and place one shared
block after the branch-specific graph allocation succeeds and before
computation. Preserve the existing allocation, compute, cleanup, and failure
behavior for both CPU and scheduled backends.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 95e6695c-8e20-49e1-a5ed-373709072048
📒 Files selected for processing (2)
include/s2_model.hsrc/s2_model.cpp
- In the GPU path of fast_decode, ggml_backend_sched_reset() was called at the end of the else block, before ggml_backend_tensor_get() read the output logits. ..This released the scheduler's buffer ownership while the output tensor was still being read. - Moved the reset to after tensor_get, guarded by !is_cpu_only since the CPU gallocr path does not use the scheduler.
- Replaced 9 sequential fast_decode() calls per generation step with a single fast_decode_batch() that reuses pre-built compute graphs. - On the CPU path (--fast-decoder-cpu or pure CPU runs), per-prefix-length graphs (n_tok = 2..10) are built once via gallocr and cached in FastGraphSlot structs. Subsequent steps replay the cached graph with only tensor_set overhead, eliminating repeated graph construction and scheduler allocation. A max-size fallback slot (n_tok = 11, all-position logits) covers the first pass while per-prefix slots are being built, then is freed once all 9 slots are ready. - On the GPU path (fast decoder on GPU), fast_decode_batch falls back to the existing sequential fast_decode() loop since the scheduler handles PCIe transfers and graph caching is not beneficial there. - Replaced RAS std::vector with a fixed-size circular buffer
1 participant
Implemented ggml_gallocr_t with proper lifecycle management.
We check if we're running purely on CPU and use the cached codepath, or fallback to old behavior when any GPU backend is involved.
This effectively eliminates redundant allocation cycles in Fast-AR decoding loop for pure CPU configs.
Before/After PR test
CPU: AMD Ryzen 7 5700X3D (8C/16T, pinned to physical cores)
Model:
s2-pro-q8_0.gguf| Backend: CPU (-ngl 0) | Allocator: mimallocPrompt: "hello world, how are you today?" (~204 tokens prefill)
My build script:
My runtime setup:
Summary by CodeRabbit
Performance
Bug Fixes