feat: KV cache reuse and prompt prefill caching - #46
Open
Skyrion9 wants to merge 5 commits into
Open
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.
📝 WalkthroughWalkthroughThis PR adds a cross-platform MappedFile abstraction and converts AudioCodec and SlowARModel to lazy, mmap-backed weight loading with GPU/CPU residency control. It adds KV-cache save/restore, prefill-state reuse in generate(), and reworks Pipeline synthesis to support VRAM swap, hot-swap offloading, and prefill caching, exposed via new CLI and server options. ChangesMmap-based lazy weight loading and VRAM swap pipeline
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related PRs
Suggested reviewers: 🚥 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)
202-225: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winFree
backend_cpuinreset_codec_impl.
load_sharedcreatesimpl_->backend_cpuat lines 958-960.reset_codec_implfreesbackend, but notbackend_cpu. The finalimpl = AudioCodec::Impl()overwrites the handle, so the CPU backend leaks on every reload and on destruction. Ifbackend_cpuis not used by any code path, remove the field instead.🔒️ Proposed fix to release the CPU backend
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 202 - 225, Update reset_codec_impl to release impl.backend_cpu with ggml_backend_free and clear it before resetting the Impl object, matching the existing backend cleanup. If backend_cpu has no consumers beyond its creation in load_shared, remove the unused field and its initialization instead.src/s2_generate.cpp (1)
192-208: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winCorrect the misleading prefill log on the supplied-state path.
src/s2_pipeline.cppalways passes a non-nullinitial_state:&cached_stateon a cache hit, and&prefill_statethat it just computed in this request (lines 1117-1118 and 1184-1185). So every Pipeline synthesis prints(prefill cached)andprefill=0here, including requests that paid the full prefill cost. The[Generate]line then contradicts the[Metrics] Synthesis: ... prefill=line.State only what this function knows.
🐛 Proposed log fix
std::cout << "[Generate] Done: " << out.n_frames << " frames generated." - << (initial_state ? " (prefill cached)" : "") << std::endl; + << (initial_state ? " (prefill supplied by caller)" : "") << std::endl;🤖 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 192 - 208, Update the verbose logging in the generation function around initial_state and prefill_ms so it does not label every supplied state as “prefill cached.” Report only that an initial state was supplied, and preserve the measured prefill value without inferring whether it came from a cache; keep the [Generate] and [Metrics] lines consistent.
🧹 Nitpick comments (5)
include/s2_model.h (1)
113-115: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winReject an out-of-range value in
set_n_past.
set_n_pastwritesn_past_without comparing it tomax_seq_len_. The pipeline sets this value from cached prefill state, so a stale or oversized value is possible.eval_cachedcatches the overflow later and returns false, which makes the cause hard to locate. Validate at the setter.♻️ Proposed refactor
- void set_n_past(int32_t n) { n_past_ = n; } + void set_n_past(int32_t n) { n_past_ = (n < 0) ? 0 : std::min(n, max_seq_len_); }🤖 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 `@include/s2_model.h` around lines 113 - 115, Update set_n_past in the model state API to validate the supplied value against max_seq_len_ before assigning n_past_. Reject values outside the valid range, including oversized cached prefill state, while preserving assignment for valid values and making the rejection explicit to callers.src/s2_codec.cpp (1)
966-966: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winAvoid copying the whole weight set in the ternary.
Model->weight_tensor_set()returns a reference, but the conditional expression materializes astd::unordered_setprvalue, so this line copies every model tensor pointer and bindsmodel_weightsto the temporary. Lifetime extension applies, so cppcheck'sdanglingTemporaryLifetimehint at line 979 is a false positive, but the copy is unnecessary and the construct is fragile. Use a function-local empty set as the fallback.♻️ Proposed refactor
- const auto & model_weights = Model ? Model->weight_tensor_set() : std::unordered_set<ggml_tensor*>(); + static const std::unordered_set<ggml_tensor*> empty_weight_set; + const std::unordered_set<ggml_tensor*> & model_weights = + Model ? Model->weight_tensor_set() : empty_weight_set;🤖 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 the weight-processing logic to use a function-local empty std::unordered_set as the fallback, then bind the reference through branches or equivalent reference-preserving logic so Model->weight_tensor_set() is not copied. Keep the existing behavior for both present and absent Model cases.Source: Linters/SAST tools
src/s2_generate.cpp (1)
38-62: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low valueAdd a size guard on the supplied
state.logits.On the cached path,
stateis fully caller-supplied. Line 81 then indexesstate.logitsup tovocab_sizewithout a size check. Current callers pass populated states, so this is defensive only.♻️ Optional guard
if (initial_state) { state = *initial_state; + if (static_cast<int32_t>(state.logits.size()) < vocab_size) { + std::cerr << "[Generate] initial_state has invalid logits size." << std::endl; + return out; + } } else {🤖 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 - 62, Add a defensive size check for caller-supplied state.logits in the initial_state branch before the later vocab_size-indexed access. Ensure the logits collection contains at least vocab_size entries; otherwise handle the invalid cached state through the existing failure path without indexing it. Keep the normal prefill path and valid cached-state behavior unchanged.src/s2_pipeline.cpp (2)
1105-1114: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winPublish
frames_availableunderdecode_mtx.The callback stores
frames_availableand callsnotify_one()without holdingdecode_mtx. The decode thread evaluates its predicate under the lock at line 1088. A notification that lands between the predicate check and the wait is lost, so that decode batch waits for the next frame or forgen_done.The final
gen_done.store(true)at line 1128 does hold the lock, so the thread always wakes and drains. The effect is added latency, not lost audio. Publish the counter the same way asgen_done.♻️ Proposed change
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); } - frames_available.store(fcd.total_frames); decode_cv.notify_one(); return 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 1105 - 1114, Update the on_frame callback in the gen_params setup to assign frames_available while holding decode_mtx, alongside the accum updates, before calling decode_cv.notify_one(). Preserve the existing atomic value and notification behavior, and do not move unrelated callback logic.
1163-1178: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExtract the hot-swap offload into one helper.
This lambda body is duplicated verbatim at lines 1220-1229 and at lines 1604-1613 in
synthesize_streaming_prompt_codes_locked. Three copies of a VRAM/RAM release sequence will drift.Each site also move-assigns into
pending_offload_thread_. A move-assign onto a joinablestd::threadcallsstd::terminate. Every current caller joins at request entry, so the target is empty today. A single helper lets you assert that invariant once.♻️ Proposed helper
// include/s2_pipeline.h, private section void spawn_hot_swap_offload_locked();// src/s2_pipeline.cpp void Pipeline::spawn_hot_swap_offload_locked() { if (pending_offload_thread_.joinable()) { pending_offload_thread_.join(); } 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..."); 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 call
spawn_hot_swap_offload_locked()at lines 1163-1178, 1216-1229, and 1599-1613.🤖 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 1163 - 1178, Extract the duplicated hot-swap VRAM/RAM release logic from the three call sites in synthesize_streaming_prompt_codes_locked into a private Pipeline helper named spawn_hot_swap_offload_locked(). Have the helper join any existing pending_offload_thread_ before releasing compute resources and spawning the background offload thread, then replace each inline lambda and move-assignment at the three sites with calls to the helper.
🤖 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 1677-1715: Update read_f32 in refresh_host_caches_from_mmap to
throw for tensor types other than GGML_TYPE_F32 and GGML_TYPE_F16, matching
tensor_to_f32 behavior, instead of returning zero-filled data. In load_vq,
validate that the loaded codebook contains at least cb_size * cb_dim elements
before computing codebook_norm, and throw on mismatch to prevent out-of-bounds
access.
- Around line 1717-1740: Fix ownership tracking across ensure_weights_loaded,
restore_weights_to_gpu, free_encoder_weights, and free_decoder_weights. Do not
treat encoder_on_gpu or decoder_on_gpu alone as proof that all_codec_weights
remain attached: restore or reload the missing partition after detachment, free
and clear model_buf before switching to partition buffers, and ensure each
free_* method releases model_buf ownership when it is the active backing
allocation before clearing tensor data.
- Around line 138-183: The allocate_codec_buffers function currently attempts to
allocate all tensors in a single ggml_backend_buft_alloc_buffer call, which can
fail on backends with allocation size limits. Follow the same chunking approach
as allocate_weight_buffers by checking ggml_backend_buft_get_max_size(buft) and
splitting tensor allocation into multiple buffers when total_bytes exceeds the
backend limit. Additionally, ensure tensor->data and tensor->buffer are reset to
initial values (likely nullptr) before attempting allocation and when allocation
or tensor placement fails, matching the model allocator's cleanup behavior.
In `@src/s2_mapped_file.cpp`:
- Around line 41-58: Update MappedFile’s open/mapping flow to keep the object
closed when no mapped view exists: store the file length in a local variable,
reject zero-length files unless an explicit open-state mechanism supports them,
and assign size_ only after MapViewOfFile or mmap succeeds. Apply the same
ordering and failure-state behavior to the corresponding path around the later
mapping logic.
- Around line 190-198: Update the Windows branch of
MappedFile::drop_page_cache() so it does not call SetProcessWorkingSetSizeEx
with -1, -1 and trim the entire process working set. Make this per-mapping
cache-drop operation a no-op on Windows, unless an existing application-wide
eviction policy explicitly requires the process-wide call.
- Around line 221-225: Update the Windows configuration before the `<windows.h>`
include so `_WIN32_WINNT` targets Windows 8 or later, enabling the
`PrefetchVirtualMemory` declaration used in the `_WIN32` branch of
`S2MappedFile`; preserve the existing call and behavior.
- Around line 204-210: Remove the conditional MADV_COLD madvise call from
warm_page_cache(), while preserving the existing MADV_WILLNEED and
MADV_SEQUENTIAL hints.
In `@src/s2_model.cpp`:
- Around line 671-702: Fix KV-cache serialization in src/s2_model.cpp lines
671-702 within SlowARModel::save_kv_state by using a per-layer layer_bytes value
based on n_positions * memory_k_->nb[2], removing the head loop, and copying one
contiguous block per layer from l * memory_k_->nb[3]. Apply the same per-layer
layout and recompute expected in src/s2_model.cpp lines 704-735 within
restore_kv_state; reject n_past values below zero or above max_seq_len_ before
any memcpy.
- Around line 1288-1320: Update SlowARModel::allocate_and_load_weights and the
shared tensor-upload path used by restore_encoder_weights,
restore_decoder_weights, and ensure_weights_loaded to validate that
gguf_data_offset_ + tensor offset + ggml_nbytes(t) stays within
mapped_gguf_.size() before calling ggml_backend_tensor_set, failing the load
when the range is invalid. Initialize the local output variables b and m to zero
before allocation.
- Around line 1254-1272: Change SlowARModel::acquire_compute_resources to return
a success boolean, validate the result of each ggml_backend_sched_new call, and
report failure when sched_ or fast_sched_ cannot be created. Update
allocate_and_load_weights to propagate that boolean and return false instead of
allowing execution with a null scheduler; preserve successful resource
initialization behavior.
In `@src/s2_pipeline.cpp`:
- Around line 1596-1622: Update the streaming cleanup block around
enable_vram_swap and the hot-swap offload thread to skip VRAM/RAM release when
params.more_segments_pending is true. Preserve cleanup for the final segment,
including synchronous releases and background offload behavior, so multi-segment
requests retain weights and mapped pages between segments.
- Around line 835-922: The vram_phase1_thread and kv_init_thread concurrently
mutate the same SlowARModel state, creating an unsafe race. In the generation
flow around vram_phase1_thread and the KV-cache initialization, ensure one
thread is joined or otherwise completed before starting the other; preserve the
existing failure checks while making
acquire_compute_resources/restore_weights_to_gpu and
init_kv_cache/reset_kv_cache execute sequentially.
- Around line 980-994: Move the asynchronous save_kv_state operation in the
prefill cache save flow below generate() completes, before prefill_cache_.valid
is set true, so KV backup serialization cannot overlap GPU computation. Update
the logic around save_n_past, kv_save_thread, and model().save_kv_state while
preserving the keep_kv_on_gpu path and existing cache status logging.
- Around line 885-947: Invalidate prefill_cache_ before every destructive KV
operation, including clear_kv_cache(), reset_kv_cache(), streaming-synthesis KV
clearing, and overwriting an existing valid entry; clear both valid and
vram_resident state. Also mark the cache invalid whenever prefill_fast or
eval_cached fails, so later requests cannot reuse stale KV data.
- Around line 305-322: Update compute_prefill_cache_key to make the cache key
cover the full prompt identity. The current implementation hashes only the first
16 frames of codebook-0 from ref_codes and omits T_prompt, causing different
prompts with identical first 16 codebook-0 codes to collide. Add num_codebooks
as a function parameter (either by passing it to the static method or converting
compute_prefill_cache_key to non-static), then update the loop that currently
iterates for only the first 16 entries to iterate through all codebooks and all
T_prompt frames instead. Include T_prompt in the key by appending it before or
after the code iteration to ensure the full prompt dimensionality is captured in
the cache lookup.
- Around line 440-483: The codec_prefers_gpu_ flag at line 473 is being set from
use_gpu_codec, which records the requested GPU placement, not whether the codec
actually ended up on GPU. When GPU codec loading fails at line 421 and falls
back to CPU at line 437, use_gpu_codec remains true but the codec resides on
CPU, causing codec_prefers_gpu_ to incorrectly reflect GPU placement and
misleading downstream weight restoration calls and the VRAM state machine
banner. Introduce a separate bool variable to track the actual achieved codec
placement (whether GPU loading succeeded or fell back to CPU), initialize it
near codec_loaded around line 363, update it based on the success or failure of
GPU vs CPU codec loading paths, and then use this achieved-placement variable
instead of use_gpu_codec when setting codec_prefers_gpu_ at line 473.
- Around line 1361-1382: Handle the boolean results from
model().restore_weights_to_gpu() and codec().restore_decoder_weights() in the
VRAM restoration block before streaming generation. If either restore fails,
stop the streaming path and propagate a descriptive error containing the restore
failure reason instead of continuing to generate or decode with incomplete GPU
state; preserve the existing successful restore flow.
---
Outside diff comments:
In `@src/s2_codec.cpp`:
- Around line 202-225: Update reset_codec_impl to release impl.backend_cpu with
ggml_backend_free and clear it before resetting the Impl object, matching the
existing backend cleanup. If backend_cpu has no consumers beyond its creation in
load_shared, remove the unused field and its initialization instead.
In `@src/s2_generate.cpp`:
- Around line 192-208: Update the verbose logging in the generation function
around initial_state and prefill_ms so it does not label every supplied state as
“prefill cached.” Report only that an initial state was supplied, and preserve
the measured prefill value without inferring whether it came from a cache; keep
the [Generate] and [Metrics] lines consistent.
---
Nitpick comments:
In `@include/s2_model.h`:
- Around line 113-115: Update set_n_past in the model state API to validate the
supplied value against max_seq_len_ before assigning n_past_. Reject values
outside the valid range, including oversized cached prefill state, while
preserving assignment for valid values and making the rejection explicit to
callers.
In `@src/s2_codec.cpp`:
- Line 966: Update the model_weights initialization near the weight-processing
logic to use a function-local empty std::unordered_set as the fallback, then
bind the reference through branches or equivalent reference-preserving logic so
Model->weight_tensor_set() is not copied. Keep the existing behavior for both
present and absent Model cases.
In `@src/s2_generate.cpp`:
- Around line 38-62: Add a defensive size check for caller-supplied state.logits
in the initial_state branch before the later vocab_size-indexed access. Ensure
the logits collection contains at least vocab_size entries; otherwise handle the
invalid cached state through the existing failure path without indexing it. Keep
the normal prefill path and valid cached-state behavior unchanged.
In `@src/s2_pipeline.cpp`:
- Around line 1105-1114: Update the on_frame callback in the gen_params setup to
assign frames_available while holding decode_mtx, alongside the accum updates,
before calling decode_cv.notify_one(). Preserve the existing atomic value and
notification behavior, and do not move unrelated callback logic.
- Around line 1163-1178: Extract the duplicated hot-swap VRAM/RAM release logic
from the three call sites in synthesize_streaming_prompt_codes_locked into a
private Pipeline helper named spawn_hot_swap_offload_locked(). Have the helper
join any existing pending_offload_thread_ before releasing compute resources and
spawning the background offload thread, then replace each inline lambda and
move-assignment at the three sites with calls to the helper.
🪄 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: d9d3df6a-f592-4fa2-8542-778904f10322
📒 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
Comment on lines
+138
to
+183
| 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 | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Compare the two weight-buffer allocators and check ggml's max-size API usage.
set -uo pipefail
rg -n -C 3 'ggml_backend_buft_get_max_size'
echo '--- codec allocator ---'
ast-grep run --pattern 'static bool allocate_codec_buffers($$$) { $$$ }' --lang cpp src/s2_codec.cppRepository: rodrigomatta/s2.cpp
Length of output: 157
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "tracked files with s2:"
git ls-files | rg '(^|/)s2_(codec|model)\.(cpp|h)$|ggml.*backend' | sed -n '1,120p'
echo
echo "search codec/weight allocators:"
rg -n -C 4 'allocate_codec_buffers|allocate_weight_buffers|ggml_backend_buft_get_max_size|restore_encoder_weights|restore_decoder_weights' . --glob '!build/**' --glob '!dist/**' || true
echo
echo "s2 codec outline:"
ast-grep outline src/s2_codec.cpp 2>/dev/null | sed -n '1,200p' || true
echo
echo "s2 codec lines 1-230:"
sed -n '1,230p' src/s2_codec.cpp
echo
echo "s2 model lines 1-230:"
sed -n '1,230p' src/s2_model.cppRepository: rodrigomatta/s2.cpp
Length of output: 28386
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "tracked files with s2:"
git ls-files | rg '(^|/)s2_(codec|model)\.(cpp|h)$|ggml.*backend' | sed -n '1,120p'
echo
echo "search codec/weight allocators:"
rg -n -C 4 'allocate_codec_buffers|allocate_weight_buffers|ggml_backend_buft_get_max_size|restore_encoder_weights|restore_decoder_weights' . --glob '!build/**' --glob '!dist/**' || true
echo
echo "s2_codec.cpp lines 100-210:"
sed -n '100,210p' src/s2_codec.cpp
echo
echo "s2_model.cpp lines 90-200:"
sed -n '90,200p' src/s2_model.cppRepository: rodrigomatta/s2.cpp
Length of output: 16850
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "codec restore/free/load restore paths:"
sed -n '1530,1760p' src/s2_codec.cpp
echo
echo "Codec impl buffer member usage:"
rg -n -C 2 'impl_->(encoder_buf|decoder_buf|model_buf)|encoder_buf|decoder_buf|model_buf|weights_allocated_|allocate_codec_buffers' src/s2_codec.cpp include/s2_codec.h
echo
echo "Static verifier: allocate_codec_buffers does not respect max_size or reset failed tensors"
python3 - <<'PY'
from pathlib import Path
p = Path('src/s2_codec.cpp')
src = p.read_text()
start = src.index('static bool allocate_codec_buffers')
end = src.index('static void reset_decode_cache', start)
body = src[start:end]
checks = {
'clears_out_buffer_at_start': 'out_buffer = nullptr;' in body[:body.index('const ggml_backend_buffer_type_t buft')],
'allocates_single_buffer': 'ggml_backend_buft_alloc_buffer(buft, total_bytes)' in body and body.count('ggml_backend_buft_alloc_buffer') == 1,
'gets_max_size': 'ggml_backend_buft_get_max_size' in body,
'chunks': 'BufferChunk' in body or 'ggml_tallocr_new' in body or '<vector>' in body and 'chunks' in body,
'clears_tensor_data_on_failure': 'tensor->data = nullptr;' in body or 'tensor->data = nullptr;' in body,
'clears_tensor_buffer_on_failure': 'tensor->buffer = nullptr;' in body,
'frees_out_buffer_on_exit_after_failure': 'ggml_backend_buffer_free' in body,
}
for k,v in checks.items():
print(f'{k}: {v}')
PYRepository: rodrigomatta/s2.cpp
Length of output: 19124
Handle per-backend buffer limits in allocate_codec_buffers.
allocate_weight_buffers splits tensors by ggml_backend_buft_get_max_size(buft), but this allocator requests total_bytes in one ggml_backend_buft_alloc_buffer call. Backends with allocation limits can fail at Restore, returning "failed to allocate codec buffer of size ...". Use the same chunking/tallocr approach, or change the storage shape accordingly.
Also reset tensor->data/tensor->buffer on allocation or tensor-placement failure; the model allocator does this before allocation and clears tensors on failure.
🤖 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, The allocate_codec_buffers function
currently attempts to allocate all tensors in a single
ggml_backend_buft_alloc_buffer call, which can fail on backends with allocation
size limits. Follow the same chunking approach as allocate_weight_buffers by
checking ggml_backend_buft_get_max_size(buft) and splitting tensor allocation
into multiple buffers when total_bytes exceeds the backend limit. Additionally,
ensure tensor->data and tensor->buffer are reset to initial values (likely
nullptr) before attempting allocation and when allocation or tensor placement
fails, matching the model allocator's cleanup behavior.
Comment on lines
+1677
to
+1715
| 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; | ||
| }; | ||
| try { | ||
| auto load_vq = [&](const std::string& prefix, int32_t in_dim, int32_t cb_dim, int32_t cb_size) -> vq_cache { | ||
| vq_cache vq; vq.input_dim = in_dim; vq.codebook_dim = cb_dim; vq.codebook_size = cb_size; | ||
| vq.in_proj_weight = read_f32(prefix + ".in_proj.weight"); vq.in_proj_bias = read_f32(prefix + ".in_proj.bias"); | ||
| vq.out_proj_weight = read_f32(prefix + ".out_proj.weight"); vq.out_proj_bias = read_f32(prefix + ".out_proj.bias"); | ||
| vq.codebook = read_f32(prefix + ".codebook.weight"); | ||
| vq.codebook_norm.resize(vq.codebook.size()); | ||
| for (int32_t c = 0; c < cb_size; ++c) { | ||
| float norm = 0.0f; const size_t base = c * cb_dim; | ||
| for (int32_t d = 0; d < cb_dim; ++d) norm += vq.codebook[base+d] * vq.codebook[base+d]; | ||
| norm = std::sqrt(std::max(norm, 1e-12f)); | ||
| for (int32_t d = 0; d < cb_dim; ++d) vq.codebook_norm[base+d] = vq.codebook[base+d] / norm; | ||
| } | ||
| return vq; | ||
| }; | ||
| impl_->semantic_vq = load_vq(impl_->tprefix + "quantizer.semantic_quantizer.quantizers.0", impl_->quantizer_input_dim, impl_->quantizer_codebook_dim, impl_->quantizer_semantic_codebook_size); | ||
| impl_->residual_vq.clear(); impl_->residual_vq.reserve(impl_->quantizer_residual_codebooks); | ||
| for (int32_t i = 0; i < impl_->quantizer_residual_codebooks; ++i) | ||
| impl_->residual_vq.push_back(load_vq(impl_->tprefix + "quantizer.quantizer.quantizers." + std::to_string(i), impl_->quantizer_input_dim, impl_->quantizer_codebook_dim, impl_->quantizer_residual_codebook_size)); | ||
| } catch (const std::exception& e) { std::cerr << "[Codec] VQ mmap load failed: " << e.what() << std::endl; return false; } | ||
| return true; | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
read_f32 returns zeros for unsupported tensor types.
The lambda handles GGML_TYPE_F32 and GGML_TYPE_F16. For any other type it returns a zero-filled vector and reports success. load_vq then builds a codebook of zeros, codebook_norm divides by the 1e-12f floor, and quantize_with_vq selects code 0 for every frame. The failure is silent and produces wrong audio codes. Throw instead, as the previous tensor_to_f32 helper at lines 599-613 does.
The loop at lines 1701-1706 also assumes vq.codebook.size() >= cb_size * cb_dim. Validate the element count before the loop, because a metadata/tensor mismatch causes an out-of-bounds read.
🐛 Proposed fix to fail 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 for " + name + ": " +
+ std::string(ggml_type_name(t->type)));
return out;- vq.codebook = read_f32(prefix + ".codebook.weight");
+ vq.codebook = read_f32(prefix + ".codebook.weight");
+ if (vq.codebook.size() < static_cast<size_t>(cb_size) * cb_dim) {
+ throw std::runtime_error("codebook too small for " + prefix);
+ }📝 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.
Suggested change
| 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; | |
| }; | |
| try { | |
| auto load_vq = [&](const std::string& prefix, int32_t in_dim, int32_t cb_dim, int32_t cb_size) -> vq_cache { | |
| vq_cache vq; vq.input_dim = in_dim; vq.codebook_dim = cb_dim; vq.codebook_size = cb_size; | |
| vq.in_proj_weight = read_f32(prefix + ".in_proj.weight"); vq.in_proj_bias = read_f32(prefix + ".in_proj.bias"); | |
| vq.out_proj_weight = read_f32(prefix + ".out_proj.weight"); vq.out_proj_bias = read_f32(prefix + ".out_proj.bias"); | |
| vq.codebook = read_f32(prefix + ".codebook.weight"); | |
| vq.codebook_norm.resize(vq.codebook.size()); | |
| for (int32_t c = 0; c < cb_size; ++c) { | |
| float norm = 0.0f; const size_t base = c * cb_dim; | |
| for (int32_t d = 0; d < cb_dim; ++d) norm += vq.codebook[base+d] * vq.codebook[base+d]; | |
| norm = std::sqrt(std::max(norm, 1e-12f)); | |
| for (int32_t d = 0; d < cb_dim; ++d) vq.codebook_norm[base+d] = vq.codebook[base+d] / norm; | |
| } | |
| return vq; | |
| }; | |
| impl_->semantic_vq = load_vq(impl_->tprefix + "quantizer.semantic_quantizer.quantizers.0", impl_->quantizer_input_dim, impl_->quantizer_codebook_dim, impl_->quantizer_semantic_codebook_size); | |
| impl_->residual_vq.clear(); impl_->residual_vq.reserve(impl_->quantizer_residual_codebooks); | |
| for (int32_t i = 0; i < impl_->quantizer_residual_codebooks; ++i) | |
| impl_->residual_vq.push_back(load_vq(impl_->tprefix + "quantizer.quantizer.quantizers." + std::to_string(i), impl_->quantizer_input_dim, impl_->quantizer_codebook_dim, impl_->quantizer_residual_codebook_size)); | |
| } catch (const std::exception& e) { std::cerr << "[Codec] VQ mmap load failed: " << e.what() << std::endl; return false; } | |
| return true; | |
| } | |
| 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 for " + name + ": " + | |
| std::string(ggml_type_name(t->type))); | |
| return out; | |
| }; | |
| try { | |
| auto load_vq = [&](const std::string& prefix, int32_t in_dim, int32_t cb_dim, int32_t cb_size) -> vq_cache { | |
| vq_cache vq; vq.input_dim = in_dim; vq.codebook_dim = cb_dim; vq.codebook_size = cb_size; | |
| vq.in_proj_weight = read_f32(prefix + ".in_proj.weight"); vq.in_proj_bias = read_f32(prefix + ".in_proj.bias"); | |
| vq.out_proj_weight = read_f32(prefix + ".out_proj.weight"); vq.out_proj_bias = read_f32(prefix + ".out_proj.bias"); | |
| vq.codebook = read_f32(prefix + ".codebook.weight"); | |
| if (vq.codebook.size() < static_cast<size_t>(cb_size) * cb_dim) { | |
| throw std::runtime_error("codebook too small for " + prefix); | |
| } | |
| vq.codebook_norm.resize(vq.codebook.size()); | |
| for (int32_t c = 0; c < cb_size; ++c) { | |
| float norm = 0.0f; const size_t base = c * cb_dim; | |
| for (int32_t d = 0; d < cb_dim; ++d) norm += vq.codebook[base+d] * vq.codebook[base+d]; | |
| norm = std::sqrt(std::max(norm, 1e-12f)); | |
| for (int32_t d = 0; d < cb_dim; ++d) vq.codebook_norm[base+d] = vq.codebook[base+d] / norm; | |
| } | |
| return vq; | |
| }; | |
| impl_->semantic_vq = load_vq(impl_->tprefix + "quantizer.semantic_quantizer.quantizers.0", impl_->quantizer_input_dim, impl_->quantizer_codebook_dim, impl_->quantizer_semantic_codebook_size); | |
| impl_->residual_vq.clear(); impl_->residual_vq.reserve(impl_->quantizer_residual_codebooks); | |
| for (int32_t i = 0; i < impl_->quantizer_residual_codebooks; ++i) | |
| impl_->residual_vq.push_back(load_vq(impl_->tprefix + "quantizer.quantizer.quantizers." + std::to_string(i), impl_->quantizer_input_dim, impl_->quantizer_codebook_dim, impl_->quantizer_residual_codebook_size)); | |
| } catch (const std::exception& e) { std::cerr << "[Codec] VQ mmap load failed: " << e.what() << std::endl; return false; } | |
| return 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_codec.cpp` around lines 1677 - 1715, Update read_f32 in
refresh_host_caches_from_mmap to throw for tensor types other than GGML_TYPE_F32
and GGML_TYPE_F16, matching tensor_to_f32 behavior, instead of returning
zero-filled data. In load_vq, validate that the loaded codebook contains at
least cb_size * cb_dim elements before computing codebook_norm, and throw on
mismatch to prevent out-of-bounds access.
Comment on lines
+1717
to
+1740
| 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
ensure_weights_loaded reports success while encoder or decoder tensors are detached.
The whole-model path and the partition paths use different buffers, but the residency flags do not record which buffer owns a tensor. This sequence produces null weight data:
ensure_weights_loadedallocatesmodel_bufforall_codec_weightsand setsencoder_on_gpu = decoder_on_gpu = true(lines 1736-1738).free_encoder_weights(lines 1555-1573) freesencoder_buf, which is stillnullptrin this state, sets every encoder tensor todata = nullptr, and clearsweights_allocated_.model_bufstays allocated, so no VRAM is released.- The next
encode()callsensure_weights_loaded.decoder_on_gpuis still true, so the function setsweights_allocated_ = trueand returns true. - The encoder graph then runs with encoder weights whose
dataisnullptr.
free_decoder_weights produces the mirror failure for decode(). restore_weights_to_gpu (lines 1543-1553) has the same gap: it clears weights_allocated_ and calls ensure_weights_loaded, which returns early when either partition flag is set, and it overwrites a non-null model_buf without freeing it.
Restore the missing partition instead of returning early, and make free_encoder_weights/free_decoder_weights release model_buf ownership before they detach tensors.
🐛 Proposed fix for the early-return path
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_->encoder_on_gpu || impl_->decoder_on_gpu) {
+ if (!impl_->encoder_on_gpu && !restore_encoder_weights()) return false;
+ if (!impl_->decoder_on_gpu && !restore_decoder_weights()) return false;
+ impl_->weights_allocated_ = true;
+ return 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_codec.cpp` around lines 1717 - 1740, Fix ownership tracking across
ensure_weights_loaded, restore_weights_to_gpu, free_encoder_weights, and
free_decoder_weights. Do not treat encoder_on_gpu or decoder_on_gpu alone as
proof that all_codec_weights remain attached: restore or reload the missing
partition after detachment, free and clear model_buf before switching to
partition buffers, and ensure each free_* method releases model_buf ownership
when it is the active backing allocation before clearing tensor data.
Comment on lines
+41
to
+58
| size_ = static_cast<size_t>(file_size.QuadPart); | ||
|
|
||
| if (size_ == 0) { | ||
| CloseHandle(fh); | ||
| return true; | ||
| } | ||
|
|
||
| HANDLE mh = CreateFileMappingW(fh, nullptr, PAGE_READONLY, 0, 0, nullptr); | ||
| if (!mh) { | ||
| CloseHandle(fh); | ||
| return false; | ||
| } | ||
|
|
||
| void* addr = MapViewOfFile(mh, FILE_MAP_READ, 0, 0, size_); | ||
| if (!addr) { | ||
| CloseHandle(mh); | ||
| CloseHandle(fh); | ||
| return false; |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Keep MappedFile closed when no view exists.
size_ is assigned before mapping. Empty files return true after the descriptor closes, while is_open() remains false. Mapping failures return false while size() retains the failed file length.
Store the length in a local variable. Assign size_ only after MapViewOfFile or mmap succeeds. Reject zero-length input unless an explicit open-state field supports it.
Also applies to: 76-87
🤖 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 41 - 58, Update MappedFile’s
open/mapping flow to keep the object closed when no mapped view exists: store
the file length in a local variable, reject zero-length files unless an explicit
open-state mechanism supports them, and assign size_ only after MapViewOfFile or
mmap succeeds. Apply the same ordering and failure-state behavior to the
corresponding path around the later mapping logic.
Comment on lines
+190
to
+198
| #elif defined(_WIN32) | ||
| (void)data_; | ||
| (void)size_; | ||
| SetProcessWorkingSetSizeEx( | ||
| GetCurrentProcess(), | ||
| static_cast<SIZE_T>(-1), | ||
| static_cast<SIZE_T>(-1), | ||
| QUOTA_LIMITS_HARDWS_MIN_DISABLE); | ||
| #endif |
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/usr/bin/env bash
set -euo pipefail
# Trace callers. Expect: process-wide trimming is not required by a per-file cache-release path.
rg -n -C 4 \
-g '*.cpp' -g '*.h' -g '*.hpp' \
'\bdrop_page_cache\s*\(|SetProcessWorkingSetSizeEx' .Repository: rodrigomatta/s2.cpp
Length of output: 4475
🌐 Web query:
Microsoft SetProcessWorkingSetSizeEx -1 minimum negative infinity maximum RemoveWorkingSetEx
💡 Result:
The function SetProcessWorkingSetSizeEx does not use "negative infinity" as an input parameter. Instead, Microsoft documentation specifies that setting both the dwMinimumWorkingSetSize and dwMaximumWorkingSetSize parameters to (SIZE_T)–1 is the mechanism used to instruct the system to remove as many pages as possible from the specified process's working set [1][2]. This operation effectively empties the working set, achieving the same result as calling the EmptyWorkingSet function [3]. There is no function named RemoveWorkingSetEx in the Windows API. It is likely that this is a misremembered reference to the functionality provided by SetProcessWorkingSetSizeEx or EmptyWorkingSet [4][3]. In summary: - To minimize a process's working set: Use SetProcessWorkingSetSizeEx (or SetProcessWorkingSetSize) with both size parameters set to (SIZE_T)–1 [1][2]. - Functionality: This forces the memory manager to trim as many pages as possible from the process [1][3]. - Terminology: The value (SIZE_T)–1 is a cast of the constant -1 to an unsigned size type (resulting in the maximum possible value for a SIZE_T), which is the documented sentinel value for this operation [1].
Citations:
- 1: https://learn.microsoft.com/en-us/windows/win32/api/memoryapi/nf-memoryapi-setprocessworkingsetsizeex
- 2: https://learn.microsoft.com/en-us/windows/win32/api/memoryapi/nf-memoryapi-setprocessworkingsetsize
- 3: https://learn.microsoft.com/en-us/windows/win32/api/psapi/nf-psapi-emptyworkingset
- 4: https://learn.microsoft.com/en-us/windows/win32/memory/working-set
Do not trim the whole process in MappedFile::drop_page_cache().
drop_page_cache() is called on individual model/codec mappings during hot-swap, but the Windows branch trims pages from the entire process. Use -1, -1 only if the application intentionally has an application-wide eviction policy, or make this branch a no-op.
🤖 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 of
MappedFile::drop_page_cache() so it does not call SetProcessWorkingSetSizeEx
with -1, -1 and trim the entire process working set. Make this per-mapping
cache-drop operation a no-op on Windows, unless an existing application-wide
eviction policy explicitly requires the process-wide call.
Comment on lines
+835
to
922
| std::thread vram_phase1_thread; | ||
| bool vram_phase1_ok = true; | ||
|
|
||
| if (params.enable_vram_swap) { | ||
| vram_phase1_thread = std::thread([this, ¶ms, &vram_phase1_ok]() { | ||
| if (model_prefers_gpu_ && !model().is_weights_on_gpu()) { | ||
| safe_print_ln("[Pipeline] Restoring Slow-AR to VRAM for generation..."); | ||
| model().acquire_compute_resources(); | ||
| if (!model().restore_weights_to_gpu()) { | ||
| safe_print_error_ln("Pipeline error: Slow-AR weight restore failed."); | ||
| vram_phase1_ok = false; | ||
| return; | ||
| } | ||
| safe_print_ln("[VRAM Diag] Post-SlowAR restore: Slow-AR=" + std::to_string(model().get_gpu_memory_usage_bytes() / 1024 / 1024) + " MB, Codec=" + std::to_string(codec().get_gpu_memory_usage_bytes() / 1024 / 1024) + " MB"); | ||
| } | ||
|
|
||
| if (!model_prefers_gpu_ && codec_prefers_gpu_ && !codec().is_decoder_on_gpu()) { | ||
| safe_print_ln("[Pipeline] Pre-loading Audio Codec decoder to VRAM (hiding behind CPU gen)..."); | ||
| codec().restore_decoder_weights(); | ||
| safe_print_ln("[VRAM Diag] Post-Decoder restore: Slow-AR=" + std::to_string(model().get_gpu_memory_usage_bytes() / 1024 / 1024) + " MB, Codec=" + std::to_string(codec().get_gpu_memory_usage_bytes() / 1024 / 1024) + " MB"); | ||
| } | ||
|
|
||
| if (model_prefers_gpu_ && codec_prefers_gpu_) { | ||
| if (codec().is_encoder_on_gpu()) { | ||
| safe_print_ln("[Pipeline] Freeing codec encoder from VRAM (not needed during generation)..."); | ||
| codec().free_encoder_weights(); | ||
| } | ||
| if (codec().is_decoder_on_gpu()) { | ||
| safe_print_ln("[Pipeline] Freeing codec decoder from VRAM (not needed during generation)..."); | ||
| codec().free_decoder_weights(); | ||
| } | ||
| if (codec().is_weights_on_gpu()) { | ||
| safe_print_ln("[Pipeline] Freeing Audio Codec from VRAM for Slow-AR generation..."); | ||
| codec().free_gpu_weights(); | ||
| } | ||
| safe_print_ln("[VRAM Diag] Post-Codec free: Slow-AR=" + std::to_string(model().get_gpu_memory_usage_bytes() / 1024 / 1024) + " MB, Codec=" + std::to_string(codec().get_gpu_memory_usage_bytes() / 1024 / 1024) + " MB"); | ||
| } | ||
|
|
||
| safe_print_ln("[VRAM Diag] End-Phase1: Slow-AR=" + | ||
| std::to_string(model().get_gpu_memory_usage_bytes() / 1024 / 1024) + " MB, Codec=" + | ||
| std::to_string(codec().get_gpu_memory_usage_bytes() / 1024 / 1024) + " MB"); | ||
| }); | ||
| } | ||
|
|
||
| const int32_t num_codebooks = model().hparams().num_codebooks; | ||
| PromptTensor prompt = build_prompt( | ||
| tokenizer(), params.text, params.prompt_text, | ||
| ref_codes, | ||
| num_codebooks, T_prompt); | ||
|
|
||
| ref_codes, num_codebooks, T_prompt); | ||
| int32_t max_seq_len = prompt.cols + params.gen.max_new_tokens; | ||
|
|
||
| const bool kv_reuse = params.enable_kv_reuse && params.is_persistent; | ||
| const bool keep_kv_on_gpu = kv_reuse && params.kv_cache_vram; | ||
|
|
||
| const bool need_fresh_kv = !kv_reuse || | ||
| model().kv_max_seq_len() < max_seq_len || | ||
| model().kv_max_seq_len() == 0; | ||
|
|
||
| if (need_fresh_kv) { | ||
| model().clear_kv_cache(); | ||
| } | ||
|
|
||
| const auto kv_t0 = std::chrono::steady_clock::now(); | ||
| if (!model().init_kv_cache(max_seq_len)) { | ||
| std::thread kv_init_thread; | ||
| bool kv_init_ok = true; | ||
|
|
||
| if (need_fresh_kv) { | ||
| kv_init_thread = std::thread([&]() { | ||
| kv_init_ok = model().init_kv_cache(max_seq_len); | ||
| }); | ||
| } else { | ||
| kv_init_thread = std::thread([&]() { | ||
| model().reset_kv_cache(); | ||
| }); | ||
| } | ||
|
|
||
| if (vram_phase1_thread.joinable()) { | ||
| vram_phase1_thread.join(); | ||
| } | ||
| if (!vram_phase1_ok) { | ||
| kv_init_thread.join(); | ||
| return false; | ||
| } | ||
|
|
||
| kv_init_thread.join(); | ||
| if (!kv_init_ok) { | ||
| safe_print_error_ln("Pipeline error: init_kv_cache failed."); | ||
| return false; | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Inspect SlowARModel KV/weight lifecycle for shared backend state and any locking.
set -euo pipefail
fd -t f 's2_model\.(h|cpp)$' src include
ast-grep outline src/s2_model.cpp --items all
for fn in init_kv_cache reset_kv_cache restore_weights_to_gpu acquire_compute_resources free_compute_buffers save_kv_state restore_kv_state set_n_past n_past kv_max_seq_len; do
echo "===== $fn ====="
rg -nP -C 12 "SlowARModel::${fn}\s*\(" src/s2_model.cpp || true
done
echo "===== locking primitives in model =====";
rg -nP '\b(std::mutex|lock_guard|unique_lock|atomic<|ggml_backend_synchronize)\b' src/s2_model.cpp include/s2_model.h || trueRepository: rodrigomatta/s2.cpp
Length of output: 9807
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "===== s2_model.h member declarations ====="
ast-grep outline include/s2_model.h --items all
for fn in SlowARModel::init_kv_cache SlowARModel::clear_kv_cache SlowARModel::restore_weights_to_gpu SlowARModel::free_gpu_weights SlowARModel::free_compute_buffers SlowARModel::acquire_compute_resources SlowARModel::get_gpu_memory_usage_bytes SlowARModel::clear_kv_cache SlowARModel::is_weights_on_gpu SlowARModel::kv_max_seq_len; do
echo "===== $fn ====="
rg -n -C 8 "$fn" include/s2_model.h src/s2_model.cpp || true
done
echo "===== allocate_and_load_weights full ====="
rg -n -C 20 "bool SlowARModel::allocate_and_load_weights" src/s2_model.cpp
echo "===== model backend initialization/backend_gpu members ====="
rg -n -C 8 "backend_gpu_|backend_cpu_|kv_buf_|memory_k_|memory_v_|original_gpu_weights_|weights_allocated_" src/s2_model.cpp include/s2_model.h
echo "===== codec gpu memory/read/write symbols in pipeline/codec files ====="
rg -n "restore_decoder_weights|free_decoder_weights|free_gpu_weights|is_decoder_on_gpu|get_gpu_memory_usage_bytes|get_decoder_memory_usage|is_encoder_on_gpu|free_encoder_weights|restore_encoder_weights|load .*backend.*gpu|allocate_and_load_weights" src include --glob '*.{cpp,h,ht' || true
echo "===== Pipeline surrounding lines ====="
sed -n '800,935p' src/s2_pipeline.cppRepository: rodrigomatta/s2.cpp
Length of output: 50375
Do not run model state changes across two threads.
When enable_vram_swap is true, vram_phase1_thread can call acquire_compute_resources() / restore_weights_to_gpu(), while kv_init_thread allocates or resets the Slow-AR KV cache on the same SlowARModel. Those paths share backend buffers and tensor states, and the model has no synchronization around them. Join or disable one thread before starting the other, or add locking/backend handles that guarantee both paths run on disjoint resources.
🤖 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 835 - 922, The vram_phase1_thread and
kv_init_thread concurrently mutate the same SlowARModel state, creating an
unsafe race. In the generation flow around vram_phase1_thread and the KV-cache
initialization, ensure one thread is joined or otherwise completed before
starting the other; preserve the existing failure checks while making
acquire_compute_resources/restore_weights_to_gpu and
init_kv_cache/reset_kv_cache execute sequentially.
Comment on lines
+885
to
+947
| const bool kv_reuse = params.enable_kv_reuse && params.is_persistent; | ||
| const bool keep_kv_on_gpu = kv_reuse && params.kv_cache_vram; | ||
|
|
||
| const bool need_fresh_kv = !kv_reuse || | ||
| model().kv_max_seq_len() < max_seq_len || | ||
| model().kv_max_seq_len() == 0; | ||
|
|
||
| if (need_fresh_kv) { | ||
| model().clear_kv_cache(); | ||
| } | ||
|
|
||
| const auto kv_t0 = std::chrono::steady_clock::now(); | ||
| if (!model().init_kv_cache(max_seq_len)) { | ||
| std::thread kv_init_thread; | ||
| bool kv_init_ok = true; | ||
|
|
||
| if (need_fresh_kv) { | ||
| kv_init_thread = std::thread([&]() { | ||
| kv_init_ok = model().init_kv_cache(max_seq_len); | ||
| }); | ||
| } else { | ||
| kv_init_thread = std::thread([&]() { | ||
| model().reset_kv_cache(); | ||
| }); | ||
| } | ||
|
|
||
| if (vram_phase1_thread.joinable()) { | ||
| vram_phase1_thread.join(); | ||
| } | ||
| if (!vram_phase1_ok) { | ||
| kv_init_thread.join(); | ||
| return false; | ||
| } | ||
|
|
||
| kv_init_thread.join(); | ||
| if (!kv_init_ok) { | ||
| safe_print_error_ln("Pipeline error: init_kv_cache failed."); | ||
| return false; | ||
| } | ||
|
|
||
| const auto kv_t1 = std::chrono::steady_clock::now(); | ||
|
|
||
| const auto gen_t0 = std::chrono::steady_clock::now(); | ||
| GenerateResult res = generate(model(), tokenizer().config(), prompt, params.gen); | ||
| const auto gen_t1 = std::chrono::steady_clock::now(); | ||
| 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
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Determine whether reset_kv_cache preserves KV contents and what clear_kv_cache frees.
set -euo pipefail
rg -nP -C 20 'SlowARModel::(reset_kv_cache|clear_kv_cache)\s*\(' src/s2_model.cpp || true
rg -nP -C 6 '\b(reset_kv_cache|clear_kv_cache|set_n_past)\s*\(' src include --type cpp || trueRepository: rodrigomatta/s2.cpp
Length of output: 12640
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== pipeline relevant sections =="
sed -n '850,980p' src/s2_pipeline.cpp
sed -n '1228,1270p' src/s2_pipeline.cpp
sed -n '1370,1412p' src/s2_pipeline.cpp
sed -n '1512,1568p' src/s2_pipeline.cpp
echo "== prefill fast definitions/usages =="
rg -n -C 4 'prefill_fast|bool prefill_fast|prefill_fast\s*(' src include --type cpp
echo "== prefill_cache declarations/assignments =="
rg -n -C 3 'prefill_cache_|struct .*Prefill|Cache .*Prefill' src/s2_pipeline.cpp include --type hpp --type h --type cppRepository: rodrigomatta/s2.cpp
Length of output: 11248
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== prefill_fast definitions/usages =="
rg -n -C 4 'prefill_fast|bool prefill_fast' src include --type cpp --type h --type hpp
echo "== prefill_cache declarations/assignments/structures =="
rg -n -C 3 'prefill_cache_|struct .*Prefill|Cache .*Prefill|valid|vram_resident' src/s2_pipeline.cpp include --type hpp --type h --type cpp
echo "== source around VRAM offload and KV-save paths =="
sed -n '940,1025p' src/s2_pipeline.cpp
sed -n '1208,1228p' src/s2_pipeline.cppRepository: rodrigomatta/s2.cpp
Length of output: 227
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n -C 4 'prefill_fast|bool prefill_fast' src include
rg -n -C 3 'prefill_cache_|struct .*Prefill|Cache .*Prefill|valid|vram_resident' src/s2_pipeline.cpp include
sed -n '940,1025p' src/s2_pipeline.cpp
sed -n '1208,1228p' src/s2_pipeline.cppRepository: rodrigomatta/s2.cpp
Length of output: 12865
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== class declaration and relevant member/function outlines =="
sed -n '118,142p' include/s2_pipeline.h
ast-grep outline src/s2_pipeline.cpp --match Pipeline::synthesize --view expanded || true
ast-grep outline src/s2_pipeline.cpp --match Pipeline::synthesize_streaming_prompt_codes_locked --view expanded || true
ast-grep outline src/s2_model.cpp --match SlowARModel::prefill_fast --view expanded || true
ast-grep outline src/s2_model.cpp --match SlowARModel::reset_kv_cache --view expanded || true
echo "== streaming function body =="
sed -n '1288,1398p' src/s2_pipeline.cppRepository: rodrigomatta/s2.cpp
Length of output: 6548
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== full synthesize flow =="
sed -n '814,994p' src/s2_pipeline.cpp
echo "== full prefill miss cache save section =="
sed -n '968,998p' src/s2_pipeline.cppRepository: rodrigomatta/s2.cpp
Length of output: 9584
Invalidate prefill_cache_ before reusing or clearing the KV cache.
The VRAM-pinned hit only calls set_n_past() and trusts the KV buffer still contains that key. prefill_cache_.valid and prefill_cache_.vram_resident should be cleared before each destructive KV operation: before clear_kv_cache(), before reset_kv_cache(), before streaming synthesis clears the KV cache, and before saving a new cache entry over an old valid entry. Mark the entry invalid on prefill_fast/eval_cached failure too. Without this, a later matching request can read another request’s KV 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_pipeline.cpp` around lines 885 - 947, Invalidate prefill_cache_ before
every destructive KV operation, including clear_kv_cache(), reset_kv_cache(),
streaming-synthesis KV clearing, and overwriting an existing valid entry; clear
both valid and vram_resident state. Also mark the cache invalid whenever
prefill_fast or eval_cached fails, so later requests cannot reuse stale KV data.
Comment on lines
+980
to
+994
| if (keep_kv_on_gpu) { | ||
| prefill_cache_.vram_resident = 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; | ||
| 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.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Check how save_kv_state reads KV tensors and whether it synchronizes the backend.
set -euo pipefail
rg -nP -C 25 'SlowARModel::save_kv_state\s*\(' src/s2_model.cpp || true
rg -nP -C 25 'SlowARModel::restore_kv_state\s*\(' src/s2_model.cpp || true
rg -nP -C 4 'ggml_backend_(tensor_get|tensor_set|synchronize|graph_compute)' src/s2_model.cpp || trueRepository: rodrigomatta/s2.cpp
Length of output: 8793
🏁 Script executed:
#!/bin/bash
# Check call sites and synchronization around save_kv_state/generate(), plus any backend synchronization in the save path.
set -euo pipefail
fd -a 'src/s2_pipeline.cpp|src/s2_model.cpp' . | sed 's#^\./##'
echo '--- s2_pipeline cache save/generate sites ---'
rg -n -C 12 'kv_save_thread|prefill_cache_.valid|generate\(|save_kv_state\(' src/s2_pipeline.cpp
echo '--- s2_model schedule fields/constructors'
rg -n -C 8 'ggml_backend_sched_|backend_gpu_|sched_|kv_save|save_kv_state\(' src/s2_model.cpp
echo '--- broader backend tensor graph calls ---'
rg -n 'ggml_backend_(sched_)?(tensor_get|tensor_set|graph_compute|synchronize|scheduler)?' src || true
echo '--- git diff stat/name/status if available ---'
git diff --stat || true
git diff -- src/s2_pipeline.cpp src/s2_model.cpp | sed -n '1,220p' || trueRepository: rodrigomatta/s2.cpp
Length of output: 40137
🌐 Web query:
ggml_backend_tensor_get documentation concurrent tensor get during graph compute
💡 Result:
The function ggml_backend_tensor_get is a synchronous operation used to copy data from a backend-allocated tensor to host memory [1][2]. Concurrent use of ggml_backend_tensor_get while a graph compute operation is running is generally unsafe and not recommended [3]. Backends, particularly hardware-accelerated ones like CUDA, operate asynchronously from the host [3][4]. Tensor memory contents are only guaranteed to be in a well-defined state between complete graph evaluations [3]. To safely access tensor data: 1. Ensure the backend has finished its computations. You can do this by calling ggml_backend_synchronize(backend) after initiating a graph compute operation [2][5]. 2. Alternatively, use the synchronous wrapper ggml_backend_graph_compute, which implicitly handles synchronization after execution is complete [2]. 3. For developers needing to inspect data during generation, ensure you synchronize the backend before calling ggml_backend_tensor_get to avoid race conditions or accessing uninitialized/in-progress memory [3][2]. Some backends support asynchronous versions of data access (e.g., get_tensor_async) for advanced use cases [6], but these still require proper synchronization primitives to ensure the graph computation has reached a point where the target tensor is ready for read access [3][2]. Always check ggml_backend_buffer_is_host(tensor->buffer) if you are attempting to read tensor data directly from memory without using the provided getter functions [3].
Citations:
- 1: https://github.com/ggml-org/ggml/blob/master/include/ggml-backend.h
- 2: https://github.com/ggml-org/ggml/blob/9be31331/src/ggml-backend.cpp
- 3: helper function to examine tensor data for any backend? ggml-org/ggml#1155
- 4: https://deepwiki.com/ggml-org/ggml/2.4-backend-interface-and-abstraction
- 5: https://github.com/ggml-org/llama.cpp/blob/f486ce9f/ggml/include/ggml-backend.h
- 6: https://github.com/ggml-org/llama.cpp/blob/0cea3622/ggml/src/ggml-backend-impl.h
Serialize the KV backup while generate() runs.
save_kv_state uses ggml_backend_tensor_get on the same GPU backend that generate() computes with. Calling an asynchronous tensor read during ggml_backend_sched_graph_compute can race with GPU work and produce a corrupted buffer or backend failure.
Move the save below generate() before prefill_cache_.valid = true, or snapshot the needed KV region before generation starts.
🤖 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 980 - 994, Move the asynchronous
save_kv_state operation in the prefill cache save flow below generate()
completes, before prefill_cache_.valid is set true, so KV backup serialization
cannot overlap GPU computation. Update the logic around save_n_past,
kv_save_thread, and model().save_kv_state while preserving the keep_kv_on_gpu
path and existing cache status logging.
Comment on lines
+1361
to
+1382
| if (params.enable_vram_swap) { | ||
| if (pending_offload_thread_.joinable()) { | ||
| pending_offload_thread_.join(); | ||
| } | ||
|
|
||
| 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 (codec_prefers_gpu_) { | ||
| if (codec().is_encoder_on_gpu()) { | ||
| codec().free_encoder_weights(); | ||
| } | ||
| if (!codec().is_decoder_on_gpu() && !codec().is_weights_on_gpu()) { | ||
| safe_print_ln("[Pipeline] Streaming: Restoring Audio Codec DECODER to VRAM..."); | ||
| codec().restore_decoder_weights(); | ||
| } | ||
| } | ||
| } | ||
|
|
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Confirm the return types of the restore APIs used by the streaming path.
set -euo pipefail
rg -nP -C 2 '\b(restore_weights_to_gpu|restore_decoder_weights|restore_encoder_weights|acquire_compute_resources)\s*\(' include/s2_model.h include/s2_codec.h || trueRepository: rodrigomatta/s2.cpp
Length of output: 1189
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "## src/s2_pipeline.cpp streaming restore path"
sed -n '1345,1390p' src/s2_pipeline.cpp | cat -n
echo
echo "## src/s2_pipeline.cpp non-streaming restore path"
sed -n '830,855p' src/s2_pipeline.cpp | cat -n
echo
echo "## s2_codec implementation restore/decode error paths"
fd -a 's2_codec' . | sed 's#^\./##'
rg -n -C 6 '\b(restore_decoder_weights|restore_encoder_weights|restore_weights_to_gpu|decode_audio|decode)\b' src include -g 's2_codec*' -g '*.cpp' -g '*.h' | head -n 220 || trueRepository: rodrigomatta/s2.cpp
Length of output: 18455
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "## s2_codec restore implementations"
rg -n -C 5 'S2Codec::restore_(weights_to_gpu|decoder_weights|encoder_weights)_to_gpu|S2Codec::restore_decoder_weights|S2Codec::restore_encoder_weights' src include -g 's2_codec*' | sed -n '1,240p'
echo
echo "## s2_model restore implementation"
rg -n -C 8 'S2Model::restore_weights_to_gpu|model_restore_weights_to_gpu' src include -g 's2_model*' | sed -n '1,240p'
echo
echo "## pipeline generate call in streaming path"
sed -n '1390,1545p' src/s2_pipeline.cpp | cat -nRepository: rodrigomatta/s2.cpp
Length of output: 193
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "## files matching s2_codec/s2_model"
git ls-files | rg '(^|/)(s2_codec|s2_model)(\.(cpp|h|hpp))?$' || true
echo
echo "## occurrences of restore decoder/encoder/weights to gpu"
rg -n -C 5 'restore_(decoder_weights|encoder_weights|weights_to_gpu)|free_decoder_weights|free_encoder_weights|free_gpu_weights' src include --glob 's2_codec.*' --glob 's2_model.*' | sed -n '1,260p' || true
echo
echo "## s2_model restore occurrences"
rg -n -C 8 'restore_weights_to_gpu|restore.*gpu|acquire_compute_resources' src include --glob 's2_model.*' | sed -n '1,260p' || true
echo
echo "## streaming generate call"
sed -n '1390,1545p' src/s2_pipeline.cpp | cat -nRepository: rodrigomatta/s2.cpp
Length of output: 19269
Handle VRAM restore failures before streaming generation.
Lines 1368-1378 discard the bool results from model().restore_weights_to_gpu() and codec().restore_decoder_weights(). If either restore fails, generate() or codec decode can run without the expected GPU weight state, and streaming reports only generic on_error values without the actual restore failure reason.
🐛 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;
+ }
}
@@
if (!codec().is_decoder_on_gpu() && !codec().is_weights_on_gpu()) {
safe_print_ln("[Pipeline] Streaming: Restoring Audio Codec DECODER to VRAM...");
- codec().restore_decoder_weights();
+ if (!codec().restore_decoder_weights()) {
+ safe_print_error_ln("Pipeline error: decoder weight restore failed.");
+ sink.on_error("Decoder 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.
Suggested change
| if (params.enable_vram_swap) { | |
| if (pending_offload_thread_.joinable()) { | |
| pending_offload_thread_.join(); | |
| } | |
| 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 (codec_prefers_gpu_) { | |
| if (codec().is_encoder_on_gpu()) { | |
| codec().free_encoder_weights(); | |
| } | |
| if (!codec().is_decoder_on_gpu() && !codec().is_weights_on_gpu()) { | |
| safe_print_ln("[Pipeline] Streaming: Restoring Audio Codec DECODER to VRAM..."); | |
| codec().restore_decoder_weights(); | |
| } | |
| } | |
| } | |
| if (params.enable_vram_swap) { | |
| if (pending_offload_thread_.joinable()) { | |
| pending_offload_thread_.join(); | |
| } | |
| 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; | |
| } | |
| } | |
| if (codec_prefers_gpu_) { | |
| if (codec().is_encoder_on_gpu()) { | |
| codec().free_encoder_weights(); | |
| } | |
| if (!codec().is_decoder_on_gpu() && !codec().is_weights_on_gpu()) { | |
| safe_print_ln("[Pipeline] Streaming: Restoring Audio Codec DECODER to VRAM..."); | |
| if (!codec().restore_decoder_weights()) { | |
| safe_print_error_ln("Pipeline error: decoder weight restore failed."); | |
| sink.on_error("Decoder 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 1361 - 1382, Handle the boolean results
from model().restore_weights_to_gpu() and codec().restore_decoder_weights() in
the VRAM restoration block before streaming generation. If either restore fails,
stop the streaming path and propagate a descriptive error containing the restore
failure reason instead of continuing to generate or decode with incomplete GPU
state; preserve the existing successful restore flow.
Comment on lines
+1596
to
+1622
|
|
||
| if (params.enable_vram_swap) { | ||
| if (params.is_persistent) { | ||
| 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_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."); | ||
| }); | ||
| pending_offload_thread_ = std::move(offload_thread); | ||
| } else { | ||
| if (codec().is_decoder_on_gpu()) codec().free_decoder_weights(); | ||
| model().free_gpu_weights(); | ||
| model().free_compute_buffers(); | ||
| } | ||
| } else { | ||
| safe_print_ln("[Pipeline] Single-shot mode: Skipping post-stream VRAM restore."); | ||
| } | ||
| } |
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟠 Major | ⚡ Quick win
Honour more_segments_pending in the streaming cleanup.
The non-streaming path gates every VRAM release on !params.more_segments_pending (lines 1148, 1164, 1216, 1233). This block ignores the flag. A multi-segment streaming request therefore frees the model weights and drops the mmap page cache after each segment, and line 1366 restores them again at the start of the next segment.
That reintroduces the per-segment restore cost the flag exists to prevent, and the vram=held metric has no streaming counterpart.
🐛 Proposed fix
if (params.enable_vram_swap) {
if (params.is_persistent) {
- if (params.enable_hot_swap) {
+ if (params.more_segments_pending) {
+ safe_print_ln("[Pipeline] Streaming: More segments pending, holding VRAM.");
+ } else if (params.enable_hot_swap) {
safe_print_ln("[Pipeline] Hot-Swap: Releasing compute resources...");📝 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.
Suggested change
| if (params.enable_vram_swap) { | |
| if (params.is_persistent) { | |
| 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_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."); | |
| }); | |
| pending_offload_thread_ = std::move(offload_thread); | |
| } else { | |
| if (codec().is_decoder_on_gpu()) codec().free_decoder_weights(); | |
| model().free_gpu_weights(); | |
| model().free_compute_buffers(); | |
| } | |
| } else { | |
| safe_print_ln("[Pipeline] Single-shot mode: Skipping post-stream VRAM restore."); | |
| } | |
| } | |
| if (params.enable_vram_swap) { | |
| if (params.is_persistent) { | |
| if (params.more_segments_pending) { | |
| safe_print_ln("[Pipeline] Streaming: More segments pending, holding VRAM."); | |
| } else 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_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."); | |
| }); | |
| pending_offload_thread_ = std::move(offload_thread); | |
| } else { | |
| if (codec().is_decoder_on_gpu()) codec().free_decoder_weights(); | |
| model().free_gpu_weights(); | |
| model().free_compute_buffers(); | |
| } | |
| } else { | |
| safe_print_ln("[Pipeline] Single-shot mode: Skipping post-stream VRAM 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_pipeline.cpp` around lines 1596 - 1622, Update the streaming cleanup
block around enable_vram_swap and the hot-swap offload thread to skip VRAM/RAM
release when params.more_segments_pending is true. Preserve cleanup for the
final segment, including synchronous releases and background offload behavior,
so multi-segment requests retain weights and mapped pages between segments.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
In upstream we do redundant prompt prefills across repeated server requests. This PR caches the KV state and generation
StepResultafter the first prefill. On cache hit, the ~190–210 ms prefill pass is skipped entirely and generation resumes directly from the cached position.This is more of an "edge" scenario in that you might want to regenerate the same input over again. For example this could be due to using a lower quant or expecting a more fitting reading of the sentence as AI can get the emotions or tonals wrong.
This PR builds upon and requires PR #44
You may revert to old behavior by using (
--no-kv-reuse) to free and reallocate the KV buffer on every request.Old behavior:
To solve this we cache the prefill result keyed by
(voice_id | prompt_text + ref_codes) + textAfter prefill, the KV state is serialized to a strided buffer in system RAM on a background thread that overlaps with generation. On cache hit, the state is restored via a single PCIe upload of the active positions.
New behavior:
When running with (
--kv-cache-vram) the KV buffer stays allocated in VRAM between requests.free_compute_buffers()is skipped whenkeep_kv_on_gpuis set, so the next request with a matching key resumes without the PCIe transfer cost.This just consumes ~168 MB VRAM which would've been stored in RAM instead between requests without this flag set.
Architecture
Any change to voice, prompt, or text invalidates the cache.
Cache entry
KV lifecycle per request
CLI flags
--no-kv-reuse--kv-cache-vramMetrics
New fields in the
[Metrics] Synthesislog line:prefill=<ms>- prefill duration (0.0 on cache hit)prefill=cached/prefill=computed- cache hit/miss indicatorNew in
[Metrics] Generate:(prefill cached)suffix on the "Done" line wheninitial_statewas usedThread safety
kv_save_threadwrites toprefill_cache_.k_data/v_dataasynchronously. It is joined aftergenerate()returns in both the overlapped and sequential decode paths, before any subsequent request can read the cache.synthesize_mutex_serializes all requests, soprefill_cache_is never accessed concurrently.save_kv_state/restore_kv_stateoperate on the KV buffer under the same mutex.Interaction with
more_segments_pendingWithin a segmented server request (multiple sentences),
more_segments_pendingis true for all segments except the last. The KV cache and prefill state are preserved across segments — the end-of-request cleanup is deferred until the final segment, avoiding per-segment prefill for multi segmented generation.Expected impact
Summary by CodeRabbit