feat: mmap lazy loading and phase-gated VRAM swap#44
Conversation
Replaced fread pipeline with a cross-platform memory-mapped file/mmap. This reduces boot times by allowing lazy weight loading for sub-models and no copy DMA. - Introducing MappedFile RAII wrapper for zero-copy memory-mapped file I/O supporting POSIX (mmap/madvise) and Windows (CreateFileMapping). - Decoupled GGUF metadata parsing from VRAM/RAM buffer allocation. - Audio Codec consumes 0 MB VRAM at init, Weight buffers are allocated on-demand, VQ codebook caches are populated directly from the mmap pointer into system RAM. - Slow-AR model weights are page-faulted from the mmap pointer to VRAM, bypassing intermediate buffers. - We've replaced read_all_tensor_data and read_tensor_data in favor of lazy loading.
… for server mode Utilizes mmap and lazy loading introduced in the previous commit to dynamically manage VRAM occupancy. Intelligently swapping in and out the required submodels depending on which phase of the processing we're at. This minimizes both peak and idle VRAM usage, allowing running larger models without OOM and increases speeds by reducing memory pressure. - Phase-Gated VRAM Swapping: Slow-AR weights and KV cache are freed immediately after generation completes, right before Audio Codec weights are restored for decode. .. We don't need the 4.2 GB (Q8_0) SlowAR to occupy VRAM as we're running inference on the Audio Codec part and possibly crash via OOM. .. Without this system, Q8_0 would hit 7-7.5 GB VRAM usage during final phase of the processing (Audio Codec) in one sentence long generation. Now it's just ~2.3 GB (Vulkan, Linux latest MESA) - CLI flags --no-vram-swap (opt out) and --hot-swap (opt in) to customize behavior. - vram-swap retains the OS page cache between requests, pagefaulting we read from RAM instead of disk, this only takes a few seconds. .. Also keeps compute buffers etc. in VRAM which are relatively small (~168 MB Vulkan) so we can immediately begin processing. .. The gguf occupies system RAM instead of VRAM, however, this occupancy is not "locked" meaning OS will free it for other applications as needed. .. This is basically tells the OS "Here's this memory pool that maps to compute buffer, keep it alive but also don't hesitate to free the memory if other apps need it." - Aggressive hot-swap mode goes a step further and explicitly instructs the kernel to reclaim memory. .. This is optimal if you want minimal, 100 MB RAM + 25 MB VRAM idles without bothering the OS and have the model on flash storage. - Backend synchronization via ggml_backend_synchronize to ensure different backends (Vulkan, CUDA, Metal, etc.) reclaim memory sooner than later to prevent PCIe thrashing. .. This is critical to reduce peak VRAM usage therefore allowing us to run larger quants without filling VRAM to the brim. Also reduces pressure on other apps and their VRAM occupancies. - Background prefetching - spawns a background thread to restore Slow-AR weights concurrently with CPU-bound voice profile loading to hide PCIe latency. - Thread safe, all synthesis entry points join pending_offload_thread_ before proceeding to prevent race conditions between background eviction and new weight restoration.
📝 WalkthroughWalkthroughThe PR adds cross-platform mmap support for GGUF files, defers model and codec weight allocation, exposes GPU residency controls, and coordinates VRAM swapping through pipeline configuration, CLI flags, and background offload threads. ChangesMapped weight lifecycle and VRAM hot-swapping
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Pipeline
participant SlowARModel
participant AudioCodec
participant MappedFile
Pipeline->>SlowARModel: restore_weights_to_gpu()
SlowARModel->>MappedFile: read mapped tensor data
Pipeline->>AudioCodec: restore_weights_to_gpu()
AudioCodec->>MappedFile: read mapped tensor data
Pipeline->>SlowARModel: free_gpu_weights()
Pipeline->>AudioCodec: free_gpu_weights()
Pipeline->>MappedFile: drop_page_cache()
Possibly related PRs
🚥 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.
Actionable comments posted: 3
🧹 Nitpick comments (6)
src/s2_codec.cpp (2)
942-942: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winBinding a
const&to a conditional with a temporary — unnecessary full-set copy (cppcheckdanglingTemporaryLifetime). Because the?:yields a prvalue, theModel->weight_tensor_set()branch is materialized into a temporary and the whole set is copied even whenModel != nullptr; the reference then binds to that (lifetime-extended) temporary. It is not actually dangling, but the copy is wasteful. Prefer a pointer to avoid the copy and silence the analyzer.♻️ Use a pointer instead of a copied reference
- const auto & model_weights = Model ? Model->weight_tensor_set() : std::unordered_set<ggml_tensor*>(); + static const std::unordered_set<ggml_tensor*> empty_weights; + const auto & model_weights = Model ? Model->weight_tensor_set() : empty_weights;🤖 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_codec.cpp` at line 942, Update the model_weights initialization in the surrounding codec flow to use a pointer to Model->weight_tensor_set() when Model is non-null, avoiding the conditional’s temporary full-set copy; provide an appropriate empty-set fallback for the null-Model case and adjust subsequent accesses to dereference the pointer while preserving existing behavior.Source: Linters/SAST tools
129-155: 🩺 Stability & Availability | 🔵 Trivial | 🏗️ Heavy lift
allocate_codec_buffersshould respect per-allocation caps. On backends like Vulkan, one largeggml_backend_buft_alloc_buffer()can fail once the request exceeds the backend’s max allocation size, even when total VRAM is available. Chunk this likeallocate_weight_buffersdoes, or guard againstggml_backend_buft_get_max_size()before a single allocation.🤖 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_codec.cpp` around lines 129 - 155, Update allocate_codec_buffers to respect the buffer type’s maximum allocation size, using ggml_backend_buft_get_max_size() and the chunking strategy established by allocate_weight_buffers. Ensure no single ggml_backend_buft_alloc_buffer call exceeds that cap while preserving alignment, total-byte accounting, error reporting, and output-buffer initialization.src/s2_model.cpp (1)
1194-1226: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueRestore path re-allocates and re-copies CPU weights unnecessarily.
free_gpu_weightsfrees onlymodel_bufs_gpu, butrestore_weights_to_gpuresetsweights_allocated_and callsallocate_and_load_weights, which unconditionally re-runs the CPU branch (allocate_weight_buffers(backend_cpu_, ...)frees and re-places all CPU tensors and re-copies them from mmap). Consider guarding the CPU allocation so only GPU weights are re-materialized on restore.🤖 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 1194 - 1226, Update allocate_and_load_weights so the restore path does not reallocate or recopy CPU weights when they remain valid; guard the backend_cpu_ allocation and CPU tensor loading using the same GPU-restore state or condition that indicates CPU weights are already materialized, while preserving initial CPU allocation and loading behavior.src/s2_pipeline.cpp (2)
732-754: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winBackground pre-fetch thread duplicated with the streaming path.
This block (condition, thread body, join-on-failure/success) is duplicated near-verbatim in
synthesize_streaming_raw(Lines 978-1002). Extracting a small helper (e.g.start_background_model_prefetch()/join_background_model_prefetch()) would avoid future logic drift between the two call sites.🤖 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_pipeline.cpp` around lines 732 - 754, Extract the duplicated VRAM pre-fetch lifecycle from the current function and synthesize_streaming_raw into shared helpers such as start_background_model_prefetch and join_background_model_prefetch. Preserve the existing enable_vram_swap/is_persistent/model_prefers_gpu_/is_weights_on_gpu condition, background acquire_compute_resources/restore_weights_to_gpu work, and joins on both failure and success paths.
893-916: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winHot-swap cleanup thread duplicated with the streaming path.
This entire persistent/hot-swap cleanup block (compute-buffer release + background offload thread that frees GPU weights and drops page caches) is duplicated near-verbatim in
synthesize_streaming_prompt_codes_locked(Lines 1264-1288). Same drift risk as the pre-fetch duplication above; consider a shared private helper.🤖 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_pipeline.cpp` around lines 893 - 916, Extract the duplicated persistent hot-swap cleanup logic from the current block and synthesize_streaming_prompt_codes_locked into a shared private helper. The helper must release compute buffers, asynchronously free model and codec GPU weights, drop both mapped-file page caches, and update pending_offload_thread_; call it from both paths while preserving existing non-hot-swap and single-shot behavior.src/main.cpp (1)
93-94: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win
--hot-swapcan silently no-op depending on flag combinations.
--hot-swapcleanup is gated onparams.is_persistent(only settruefor--server, Line 316), and the entire hot-swap block ins2_pipeline.cppis additionally nested underparams.enable_vram_swap. So--hot-swapalone in CLI/file mode does nothing, and--hot-swap --no-vram-swaptogether also silently disables hot-swap. Consider warning the user in these cases, similar to the existing warnings pattern (e.g. Lines 271-277).💡 Suggested warning
+ if (params.enable_hot_swap && !use_server) { + safe_print_error_ln("Warning: --hot-swap has no effect outside --server mode.\n"); + } + if (params.enable_hot_swap && !params.enable_vram_swap) { + safe_print_error_ln("Warning: --hot-swap has no effect when combined with --no-vram-swap.\n"); + }Also applies to: 191-202
🤖 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/main.cpp` around lines 93 - 94, Add validation warnings in the argument/configuration handling near the existing warnings and the --hot-swap option: warn when --hot-swap is enabled without persistent server mode, and when it is combined with --no-vram-swap, since the s2_pipeline.cpp cleanup path requires both conditions. Keep the existing behavior unchanged while clearly informing users that hot-swap will be inactive.
🤖 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.
Inline comments:
In `@src/s2_codec.cpp`:
- Around line 934-936: Update reset_codec_impl to release impl_->backend_cpu
during reset and destruction, matching the existing cleanup for impl_->backend,
impl_->ctx_w, and impl_->model_buf. Ensure the pointer is cleared after freeing,
and preserve the current initialization behavior in the backend_cpu setup block.
- Around line 1528-1544: Update read_f32 in refresh_host_caches_from_mmap to
throw an error for tensor types other than GGML_TYPE_F32 and GGML_TYPE_F16,
preventing unsupported data from being returned as zero-filled values. Also
update the base calculation in the same function to widen c/code before
multiplication, avoiding int32_t overflow while preserving the existing indexing
behavior.
In `@src/s2_pipeline.cpp`:
- Around line 442-455: Update the codec state assignment in the initialization
flow so codec_prefers_gpu_ reflects the codec’s actual post-load backend, not
the pre-fallback use_gpu_codec request. Reuse the codec’s existing GPU-residency
query or equivalent state established after the GPU load/fallback logic, then
keep the VRAM State Machine diagnostics and downstream swap decisions based on
that corrected flag.
---
Nitpick comments:
In `@src/main.cpp`:
- Around line 93-94: Add validation warnings in the argument/configuration
handling near the existing warnings and the --hot-swap option: warn when
--hot-swap is enabled without persistent server mode, and when it is combined
with --no-vram-swap, since the s2_pipeline.cpp cleanup path requires both
conditions. Keep the existing behavior unchanged while clearly informing users
that hot-swap will be inactive.
In `@src/s2_codec.cpp`:
- Line 942: Update the model_weights initialization in the surrounding codec
flow to use a pointer to Model->weight_tensor_set() when Model is non-null,
avoiding the conditional’s temporary full-set copy; provide an appropriate
empty-set fallback for the null-Model case and adjust subsequent accesses to
dereference the pointer while preserving existing behavior.
- Around line 129-155: Update allocate_codec_buffers to respect the buffer
type’s maximum allocation size, using ggml_backend_buft_get_max_size() and the
chunking strategy established by allocate_weight_buffers. Ensure no single
ggml_backend_buft_alloc_buffer call exceeds that cap while preserving alignment,
total-byte accounting, error reporting, and output-buffer initialization.
In `@src/s2_model.cpp`:
- Around line 1194-1226: Update allocate_and_load_weights so the restore path
does not reallocate or recopy CPU weights when they remain valid; guard the
backend_cpu_ allocation and CPU tensor loading using the same GPU-restore state
or condition that indicates CPU weights are already materialized, while
preserving initial CPU allocation and loading behavior.
In `@src/s2_pipeline.cpp`:
- Around line 732-754: Extract the duplicated VRAM pre-fetch lifecycle from the
current function and synthesize_streaming_raw into shared helpers such as
start_background_model_prefetch and join_background_model_prefetch. Preserve the
existing enable_vram_swap/is_persistent/model_prefers_gpu_/is_weights_on_gpu
condition, background acquire_compute_resources/restore_weights_to_gpu work, and
joins on both failure and success paths.
- Around line 893-916: Extract the duplicated persistent hot-swap cleanup logic
from the current block and synthesize_streaming_prompt_codes_locked into a
shared private helper. The helper must release compute buffers, asynchronously
free model and codec GPU weights, drop both mapped-file page caches, and update
pending_offload_thread_; call it from both paths while preserving existing
non-hot-swap and single-shot behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: bc17f0ca-6456-4aec-9e41-64291ff19de0
📒 Files selected for processing (10)
CMakeLists.txtinclude/s2_codec.hinclude/s2_mapped_file.hinclude/s2_model.hinclude/s2_pipeline.hsrc/main.cppsrc/s2_codec.cppsrc/s2_mapped_file.cppsrc/s2_model.cppsrc/s2_pipeline.cpp
| if (!impl_->backend_cpu) { | ||
| impl_->backend_cpu = ggml_backend_cpu_init(); | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
backend_cpu is initialized but never freed (and appears unused). reset_codec_impl (Lines 193-208) frees impl.backend, impl.ctx_w, and impl.model_buf, but not impl.backend_cpu, so this ggml_backend_cpu_init() leaks on every reset/destruction. It also does not appear to be used by the allocation/compute paths (those use impl_->backend). Either wire backend_cpu into the lifecycle and free it in reset_codec_impl, or drop the field.
🛠️ Free backend_cpu in reset_codec_impl
if (impl.backend) {
ggml_backend_free(impl.backend);
impl.backend = nullptr;
}
+ if (impl.backend_cpu) {
+ ggml_backend_free(impl.backend_cpu);
+ impl.backend_cpu = nullptr;
+ }
impl = AudioCodec::Impl();🤖 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_codec.cpp` around lines 934 - 936, Update reset_codec_impl to release
impl_->backend_cpu during reset and destruction, matching the existing cleanup
for impl_->backend, impl_->ctx_w, and impl_->model_buf. Ensure the pointer is
cleared after freeing, and preserve the current initialization behavior in the
backend_cpu setup block.
| bool AudioCodec::refresh_host_caches_from_mmap() { | ||
| if (!impl_ || !impl_->mapped_gguf_.is_open()) return false; | ||
| auto read_f32 = [&](const std::string& name) -> std::vector<float> { | ||
| ggml_tensor* t = ggml_get_tensor(impl_->ctx_w, name.c_str()); | ||
| if (!t) throw std::runtime_error("missing vq tensor: " + name); | ||
| auto it = impl_->tensor_offsets.find(t); | ||
| if (it == impl_->tensor_offsets.end()) throw std::runtime_error("missing offset"); | ||
| const size_t n = ggml_nelements(t); | ||
| std::vector<float> out(n); | ||
| const uint8_t* src = impl_->mapped_gguf_.data() + impl_->gguf_data_offset + it->second; | ||
| if (t->type == GGML_TYPE_F32) std::memcpy(out.data(), src, n * sizeof(float)); | ||
| else if (t->type == GGML_TYPE_F16) { | ||
| const ggml_fp16_t* tmp = reinterpret_cast<const ggml_fp16_t*>(src); | ||
| for (size_t i = 0; i < n; ++i) out[i] = ggml_fp16_to_fp32(tmp[i]); | ||
| } | ||
| return out; | ||
| }; |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
refresh_host_caches_from_mmap silently zero-fills unsupported tensor types. read_f32 handles only GGML_TYPE_F32/GGML_TYPE_F16; any other type falls through with out left zero-initialized, producing silently-wrong VQ codebooks instead of a hard failure. The prior tensor_to_f32 threw on unsupported types. Also, base = c * cb_dim (Line 1553) multiplies two int32_t before widening to size_t, risking overflow for large codebooks; cast first as done elsewhere (static_cast<size_t>(code) * codebook_dim).
🛡️ Fail loudly on unsupported types
if (t->type == GGML_TYPE_F32) std::memcpy(out.data(), src, n * sizeof(float));
else if (t->type == GGML_TYPE_F16) {
const ggml_fp16_t* tmp = reinterpret_cast<const ggml_fp16_t*>(src);
for (size_t i = 0; i < n; ++i) out[i] = ggml_fp16_to_fp32(tmp[i]);
- }
+ } else {
+ throw std::runtime_error("unsupported vq tensor type: " + name);
+ }📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| bool AudioCodec::refresh_host_caches_from_mmap() { | |
| if (!impl_ || !impl_->mapped_gguf_.is_open()) return false; | |
| auto read_f32 = [&](const std::string& name) -> std::vector<float> { | |
| ggml_tensor* t = ggml_get_tensor(impl_->ctx_w, name.c_str()); | |
| if (!t) throw std::runtime_error("missing vq tensor: " + name); | |
| auto it = impl_->tensor_offsets.find(t); | |
| if (it == impl_->tensor_offsets.end()) throw std::runtime_error("missing offset"); | |
| const size_t n = ggml_nelements(t); | |
| std::vector<float> out(n); | |
| const uint8_t* src = impl_->mapped_gguf_.data() + impl_->gguf_data_offset + it->second; | |
| if (t->type == GGML_TYPE_F32) std::memcpy(out.data(), src, n * sizeof(float)); | |
| else if (t->type == GGML_TYPE_F16) { | |
| const ggml_fp16_t* tmp = reinterpret_cast<const ggml_fp16_t*>(src); | |
| for (size_t i = 0; i < n; ++i) out[i] = ggml_fp16_to_fp32(tmp[i]); | |
| } | |
| return out; | |
| }; | |
| bool AudioCodec::refresh_host_caches_from_mmap() { | |
| if (!impl_ || !impl_->mapped_gguf_.is_open()) return false; | |
| auto read_f32 = [&](const std::string& name) -> std::vector<float> { | |
| ggml_tensor* t = ggml_get_tensor(impl_->ctx_w, name.c_str()); | |
| if (!t) throw std::runtime_error("missing vq tensor: " + name); | |
| auto it = impl_->tensor_offsets.find(t); | |
| if (it == impl_->tensor_offsets.end()) throw std::runtime_error("missing offset"); | |
| const size_t n = ggml_nelements(t); | |
| std::vector<float> out(n); | |
| const uint8_t* src = impl_->mapped_gguf_.data() + impl_->gguf_data_offset + it->second; | |
| if (t->type == GGML_TYPE_F32) std::memcpy(out.data(), src, n * sizeof(float)); | |
| else if (t->type == GGML_TYPE_F16) { | |
| const ggml_fp16_t* tmp = reinterpret_cast<const ggml_fp16_t*>(src); | |
| for (size_t i = 0; i < n; ++i) out[i] = ggml_fp16_to_fp32(tmp[i]); | |
| } else { | |
| throw std::runtime_error("unsupported vq tensor type: " + name); | |
| } | |
| return out; | |
| }; |
🤖 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_codec.cpp` around lines 1528 - 1544, Update read_f32 in
refresh_host_caches_from_mmap to throw an error for tensor types other than
GGML_TYPE_F32 and GGML_TYPE_F16, preventing unsupported data from being returned
as zero-filled values. Also update the base calculation in the same function to
widen c/code before multiplication, avoiding int32_t overflow while preserving
the existing indexing behavior.
| initialized_ = true; | ||
|
|
||
| model_prefers_gpu_ = model().is_weights_on_gpu(); | ||
| codec_prefers_gpu_ = use_gpu_codec; | ||
|
|
||
| if (model_prefers_gpu_ && codec_prefers_gpu_) { | ||
| safe_print_ln("[Pipeline] VRAM State Machine: Case 1 (Both prefer GPU) - Codec is lazily allocated on demand."); | ||
| } else if (model_prefers_gpu_ && !codec_prefers_gpu_) { | ||
| safe_print_ln("[Pipeline] VRAM State Machine: Case 2 (Slow-AR GPU, Codec CPU) - Ready."); | ||
| } else if (!model_prefers_gpu_ && codec_prefers_gpu_) { | ||
| safe_print_ln("[Pipeline] VRAM State Machine: Case 3 (Slow-AR CPU, Codec GPU) - Codec is lazily allocated on demand."); | ||
| } else { | ||
| safe_print_ln("[Pipeline] VRAM State Machine: Case 4 (All CPU) - Ready."); | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
codec_prefers_gpu_ records intent, not the actual post-fallback backend.
model_prefers_gpu_ is derived from the actual result via model().is_weights_on_gpu() (Line 444), but codec_prefers_gpu_ is assigned use_gpu_codec (Line 445) — a value computed before the GPU load attempt/fallback at Lines 391-417. If the codec's GPU load fails and it falls back to CPU (Lines 402-409), codec_prefers_gpu_ stays true even though the codec is actually CPU-backed. This corrupts the "VRAM State Machine" diagnostics (wrong case printed) and feeds every downstream swap decision (Lines 815, 821, 859, 873, 911, 1045, 1283) with a flag that doesn't reflect reality. It happens to be masked today only because AudioCodec::restore_weights_to_gpu()/free_gpu_weights() no-op when there are no GPU weights, but the state is still wrong and fragile.
🐛 Proposed fix
- codec_prefers_gpu_ = use_gpu_codec;
+ codec_prefers_gpu_ = codec().is_weights_on_gpu();📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| initialized_ = true; | |
| model_prefers_gpu_ = model().is_weights_on_gpu(); | |
| codec_prefers_gpu_ = use_gpu_codec; | |
| if (model_prefers_gpu_ && codec_prefers_gpu_) { | |
| safe_print_ln("[Pipeline] VRAM State Machine: Case 1 (Both prefer GPU) - Codec is lazily allocated on demand."); | |
| } else if (model_prefers_gpu_ && !codec_prefers_gpu_) { | |
| safe_print_ln("[Pipeline] VRAM State Machine: Case 2 (Slow-AR GPU, Codec CPU) - Ready."); | |
| } else if (!model_prefers_gpu_ && codec_prefers_gpu_) { | |
| safe_print_ln("[Pipeline] VRAM State Machine: Case 3 (Slow-AR CPU, Codec GPU) - Codec is lazily allocated on demand."); | |
| } else { | |
| safe_print_ln("[Pipeline] VRAM State Machine: Case 4 (All CPU) - Ready."); | |
| } | |
| initialized_ = true; | |
| model_prefers_gpu_ = model().is_weights_on_gpu(); | |
| codec_prefers_gpu_ = codec().is_weights_on_gpu(); | |
| if (model_prefers_gpu_ && codec_prefers_gpu_) { | |
| safe_print_ln("[Pipeline] VRAM State Machine: Case 1 (Both prefer GPU) - Codec is lazily allocated on demand."); | |
| } else if (model_prefers_gpu_ && !codec_prefers_gpu_) { | |
| safe_print_ln("[Pipeline] VRAM State Machine: Case 2 (Slow-AR GPU, Codec CPU) - Ready."); | |
| } else if (!model_prefers_gpu_ && codec_prefers_gpu_) { | |
| safe_print_ln("[Pipeline] VRAM State Machine: Case 3 (Slow-AR CPU, Codec GPU) - Codec is lazily allocated on demand."); | |
| } else { | |
| safe_print_ln("[Pipeline] VRAM State Machine: Case 4 (All CPU) - Ready."); | |
| } |
🤖 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_pipeline.cpp` around lines 442 - 455, Update the codec state
assignment in the initialization flow so codec_prefers_gpu_ reflects the codec’s
actual post-load backend, not the pre-fallback use_gpu_codec request. Reuse the
codec’s existing GPU-residency query or equivalent state established after the
GPU load/fallback logic, then keep the VRAM State Machine diagnostics and
downstream swap decisions based on that corrected flag.
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
src/s2_codec.cpp (1)
942-942: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueAvoid copying the entire model weight set.
The ternary mixes an lvalue reference (
Model->weight_tensor_set()) with a prvalue temporary, so the result is a prvalue andmodel_weightsbinds to a full copy of the set (this is also why cppcheck flags line 955 — it's a false positive since lifetime extension applies). Prefer a pointer to avoid the copy; it also makes theModel &&guard meaningful.♻️ Suggested change
- const auto & model_weights = Model ? Model->weight_tensor_set() : std::unordered_set<ggml_tensor*>(); + const std::unordered_set<ggml_tensor*> * model_weights = + Model ? &Model->weight_tensor_set() : nullptr;and at line 955:
- if (Model && model_weights.find(t) != model_weights.end()) continue; + if (model_weights && model_weights->find(t) != model_weights->end()) continue;🤖 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_codec.cpp` at line 942, Update the model_weights initialization near weight_tensor_set() to use a pointer, selecting the model’s weight set only when Model is non-null and otherwise storing nullptr. Adjust its use near line 955 to dereference the pointer only under the existing Model guard, preserving lifetime safety without copying the entire set.
🤖 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.
Inline comments:
In `@src/s2_mapped_file.cpp`:
- Around line 158-171: Update MappedFile::drop_page_cache so the madvise call
remains available on both Linux and macOS, while the fd_ posix_fadvise call and
POSIX_FADV_DONTNEED usage are compiled only under __linux__. Keep the existing
Windows VirtualUnlock branch unchanged.
In `@src/s2_pipeline.cpp`:
- Around line 895-909: Reorder the offline hot-swap flow around the
pending_offload_thread_ creation: perform clear_kv_cache(), Post-Phase3
diagnostics, and all model()/codec() GPU-memory or residency reads before
starting the background thread. Spawn the offload_thread only as the final
action so its free_gpu_weights() and drop_page_cache() calls cannot race with
remaining state access.
---
Nitpick comments:
In `@src/s2_codec.cpp`:
- Line 942: Update the model_weights initialization near weight_tensor_set() to
use a pointer, selecting the model’s weight set only when Model is non-null and
otherwise storing nullptr. Adjust its use near line 955 to dereference the
pointer only under the existing Model guard, preserving lifetime safety without
copying the entire set.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: bc17f0ca-6456-4aec-9e41-64291ff19de0
📒 Files selected for processing (10)
CMakeLists.txtinclude/s2_codec.hinclude/s2_mapped_file.hinclude/s2_model.hinclude/s2_pipeline.hsrc/main.cppsrc/s2_codec.cppsrc/s2_mapped_file.cppsrc/s2_model.cppsrc/s2_pipeline.cpp
| void MappedFile::drop_page_cache() { | ||
| #if defined(__linux__) || defined(__APPLE__) | ||
| if (data_ && data_ != MAP_FAILED && size_ > 0) { | ||
| ::madvise(data_, size_, MADV_DONTNEED); | ||
| } | ||
| if (fd_ >= 0) { | ||
| ::posix_fadvise(fd_, 0, 0, POSIX_FADV_DONTNEED); | ||
| } | ||
| #elif defined(_WIN32) | ||
| if (data_ && size_ > 0) { | ||
| ::VirtualUnlock(data_, size_); | ||
| } | ||
| #endif | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🌐 Web query:
Is posix_fadvise available on macOS Darwin
💡 Result:
No, the posix_fadvise function is not available on macOS (Darwin) [1][2][3]. While posix_fadvise is a standardized function for providing file access pattern advice [4], it is not implemented on the macOS platform [1][2]. Developers working on cross-platform software often encounter compilation errors when attempting to use it on macOS [3] and must use platform-specific alternatives to achieve similar functionality [1][5]. Common alternatives on macOS include: 1. F_RDADVISE: Used with fcntl to provide prefetching hints [1][5]. It requires a radvisory structure specifying the offset and length of the data to be prefetched [5]. 2. F_NOCACHE: Used with fcntl to disable caching for a file, which can serve as an alternative to advice patterns like POSIX_FADV_DONTNEED or POSIX_FADV_NOREUSE when processing large files to avoid flushing hot pages from the buffer cache [5]. Code designed to be portable across POSIX-compliant systems typically uses preprocessor directives (such as #ifdef or feature detection) to exclude or replace calls to posix_fadvise when running on macOS [1][3][6].
Citations:
- 1: https://www.postgresql.org/message-id/E1sjBKh-001noq-Cf%40gemulon.postgresql.org
- 2: Improve flushing pagecache on darwin axboe/fio#1883
- 3: Build fails on Mac OSX: unavailable posix_fadvise() MikePopoloski/slang#168
- 4: http://pubs.opengroup.org/onlinepubs/9699919799/functions/posix_fadvise.html
- 5: feat(fast_io): apply F_NOCACHE hint on macOS source-file reads oferchen/rsync#4068
- 6: https://www.postgresql.org/message-id/attachment/164094/v2-0001-Add-prefetching-support-on-macOS.patch
🏁 Script executed:
set -euo pipefail
printf '\n## src/s2_mapped_file.cpp around reported lines\n'
sed -n '140,190p' src/s2_mapped_file.cpp
printf '\n## occurrences of posix_fadvise and __APPLE__\n'
rg -n --hidden --glob '!**/.git/**' 'posix_fadvise|POSIX_FADV_DONTNEED|__APPLE__|macOS|Darwin|Metal' .
printf '\n## file list for build/config hints\n'
git ls-files | rg '(^|/)(CMakeLists\.txt|Makefile|README|README\.md|LICENSE|\.github|cmake|build|src/.*\.(cpp|h|hpp))$'Repository: rodrigomatta/s2.cpp
Length of output: 2026
Guard posix_fadvise behind __linux__ only posix_fadvise / POSIX_FADV_DONTNEED aren’t available on macOS, so this __APPLE__ path won’t compile. Keep madvise(..., MADV_DONTNEED) for Apple, but move the posix_fadvise call into the Linux-only block.
🤖 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_mapped_file.cpp` around lines 158 - 171, Update
MappedFile::drop_page_cache so the madvise call remains available on both Linux
and macOS, while the fd_ posix_fadvise call and POSIX_FADV_DONTNEED usage are
compiled only under __linux__. Keep the existing Windows VirtualUnlock branch
unchanged.
| if (params.enable_hot_swap) { | ||
| safe_print_ln("[Pipeline] Hot-Swap: Releasing compute resources..."); | ||
| model().free_compute_buffers(); | ||
|
|
||
| safe_print_ln("[Pipeline] Hot-Swap: Spawning background thread to free VRAM & RAM..."); | ||
| std::thread offload_thread([this]() { | ||
| if (model().is_weights_on_gpu()) model().free_gpu_weights(); | ||
| if (codec().is_weights_on_gpu()) codec().free_gpu_weights(); | ||
|
|
||
| model().mapped_file().drop_page_cache(); | ||
| codec().mapped_file().drop_page_cache(); | ||
|
|
||
| safe_print_ln("[Pipeline] Hot-Swap: Background VRAM & RAM free complete."); | ||
| }); | ||
| pending_offload_thread_ = std::move(offload_thread); |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | 🏗️ Heavy lift
Data race: the hot-swap offload thread frees GPU weights while the main thread still reads them.
In the offline persistent hot-swap path the background thread is spawned here and immediately begins model().free_gpu_weights() / codec().free_gpu_weights() (both call free_backend_buffers(...), clearing weights_.model_bufs_gpu and resetting residency flags). But the main thread keeps running after the spawn and reads the very same shared state:
- Line 923:
model().clear_kv_cache() - Lines 955–959:
model().get_gpu_memory_usage_bytes()/codec().get_gpu_memory_usage_bytes(), which iteratemodel_bufs_gpu/ readmodel_bufconcurrently with the thread clearing them.
This is undefined behavior (iterating a vector being cleared, reading freed buffer sizes) and can crash. The streaming variant is safe because it spawns the offload thread after all diagnostics; the offline path emits Post-Phase3 after the spawn. Move the post-decode diagnostics and any model/codec state reads before spawning the offload thread (or snapshot the values first).
🔒 Suggested direction
Emit the Post-Phase3 VRAM diagnostic and finish all get_gpu_memory_usage_bytes() reads (and clear_kv_cache()) before line 900, then spawn offload_thread as the last action so the background thread has exclusive access to the weight/buffer state it mutates.
🤖 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_pipeline.cpp` around lines 895 - 909, Reorder the offline hot-swap
flow around the pending_offload_thread_ creation: perform clear_kv_cache(),
Post-Phase3 diagnostics, and all model()/codec() GPU-memory or residency reads
before starting the background thread. Spawn the offload_thread only as the
final action so its free_gpu_weights() and drop_page_cache() calls cannot race
with remaining state access.
Replace MADV_RANDOM with MADV_SEQUENTIAL. Weight loading iterates tensors in sequential file order, so disabling readahead forced ~1.3M individual 4 KB I/O syscalls on ..cold reads after drop_page_cache(). MADV_SEQUENTIAL enables aggressive kernel readahead from byte 0, reducing syscall count. - Add MADV_SEQUENTIAL (Linux+macOS), FILE_FLAG_SEQUENTIAL_SCAN (Windows) as the equivalent hint. - Add MADV_HUGEPAGE (Linux) to reduce TLB pressure during multi-GB loads - Add MADV_DONTDUMP (Linux) to exclude the mapping from core dumps
1 participant
What this solves
In the final stages when it's time to process Audio Codec, Slow-AR isn't evicted from VRAM despite being no longer needed. So our VRAM footprint becomes
Slow-AR + compute buffers + Audio Codec + compute buffers + KV cachespeaking at6-7 GB on q8_0with transient spikes during buffer allocations (possibly due to fragmentaton?). In effect this causes high memory pressure and leads to OOMs if you were nearing hardware limit already.To solve this we selectively load the submodels and evict them from VRAM once they finish. By doing this we achieve 2.2 GB VRAM usage during Audio Codec phase of processing, and we also only load the Audio Codec when it's needed so the initial VRAM load is also lowered by ~1 GB.
With memory-mapped page cache sitting in RAM we can near-instantly fetch them when needed, and we hide the latency of this (mostly PCIe) fetch by scheduling it in background thread as we finish other sequential processes.
We've two atomic commits for easier review:
The new default behavior after the second commit can be categorized by opt-in and opt-out:
Opt out vram-swap
Unless you opt out, we evict already-processed submodels from VRAM, and lazily loads them as needed. Because we're using mmap these loads are near-instant since we fetch them from RAM instead of disk. We also keep compute buffers in VRAM as they're relatively small but take longer to re-initiate (~170 MB in vulkan). Your OS might (partially) reclaim the RAM page cache incurring some disk IO for the reclaimed pages, but it's generally as fast as keeping the models loaded in VRAM all the time. This depends on page faults. It's also more efficient than fread (double copy) since we don't load all tensors at once and use zero-copy DMA.
--no-vram-swapOpt in hot-swap
With the opt-in hot-swap behavior, we free anything from not only VRAM but also RAM. This includes compute buffers and the mmap page cache, landing us at a 100 MB RAM + 25 MB VRAM footprint. This is the "I don't trust my drivers & OS to free memory in a timely fashion and want to evict them aggressively, I need my VRAM back immediately for other tasks" option. This also ensures you'll need to read from disk each request, so I advise only using this if you absolutely need lowest memory footprint or the model is in m.2 SSD. It's a rather "dumb" solution compared to vram-swap. I advise against using this unless you know what you're doing as this can take 4x longer to process compared to vram-swap only. You're in effect doing cold boots every request and purely limited by your IO speed (mines a slow 200 MB/s SATA SSD)
--hot-swapKey Features
1. Zero-Copy
mmapArchitectureMappedFileRAII wrapper (POSIXmmap/ WindowsCreateFileMapping).2. Lazy Weight Allocation
encode()ordecode()is called viaensure_weights_loaded().mmappointer into system RAM.3. Phase-Gated VRAM Swapping
ggml_backend_synchronizeto nudge drivers to free VRAM instead of deferring & thrashing.4. Hot-Swap Mode (
--hot-swap)madvise(MADV_DONTNEED)/posix_fadvise(DONTNEED)to evict the 5.3 GB GGUF file from the OS page cache.Practical usage/footprint
--hot-swap--no-vram-swapSome metrics
Read me!:
--no-vram-swapflag.Notes
Summary by CodeRabbit
--no-vram-swapand--hot-swap, plus persistent pipeline mode for server runs.