feat: KV cache quantization - #47
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.
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
…ipeline Extended the phase-gated VRAM swap with finer-grained codec weight management, and concurrent decode threading. - Codec encoder/decoder granular split: codec weights are classified at load time into encoder and decoder groups via tensor name prefixes. - New methods (free/restore/is_on_gpu/get_bytes for each group) allowing the pipeline to load only the encoder for reference audio encoding, free it before generation, ..load only the decoder for the decode phase. - Minimizing VRAM footprint at each step while catering to lower latency by utilizing background threads to lazy load as needed. - The priority is to reduce Slow-AR processing stage's VRAM footprint as this is the heaviest model in the pipeline, so we never keep another (unnecessary) model loaded when that's on. And we free it before loading any new models. In effect, OOM is far less likely and as long as your GPU can fit the Slow-AR and its buffers it'll run the whole pipeline without issue. - Deferred weight loading: when VRAM swap is active and the model prefers GPU, init() skips Slow-AR weight allocation entirely and calls warm_page_cache() to pre-fault mmap pages via ..MADV_WILLNEED + MADV_COLD (Linux), fcntl F_RDAHEAD (macOS), or PrefetchVirtualMemory ..(Windows). First-request restore hits warm RAM instead of cold disk. Replacing the old incorrect VirtualUnlock. - prefers_gpu() provides an intent-based check that works before weights are loaded, replacing is_weights_on_gpu() for init decisions to handle edge cases better. - Codec eviction + Slow-AR restore and KV cache init now runs in background threads, hiding PCIe latency behind the CPU bound prompt construction. - Overlapped decode path: when Slow-AR is on GPU and codec is on CPU, a producer-consumer thread pair decodes audio frames concurrently with generation via mutex, condition variable, and atomic frame counters. This reduces Total RTF by starting the CPU codec work as early as we can instead of waiting for GPU to finish Slow-AR processing in its entirety. - server-aware swapping: more_segments_pending server sets this flag on all sentence segments except the last within a single request, keeping Slow-AR resident in VRAM across segments and eliminating per-segment restore overhead. - Streaming path: granular codec management (free encoder, restore decoder only), Slow-AR freed in non-hot-swap persistent mode (fixes VRAM leak where Slow-AR was never freed after streaming requests). - pre_restore_thread removed from synthesize_raw/synthesize_streaming_raw; replaced with pending_offload_thread_ join before encoding to prevent races between background eviction and encoder weight restoration. - --fast-decoder-cpu / --codebook-cpu CLI flags force specific tensor groups onto CPU, saving ~200-400 MB / ~56 MB VRAM respectively at the cost of PCIe transfers. Previously, any --gpu-layers value also offloaded fast-decoder and codebook tensors. These flags decouple that decision for finer-grained VRAM control. - allocate_weight_buffers nulls stale tensor data/buffer pointers after freeing, preventing use-after-free when weights are re-allocated after a free/restore cycle. - MappedFile::open uses CreateFileW with UTF-8 -> UTF-16 conversion, fixing model loading on Windows paths containing non-ASCII characters.
Caches the prompt prefill result across requests so repeated synthesis with the same voice and text prefix skipping the ~190-210 ms prefill pass. The KV buffer is reused when its capacity covers the new request, replacing per-request free/realloc with a memset reset. - PrefillCacheEntry stores the prefill StepResult, n_past position, and KV cache contents keyed by voice_id (or prompt text + first 16 reference codes) concatenated with the synthesis text. Cache is invalidated on key mismatch or insufficient max_seq_len. - Two residency modes controlled by --kv-cache-vram: .. Default (system RAM): KV state is serialized to a compact strided buffer in system RAM via save_kv_state() after prefill, running on a background thread that overlaps with generation. Restored via restore_kv_state() on cache hit, takes a PCIe transfer instead of a full re-prefill. .. --kv-cache-vram (VRAM pin flag): KV buffer stays allocated in VRAM between requests. free_compute_buffers() is skipped when keep_kv_on_gpu is set by this, so the next request with a matching key resumes generation immediately without PCIe transfer latency. - generate() accepts an optional initial_state pointer; when provided, the internal prefill is skipped and the cached StepResult (hidden state + logits + n_past) is used directly. The pipeline performs the prefill itself on cache miss and passes the resulting state. - reset_kv_cache() memsets the existing KV buffer to zero without freeing or reallocating it, replacing the old clear + init_kv_cache cycle when the buffer is already large enough for the new request. - KV lifecycle respects more_segments_pending: within a segmented server request, the KV buffer and prefill cache are preserved across sentence segments so we don't incur transfer penalties knowing the full prompt is still processing. - CLI flags: --no-kv-reuse (opt out, restores old per-request free/realloc behavior), --kv-cache-vram (pin prefill cache in VRAM between requests instead of serializing to system RAM). - Metrics: prefill_ms and prefill=cached/computed added to the synthesis log line.
…e-v) Allows configuring the quantization of KV caches (F16 default as before.) - New CLI flags --cache-type-k <type> and --cache-type-v <type>, accepting f32, f16, bf16, q8_0, q4_0, q4_1, iq4_nl, q5_0, q5_1. (incorrect values fall back to f16 with a warning.) - PipelineParams gains kv_cache_type_k / kv_cache_type_v strings, parsed in Pipeline::init via parse_kv_cache_type and applied through SlowARModel::set_kv_cache_types before any cache allocation. - SlowARModel stores kv_cache_type_k_ / kv_cache_type_v_ (default GGML_TYPE_F16) and init_kv_cache now allocates memory_k_ / memory_v_ with the configured types instead of hardcoded F16. The allocation log line now reports both types. - Fix save_kv_state / restore_kv_state to use per-tensor strides. Previously both functions derived pos_bytes, total size, and layer/head offsets from memory_k_->nb for K and V alike, which is ..only correct when K and V share a type. With mixed types (e.g. K=f16, V=q8_0) the V copies read/wrote the wrong byte ranges, corrupting the restored cache on a prefill-cache hit. Both functions now compute ..k_pos_bytes/v_pos_bytes, k_full_bytes/v_full_bytes, and separate k_src/v_src (and k_dst/v_dst) offsets from memory_k_ and memory_v_ respectively. Quantize V first. V-cache quantization is more tolerable because V is linearly combined into the attention output, so its error stays bounded. K-cache quantization perturbs the Q·K dot products that feed softmax, and the exponential amplifies the error, which can put the model in an infinite generation loop (even in q8_0), never hitting EOS. The --cache-type-k help text carries a warning to that effect. All V quants were tested on Vulkan s2-pro-q4_k_m: V=q8_0 saves ~23% KV VRAM and is perceptually identical to f16; more aggressive V types save up to ~36% but trade audible fidelity (e.g. q5_0 adds a sharp hiss on higher frequencies) with no generation-speed benefit. YMMV by hardware and backend.
🚥 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: 17
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
src/s2_codec.cpp (1)
204-224: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winFree
impl_->backend_cpu, or remove the field.
load_sharedinitializesimpl_->backend_cpu, butreset_codec_implonly freesimpl_->backendbefore overwriting*impl_with a freshImpl. This leaks one CPU backend per shared reload, and the destructor also leaves the last initialized CPU backend unfreed. If the codec does not use a separate CPU backend, remove the field and initialization.🐛 Proposed fix
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 204 - 224, Update reset_codec_impl to free impl.backend_cpu and clear it before reinitializing the Impl, ensuring both reloads and destruction release the CPU backend. If the separate CPU backend is unused, instead remove the backend_cpu field and its initialization from load_shared.src/s2_server.cpp (1)
229-248: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick winA cancelled or failed segment leaves VRAM held.
Line 237 sets
more_segments_pending = truefor every segment except the last.src/s2_pipeline.cppgates its VRAM and KV teardown on!params.more_segments_pendingat lines 1251, 1269, and 1182, so a segment with the flag set deliberately keeps the Slow-AR weights and the KV cache resident.The loop returns early at line 231 on cancellation and at line 247 on segment failure. The last segment issued still had the flag set, so nothing runs the teardown. The weights stay in VRAM until the next synthesis request arrives. On an idle server after a cancelled request that holds the full Slow-AR VRAM footprint, which is the cost this PR reduces.
Run one final synthesis pass with the flag cleared, or add a pipeline method that releases held VRAM, and call it on the abort paths.
Do you want me to open an issue to track a
Pipeline::release_held_vram()entry point for the abort paths?🤖 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_server.cpp` around lines 229 - 248, Ensure the segment loop releases deferred VRAM and KV resources when cancellation or synthesis failure occurs. Before returning from the abort paths around sink.is_cancelled() and the !ok check in the synthesis loop, run a final cleanup pass with more_segments_pending cleared or invoke a dedicated pipeline release method, while preserving the existing deferred teardown for successful multi-segment synthesis.
🧹 Nitpick comments (11)
src/s2_mapped_file.cpp (2)
43-46: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winReturn
falsefor a zero-size file.
open()returnstruefor an empty file, butdata_staysnullptr, sois_open()returnsfalse. The two results disagree. Current callers testis_open(), so they behave correctly today. A caller that trusts the return value would proceed with an unmapped file.♻️ Proposed change for the POSIX branch
if (size_ == 0) { ::close(fd); - return true; + return false; }Apply the same change to the Windows branch at lines 43-46.
Also applies to: 78-81
🤖 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 43 - 46, Update the zero-size file handling in the Windows branch of open() to close the handle and return false, keeping the result consistent with is_open() when data_ remains nullptr.
190-198: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueNarrow the Windows page-cache drop, or document its process-wide effect.
SetProcessWorkingSetSizeExtrims the working set of the entire process. It does not target the pages of this mapping. The call also changes a process-wide quota flag and its result is discarded. ConsiderVirtualUnlockon the mapped range, or add a comment that states the process-wide scope.🤖 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 190 - 198, Update the Windows branch in the mapped-file cache-drop implementation to avoid using process-wide SetProcessWorkingSetSizeEx for a single mapping; use VirtualUnlock with the mapped range and size represented by data_ and size_. Preserve the existing no-op behavior only if the platform-specific operation cannot target the mapping, and handle the operation result consistently with surrounding code.src/s2_codec.cpp (3)
1742-1742: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low valueGuard
impl_inmapped_file().Every other new accessor tests
impl_first.mapped_file()dereferences it directly, andimpl_is null after~AudioCodecruns or after a failed construction. Return a reference only whenimpl_is valid, or assert the precondition.🤖 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 1742, Update AudioCodec::mapped_file() to guard impl_ before dereferencing it, matching the other accessors; assert or otherwise enforce that impl_ is valid before returning mapped_gguf_, while preserving the reference-returning API.
768-773: 🎯 Functional Correctness | 🔵 Trivial | 💤 Low valueMatch the tensor-name prefix instead of a substring.
find(...) != std::string::nposmatches anywhere in the name. With an emptytprefix(architecturefish-speech-codec, line 901), any name that containsencoder.at any position classifies as an encoder tensor. Lines 985-989 place each tensor in exactly one group, so a wrong match assigns a tensor to the wrong residency group.rfind(marker, 0) == 0states the intent exactly.♻️ Proposed refactor
static bool is_encoder_tensor(const std::string & name, const std::string & tprefix) { - if (name.find(tprefix + "encoder.") != std::string::npos) return true; - if (name.find(tprefix + "quantizer.pre_module.") != std::string::npos) return true; - if (name.find(tprefix + "quantizer.downsample.") != std::string::npos) return true; + if (name.rfind(tprefix + "encoder.", 0) == 0) return true; + if (name.rfind(tprefix + "quantizer.pre_module.", 0) == 0) return true; + if (name.rfind(tprefix + "quantizer.downsample.", 0) == 0) return true; return false; }🤖 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 768 - 773, Update is_encoder_tensor to require each encoder or quantizer marker to occur at the beginning of name, using a prefix check equivalent to rfind(..., 0) == 0 instead of substring find checks. Preserve the existing markers and boolean classification behavior.
966-966: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winAvoid the hidden copy of the model weight set.
The conditional operator mixes a
const std::unordered_set&branch with a prvalue branch, so the common type is a prvalue. WhenModelis non-null, the whole weight set is copied. Lifetime extension keeps the reference valid, so thefindat line 979 is safe, but Cppcheck reportsdanglingTemporaryLifetimeon this pattern and the copy is avoidable.♻️ Proposed refactor
- const auto & model_weights = Model ? Model->weight_tensor_set() : std::unordered_set<ggml_tensor*>(); + static const std::unordered_set<ggml_tensor*> empty_model_weights; + const std::unordered_set<ggml_tensor*> & model_weights = + Model ? Model->weight_tensor_set() : empty_model_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 966, Update the model_weights initialization near weight_tensor_set() to avoid the conditional operator’s prvalue temporary and hidden set copy. Preserve a reference to Model->weight_tensor_set() when Model exists, and use a separate stable empty-set object for the null-Model case so the subsequent find remains safe without copying.Source: Linters/SAST tools
src/s2_model.cpp (2)
1284-1287: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winCheck the scheduler creation result.
ggml_backend_sched_newcan return null.eval_cachedthen callsggml_backend_sched_reset(sched_)at line 1027 with a null scheduler.acquire_compute_resourcesreturnsvoid, so a failure is invisible toallocate_and_load_weightsat line 1334. Log the failure, or change the signature to report it.🤖 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 1284 - 1287, Validate the results of both ggml_backend_sched_new calls in the scheduler initialization path, including sched_ and the conditional fast_sched_. On failure, log a clear error and propagate failure through acquire_compute_resources so allocate_and_load_weights cannot continue with a null scheduler.
688-709: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winRead only the saved region instead of staging both full tensors.
Lines 690-693 allocate host buffers for the complete K and V caches and transfer every byte from the backend, even when
n_positionsis far belowmax_seq_len_. Each (layer, head) slice is already contiguous, so it can be read directly.♻️ Proposed refactor
- const size_t k_full_bytes = ggml_nbytes(memory_k_); - const size_t v_full_bytes = ggml_nbytes(memory_v_); - std::vector<uint8_t> k_full(k_full_bytes); - std::vector<uint8_t> v_full(v_full_bytes); - ggml_backend_tensor_get(memory_k_, k_full.data(), 0, k_full_bytes); - ggml_backend_tensor_get(memory_v_, v_full.data(), 0, v_full_bytes); - k_out.resize(k_out_bytes); v_out.resize(v_out_bytes); size_t k_off = 0, v_off = 0; for (int32_t l = 0; l < n_layer; ++l) { for (int32_t h = 0; h < n_head_kv; ++h) { const size_t k_src = static_cast<size_t>(l) * memory_k_->nb[3] + static_cast<size_t>(h) * memory_k_->nb[2]; const size_t v_src = static_cast<size_t>(l) * memory_v_->nb[3] + static_cast<size_t>(h) * memory_v_->nb[2]; - std::memcpy(k_out.data() + k_off, k_full.data() + k_src, k_pos_bytes); - std::memcpy(v_out.data() + v_off, v_full.data() + v_src, v_pos_bytes); + ggml_backend_tensor_get(memory_k_, k_out.data() + k_off, k_src, k_pos_bytes); + ggml_backend_tensor_get(memory_v_, v_out.data() + v_off, v_src, v_pos_bytes); k_off += k_pos_bytes; v_off += v_pos_bytes; } }Confirm the stride mapping first. This refactor keeps whichever layout the fix at lines 621-623 settles on.
🤖 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 688 - 709, Update the cache extraction logic around the K/V transfer loop to avoid allocating or downloading full memory_k_ and memory_v_ tensors. After confirming the layout and stride mapping established near the existing offset calculations, read each contiguous layer/head saved slice directly from the backend into k_out and v_out, preserving the current output ordering and n_positions-limited byte counts.src/s2_pipeline.cpp (4)
792-797: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low valueJoin the pending offload thread unconditionally.
The join is gated on
params.enable_vram_swap. A previous request can spawnpending_offload_thread_while the flag is true. If a later call arrives with the flag false, this join is skipped and the background thread can free GPU weights while the new request generates. The thread is joined only in the destructor.The thread exists only when a previous request created it, so joining it does not depend on the current request's flag.
♻️ Proposed refactor
- if (params.enable_vram_swap) { - if (pending_offload_thread_.joinable()) { - pending_offload_thread_.join(); - } + if (pending_offload_thread_.joinable()) { + pending_offload_thread_.join(); }🤖 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 792 - 797, Update the pending offload synchronization in the surrounding request flow to check and join pending_offload_thread_ unconditionally, removing the params.enable_vram_swap guard while retaining the joinable check. Ensure every call waits for an existing offload thread regardless of the current request’s flag.
583-604: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low valueFree the encoder weights through a scope guard.
need_encoder_vramgates both the restore at line 588 and the free at line 602, so the pair is balanced on the normal path. Ifcodec().encodeat line 592 throws, line 602 does not run and the encoder stays resident in VRAM for the rest of the process. The repository convention allowsstd::runtime_errorfor hard failures, so a throw from the codec is possible.Bind the free to scope exit.
As per coding guidelines: "Mix bool returns for recoverable operations and std::runtime_error for hard failures in C++".
🤖 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 583 - 604, Use a scope guard in the encode flow around need_encoder_vram so codec().free_encoder_weights() always runs when encoder weights were restored, including if codec().encode throws. Remove the normal-path-only cleanup and preserve the existing conditional VRAM behavior.Source: Coding guidelines
1186-1196: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExtract the duplicated hot-swap offload body.
The same offload lambda body appears three times: lines 1186-1194, lines 1238-1246, and lines 1622-1630. All three free the model GPU weights, the codec decoder, the codec encoder, the codec GPU weights, and then drop both mapped-file page caches. Any future change to the eviction order must be applied in three places.
Extract one private member function and spawn the thread from it.
♻️ Proposed refactor
+void Pipeline::spawn_hot_swap_offload() { + safe_print_ln("[Pipeline] Hot-Swap: Spawning background thread to free VRAM & RAM..."); + pending_offload_thread_ = std::thread([this]() { + if (model().is_weights_on_gpu()) model().free_gpu_weights(); + if (codec().is_decoder_on_gpu()) codec().free_decoder_weights(); + if (codec().is_encoder_on_gpu()) codec().free_encoder_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."); + }); +}Then replace each of the three sites:
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]() { - ... - }); - pending_offload_thread_ = std::move(offload_thread); + spawn_hot_swap_offload();🤖 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 1186 - 1196, Extract the shared hot-swap offload lambda body into one private member function on the pipeline class, preserving the existing free order, page-cache drops, and completion message. Update each of the three duplicated thread-spawn sites to launch that member function instead, including the existing pending-thread assignment behavior.
138-150: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueIdentify which cache-type flag failed.
Lines 484-486 call
parse_kv_cache_typefor both K and V. The warning text names--cache-type, so a user who mistypes one of the two flags cannot tell which value was rejected. Pass the flag name to the helper.♻️ Proposed refactor
-static ggml_type parse_kv_cache_type(const std::string & type_str) { +static ggml_type parse_kv_cache_type(const std::string & type_str, const char * flag_name) { if (type_str == "f32") return GGML_TYPE_F32; @@ - safe_print_warn_ln("Warning: unknown --cache-type value '" + type_str + "', defaulting to f16."); + safe_print_warn_ln(std::string("Warning: unknown ") + flag_name + " value '" + type_str + + "', defaulting to f16."); return GGML_TYPE_F16; }model().set_kv_cache_types( - parse_kv_cache_type(params.kv_cache_type_k), - parse_kv_cache_type(params.kv_cache_type_v)); + parse_kv_cache_type(params.kv_cache_type_k, "--cache-type-k"), + parse_kv_cache_type(params.kv_cache_type_v, "--cache-type-v"));🤖 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 138 - 150, Update parse_kv_cache_type to accept the relevant flag name and include it in the unknown-value warning, then pass the K and V cache-type flag names from both call sites so the warning identifies which setting was rejected.
🤖 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 1679-1693: Update the read_f32 lambda to throw an error when
t->type is neither GGML_TYPE_F32 nor GGML_TYPE_F16, matching tensor_to_f32
behavior; do not return the zero-initialized vector for unsupported tensor
types.
- Around line 1717-1740: Update AudioCodec::ensure_weights_loaded() to track and
restore residency independently for encoder, decoder, and shared/combined codec
tensor groups instead of returning success solely from the global
weights_allocated_ flag. Ensure each group needed by encode() or decode() is
allocated and loaded before use, including when only some groups remain resident
after freeing.
- Around line 138-183: Update allocate_codec_buffers to preserve any existing
out_buffer handle instead of unconditionally resetting it, and reuse the
established ggml_tallocr-based tensor allocation path used by the model
allocator. Replace manual tensor->data and tensor->buffer assignments with
backend-aware allocation through ggml_tallocr, while retaining the existing
buffer ownership and error handling behavior.
In `@src/s2_generate.cpp`:
- Around line 38-42: Validate the cached state in the initial_state branch
before the first sampling call: ensure initial_state->logits contains at least
vocab_size entries, and reject undersized or empty cached states before
assigning or sampling from state. Keep the existing fresh prefill path
unchanged, and use the surrounding generate function’s established
error-handling behavior for the rejection.
In `@src/s2_mapped_file.cpp`:
- Around line 204-210: Remove the MADV_COLD block from warm_page_cache, leaving
MADV_WILLNEED and MADV_SEQUENTIAL unchanged. If cooling behavior is required,
move the MADV_COLD advice to the drop_page_cache implementation instead.
- Around line 221-235: Define the Windows SDK target as Windows 8 for the target
containing warm_page_cache, using _WIN32_WINNT_WIN8 or an equivalent
target-specific _WIN32_WINNT compile definition in CMakeLists.txt. Ensure
PrefetchVirtualMemory and WIN32_MEMORY_RANGE_ENTRY are declared without
affecting non-Windows builds.
In `@src/s2_model.cpp`:
- Around line 1320-1330: Unchecked GGUF offsets and tensor sizes can read beyond
the mmap, while missing offsets are treated as success. In the class containing
the weight-loading loops around original_gpu_weights_ and original_cpu_weights_,
add a bounds-checked upload helper that validates gguf_data_offset_ + tensor
offset + ggml_nbytes(t) against mapped_gguf_.size(), returns false for missing
offsets, and route both loops through it before marking weights allocated. Apply
the same helper/check in src/s2_codec.cpp at lines 1729-1734, and use it for
uploads in restore_encoder_weights (1587-1592), restore_decoder_weights
(1634-1639), and the read_f32 lambda in refresh_host_caches_from_mmap (1686).
- Around line 1290-1302: Update SlowARModel::get_gpu_memory_usage_bytes so
kv_buf_ contributes to the total only when the KV cache was initialized on the
GPU backend, matching the backend selection in init_kv_cache. Keep counting
model GPU buffers unchanged and exclude CPU-backed KV cache memory from the
reported GPU usage.
- Around line 621-623: Restore the tensor dimensions in the KV cache allocations
in the memory_k_ and memory_v_ initialization so they use head_dim, n_head_kv,
max_seq_len, n_layer order expected by eval_cached. Keep the existing tensor
types and allocation context unchanged, ensuring nb[2] remains the per-token
stride and nb[1] the per-head stride.
In `@src/s2_pipeline.cpp`:
- Around line 447-452: Track the codec’s achieved placement separately from the
requested use_gpu_codec flag across the codec-loading flow. Update the
successful GPU and CPU load paths near codec_loaded and the assignments around
use_gpu_codec so codec_prefers_gpu_ reflects the backend that actually loaded,
including the CPU fallback after GPU failure; use this achieved-placement state
for subsequent VRAM, weight-management, restore, and overlap decisions.
- Around line 1384-1388: Update the streaming restore block around
model().restore_weights_to_gpu() to check its boolean result; when restoration
fails, report the failure through the request’s sink using the existing
error-reporting pattern and return false before init_kv_cache or generate can
run. Preserve the successful restoration path unchanged.
- Around line 856-940: Serialize the SlowARModel initialization paths used by
vram_phase1_thread and kv_init_thread so model().acquire_compute_resources(),
model().restore_weights_to_gpu(), model().init_kv_cache(), and
model().reset_kv_cache() cannot run concurrently. Add synchronization around the
shared model/backend allocation operations, or reorder the thread startup and
joins to prevent overlap while preserving the existing success and failure
handling.
- Around line 1294-1302: Rename the prefill status field appended in the metrics
construction near can_overlap_decode and prefill_hit from “prefill” to a
distinct key such as “prefill_mode,” while leaving the earlier prefill duration
field unchanged.
- Around line 944-965: Invalidate the VRAM-resident prefill cache entry whenever
the KV cache is cleared or reinitialized in the fresh-KV path around
clear_kv_cache() and init_kv_cache(). Ensure prefill_cache_.valid and its
VRAM-resident state cannot remain true after reallocation, so the lookup in the
prefill-hit branch does not reuse stale n_past or state.
- Around line 1002-1013: Update the prefill-cache save flow around the
kv_save_thread creation so prefill_cache_.valid remains false while
save_kv_state is writing; set it only after the asynchronous save completes,
either inside the save thread after the write or immediately after the thread is
joined. Keep the existing cache publication and logging behavior otherwise
unchanged.
- Around line 1102-1149: Add a small RAII joining guard near the top of the
file, then bind it to both locally owned threads, decode_thread and
kv_save_thread, immediately after their creation. The guard destructor must join
any still-joinable thread so exceptions from generate or codec().decode,
including std::bad_alloc and std::runtime_error, cannot trigger std::terminate;
retain the existing explicit join() calls for the normal path.
- Around line 319-336: Update Pipeline::compute_prefill_cache_key to
length-prefix every variable-length request component, especially
params.prompt_text and params.text, before concatenating them; ensure delimiter
characters within either value cannot create equivalent keys for different
inputs. Preserve the existing voice/ref-code/noprompt distinctions and ref-code
content while applying unambiguous length encoding consistently.
---
Outside diff comments:
In `@src/s2_codec.cpp`:
- Around line 204-224: Update reset_codec_impl to free impl.backend_cpu and
clear it before reinitializing the Impl, ensuring both reloads and destruction
release the CPU backend. If the separate CPU backend is unused, instead remove
the backend_cpu field and its initialization from load_shared.
In `@src/s2_server.cpp`:
- Around line 229-248: Ensure the segment loop releases deferred VRAM and KV
resources when cancellation or synthesis failure occurs. Before returning from
the abort paths around sink.is_cancelled() and the !ok check in the synthesis
loop, run a final cleanup pass with more_segments_pending cleared or invoke a
dedicated pipeline release method, while preserving the existing deferred
teardown for successful multi-segment synthesis.
---
Nitpick comments:
In `@src/s2_codec.cpp`:
- Line 1742: Update AudioCodec::mapped_file() to guard impl_ before
dereferencing it, matching the other accessors; assert or otherwise enforce that
impl_ is valid before returning mapped_gguf_, while preserving the
reference-returning API.
- Around line 768-773: Update is_encoder_tensor to require each encoder or
quantizer marker to occur at the beginning of name, using a prefix check
equivalent to rfind(..., 0) == 0 instead of substring find checks. Preserve the
existing markers and boolean classification behavior.
- Line 966: Update the model_weights initialization near weight_tensor_set() to
avoid the conditional operator’s prvalue temporary and hidden set copy. Preserve
a reference to Model->weight_tensor_set() when Model exists, and use a separate
stable empty-set object for the null-Model case so the subsequent find remains
safe without copying.
In `@src/s2_mapped_file.cpp`:
- Around line 43-46: Update the zero-size file handling in the Windows branch of
open() to close the handle and return false, keeping the result consistent with
is_open() when data_ remains nullptr.
- Around line 190-198: Update the Windows branch in the mapped-file cache-drop
implementation to avoid using process-wide SetProcessWorkingSetSizeEx for a
single mapping; use VirtualUnlock with the mapped range and size represented by
data_ and size_. Preserve the existing no-op behavior only if the
platform-specific operation cannot target the mapping, and handle the operation
result consistently with surrounding code.
In `@src/s2_model.cpp`:
- Around line 1284-1287: Validate the results of both ggml_backend_sched_new
calls in the scheduler initialization path, including sched_ and the conditional
fast_sched_. On failure, log a clear error and propagate failure through
acquire_compute_resources so allocate_and_load_weights cannot continue with a
null scheduler.
- Around line 688-709: Update the cache extraction logic around the K/V transfer
loop to avoid allocating or downloading full memory_k_ and memory_v_ tensors.
After confirming the layout and stride mapping established near the existing
offset calculations, read each contiguous layer/head saved slice directly from
the backend into k_out and v_out, preserving the current output ordering and
n_positions-limited byte counts.
In `@src/s2_pipeline.cpp`:
- Around line 792-797: Update the pending offload synchronization in the
surrounding request flow to check and join pending_offload_thread_
unconditionally, removing the params.enable_vram_swap guard while retaining the
joinable check. Ensure every call waits for an existing offload thread
regardless of the current request’s flag.
- Around line 583-604: Use a scope guard in the encode flow around
need_encoder_vram so codec().free_encoder_weights() always runs when encoder
weights were restored, including if codec().encode throws. Remove the
normal-path-only cleanup and preserve the existing conditional VRAM behavior.
- Around line 1186-1196: Extract the shared hot-swap offload lambda body into
one private member function on the pipeline class, preserving the existing free
order, page-cache drops, and completion message. Update each of the three
duplicated thread-spawn sites to launch that member function instead, including
the existing pending-thread assignment behavior.
- Around line 138-150: Update parse_kv_cache_type to accept the relevant flag
name and include it in the unknown-value warning, then pass the K and V
cache-type flag names from both call sites so the warning identifies which
setting was rejected.
🪄 Autofix
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: b31d1e0a-bb5d-498e-820c-7c4c057ecd6d
📒 Files selected for processing (13)
CMakeLists.txtinclude/s2_codec.hinclude/s2_generate.hinclude/s2_mapped_file.hinclude/s2_model.hinclude/s2_pipeline.hsrc/main.cppsrc/s2_codec.cppsrc/s2_generate.cppsrc/s2_mapped_file.cppsrc/s2_model.cppsrc/s2_pipeline.cppsrc/s2_server.cpp
| static bool allocate_codec_buffers(ggml_backend_t backend, | ||
| const std::vector<ggml_tensor *> & tensors, | ||
| ggml_backend_buffer_t & out_buffer, | ||
| size_t & total_bytes, | ||
| std::string & error_message) { | ||
| out_buffer = nullptr; | ||
| total_bytes = 0; | ||
| error_message.clear(); | ||
| if (backend == nullptr || tensors.empty()) return true; | ||
|
|
||
| const ggml_backend_buffer_type_t buft = ggml_backend_get_default_buffer_type(backend); | ||
| const size_t alignment = ggml_backend_buft_get_alignment(buft); | ||
|
|
||
| for (ggml_tensor * tensor : tensors) { | ||
| const size_t alloc_size = ggml_backend_buft_get_alloc_size(buft, tensor); | ||
| const size_t rem = total_bytes % alignment; | ||
| if (rem != 0) total_bytes += (alignment - rem); | ||
| total_bytes += alloc_size; | ||
| } | ||
|
|
||
| if (total_bytes == 0) return true; | ||
|
|
||
| out_buffer = ggml_backend_buft_alloc_buffer(buft, total_bytes); | ||
| if (!out_buffer) { | ||
| error_message = "failed to allocate codec buffer of size " + std::to_string(total_bytes); | ||
| return false; | ||
| } | ||
|
|
||
| ggml_backend_buffer_set_usage(out_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); | ||
|
|
||
| char * base_ptr = static_cast<char *>(ggml_backend_buffer_get_base(out_buffer)); | ||
| size_t current_offset = 0; | ||
|
|
||
| for (ggml_tensor * tensor : tensors) { | ||
| const size_t alloc_size = ggml_backend_buft_get_alloc_size(buft, tensor); | ||
| const size_t rem = current_offset % alignment; | ||
| if (rem != 0) current_offset += (alignment - rem); | ||
|
|
||
| tensor->data = base_ptr + current_offset; | ||
| tensor->buffer = out_buffer; | ||
|
|
||
| current_offset += alloc_size; | ||
| } | ||
|
|
||
| return true; | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Compare codec allocation with the model allocator path and locate init_tensor use.
set -euo pipefail
rg -n -C 3 'ggml_tallocr_new|ggml_tallocr_alloc|ggml_backend_buffer_init_tensor' --type=cpp
rg -n -C 3 'allocate_codec_buffers\(' src/s2_codec.cppRepository: rodrigomatta/s2.cpp
Length of output: 157
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf 'repo files:\n'
git ls-files | sed -n '1,80p'
printf '\nSearch all tracked files for ggml_tallocr and codec buffer functions:\n'
rg -n -C 3 'ggml_tallocr_new|ggml_tallocr_alloc|ggml_backend_buffer_init_tensor|allocate_codec_buffers\(|ensure_weights_loaded|free_encoder_weights|model_buf|encoder_buf|decoder_buf' .
printf '\nInspect src/s2_codec.cpp around relevant symbols:\n'
wc -l src/s2_codec.cpp
sed -n '1,230p' src/s2_codec.cpp
sed -n '1520,1590p' src/s2_codec.cpp
sed -n '1700,1745p' src/s2_codec.cppRepository: rodrigomatta/s2.cpp
Length of output: 31684
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf 'Inspect src/s2_codec.cpp restore paths:\n'
sed -n '1536,1642p' src/s2_codec.cpp
sed -n '1717,1742p' src/s2_codec.cpp
printf '\nInspect CMake options and ggml usage flags:\n'
sed -n '1,220p' CMakeLists.txt
rg -n -C 2 'GGML_USE_(CUDA|METAL|VULKAN)|add_subdirectory\(ggml|project\(|set\((GLOB|GGML)' CMakeLists.txt
printf '\nCheck if ggml repository sources include allocator/init_tensor implementations:\n'
git ls-files ggml | rg '/(allocator|backend|cuda|metal)/|ggml-backend|ggml-tallocr' | sed -n '1,120p'
rg -n -C 3 'ggml_tallocr_new|ggml_tallocr_alloc|ggml_backend_buffer_init_tensor|init_tensor' ggml include common 2>/dev/null | sed -n '1,220p' || trueRepository: rodrigomatta/s2.cpp
Length of output: 12484
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf 'Find all files with allocator/init_tensor symbols, including submodules:\n'
git submodule status --recursive || true
files=$(git ls-files && git ls-tree -r --name-only HEAD ggml 2>/dev/null || true)
printf '%s\n' "$files" | rg -n -C 3 'ggml_tallocr_new|ggml_tallocr_alloc|ggml_backend_buffer_init_tensor' || true
printf '\nSearch for tensor buffer/data init calls in ggml and local ggml utils:\n'
rg -n -C 3 'tensor->buffer|backend_buffer_init_tensor|ggml_backend_buffer_init_tensor|ggml_backend_buffer_get_base' ggml src include 2>/dev/null | sed -n '1,240p' || true
printf '\nSearch project-wide for any free_codec_buffers or reset_codec_impl model_buf handling:\n'
rg -n -C 3 'reset_codec_impl|ensure_weights_loaded|weights_allocated_|restore_weights_to_gpu|restore_encoder_weights|restore_decoder_weights|free_gpu_weights|free_encoder_weights|free_decoder_weights' src include CMakeLists.txt | sed -n '1,320p'Repository: rodrigomatta/s2.cpp
Length of output: 19274
🌐 Web query:
ggml tallocr API ggml_tallocr_new ggml_tallocr_alloc ggml_backend_buffer_init_tensor documentation
💡 Result:
In the GGML library, ggml_tallocr is a lightweight, linear (bump-pointer) allocator designed for assigning individual tensors into a pre-existing backend buffer [1][2][3]. It is intended for manual, one-off allocations where you have control over the buffer and want to place tensors at specific offsets [1]. Key API components: 1. ggml_tallocr_new: - Signature: struct ggml_tallocr ggml_tallocr_new(ggml_backend_buffer_t buffer) - Function: Initializes a new ggml_tallocr struct using an existing backend buffer [1]. - Behavior: It retrieves the base address and alignment requirements of the provided buffer [4][5]. It does not perform heap allocations and does not take ownership of the buffer; you remain responsible for freeing the buffer [1]. 2. ggml_tallocr_alloc: - Signature: enum ggml_status ggml_tallocr_alloc(struct ggml_tallocr * talloc, struct ggml_tensor * tensor) - Function: Allocates space for a single tensor within the allocator's buffer [1]. - Behavior: It calculates the required memory size for the tensor (including necessary alignment padding), checks if sufficient space remains in the buffer, updates the allocator's internal offset, and finally calls ggml_backend_tensor_alloc to finalize the tensor's memory assignment [4][5][2]. 3. ggml_backend_buffer_init_tensor: - Signature: enum ggml_status ggml_backend_buffer_init_tensor(ggml_backend_buffer_t buffer, struct ggml_tensor * tensor) - Function: An internal or backend-specific function used to initialize a tensor within a buffer [6][7][8]. - Behavior: While standard tensor allocation is often handled by higher-level functions, this interface allows the backend to perform hardware-specific initialization (e.g., setting up metadata or custom resources) if the backend supports it [9][7]. It is often called during the allocation process to ensure the tensor is correctly prepared for the specific backend device [9]. In modern GGML usage, for complex computation graphs, the library provides ggml_gallocr, which performs automatic lifetime analysis and memory reuse [1][2]. Use ggml_tallocr only when precise, manual control over individual tensor placement is required [1][3].
Citations:
- 1: https://mintlify.wiki/ggml-org/ggml/api/allocator
- 2: https://deepwiki.com/ggml-org/ggml/2.3-memory-management-and-allocation
- 3: ggml-alloc v3 ggml-org/ggml#727
- 4: https://github.com/ggml-org/ggml/blob/c044a8ee/src/ggml-alloc.c
- 5: https://github.com/ggml-org/ggml/blob/3f3fe907/src/ggml-alloc.c
- 6: https://github.com/ggml-org/ggml/blob/3f3fe907/include/ggml-backend.h
- 7: https://github.com/pudepiedj/llama.cpp/blob/5b09797321430f08caf0473143a962916ab2ea89/ggml-backend.c
- 8: https://github.com/pudepiedj/llama.cpp/blob/8a3012a4ad08112bb3dc3f1399afec4e93780c44/ggml-backend.h
- 9: How to implement accelerator support as a backend? ggml-org/llama.cpp#11968
Use ggml_tallocr for codec tensor placement and do not discard the existing buffer handle.
allocate_codec_buffers receives impl_->model_buf, impl_->encoder_buf, or impl_->decoder_buf by reference. Setting out_buffer = nullptr abandons any existing handle and can leak a GPU buffer when restore_*_weights() calls this path after freeing only the tensor contents. In addition, assigning tensor->data and tensor->buffer directly bypasses the backend tensor allocation call used by the model allocator, so backends can miss required per-tensor initialization.
🤖 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 138 - 183, Update allocate_codec_buffers to
preserve any existing out_buffer handle instead of unconditionally resetting it,
and reuse the established ggml_tallocr-based tensor allocation path used by the
model allocator. Replace manual tensor->data and tensor->buffer assignments with
backend-aware allocation through ggml_tallocr, while retaining the existing
buffer ownership and error handling behavior.
| 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
Throw for unsupported tensor types in read_f32.
read_f32 handles GGML_TYPE_F32 and GGML_TYPE_F16 only. For any other type it returns a zero-filled vector and reports no error. tensor_to_f32 at lines 599-613 throws in that case, so this is a regression on the mmap path.
With a zero-filled codebook, codebook_norm becomes 0 / sqrt(1e-12), quantize_with_vq selects code 0 for every frame, and encode produces meaningless codes while returning success.
🐛 Proposed fix
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: " +
+ std::string(ggml_type_name(t->type)));
+ }
return out;📝 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.
| 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; | |
| }; | |
| 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: " + | |
| std::string(ggml_type_name(t->type))); | |
| } | |
| 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 1679 - 1693, Update the read_f32 lambda to
throw an error when t->type is neither GGML_TYPE_F32 nor GGML_TYPE_F16, matching
tensor_to_f32 behavior; do not return the zero-initialized vector for
unsupported tensor types.
| bool AudioCodec::ensure_weights_loaded() { | ||
| if (!impl_ || impl_->weights_allocated_) return true; | ||
| if (impl_->decoder_on_gpu || impl_->encoder_on_gpu) { | ||
| impl_->weights_allocated_ = true; | ||
| return true; | ||
| } | ||
| if (!impl_->mapped_gguf_.is_open()) return false; | ||
| S2_LOG_INFO_STREAM("[Codec] >>> Allocating and loading Audio Codec weights on demand..." << std::endl); | ||
| size_t b = 0; std::string e; | ||
| if (!allocate_codec_buffers(impl_->backend, impl_->all_codec_weights, impl_->model_buf, b, e)) { | ||
| std::cerr << "[Codec] Alloc failed: " << e << std::endl; return false; | ||
| } | ||
| const uint8_t* base = impl_->mapped_gguf_.data(); | ||
| for (ggml_tensor * t : impl_->all_codec_weights) { | ||
| auto it = impl_->tensor_offsets.find(t); | ||
| if (it != impl_->tensor_offsets.end()) | ||
| ggml_backend_tensor_set(t, base + impl_->gguf_data_offset + it->second, 0, ggml_nbytes(t)); | ||
| } | ||
| impl_->weights_allocated_ = true; | ||
| impl_->weights_on_gpu = !ggml_backend_is_cpu(impl_->backend); | ||
| impl_->encoder_on_gpu = impl_->weights_on_gpu; | ||
| impl_->decoder_on_gpu = impl_->weights_on_gpu; | ||
| return true; | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Trace codec residency transitions around encode and decode in the pipeline.
set -euo pipefail
rg -n -C 6 'free_encoder_weights|restore_encoder_weights|free_decoder_weights|restore_decoder_weights' src/s2_pipeline.cpp
rg -n -C 4 'codec\(\)\.encode|encode_prompt_audio' src/s2_pipeline.cppRepository: rodrigomatta/s2.cpp
Length of output: 11274
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== src/s2_codec.cpp relevant lines =="
sed -n '1520,1785p' src/s2_codec.cpp | cat -n | sed 's/^/src\/s2_codec.cpp:/'
echo
echo "== relevant residency flag usages =="
rg -n 'weights_on_gpu|encoder_on_gpu|decoder_on_gpu|weights_allocated_|weights|ensure_weights_loaded|free_encoder_weights|restore_encoder_weights|free_decoder_weights|restore_decoder_weights|free_gpu_weights|restore_weights_to_gpu|ggml_backend_tensor_set|ggml_graph_compute|ggml_backend_tensor_get|backend_tensor_get' src/s2_codec.cpp src/s2_pipeline.cpp | sed -n '1,240p'Repository: rodrigomatta/s2.cpp
Length of output: 30610
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== encode/decode implementation outlines =="
ast-grep outline src/s2_codec.cpp --match encode --view expanded || true
ast-grep outline src/s2_codec.cpp --match decode --view expanded || true
echo
echo "== encode/decode implementation lines =="
sed -n '900,1155p' src/s2_codec.cpp | cat -n | sed 's/^/src\/s2_codec.cpp:/'
sed -n '1360,1520p' src/s2_codec.cpp | cat -n | sed 's/^/src\/s2_codec.cpp:/'
echo
echo "== encode usage call-sites =="
rg -n -C 3 '\bencode\(' src/s2_codec.cpp src/s2_pipeline.h src/s2_pipeline.cpp | sed -n '1,220p'Repository: rodrigomatta/s2.cpp
Length of output: 33437
Ensure each codec group is allocated before encoding or decoding.
weights_allocated_ is a single boolean for encoder, decoder, and shared/combined codec tensors. If one group is on GPU and another is freed, later calls to encode() or decode() return early here and then use tensors with null data. Gate the allocation/restoration per group instead of treating partial residency as success.
🤖 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 1717 - 1740, Update
AudioCodec::ensure_weights_loaded() to track and restore residency independently
for encoder, decoder, and shared/combined codec tensor groups instead of
returning success solely from the global weights_allocated_ flag. Ensure each
group needed by encode() or decode() is allocated and loaded before use,
including when only some groups remain resident after freeing.
| StepResult state; | ||
| if (params.verbose && log_enabled(LogLevel::Info)) { | ||
| std::cout << "[Generate] Prefilling " << prompt.cols << " tokens..." << std::endl; | ||
| } | ||
| const auto prefill_t0 = std::chrono::steady_clock::now(); | ||
| if (!model.prefill_fast(prompt_tm, prompt.cols, params.n_threads, state)) { | ||
| std::cerr << "[Generate] Prefill failed." << std::endl; | ||
| return out; | ||
| double prefill_ms = 0.0; | ||
|
|
||
| if (initial_state) { | ||
| state = *initial_state; |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Validate initial_state->logits before the first sample.
Line 81 calls apply_mask_and_sample(state.logits, ...), and the lambda at line 68 reads logits[i] for every i < vocab_size with no bounds check. On the fresh path prefill_fast guarantees that state.logits has vocab_size entries. On the cached path the size depends entirely on the caller-supplied StepResult. A caller that supplies a default-constructed or truncated state causes an out-of-bounds read on the first sample.
Reject an undersized cached state instead of reading past the end.
🛡️ Proposed guard
if (initial_state) {
+ if (static_cast<int32_t>(initial_state->logits.size()) < vocab_size) {
+ std::cerr << "[Generate] Invalid initial_state: logits size "
+ << initial_state->logits.size() << " < vocab_size " << vocab_size << std::endl;
+ return out;
+ }
state = *initial_state;
} else {📝 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.
| StepResult state; | |
| if (params.verbose && log_enabled(LogLevel::Info)) { | |
| std::cout << "[Generate] Prefilling " << prompt.cols << " tokens..." << std::endl; | |
| } | |
| const auto prefill_t0 = std::chrono::steady_clock::now(); | |
| if (!model.prefill_fast(prompt_tm, prompt.cols, params.n_threads, state)) { | |
| std::cerr << "[Generate] Prefill failed." << std::endl; | |
| return out; | |
| double prefill_ms = 0.0; | |
| if (initial_state) { | |
| state = *initial_state; | |
| StepResult state; | |
| double prefill_ms = 0.0; | |
| if (initial_state) { | |
| if (static_cast<int32_t>(initial_state->logits.size()) < vocab_size) { | |
| std::cerr << "[Generate] Invalid initial_state: logits size " | |
| << initial_state->logits.size() << " < vocab_size " << vocab_size << std::endl; | |
| return out; | |
| } | |
| state = *initial_state; |
🤖 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_generate.cpp` around lines 38 - 42, Validate the cached state in the
initial_state branch before the first sampling call: ensure
initial_state->logits contains at least vocab_size entries, and reject
undersized or empty cached states before assigning or sampling from state. Keep
the existing fresh prefill path unchanged, and use the surrounding generate
function’s established error-handling behavior for the rejection.
| #ifdef __linux__ | ||
| ::madvise(data_, size_, MADV_WILLNEED); | ||
| ::madvise(data_, size_, MADV_SEQUENTIAL); | ||
|
|
||
| #ifdef MADV_COLD | ||
| ::madvise(data_, size_, MADV_COLD); | ||
| #endif |
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🌐 Web query:
Linux madvise MADV_COLD semantics inactive list versus MADV_WILLNEED
💡 Result:
In the Linux kernel, MADV_COLD and MADV_WILLNEED serve opposing roles in memory management, representing non-destructive hints regarding the expected future usage of memory pages [1][2]. MADV_COLD Introduced in Linux 5.4, MADV_COLD is a hint indicating that the specified range of pages is not expected to be used in the near future [1][3]. Its semantics are centered on page deactivation: - Deactivation: It moves the targeted pages from the active LRU (Least Recently Used) list to the inactive LRU list [3][4]. This makes them more probable candidates for reclamation should the system experience memory pressure [1][3]. - Non-destructive: Unlike MADV_DONTNEED or MADV_FREE, MADV_COLD is a non-destructive operation [1][2]. It does not discard the page contents [3][4]. If the pages are accessed again, they are simply faulted back in (if they were reclaimed) or remain available in memory, and the kernel will eventually promote them back to the active list [3][4]. MADV_WILLNEED In contrast, MADV_WILLNEED is a proactive hint indicating that the application expects to access the specified memory range in the near future [5][6]. - Pre-faulting/Read-ahead: The kernel uses this hint to trigger background activity to populate the page cache [6][7]. Specifically, it calls the kernel's read-ahead mechanism to bring the specified file-backed pages into memory before the application actually touches them [7]. - Performance optimization: It is intended to hide I/O latency by warming up the cache [5][6]. Summary of Differences MADV_COLD helps the kernel manage memory pressure by "cooling down" pages, making them easier to evict without data loss [1][2]. MADV_WILLNEED helps the application by "warming up" memory, proactively pulling data into the cache to avoid future page faults [5][7]. Top results: [1][3][6][7][2]
Citations:
- 1: https://man7.org/linux/man-pages/man2/madvise.2.html
- 2: https://lore.kernel.org/lkml/20190627115405.255259-1-minchan@kernel.org/
- 3: https://lists.openwall.net/linux-kernel/2019/07/11/24
- 4: https://lists.openwall.net/linux-kernel/2019/07/11/525
- 5: https://man7.org/linux/man-pages/man3/posix_madvise.3.html
- 6: https://github.com/torvalds/linux/blob/0adb32858b0bddf4ada5f364a84ed60b196dbcda/mm/madvise.c
- 7: https://stackoverflow.com/questions/31215250/is-mmap-madvise-really-a-form-of-async-i-o
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf 'Repo files matching s2_mapped_file:\n'
fd -a 's2_mapped_file\.(cpp|h|hpp)$' . || true
printf '\nRelevant section if present:\n'
file="$(fd 's2_mapped_file\.cpp$' . | head -n 1 || true)"
if [ -n "${file:-}" ]; then
wc -l "$file"
sed -n '1,260p' "$file" | nl -ba | sed -n '170,250p'
fi
printf '\nSearch for warm/drop/cold usages:\n'
rg -n "warm_page_cache|drop_page_cache|MADV_(WILLNEED|SEQUENTIAL|COLD|FREE)" .Repository: rodrigomatta/s2.cpp
Length of output: 379
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf 'src/s2_mapped_file.cpp section lines 170-250:\n'
sed -n '170,250p' src/s2_mapped_file.cpp
printf '\nHeader class outline / relevant declarations:\n'
sed -n '1,140p' include/s2_mapped_file.h
printf '\nmadvice usage references:\n'
rg -n "warm_page_cache|drop_page_cache|MADV_(WILLNEED|SEQUENTIAL|COLD|FREE)" src include || trueRepository: rodrigomatta/s2.cpp
Length of output: 4165
🌐 Web query:
Linux madvise MADV_COLD reads pages first before moving inactive LRU madvise.c
💡 Result:
The Linux kernel's implementation of MADV_COLD does not read pages before moving them to the inactive LRU list [1][2]. MADV_COLD is a hint to the kernel that a specified range of memory is not expected to be used in the near future, allowing the kernel to deactivate those pages and make them easier to reclaim during memory pressure [3][2]. The implementation, located in mm/madvise.c (specifically via madvise_cold and the underlying page walk functions), works by iterating through the page table entries (PTEs) of the specified range [1]. For each present page, it performs a check (such as verifying it is an LRU page that is currently active and not unevictable) and then directly moves it to the inactive list [2]. This process involves clearing page flags (like the referenced bit) and updating the page's status on the LRU lists, without performing any data read operations on the pages themselves [2]. In contrast to operations like MADV_WILLNEED, which might involve prefaulting or reading memory to prepare it, MADV_COLD is intended as a non-destructive, lightweight operation that simply adjusts the kernel's internal page management status to favor future reclamation [2][4].
Citations:
- 1: https://codebrowser.dev/linux/linux/mm/madvise.c.html
- 2: https://lists.openwall.net/linux-kernel/2019/07/11/24
- 3: https://man7.org/linux/man-pages/man2/madvise.2.html
- 4: https://lore.kernel.org/lkml/20190603053655.127730-1-minchan@kernel.org/
Remove MADV_COLD from warm_page_cache.
MADV_WILLNEED requests read-ahead for the mapping. MADV_COLD then marks pages as not needed soon and moves them onto the inactive list, so Linux can reclaim them first. Apply it to drop_page_cache if the mapping should be cooled after use.
🤖 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 204 - 210, Remove the MADV_COLD block
from warm_page_cache, leaving MADV_WILLNEED and MADV_SEQUENTIAL unchanged. If
cooling behavior is required, move the MADV_COLD advice to the drop_page_cache
implementation instead.
| const std::string cache_key = compute_prefill_cache_key(params, ref_codes, T_prompt); | ||
| bool prefill_hit = false; | ||
| StepResult cached_state; | ||
|
|
||
| if (res.n_frames == 0) { | ||
| safe_print_error_ln("Pipeline error: generation produced no frames."); | ||
| return false; | ||
| if (kv_reuse && prefill_cache_.valid && | ||
| prefill_cache_.cache_key == cache_key && | ||
| prefill_cache_.max_seq_len >= max_seq_len) | ||
| { | ||
| if (keep_kv_on_gpu && prefill_cache_.vram_resident) { | ||
| model().set_n_past(prefill_cache_.n_past); | ||
| cached_state = prefill_cache_.state; | ||
| prefill_hit = true; | ||
| safe_print_ln("[Pipeline] Prefill cache HIT (VRAM-pinned, key=" + cache_key + ")"); | ||
| } else if (!keep_kv_on_gpu && !prefill_cache_.k_data.empty()) { | ||
| if (model().restore_kv_state(prefill_cache_.k_data, prefill_cache_.v_data, | ||
| prefill_cache_.n_past)) { | ||
| cached_state = prefill_cache_.state; | ||
| prefill_hit = true; | ||
| safe_print_ln("[Pipeline] Prefill cache HIT (system RAM, key=" + cache_key + ")"); | ||
| } | ||
| } | ||
| } |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🔴 Critical | ⚡ Quick win
Also compare n_past consistency on a VRAM-pinned hit.
On the VRAM-pinned branch at lines 952-956 the code calls set_n_past(prefill_cache_.n_past) and reuses prefill_cache_.state without checking that the KV cache still holds that prefix. Lines 906-912 can have run clear_kv_cache() and init_kv_cache() on this same request when model().kv_max_seq_len() < max_seq_len, which discards the pinned KV contents while prefill_cache_.vram_resident stays true.
Trace: need_fresh_kv true → clear_kv_cache() at 911 → init_kv_cache() at 920 → the cache is empty. The lookup at 948 still passes because prefill_cache_.max_seq_len records the old kv_max_seq_len(), and valid and cache_key are untouched. Generation then continues from n_past over uninitialized KV data.
Invalidate the VRAM-resident entry whenever the KV cache is reallocated.
🐛 Proposed fix
if (need_fresh_kv) {
model().clear_kv_cache();
+ if (prefill_cache_.vram_resident) {
+ prefill_cache_.valid = false;
+ prefill_cache_.vram_resident = false;
+ }
}🤖 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 944 - 965, Invalidate the VRAM-resident
prefill cache entry whenever the KV cache is cleared or reinitialized in the
fresh-KV path around clear_kv_cache() and init_kv_cache(). Ensure
prefill_cache_.valid and its VRAM-resident state cannot remain true after
reallocation, so the lookup in the prefill-hit branch does not reuse stale
n_past or state.
| kv_save_thread = std::thread([this, save_n_past]() { | ||
| model().save_kv_state(prefill_cache_.k_data, | ||
| prefill_cache_.v_data, | ||
| save_n_past); | ||
| }); | ||
| } | ||
| prefill_cache_.valid = true; | ||
| safe_print_ln("[Pipeline] Prefill cache SAVED (key=" + cache_key + | ||
| ", n_past=" + std::to_string(prefill_cache_.n_past) + | ||
| (keep_kv_on_gpu ? ", VRAM)" : ", RAM, async)")); | ||
| } | ||
| } |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Do not publish prefill_cache_.valid before the async save completes.
Line 1002 starts kv_save_thread, which writes prefill_cache_.k_data and prefill_cache_.v_data. Line 1008 sets prefill_cache_.valid = true while that write is in flight.
Within this call the thread is joined at line 1141 or 1208, and synthesize_mutex_ serializes requests, so a later request cannot read a partial buffer on the normal path. The exposure is the abnormal path: if the function unwinds or returns before the join, valid stays true with a half-written k_data. The next request then reaches line 957 and calls restore_kv_state with truncated data.
Set valid in the save thread after the write, or set it after the join.
🐛 Proposed fix
if (keep_kv_on_gpu) {
prefill_cache_.vram_resident = true;
+ prefill_cache_.valid = true;
} else {
const int32_t save_n_past = prefill_cache_.n_past;
kv_save_thread = std::thread([this, save_n_past]() {
model().save_kv_state(prefill_cache_.k_data,
prefill_cache_.v_data,
save_n_past);
+ prefill_cache_.valid = true;
});
}
- prefill_cache_.valid = true;🤖 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 1002 - 1013, Update the prefill-cache save
flow around the kv_save_thread creation so prefill_cache_.valid remains false
while save_kv_state is writing; set it only after the asynchronous save
completes, either inside the save thread after the write or immediately after
the thread is joined. Keep the existing cache publication and logging behavior
otherwise unchanged.
| std::thread decode_thread([&]() { | ||
| int32_t last_committed = 0; | ||
| while (true) { | ||
| std::unique_lock<std::mutex> lock(decode_mtx); | ||
| decode_cv.wait(lock, [&]() { | ||
| return frames_available.load() > last_committed | ||
| || gen_done.load(); | ||
| }); | ||
| const int32_t avail = frames_available.load(); | ||
| const bool done = gen_done.load(); | ||
| lock.unlock(); | ||
| if (avail <= last_committed && done) | ||
| break; | ||
| if (!decode_window(avail, done)) { | ||
| decode_failed = true; | ||
| break; | ||
| } | ||
| last_committed = committed_frames; | ||
| } | ||
| }); | ||
|
|
||
| gen_params.on_frame = [&](const FrameCallbackData & fcd) -> bool { | ||
| { | ||
| std::lock_guard<std::mutex> lock(decode_mtx); | ||
| for (int32_t cb = 0; cb < fcd.num_codebooks; ++cb) | ||
| accum[cb].push_back(fcd.codes[cb]); | ||
| } | ||
| frames_available.store(fcd.total_frames); | ||
| decode_cv.notify_one(); | ||
| return true; | ||
| }; | ||
|
|
||
| const auto gen_t0 = std::chrono::steady_clock::now(); | ||
| res = generate(model(), tokenizer().config(), prompt, gen_params, | ||
| prefill_hit ? &cached_state : &prefill_state); | ||
| const auto gen_t1 = std::chrono::steady_clock::now(); | ||
| gen_ms = std::chrono::duration<double, std::milli>(gen_t1 - gen_t0).count(); | ||
|
|
||
| if (kv_save_thread.joinable()) { | ||
| kv_save_thread.join(); | ||
| } | ||
|
|
||
| { | ||
| std::lock_guard<std::mutex> lock(decode_mtx); | ||
| gen_done.store(true); | ||
| } | ||
| decode_cv.notify_one(); | ||
| decode_thread.join(); |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | 🏗️ Heavy lift
A throw between thread creation and join() calls std::terminate.
decode_thread is created at line 1102 and joined at line 1149. kv_save_thread is created at line 1002 and joined at line 1141. Between those points the code calls generate(...) at line 1135 and, inside the decode lambda, codec().decode(...) at line 1076.
generate allocates out.codes with num_cb * max_new_tokens elements and several per-step std::vector<float> buffers, so std::bad_alloc is reachable. The repository convention also allows std::runtime_error for hard failures in these components. If any of those calls throws, the stack unwinds and ~std::thread runs on a joinable thread, which calls std::terminate. On the server path that kills the whole process instead of failing one request.
Wrap each thread in a joining guard so unwinding joins instead of terminating.
As per coding guidelines: "Mix bool returns for recoverable operations and std::runtime_error for hard failures in C++".
🛡️ Proposed fix
Add a small joining guard near the top of the file:
struct ThreadJoiner {
explicit ThreadJoiner(std::thread & t) : t_(t) {}
~ThreadJoiner() { if (t_.joinable()) t_.join(); }
std::thread & t_;
};Then bind every locally owned thread:
std::thread kv_save_thread;
+ ThreadJoiner kv_save_joiner(kv_save_thread);
@@
std::thread decode_thread([&]() {
+ ThreadJoiner decode_joiner(decode_thread);Keep the explicit join() calls; the guard only covers the unwinding path.
🤖 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 1102 - 1149, Add a small RAII joining guard
near the top of the file, then bind it to both locally owned threads,
decode_thread and kv_save_thread, immediately after their creation. The guard
destructor must join any still-joinable thread so exceptions from generate or
codec().decode, including std::bad_alloc and std::runtime_error, cannot trigger
std::terminate; retain the existing explicit join() calls for the normal path.
Source: Coding guidelines
| " ms, prefill=" + std::to_string(prefill_ms) + | ||
| " ms, generate=" + std::to_string(gen_ms) + | ||
| " ms, decode=" + std::to_string(decode_ms) + | ||
| " ms, decode_wall=" + std::to_string(decode_wall_ms) + | ||
| " ms, decode_batches=" + std::to_string(decode_batches) + | ||
| ", decode_stride=" + std::to_string(offline_decode_stride_frames) + | ||
| " frames, total=" + std::to_string(total_ms) + | ||
| " frames" + | ||
| (can_overlap_decode ? ", decode_mode=overlapped" : ", decode_mode=sequential") + | ||
| (prefill_hit ? ", prefill=cached" : ", prefill=computed") + |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
The metrics line emits prefill= twice with different meanings.
Line 1294 emits prefill=<duration> and line 1302 emits prefill=cached or prefill=computed. A consumer that parses the metrics line by key reads an ambiguous prefill field. Rename the status key.
🐛 Proposed fix
- (prefill_hit ? ", prefill=cached" : ", prefill=computed") +
+ (prefill_hit ? ", prefill_source=cached" : ", prefill_source=computed") +📝 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.
| " ms, prefill=" + std::to_string(prefill_ms) + | |
| " ms, generate=" + std::to_string(gen_ms) + | |
| " ms, decode=" + std::to_string(decode_ms) + | |
| " ms, decode_wall=" + std::to_string(decode_wall_ms) + | |
| " ms, decode_batches=" + std::to_string(decode_batches) + | |
| ", decode_stride=" + std::to_string(offline_decode_stride_frames) + | |
| " frames, total=" + std::to_string(total_ms) + | |
| " frames" + | |
| (can_overlap_decode ? ", decode_mode=overlapped" : ", decode_mode=sequential") + | |
| (prefill_hit ? ", prefill=cached" : ", prefill=computed") + | |
| " ms, prefill=" + std::to_string(prefill_ms) + | |
| " ms, generate=" + std::to_string(gen_ms) + | |
| " ms, decode=" + std::to_string(decode_ms) + | |
| " ms, decode_wall=" + std::to_string(decode_wall_ms) + | |
| " ms, decode_batches=" + std::to_string(decode_batches) + | |
| ", decode_stride=" + std::to_string(offline_decode_stride_frames) + | |
| " frames" + | |
| (can_overlap_decode ? ", decode_mode=overlapped" : ", decode_mode=sequential") + | |
| (prefill_hit ? ", prefill_source=cached" : ", prefill_source=computed") + |
🤖 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 1294 - 1302, Rename the prefill status
field appended in the metrics construction near can_overlap_decode and
prefill_hit from “prefill” to a distinct key such as “prefill_mode,” while
leaving the earlier prefill duration field unchanged.
| if (model_prefers_gpu_ && !model().is_weights_on_gpu()) { | ||
| safe_print_ln("[Pipeline] Streaming: Restoring Slow-AR to VRAM..."); | ||
| model().acquire_compute_resources(); | ||
| model().restore_weights_to_gpu(); | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Check the result of restore_weights_to_gpu() in the streaming path.
Line 1387 discards the returned bool. The non-streaming path at lines 861-865 checks the same call, prints an error, and aborts the request. Here a failed restore continues into init_kv_cache at line 1424 and generate at line 1553 with the Slow-AR weights missing or partially resident. The request then produces incorrect audio or crashes, and the sink never receives an error.
Report the failure to the sink and return false.
🐛 Proposed fix
if (model_prefers_gpu_ && !model().is_weights_on_gpu()) {
safe_print_ln("[Pipeline] Streaming: Restoring Slow-AR to VRAM...");
model().acquire_compute_resources();
- model().restore_weights_to_gpu();
+ if (!model().restore_weights_to_gpu()) {
+ safe_print_error_ln("Pipeline error: Slow-AR weight restore failed.");
+ sink.on_error("Slow-AR weight restore failed");
+ return false;
+ }
}📝 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.
| if (model_prefers_gpu_ && !model().is_weights_on_gpu()) { | |
| safe_print_ln("[Pipeline] Streaming: Restoring Slow-AR to VRAM..."); | |
| model().acquire_compute_resources(); | |
| model().restore_weights_to_gpu(); | |
| } | |
| if (model_prefers_gpu_ && !model().is_weights_on_gpu()) { | |
| safe_print_ln("[Pipeline] Streaming: Restoring Slow-AR to VRAM..."); | |
| model().acquire_compute_resources(); | |
| if (!model().restore_weights_to_gpu()) { | |
| safe_print_error_ln("Pipeline error: Slow-AR weight restore failed."); | |
| sink.on_error("Slow-AR weight restore failed"); | |
| return false; | |
| } | |
| } |
🤖 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 1384 - 1388, Update the streaming restore
block around model().restore_weights_to_gpu() to check its boolean result; when
restoration fails, report the failure through the request’s sink using the
existing error-reporting pattern and return false before init_kv_cache or
generate can run. Preserve the successful restoration path unchanged.
1 participant
This PR introduces new parameters to enable KV cache quantization. By quantizing KV cache you can reduce both VRAM occupied space and memory bandwidth bottleneck. It's usually utilized by LLMs for long context generations, as for s2.cpp using
q8_0toq5_0for V cache gives 23% to 36% KV size reduction.This PR builds upon and requires PR #46
Changes
--cache-type-k <type>and--cache-type-v <type>.f32,f16,bf16,q8_0,q4_0,q4_1,iq4_nl,q5_0,q5_1.f16. Incorrect values fall back tof16with a warning.PipelineParamsnow carries the string types, which are parsed inPipeline::initviaparse_kv_cache_type()and applied toSlowARModelviaset_kv_cache_types()before any cache allocation.init_kv_cachenow allocatesmemory_k_andmemory_v_using the configuredggml_typeinstead of hardcodedGGML_TYPE_F16. The allocation log line now reports both types.save_kv_stateandrestore_kv_state. Previously, these functions derivedpos_bytes, total size, and layer/head offsets frommemory_k_->nbfor both K and V. This is only correct when K and V share the exact same type. With mixed types (e.g., K=f16, V=q8_0), the memory copies would read/write incorrect byte ranges, corrupting the restored cache on a prefill cache hit. Now they compute independent strides (k_pos_bytes,v_pos_bytes), total sizes, and source, destination offsets.Usage
Read caveats below first
Caveats
1. Quantize V first:
V cache quantization is more tolerable because V is linearly combined into the attention output.
2. K cache quantization Warning (Infinite generation):$Q \cdot K$ dot products that feed into the softmax. The exponential nature of softmax amplifies these quantization errors, which then puts the model in an infinite generation loop by never hitting EOS, even at
K cache quantization affects the
q8_0. The--cache-type-khelp text has a warning about this.Included K cache quantization because user preference should come first with warnings and some higher quant ggufs like Q8 might run quantized K cache without generating garbage output etc.
3. Which cache quant to use?
q8_0is recommended for perceivably same quality asf16.q5_0is where degradation is more obvious as breathiness/higher vocal frequencies get a sharper/synthetic sound. I recommend not going belowq5_1, but YMMV between different voice profiles and base model quants (I've used q4_k_m for testing this PR specifically - my performance tests in other PRs are the Q8 model.)You can look at llama.cpp perplexity loss numbers for each quant in some of their issues/PRs to get a better idea of each type and their quality, speed and size tradeoffs.
4. Quality vs. VRAM Tradeoffs:
Tested on Vulkan backend with
s2-pro-q4_k_mV=q8_0: Saves ~23% KV VRAM. Perceptually identical tof16. (Highly Recommended)V=q5_0/q4_0: Saves up to ~36% KV VRAM. Trades audible fidelity (e.g., introduces a sharp hiss on higher frequencies) with no generation-speed benefit. YMMV by hardware and backend. If you must use lower thanq8_0, test for speed and useq5_1where possible. It's generally the better trade-off afterq8_0Summary by CodeRabbit
New Features
Improvements