From ad61fbbf419e4b4cf06cbe45672be2873b9b5d47 Mon Sep 17 00:00:00 2001 From: Skyrion9 Date: Sun, 19 Jul 2026 22:49:31 +0300 Subject: [PATCH 1/5] feat: mmap-based lazy weight loading implementation 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. --- CMakeLists.txt | 1 + include/s2_codec.h | 21 ++- include/s2_mapped_file.h | 40 +++++ include/s2_model.h | 30 +++- src/s2_codec.cpp | 290 ++++++++++++++++++++++++++---------- src/s2_mapped_file.cpp | 173 ++++++++++++++++++++++ src/s2_model.cpp | 309 ++++++++++++++++++--------------------- src/s2_pipeline.cpp | 121 +++------------ 8 files changed, 637 insertions(+), 348 deletions(-) create mode 100644 include/s2_mapped_file.h create mode 100644 src/s2_mapped_file.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index f6d8a51..c4257f7 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -89,6 +89,7 @@ set(S2_CORE_SOURCES src/s2_pipeline.cpp src/s2_server.cpp src/s2_voice.cpp + src/s2_mapped_file.cpp ) function(s2_configure_target target_name) diff --git a/include/s2_codec.h b/include/s2_codec.h index 0781624..8a7e791 100644 --- a/include/s2_codec.h +++ b/include/s2_codec.h @@ -1,6 +1,7 @@ #pragma once #include "s2_backend.h" +#include "s2_mapped_file.h" #include "ggml.h" #include "ggml-alloc.h" #include "ggml-backend.h" @@ -10,6 +11,8 @@ #include #include #include +#include +#include #include "s2_model.h" @@ -24,10 +27,6 @@ class AudioCodec { bool load_shared(SlowARModel* Model, gguf_context * gguf_ctx, const std::string & gguf_path, int32_t gpu_device = -1, BackendType backend_type = BackendType::CPU); - bool read_tensor_data(const std::string & gguf_path, gguf_context * gguf_ctx); - - bool refresh_host_caches(); - ggml_context * weights_ctx() const; bool encode(const float * audio, int32_t n_samples, int32_t n_threads, @@ -38,6 +37,20 @@ class AudioCodec { void clear_decode_cache(); + MappedFile& mapped_file(); + + bool restore_weights_to_gpu(); + + bool free_gpu_weights(); + + bool is_weights_on_gpu() const; + + bool refresh_host_caches_from_mmap(); + + bool ensure_weights_loaded(); + + size_t get_gpu_memory_usage_bytes() const; + int32_t sample_rate() const { return sample_rate_; } int32_t hop_length() const { return hop_length_; } int32_t num_codebooks() const { return num_codebooks_; } diff --git a/include/s2_mapped_file.h b/include/s2_mapped_file.h new file mode 100644 index 0000000..c355a6d --- /dev/null +++ b/include/s2_mapped_file.h @@ -0,0 +1,40 @@ +#pragma once + +#include +#include +#include + +namespace s2 { + +class MappedFile { +public: + MappedFile() = default; + ~MappedFile() { close(); } + + MappedFile(const MappedFile&) = delete; + MappedFile& operator=(const MappedFile&) = delete; + + MappedFile(MappedFile&& other) noexcept; + MappedFile& operator=(MappedFile&& other) noexcept; + + bool open(const std::string& path); + void close(); + void drop_page_cache(); + + bool is_open() const { return data_ != nullptr; } + const uint8_t* data() const { return static_cast(data_); } + size_t size() const { return size_; } + +private: + void* data_ = nullptr; + size_t size_ = 0; + +#ifdef _WIN32 + void* file_handle_ = nullptr; + void* mapping_handle_ = nullptr; +#else + int fd_ = -1; +#endif +}; + +} diff --git a/include/s2_model.h b/include/s2_model.h index eb36a6e..123f73d 100644 --- a/include/s2_model.h +++ b/include/s2_model.h @@ -1,5 +1,6 @@ #pragma once +#include "s2_mapped_file.h" #include "s2_backend.h" #include "ggml.h" #include "ggml-alloc.h" @@ -19,6 +20,7 @@ #include #include #include +#include #include namespace s2 { @@ -98,8 +100,6 @@ class SlowARModel { bool load_shared(gguf_context * gguf_ctx, const std::string & gguf_path, int32_t gpu_device = -1, BackendType backend_type = BackendType::CPU, int32_t n_gpu_layers = -1); - bool read_tensor_data(const std::string & gguf_path, gguf_context * gguf_ctx); - ggml_context * weights_ctx() { return weights_.ctx_w; } const std::unordered_set & weight_tensor_set() const { return weight_tensor_set_; } @@ -109,6 +109,22 @@ class SlowARModel { void clear_kv_cache(); + MappedFile& mapped_file() { return mapped_gguf_; } + + bool allocate_and_load_weights(); + + bool restore_weights_to_gpu(); + + bool free_gpu_weights(); + + void free_compute_buffers(); + + void acquire_compute_resources(); + + size_t get_gpu_memory_usage_bytes() const; + + bool is_weights_on_gpu() const { return weights_on_gpu_; } + private: bool eval_cached(const std::vector & flat_tokens, int32_t n_tokens, int32_t n_threads, @@ -152,6 +168,16 @@ class SlowARModel { std::unordered_set weight_tensor_set_; + std::string gguf_path_; + size_t gguf_data_offset_ = 0; + std::unordered_map tensor_offsets_; + std::vector original_gpu_weights_; + std::vector original_cpu_weights_; + bool weights_on_gpu_ = false; + bool weights_allocated_ = false; + + MappedFile mapped_gguf_; + }; } diff --git a/src/s2_codec.cpp b/src/s2_codec.cpp index e4dc3bc..fae1ced 100755 --- a/src/s2_codec.cpp +++ b/src/s2_codec.cpp @@ -20,6 +20,7 @@ #include #include #include +#include namespace s2 { @@ -61,9 +62,10 @@ struct codec_decode_cache { }; struct AudioCodec::Impl { - ggml_backend_t backend = nullptr; - ggml_context * ctx_w = nullptr; - ggml_backend_buffer_t model_buf = nullptr; + ggml_backend_t backend = nullptr; + ggml_backend_t backend_cpu = nullptr; + ggml_context * ctx_w = nullptr; + ggml_backend_buffer_t model_buf = nullptr; std::string tprefix; int32_t sample_rate = 0; @@ -103,6 +105,15 @@ struct AudioCodec::Impl { std::vector residual_vq; codec_decode_cache decode_cache; + std::string gguf_path; + size_t gguf_data_offset = 0; + std::unordered_map tensor_offsets; + std::vector original_gpu_weights; + std::vector original_cpu_weights; + std::vector all_codec_weights; + bool weights_on_gpu = false; + MappedFile mapped_gguf_; + bool weights_allocated_ = false; }; static const char * backend_type_name(BackendType backend_type) { @@ -115,6 +126,53 @@ static const char * backend_type_name(BackendType backend_type) { return "Unknown"; } +static bool allocate_codec_buffers(ggml_backend_t backend, + const std::vector & 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(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; +} + static void reset_decode_cache(codec_decode_cache & cache, bool preserve_failed_n_frames = true) { if (cache.allocr) { ggml_gallocr_free(cache.allocr); @@ -873,103 +931,76 @@ bool AudioCodec::load_shared(SlowARModel* Model, gguf_context * shared_gguf_ctx, streaming_history_frames_ = 160; } - impl_->model_buf = ggml_backend_alloc_ctx_tensors(impl_->ctx_w, impl_->backend); - if (!impl_->model_buf) throw std::runtime_error("ggml_backend_alloc_ctx_tensors() failed"); + if (!impl_->backend_cpu) { + impl_->backend_cpu = ggml_backend_cpu_init(); + } - impl_->semantic_vq = vq_cache(); - impl_->residual_vq.clear(); - } catch (const std::exception & e) { - std::cerr << "[Codec] " << e.what() << std::endl; - reset_codec_impl(*impl_); - return false; - } - S2_LOG_INFO_STREAM("[Codec] Backend: " << backend_name() << std::endl); - return true; -} + impl_->gguf_path = gguf_path; + impl_->gguf_data_offset = gguf_get_data_offset(shared_gguf_ctx); -bool AudioCodec::refresh_host_caches() { - if (!impl_ || !impl_->ctx_w) { - return false; - } + const int64_t n_tensors = gguf_get_n_tensors(shared_gguf_ctx); + const auto & model_weights = Model ? Model->weight_tensor_set() : std::unordered_set(); - try { - impl_->semantic_vq = load_vq_cache(impl_->ctx_w, - impl_->tprefix + "quantizer.semantic_quantizer.quantizers.0", - impl_->quantizer_input_dim, impl_->quantizer_codebook_dim, - impl_->quantizer_semantic_codebook_size); + impl_->all_codec_weights.clear(); + impl_->tensor_offsets.clear(); + impl_->original_gpu_weights.clear(); + impl_->original_cpu_weights.clear(); - 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_cache(impl_->ctx_w, - impl_->tprefix + "quantizer.quantizer.quantizers." + std::to_string(i), - impl_->quantizer_input_dim, impl_->quantizer_codebook_dim, - impl_->quantizer_residual_codebook_size)); + for (int64_t ti = 0; ti < n_tensors; ++ti) { + const char * tname = gguf_get_tensor_name(shared_gguf_ctx, ti); + ggml_tensor * t = ggml_get_tensor(impl_->ctx_w, tname); + if (!t) continue; + + // Skip tensors that belong to the SlowAR Model + if (Model && model_weights.find(t) != model_weights.end()) continue; + + impl_->all_codec_weights.push_back(t); + impl_->tensor_offsets[t] = gguf_get_tensor_offset(shared_gguf_ctx, ti); + + if (!ggml_backend_is_cpu(impl_->backend)) { + impl_->original_gpu_weights.push_back(t); + } else { + impl_->original_cpu_weights.push_back(t); + } } - } catch (const std::exception & e) { - std::cerr << "[Codec] Failed to refresh VQ caches: " << e.what() << std::endl; - impl_->semantic_vq = vq_cache(); - impl_->residual_vq.clear(); - return false; - } - return true; -} + impl_->weights_on_gpu = false; -bool AudioCodec::read_tensor_data(const std::string & gguf_path, gguf_context * gguf_ctx) { - if (!impl_ || !impl_->ctx_w) return false; + impl_->mapped_gguf_.open(gguf_path); + if (!impl_->mapped_gguf_.is_open()) { + throw std::runtime_error("Failed to mmap " + gguf_path); + } - const size_t data_offset = gguf_get_data_offset(gguf_ctx); - const int64_t n_tensors = gguf_get_n_tensors(gguf_ctx); + impl_->weights_allocated_ = false; - std::FILE * f = std::fopen(gguf_path.c_str(), "rb"); - if (!f) { - std::cerr << "[Codec] Failed to reopen " << gguf_path << " for data loading." << std::endl; + impl_->semantic_vq = vq_cache(); + impl_->residual_vq.clear(); + } catch (const std::exception & e) { + std::cerr << "[Codec] " << e.what() << std::endl; + reset_codec_impl(*impl_); return false; } - for (int64_t ti = 0; ti < n_tensors; ++ti) { - const char * name = gguf_get_tensor_name(gguf_ctx, ti); - ggml_tensor * t = ggml_get_tensor(impl_->ctx_w, name); - if (!t) continue; - const size_t off = data_offset + gguf_get_tensor_offset(gguf_ctx, ti); - const size_t nbytes = ggml_nbytes(t); - std::vector tmp(nbytes); -#ifdef _WIN32 - _fseeki64(f, (int64_t)off, SEEK_SET); -#else - fseeko(f, (off_t)off, SEEK_SET); -#endif - if (std::fread(tmp.data(), 1, nbytes, f) != nbytes) { - std::fclose(f); - std::cerr << "[Codec] Failed to read tensor: " << name << std::endl; - return false; - } - ggml_backend_tensor_set(t, tmp.data(), 0, nbytes); - } - std::fclose(f); - return refresh_host_caches(); + S2_LOG_INFO_STREAM("[Codec] Backend: " << backend_name() << std::endl); + return true; } bool AudioCodec::load(const std::string & gguf_path, int32_t gpu_device, BackendType backend_type) { - struct gguf_init_params params = { true, nullptr }; gguf_context * ctx_gguf = gguf_init_from_file(gguf_path.c_str(), params); if (!ctx_gguf) { std::cerr << "[Codec] Failed to open " << gguf_path << std::endl; return false; } - if (!load_shared(nullptr, ctx_gguf, gguf_path, gpu_device, backend_type)) { gguf_free(ctx_gguf); return false; } - - if (!read_tensor_data(gguf_path, ctx_gguf)) { - gguf_free(ctx_gguf); + gguf_free(ctx_gguf); + + if (!refresh_host_caches_from_mmap()) { + std::cerr << "[Codec] Failed to refresh VQ caches from mmap." << std::endl; return false; } - - gguf_free(ctx_gguf); return true; } @@ -980,6 +1011,8 @@ ggml_context * AudioCodec::weights_ctx() const { bool AudioCodec::encode(const float * audio, int32_t n_samples, int32_t n_threads, std::vector & codes_out, int32_t & n_frames_out) { + if (!ensure_weights_loaded()) return false; + const int32_t frame_length = (impl_->frame_length > 0) ? impl_->frame_length : 512; const int32_t padded = ((n_samples + frame_length - 1) / frame_length) * frame_length; std::vector audio_padded(padded, 0.0f); @@ -1298,6 +1331,7 @@ bool AudioCodec::decode(const int32_t * codes, int32_t n_frames, int32_t n_threa std::vector & audio_out) { if (n_frames <= 0) return false; if (!impl_ || !impl_->backend) return false; + if (!ensure_weights_loaded()) return false; if (!ggml_backend_is_cpu(impl_->backend) && run_cached_decode_graph(*impl_, codes, n_frames, n_threads, audio_out)) { @@ -1444,4 +1478,112 @@ bool AudioCodec::decode(const int32_t * codes, int32_t n_frames, int32_t n_threa return true; } +bool AudioCodec::is_weights_on_gpu() const { + return impl_ ? impl_->weights_on_gpu : false; +} + +bool AudioCodec::free_gpu_weights() { + if (!impl_ || !impl_->weights_on_gpu) return true; + + S2_LOG_INFO_STREAM("[Codec] >>> FREEING Audio Codec GPU weights..." << std::endl); + const auto t0 = std::chrono::steady_clock::now(); + + ggml_backend_synchronize(impl_->backend); + + if (impl_->model_buf) { + ggml_backend_buffer_free(impl_->model_buf); + impl_->model_buf = nullptr; + } + + for (ggml_tensor * t : impl_->all_codec_weights) { + if (t) { t->data = nullptr; t->buffer = nullptr; } + } + + impl_->weights_allocated_ = false; + impl_->weights_on_gpu = false; + + const auto t1 = std::chrono::steady_clock::now(); + const double free_ms = std::chrono::duration(t1 - t0).count(); + S2_LOG_INFO_STREAM("[Codec] <<< Audio Codec GPU weights FREED in " << free_ms << " ms" << std::endl); + return true; +} + +bool AudioCodec::restore_weights_to_gpu() { + if (!impl_ || impl_->weights_on_gpu || impl_->original_gpu_weights.empty()) return true; + S2_LOG_INFO_STREAM("[Codec] >>> RESTORING Audio Codec weights from mmap to GPU..." << std::endl); + const auto t0 = std::chrono::steady_clock::now(); + impl_->weights_allocated_ = false; + if (!ensure_weights_loaded()) return false; + const auto t1 = std::chrono::steady_clock::now(); + const double restore_ms = std::chrono::duration(t1 - t0).count(); + S2_LOG_INFO_STREAM("[Codec] <<< Audio Codec weights RESTORED in " << restore_ms << " ms" << std::endl); + return true; +} + +size_t AudioCodec::get_gpu_memory_usage_bytes() const { + if (!impl_ || !impl_->model_buf || !impl_->weights_on_gpu) return 0; + return ggml_backend_buffer_get_size(impl_->model_buf); +} + +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 { + 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 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(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::ensure_weights_loaded() { + if (!impl_ || impl_->weights_allocated_) 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); + return true; +} + +MappedFile& AudioCodec::mapped_file() { return impl_->mapped_gguf_; } + } diff --git a/src/s2_mapped_file.cpp b/src/s2_mapped_file.cpp new file mode 100644 index 0000000..e0c3021 --- /dev/null +++ b/src/s2_mapped_file.cpp @@ -0,0 +1,173 @@ +#include "../include/s2_mapped_file.h" + +#ifdef _WIN32 +#include +#else +#include +#include +#include +#include +#include +#endif + +namespace s2 { + +bool MappedFile::open(const std::string& path) { + close(); + +#ifdef _WIN32 + HANDLE fh = CreateFileA(path.c_str(), GENERIC_READ, FILE_SHARE_READ, + nullptr, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, nullptr); + if (fh == INVALID_HANDLE_VALUE) { + return false; + } + + LARGE_INTEGER file_size; + if (!GetFileSizeEx(fh, &file_size)) { + CloseHandle(fh); + return false; + } + size_ = static_cast(file_size.QuadPart); + + if (size_ == 0) { + CloseHandle(fh); + return true; + } + + HANDLE mh = CreateFileMappingA(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; + } + + data_ = addr; + file_handle_ = static_cast(fh); + mapping_handle_ = static_cast(mh); + +#else + int fd = ::open(path.c_str(), O_RDONLY); + if (fd < 0) { + return false; + } + + struct stat st; + if (fstat(fd, &st) < 0) { + ::close(fd); + return false; + } + size_ = static_cast(st.st_size); + + if (size_ == 0) { + ::close(fd); + return true; + } + + void* addr = ::mmap(nullptr, size_, PROT_READ, MAP_PRIVATE, fd, 0); + if (addr == MAP_FAILED) { + ::close(fd); + return false; + } + + ::madvise(addr, size_, MADV_RANDOM); + + data_ = addr; + fd_ = fd; +#endif + + return true; +} + +void MappedFile::close() { +#ifdef _WIN32 + if (data_) { + UnmapViewOfFile(data_); + data_ = nullptr; + } + if (mapping_handle_) { + CloseHandle(static_cast(mapping_handle_)); + mapping_handle_ = nullptr; + } + if (file_handle_) { + CloseHandle(static_cast(file_handle_)); + file_handle_ = nullptr; + } +#else + if (data_ && data_ != MAP_FAILED) { + ::munmap(data_, size_); + data_ = nullptr; + } + if (fd_ >= 0) { + ::close(fd_); + fd_ = -1; + } +#endif + size_ = 0; +} + +MappedFile::MappedFile(MappedFile&& other) noexcept { +#ifdef _WIN32 + data_ = other.data_; + other.data_ = nullptr; + size_ = other.size_; + other.size_ = 0; + file_handle_ = other.file_handle_; + other.file_handle_ = nullptr; + mapping_handle_ = other.mapping_handle_; + other.mapping_handle_ = nullptr; +#else + data_ = other.data_; + other.data_ = nullptr; + size_ = other.size_; + other.size_ = 0; + fd_ = other.fd_; + other.fd_ = -1; +#endif +} + +MappedFile& MappedFile::operator=(MappedFile&& other) noexcept { + if (this != &other) { + close(); +#ifdef _WIN32 + data_ = other.data_; + other.data_ = nullptr; + size_ = other.size_; + other.size_ = 0; + file_handle_ = other.file_handle_; + other.file_handle_ = nullptr; + mapping_handle_ = other.mapping_handle_; + other.mapping_handle_ = nullptr; +#else + data_ = other.data_; + other.data_ = nullptr; + size_ = other.size_; + other.size_ = 0; + fd_ = other.fd_; + other.fd_ = -1; +#endif + } + return *this; +} + +void MappedFile::drop_page_cache() { +#if defined(__linux__) || defined(__APPLE__) + if (data_ && data_ != MAP_FAILED && size_ > 0) { + ::madvise(data_, size_, MADV_DONTNEED); + } + if (fd_ >= 0) { + ::posix_fadvise(fd_, 0, 0, POSIX_FADV_DONTNEED); + } +#elif defined(_WIN32) + if (data_ && size_ > 0) { + ::VirtualUnlock(data_, size_); + } +#endif +} + +} diff --git a/src/s2_model.cpp b/src/s2_model.cpp index 9ced623..992b398 100755 --- a/src/s2_model.cpp +++ b/src/s2_model.cpp @@ -1,5 +1,6 @@ #include "../include/s2_model.h" #include "../include/s2_log.h" +#include "../include/s2_mapped_file.h" #include "s2_ggml_utils.h" #include #include @@ -217,6 +218,8 @@ bool SlowARModel::load_shared(gguf_context * ctx_gguf, const std::string & gguf_ } gguf_free(local_gguf); + gguf_path_ = gguf_path; + gguf_data_offset_ = gguf_get_data_offset(ctx_gguf); S2_LOG_INFO_STREAM("[Model] Reading metadata from " << gguf_path << std::endl); @@ -478,6 +481,15 @@ bool SlowARModel::load_shared(gguf_context * ctx_gguf, const std::string & gguf_ weight_tensor_set_.insert(weight_tensors.begin(), weight_tensors.end()); + const int64_t n_tensors_gguf = gguf_get_n_tensors(ctx_gguf); + for (int64_t ti = 0; ti < n_tensors_gguf; ++ti) { + const char * tname = gguf_get_tensor_name(ctx_gguf, ti); + ggml_tensor * t = ggml_get_tensor(weights_.ctx_w, tname); + if (t && weight_tensor_set_.find(t) != weight_tensor_set_.end()) { + tensor_offsets_[t] = gguf_get_tensor_offset(ctx_gguf, ti); + } + } + const bool full_model_offload = backend_gpu_ != nullptr && backend_type != BackendType::CUDA && @@ -521,196 +533,39 @@ bool SlowARModel::load_shared(gguf_context * ctx_gguf, const std::string & gguf_ } } - std::vector gpu_buffers; - std::vector cpu_buffers; - - if (!gpu_weight_tensors.empty() && backend_gpu_) { - size_t gpu_bytes = 0; - size_t gpu_max_buffer = 0; - std::string gpu_alloc_error; - if (!allocate_weight_buffers( - backend_gpu_, - gpu_weight_tensors, - gpu_buffers, - gpu_bytes, - gpu_max_buffer, - gpu_alloc_error)) { - const double requested_mb = gpu_bytes / (1024.0 * 1024.0); - const double max_chunk_mb = - gpu_max_buffer == static_cast(-1) - ? 0.0 - : gpu_max_buffer / (1024.0 * 1024.0); - std::cerr << "[Model] GPU weight buffer allocation failed." << std::endl; - std::cerr << "[Model] Requested GPU memory: " << requested_mb << " MB for " - << gpu_weight_tensors.size() << " tensors (" << n_gpu_layers_ - << " layers)" << std::endl; - if (gpu_max_buffer != static_cast(-1)) { - std::cerr << "[Model] Backend max buffer size: " << max_chunk_mb - << " MB per allocation." << std::endl; - } - if (!gpu_alloc_error.empty()) { - std::cerr << "[Model] " << gpu_alloc_error << std::endl; - } - std::cerr << "[Model] Suggest using a lower --gpu-layers value (e.g., --gpu-layers " - << std::max(1, n_gpu_layers_ / 2) << ")" << std::endl; - return false; - } - } - - if (!cpu_weight_tensors.empty()) { - size_t cpu_bytes = 0; - size_t cpu_max_buffer = 0; - std::string cpu_alloc_error; - if (!allocate_weight_buffers( - backend_cpu_, - cpu_weight_tensors, - cpu_buffers, - cpu_bytes, - cpu_max_buffer, - cpu_alloc_error)) { - free_backend_buffers(gpu_buffers); - std::cerr << "[Model] Failed to allocate CPU weight buffer." << std::endl; - if (!cpu_alloc_error.empty()) { - std::cerr << "[Model] " << cpu_alloc_error << std::endl; - } - return false; - } - } - - weights_.model_bufs_gpu = std::move(gpu_buffers); - weights_.model_bufs_cpu = std::move(cpu_buffers); - - { - ggml_backend_t backends[2]; - int n_backends; - if (backend_gpu_) { - backends[0] = backend_gpu_; - backends[1] = backend_cpu_; - n_backends = 2; - } else { - backends[0] = backend_cpu_; - n_backends = 1; - } - - sched_ = ggml_backend_sched_new(backends, NULL, n_backends, 32768, false, true); - if (!sched_) { - std::cerr << "[Model] Failed to create Slow-AR scheduler." << std::endl; - return false; - } - - if (hparams_.has_fast_decoder) { - fast_sched_ = ggml_backend_sched_new(backends, NULL, n_backends, 16384, false, true); - if (!fast_sched_) { - std::cerr << "[Model] Failed to create Fast-AR scheduler." << std::endl; - return false; - } - } - } - - if (n_gpu_layers_ > 0 && backend_gpu_) { - const int32_t first_cpu_layer = n_gpu_layers_; - const int32_t last_gpu_layer = n_gpu_layers_ - 1; - if (first_cpu_layer < hparams_.block_count) { - S2_LOG_INFO_STREAM("[Model] Layers 0-" << last_gpu_layer << " on " - << ggml_backend_name(backend_gpu_) - << ", " << first_cpu_layer << "-" << (hparams_.block_count - 1) << " on CPU" - << std::endl); - } else { - S2_LOG_INFO_STREAM("[Model] Layers 0-" << (hparams_.block_count - 1) << " on " - << ggml_backend_name(backend_gpu_) << " (all)" << std::endl); - } - } else { - S2_LOG_INFO_STREAM("[Model] All " << hparams_.block_count << " layers on CPU" << std::endl); - } - - const size_t gpu_weight_bytes = total_backend_buffer_bytes(weights_.model_bufs_gpu); - if (!weights_.model_bufs_gpu.empty()) { - S2_LOG_INFO_STREAM("[Model] GPU weight buffers: " << weights_.model_bufs_gpu.size() - << " chunk(s), " << (gpu_weight_bytes / 1024.0 / 1024.0) - << " MB total" << std::endl); - } - const size_t cpu_weight_bytes = total_backend_buffer_bytes(weights_.model_bufs_cpu); - if (!weights_.model_bufs_cpu.empty()) { - S2_LOG_INFO_STREAM("[Model] CPU weight buffers: " << weights_.model_bufs_cpu.size() - << " chunk(s), " << (cpu_weight_bytes / 1024.0 / 1024.0) - << " MB total" << std::endl); - } - const size_t total_bytes = gpu_weight_bytes + cpu_weight_bytes; - S2_LOG_INFO_STREAM("[Model] Total model size: " - << (total_bytes / 1024.0 / 1024.0) << " MB" << std::endl); - - S2_LOG_INFO_STREAM("[Model] KV cache: " << (n_gpu_layers_ > 0 && backend_gpu_ ? "GPU" : "CPU") - << ", n_gpu_layers=" << n_gpu_layers_ << std::endl); - - if (backend_type == BackendType::CUDA && - backend_gpu_ && - ggml_is_quantized(weights_.embeddings->type)) { - S2_LOG_INFO_STREAM("[Model] Keeping quantized embedding tables on CPU for CUDA stability." - << std::endl); - } - - return true; -} - -bool SlowARModel::read_tensor_data(const std::string & gguf_path, gguf_context * ctx_gguf) { - const size_t data_offset = gguf_get_data_offset(ctx_gguf); - const int64_t n_tensors = gguf_get_n_tensors(ctx_gguf); + original_gpu_weights_ = gpu_weight_tensors; + original_cpu_weights_ = cpu_weight_tensors; - std::FILE * f = std::fopen(gguf_path.c_str(), "rb"); - if (!f) { - std::cerr << "[Model] Cannot reopen " << gguf_path << " for data loading." << std::endl; + mapped_gguf_.open(gguf_path); + if (!mapped_gguf_.is_open()) { + std::cerr << "[Model] Failed to mmap " << gguf_path << std::endl; return false; } - std::vector tmp; - for (int64_t ti = 0; ti < n_tensors; ++ti) { - const char * tname = gguf_get_tensor_name(ctx_gguf, ti); - ggml_tensor * t = ggml_get_tensor(weights_.ctx_w, tname); - if (!t || weight_tensor_set_.find(t) == weight_tensor_set_.end()) continue; + weights_allocated_ = false; + weights_on_gpu_ = false; + - const size_t toff = data_offset + gguf_get_tensor_offset(ctx_gguf, ti); - const size_t tsize = ggml_nbytes(t); - if (tmp.size() < tsize) tmp.resize(tsize); -#ifdef _WIN32 - _fseeki64(f, (int64_t)toff, SEEK_SET); -#else - fseeko(f, (off_t)toff, SEEK_SET); -#endif - if (std::fread(tmp.data(), 1, tsize, f) != tsize) { - std::cerr << "[Model] Failed to read tensor: " << tname << std::endl; - std::fclose(f); - return false; - } - ggml_backend_tensor_set(t, tmp.data(), 0, tsize); - } - tmp.clear(); - tmp.shrink_to_fit(); - std::fclose(f); - - S2_LOG_INFO_STREAM("[Model] Weights loaded. Total tensors: " << n_tensors << std::endl); return true; } bool SlowARModel::load(const std::string & gguf_path, int32_t gpu_device, BackendType backend_type, int32_t n_gpu_layers) { - struct gguf_init_params params = { true, nullptr }; gguf_context * ctx_gguf = gguf_init_from_file(gguf_path.c_str(), params); if (!ctx_gguf) { std::cerr << "[Model] Failed to load GGUF from " << gguf_path << std::endl; return false; } - if (!load_shared(ctx_gguf, gguf_path, gpu_device, backend_type, n_gpu_layers)) { gguf_free(ctx_gguf); return false; } - - if (!read_tensor_data(gguf_path, ctx_gguf)) { - gguf_free(ctx_gguf); + gguf_free(ctx_gguf); + + if (!allocate_and_load_weights()) { + std::cerr << "[Model] Failed to allocate and load weights from mmap." << std::endl; return false; } - - gguf_free(ctx_gguf); return true; } @@ -1254,4 +1109,120 @@ bool SlowARModel::fast_decode(const std::vector & hidden_in, return true; } +bool SlowARModel::restore_weights_to_gpu() { + if (!backend_gpu_ || weights_on_gpu_) return true; + S2_LOG_INFO_STREAM("[Model] >>> RESTORING Slow-AR weights from mmap to GPU..." << std::endl); + const auto t0 = std::chrono::steady_clock::now(); + weights_allocated_ = false; + if (!allocate_and_load_weights()) return false; + const auto t1 = std::chrono::steady_clock::now(); + const double restore_ms = std::chrono::duration(t1 - t0).count(); + S2_LOG_INFO_STREAM("[Model] <<< Slow-AR weights RESTORED in " << restore_ms << " ms" << std::endl); + return true; +} + +bool SlowARModel::free_gpu_weights() { + if (!backend_gpu_ || !weights_on_gpu_) return true; + + S2_LOG_INFO_STREAM("[Model] >>> FREEING Slow-AR GPU weights" << std::endl); + const auto t0 = std::chrono::steady_clock::now(); + + ggml_backend_synchronize(backend_gpu_); + + free_backend_buffers(weights_.model_bufs_gpu); + weights_.model_bufs_gpu.clear(); + + for (ggml_tensor * t : original_gpu_weights_) { + if (t) { t->data = nullptr; t->buffer = nullptr; } + } + + weights_allocated_ = false; + weights_on_gpu_ = false; + const auto t1 = std::chrono::steady_clock::now(); + const double free_ms = std::chrono::duration(t1 - t0).count(); + S2_LOG_INFO_STREAM("[Model] <<< Slow-AR GPU weights FREED in " << free_ms << " ms" << std::endl); + return true; +} + +void SlowARModel::free_compute_buffers() { + if (backend_gpu_) ggml_backend_synchronize(backend_gpu_); + clear_kv_cache(); + if (sched_) { + ggml_backend_sched_free(sched_); + sched_ = nullptr; + } + if (fast_sched_) { + ggml_backend_sched_free(fast_sched_); + fast_sched_ = nullptr; + } +} + +void SlowARModel::acquire_compute_resources() { + if (sched_) return; + + ggml_backend_t backends[2]; + int n_backends; + if (backend_gpu_) { + backends[0] = backend_gpu_; + backends[1] = backend_cpu_; + n_backends = 2; + } else { + backends[0] = backend_cpu_; + n_backends = 1; + } + + sched_ = ggml_backend_sched_new(backends, NULL, n_backends, 32768, false, true); + if (hparams_.has_fast_decoder) { + fast_sched_ = ggml_backend_sched_new(backends, NULL, n_backends, 16384, false, true); + } +} + +size_t SlowARModel::get_gpu_memory_usage_bytes() const { + size_t total = 0; + + for (const auto & buf : weights_.model_bufs_gpu) { + if (buf) total += ggml_backend_buffer_get_size(buf); + } + + if (kv_buf_) { + total += ggml_backend_buffer_get_size(kv_buf_); + } + + return total; +} + +bool SlowARModel::allocate_and_load_weights() { + if (weights_allocated_) return true; + if (!mapped_gguf_.is_open()) return false; + + size_t b, m; std::string e; + if (!original_gpu_weights_.empty() && backend_gpu_) { + if (!allocate_weight_buffers(backend_gpu_, original_gpu_weights_, weights_.model_bufs_gpu, b, m, e)) { + std::cerr << "[Model] GPU alloc failed: " << e << std::endl; return false; + } + } + if (!original_cpu_weights_.empty()) { + if (!allocate_weight_buffers(backend_cpu_, original_cpu_weights_, weights_.model_bufs_cpu, b, m, e)) { + std::cerr << "[Model] CPU alloc failed: " << e << std::endl; return false; + } + } + + const uint8_t* base = mapped_gguf_.data(); + for (ggml_tensor * t : original_gpu_weights_) { + auto it = tensor_offsets_.find(t); + if (it != tensor_offsets_.end()) + ggml_backend_tensor_set(t, base + gguf_data_offset_ + it->second, 0, ggml_nbytes(t)); + } + for (ggml_tensor * t : original_cpu_weights_) { + auto it = tensor_offsets_.find(t); + if (it != tensor_offsets_.end()) + ggml_backend_tensor_set(t, base + gguf_data_offset_ + it->second, 0, ggml_nbytes(t)); + } + + weights_allocated_ = true; + weights_on_gpu_ = !original_gpu_weights_.empty(); + acquire_compute_resources(); + return true; +} + } diff --git a/src/s2_pipeline.cpp b/src/s2_pipeline.cpp index 0dd5f01..be22a63 100644 --- a/src/s2_pipeline.cpp +++ b/src/s2_pipeline.cpp @@ -303,91 +303,6 @@ static void sync_tokenizer_config_from_model(Tokenizer& tokenizer, const SlowARM Pipeline::Pipeline() {} Pipeline::~Pipeline() {} -static bool read_all_tensor_data( - const std::string & gguf_path, - gguf_context * gguf_ctx, - s2::SlowARModel & model, - s2::AudioCodec & codec) -{ - const size_t data_offset = gguf_get_data_offset(gguf_ctx); - const int64_t n_tensors = gguf_get_n_tensors(gguf_ctx); - - std::FILE * f = std::fopen(gguf_path.c_str(), "rb"); - if (!f) { - std::cerr << "[Pipeline] Cannot reopen " << gguf_path << " for data loading." << std::endl; - return false; - } - - const auto & model_weights = model.weight_tensor_set(); - ggml_context * codec_ctx = codec.weights_ctx(); - std::vector tmp; - - for (int64_t ti = 0; ti < n_tensors; ++ti) { - const char * tname = gguf_get_tensor_name(gguf_ctx, ti); - const size_t toff = data_offset + gguf_get_tensor_offset(gguf_ctx, ti); - - ggml_tensor * t = ggml_get_tensor(model.weights_ctx(), tname); - if (t && model_weights.find(t) != model_weights.end()) { - const size_t tsize = ggml_nbytes(t); - if (tmp.size() < tsize) tmp.resize(tsize); -#ifdef _WIN32 - _fseeki64(f, (int64_t)toff, SEEK_SET); -#else - fseeko(f, (off_t)toff, SEEK_SET); -#endif - if (std::fread(tmp.data(), 1, tsize, f) != tsize) { - std::cerr << "[Pipeline] Failed to read tensor: " << tname << std::endl; - std::fclose(f); - return false; - } - ggml_backend_tensor_set(t, tmp.data(), 0, tsize); - continue; - } - - if (codec_ctx) { - t = ggml_get_tensor(codec_ctx, tname); - if (t) { - const size_t tsize = ggml_nbytes(t); - if (tmp.size() < tsize) tmp.resize(tsize); -#ifdef _WIN32 - _fseeki64(f, (int64_t)toff, SEEK_SET); -#else - fseeko(f, (off_t)toff, SEEK_SET); -#endif - if (std::fread(tmp.data(), 1, tsize, f) != tsize) { - std::cerr << "[Pipeline] Failed to read tensor: " << tname << std::endl; - std::fclose(f); - return false; - } - ggml_backend_tensor_set(t, tmp.data(), 0, tsize); - continue; - } - } - - } - tmp.clear(); - tmp.shrink_to_fit(); - std::fclose(f); - - if (!codec.refresh_host_caches()) { - std::cerr << "[Pipeline] Failed to refresh codec host caches after weight load." << std::endl; - return false; - } - -#ifdef __linux__ - { - int fd = ::open(gguf_path.c_str(), O_RDONLY); - if (fd >= 0) { - ::posix_fadvise(fd, 0, 0, POSIX_FADV_DONTNEED); - ::close(fd); - } - } -#endif - - S2_LOG_INFO_STREAM("[Model] Weights loaded. Total tensors: " << n_tensors << std::endl); - return true; -} - bool Pipeline::init(const PipelineParams & params) { tokenizer_ref_ = &owned_tokenizer_; model_ref_ = &owned_model_; @@ -477,45 +392,51 @@ bool Pipeline::init(const PipelineParams & params) { safe_print_ln("Pipeline: loading codec on " + backend_label + " device " + std::to_string(codec_gpu_device) + "..."); } + codec_loaded = codec().load_shared(&model(), shared_gguf, params.model_path, codec_gpu_device, codec_backend_type); + if (!codec_loaded) { if (codec_backend_type == BackendType::Metal) { - safe_print_warn_ln( - "Pipeline warning: codec " + backend_label + - " load failed, falling back to CPU."); + safe_print_warn_ln("Pipeline warning: codec " + backend_label + " load failed, falling back to CPU."); } else { - safe_print_warn_ln( - "Pipeline warning: codec " + backend_label + - " load failed on device " + std::to_string(codec_gpu_device) + - ", falling back to CPU."); + safe_print_warn_ln("Pipeline warning: codec " + backend_label + " load failed on device " + + std::to_string(codec_gpu_device) + ", falling back to CPU."); } } } + if (!codec_loaded) { if (!use_gpu_codec) { safe_print_ln("Pipeline: loading codec on CPU."); } codec_loaded = codec().load_shared(&model(), shared_gguf, params.model_path, -1, BackendType::CPU); } + if (!codec_loaded) { safe_print_error_ln("Pipeline error: could not load codec from " + params.model_path); gguf_free(shared_gguf); return false; } - if (!read_all_tensor_data(params.model_path, shared_gguf, model(), codec())) { - safe_print_error_ln("Pipeline error: failed to read tensor data from " + params.model_path); - gguf_free(shared_gguf); - return false; - } - gguf_free(shared_gguf); + if (!codec().refresh_host_caches_from_mmap()) { + safe_print_error_ln("Pipeline error: failed to refresh VQ caches from mmap"); + return false; + } const auto codec_t1 = std::chrono::steady_clock::now(); + const auto model_weights_t0 = std::chrono::steady_clock::now(); + if (!model().allocate_and_load_weights()) { + safe_print_error_ln("Pipeline error: failed to allocate and load Slow-AR weights"); + return false; + } + const auto model_weights_t1 = std::chrono::steady_clock::now(); + sync_tokenizer_config_from_model(tokenizer(), model()); initialized_ = true; + const auto init_t1 = std::chrono::steady_clock::now(); safe_print_ln( "[Metrics] Init: tokenizer=" + @@ -524,7 +445,9 @@ bool Pipeline::init(const PipelineParams & params) { std::to_string(std::chrono::duration(model_t1 - model_t0).count()) + " ms, codec=" + std::to_string(std::chrono::duration(codec_t1 - codec_t0).count()) + - " ms (" + codec().backend_name() + "), total=" + + " ms (" + codec().backend_name() + "), model_weights=" + + std::to_string(std::chrono::duration(model_weights_t1 - model_weights_t0).count()) + + " ms, total=" + std::to_string(std::chrono::duration(init_t1 - init_t0).count()) + " ms, max_rss=" + std::to_string(get_max_rss_mb()) + " MB"); From c6680adcd14bf5b49c1ee6574cbb8b95603f3784 Mon Sep 17 00:00:00 2001 From: Skyrion9 Date: Sun, 19 Jul 2026 23:40:41 +0300 Subject: [PATCH 2/5] feat: phase-gated VRAM swap and hot-swap state machine implementation 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. --- include/s2_pipeline.h | 7 ++ src/main.cpp | 25 +++--- src/s2_pipeline.cpp | 195 +++++++++++++++++++++++++++++++++++++++++- 3 files changed, 213 insertions(+), 14 deletions(-) diff --git a/include/s2_pipeline.h b/include/s2_pipeline.h index 19b3c3d..de276e3 100644 --- a/include/s2_pipeline.h +++ b/include/s2_pipeline.h @@ -11,6 +11,7 @@ #include #include #include +#include namespace s2 { @@ -46,6 +47,9 @@ struct PipelineParams { std::string voice_id; bool save_voice = false; std::string voice_storage_dir = "./voices"; + bool enable_vram_swap = true; + bool enable_hot_swap = false; + bool is_persistent = false; }; class Pipeline { @@ -108,8 +112,11 @@ class Pipeline { Tokenizer* tokenizer_ref_ = &owned_tokenizer_; SlowARModel* model_ref_ = &owned_model_; AudioCodec* codec_ref_ = &owned_codec_; + std::thread pending_offload_thread_; mutable std::mutex synthesize_mutex_; bool initialized_ = false; + bool model_prefers_gpu_ = false; + bool codec_prefers_gpu_ = false; }; } diff --git a/src/main.cpp b/src/main.cpp index d3e0b49..eff3531 100644 --- a/src/main.cpp +++ b/src/main.cpp @@ -90,6 +90,8 @@ void print_uso() { safe_print(" --codec-auto Benchmark codec backends and keep the fastest (default)\n"); safe_print(" --codec-follow-backend Force codec to follow the selected GPU backend\n"); safe_print(" --codec-cpu Force codec on CPU even when model uses GPU\n"); + safe_print(" --no-vram-swap Disable phase-gated VRAM swapping between Slow-AR and Codec (keeps both in VRAM simultaneously instead)\n"); + safe_print(" --hot-swap Aggressively evict weights and OS page cache after each server request (~100 MB RAM, ~25 MB VRAM idle)\n"); safe_print(" --stream-file Write output WAV through the streaming path\n"); safe_print(" --stream-decode-stride Decode cadence in frames (0 = auto: server 4, file/offline 16)\n"); safe_print(" --codec-context-frames Override codec decode history (lower uses less VRAM, default: auto)\n"); @@ -186,16 +188,18 @@ int main(int argc, char** argv) { else if (arg == "-temp" || arg == "--temp" || arg == "--temperature") { if (i+1 < argc) { try { params.gen.temperature = std::stof(argv[++i]); } catch(...) {} } } else if (arg == "-top-p" || arg == "--top-p") { if (i+1 < argc) { try { params.gen.top_p = std::stof(argv[++i]); } catch(...) {} } } else if (arg == "-top-k" || arg == "--top-k") { if (i+1 < argc) { try { params.gen.top_k = std::stoi(argv[++i]); } catch(...) {} } } - else if (arg == "--dynamic-normalize") { params.normalize_dynamic = true; } - else if (arg == "--no-dynamic-normalize") { params.normalize_dynamic = false; } - else if (arg == "--no-trim-silence") { params.trim_silence = false; } - else if (arg == "--trim-silence") { params.trim_silence = true; } - else if (arg == "--no-normalize") { params.normalize_output = false; } - else if (arg == "--normalize") { params.normalize_output = true; } - else if (arg == "--codec-auto") { params.codec_auto_backend = true; params.codec_follow_backend = true; } - else if (arg == "--codec-follow-backend") { params.codec_auto_backend = false; params.codec_follow_backend = true; } - else if (arg == "--codec-cpu") { params.codec_auto_backend = false; params.codec_follow_backend = false; } - else if (arg == "--stream-file") { use_stream_file = true; } + else if (arg == "--dynamic-normalize") { params.normalize_dynamic = true; } + else if (arg == "--no-dynamic-normalize") { params.normalize_dynamic = false; } + else if (arg == "--no-trim-silence") { params.trim_silence = false; } + else if (arg == "--trim-silence") { params.trim_silence = true; } + else if (arg == "--no-normalize") { params.normalize_output = false; } + else if (arg == "--normalize") { params.normalize_output = true; } + else if (arg == "--codec-auto") { params.codec_auto_backend = true; params.codec_follow_backend = true; } + else if (arg == "--codec-follow-backend") { params.codec_auto_backend = false; params.codec_follow_backend = true; } + else if (arg == "--codec-cpu") { params.codec_auto_backend = false; params.codec_follow_backend = false; } + else if (arg == "--no-vram-swap") { params.enable_vram_swap = false; } + else if (arg == "--hot-swap") { params.enable_hot_swap = true; } + else if (arg == "--stream-file") { use_stream_file = true; } else if (arg == "--stream-decode-stride") { if (i+1 < argc) { try { params.stream_decode_stride_frames = std::stoi(argv[++i]); } catch(...) {} @@ -309,6 +313,7 @@ int main(int argc, char** argv) { if (use_server) { serverParams.pipeline = params; + serverParams.pipeline.is_persistent = true; s2::Server server; if (!server.serve(serverParams)) { safe_print_error("Server initialization failed.\n"); diff --git a/src/s2_pipeline.cpp b/src/s2_pipeline.cpp index be22a63..d194c63 100644 --- a/src/s2_pipeline.cpp +++ b/src/s2_pipeline.cpp @@ -301,7 +301,11 @@ static void sync_tokenizer_config_from_model(Tokenizer& tokenizer, const SlowARM } Pipeline::Pipeline() {} -Pipeline::~Pipeline() {} +Pipeline::~Pipeline() { + if (pending_offload_thread_.joinable()) { + pending_offload_thread_.join(); + } +} bool Pipeline::init(const PipelineParams & params) { tokenizer_ref_ = &owned_tokenizer_; @@ -436,6 +440,19 @@ bool Pipeline::init(const PipelineParams & params) { sync_tokenizer_config_from_model(tokenizer(), model()); initialized_ = true; + + model_prefers_gpu_ = model().is_weights_on_gpu(); + codec_prefers_gpu_ = use_gpu_codec; + + if (model_prefers_gpu_ && codec_prefers_gpu_) { + safe_print_ln("[Pipeline] VRAM State Machine: Case 1 (Both prefer GPU) - Codec is lazily allocated on demand."); + } else if (model_prefers_gpu_ && !codec_prefers_gpu_) { + safe_print_ln("[Pipeline] VRAM State Machine: Case 2 (Slow-AR GPU, Codec CPU) - Ready."); + } else if (!model_prefers_gpu_ && codec_prefers_gpu_) { + safe_print_ln("[Pipeline] VRAM State Machine: Case 3 (Slow-AR CPU, Codec GPU) - Codec is lazily allocated on demand."); + } else { + safe_print_ln("[Pipeline] VRAM State Machine: Case 4 (All CPU) - Ready."); + } const auto init_t1 = std::chrono::steady_clock::now(); safe_print_ln( @@ -712,11 +729,30 @@ bool Pipeline::synthesize_raw(const PipelineParams & params, AudioData & ref_aud return false; } + std::thread pre_restore_thread; + bool pre_restore_started = false; + + if (params.enable_vram_swap && params.is_persistent && + model_prefers_gpu_ && !model().is_weights_on_gpu()) { + safe_print_ln("[Pipeline] Hot-Swap: Pre-fetching Slow-AR to VRAM in background..."); + pre_restore_started = true; + pre_restore_thread = std::thread([this]() { + model().acquire_compute_resources(); + model().restore_weights_to_gpu(); + }); + } + if (!resolve_reference_prompt_locked(params, ref_audio, ref_codes, T_prompt, effective_prompt_text, ref_encode_ms)) { + if (pre_restore_started) pre_restore_thread.join(); return false; } + if (pre_restore_started) { + pre_restore_thread.join(); + safe_print_ln("[Pipeline] Hot-Swap: Background pre-fetch complete."); + } + PipelineParams effective_params = params; effective_params.prompt_text = std::move(effective_prompt_text); @@ -756,12 +792,43 @@ bool Pipeline::synthesize_prompt_codes_locked(const PipelineParams & params, con return false; } + if (params.enable_vram_swap) { + if (pending_offload_thread_.joinable()) { + pending_offload_thread_.join(); + } + } + CodecDecodeCacheScope codec_decode_cache_scope(codec()); model().clear_kv_cache(); safe_print_ln("--- Pipeline Synthesize ---"); safe_print_ln("Text: " + params.text); + if (params.enable_vram_swap) { + if (model_prefers_gpu_ && !model().is_weights_on_gpu()) { + safe_print_ln("[Pipeline] Restoring Slow-AR to VRAM for generation..."); + model().acquire_compute_resources(); + model().restore_weights_to_gpu(); + 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_weights_on_gpu()) { + safe_print_ln("[Pipeline] Pre-loading Audio Codec to VRAM (hiding behind CPU gen)..."); + codec().restore_weights_to_gpu(); + safe_print_ln("[VRAM Diag] Post-Codec 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_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( @@ -786,19 +853,72 @@ bool Pipeline::synthesize_prompt_codes_locked(const PipelineParams & params, con return false; } + if (params.enable_vram_swap) { + if (model_prefers_gpu_ && model().is_weights_on_gpu()) { + if (params.is_persistent) { + if (codec_prefers_gpu_) { + safe_print_ln("[Pipeline] Freeing Slow-AR from VRAM to make room for GPU Audio Codec..."); + model().free_gpu_weights(); + model().clear_kv_cache(); + safe_print_ln("[VRAM Diag] Post-SlowAR 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"); + } + } else { + safe_print_ln("[Pipeline] Single-shot: Freeing Slow-AR from VRAM..."); + model().free_gpu_weights(); + model().free_compute_buffers(); + safe_print_ln("[VRAM Diag] Post-SlowAR 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"); + } + } + + if (codec_prefers_gpu_ && !codec().is_weights_on_gpu()) { + safe_print_ln("[Pipeline] Restoring Audio Codec to VRAM for decode..."); + codec().restore_weights_to_gpu(); + safe_print_ln("[VRAM Diag] Post-Codec 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"); + } + } + const int32_t offline_decode_stride_frames = params.stream_decode_stride_frames > 0 ? params.stream_decode_stride_frames : 16; double decode_ms = 0.0; int32_t decode_batches = 0; const auto decode_t0 = std::chrono::steady_clock::now(); - if (!decode_codes_windowed(codec(), res.codes.data(), res.n_frames, num_codebooks, + + bool decode_ok = decode_codes_windowed(codec(), res.codes.data(), res.n_frames, num_codebooks, params.gen.n_threads, offline_decode_stride_frames, params.codec_decode_context_frames, - audio_out, &decode_ms, &decode_batches)) { + audio_out, &decode_ms, &decode_batches); + + const auto decode_t1 = std::chrono::steady_clock::now(); + + 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_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_weights_on_gpu()) codec().free_gpu_weights(); + } + } else { + safe_print_ln("[Pipeline] Single-shot mode: Skipping post-decode VRAM restore."); + } + } + + if (!decode_ok) { safe_print_error_ln("Pipeline error: decode failed."); return false; } - const auto decode_t1 = std::chrono::steady_clock::now(); model().clear_kv_cache(); const auto synth_t1 = std::chrono::steady_clock::now(); @@ -831,6 +951,13 @@ bool Pipeline::synthesize_prompt_codes_locked(const PipelineParams & params, con " ms/frame, gen_rtf=" + std::to_string(gen_rtf) + ", total_rtf=" + std::to_string(total_rtf) + ", max_rss=" + std::to_string(get_max_rss_mb()) + " MB"); + + if (params.enable_vram_swap) { + safe_print_ln("[VRAM Diag] Post-Phase3: 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"); + } + return true; } @@ -848,12 +975,31 @@ bool Pipeline::synthesize_streaming_raw(const PipelineParams & params, AudioData return false; } + std::thread pre_restore_thread; + bool pre_restore_started = false; + + if (params.enable_vram_swap && params.is_persistent && + model_prefers_gpu_ && !model().is_weights_on_gpu()) { + safe_print_ln("[Pipeline] Hot-Swap: Pre-fetching Slow-AR to VRAM in background..."); + pre_restore_started = true; + pre_restore_thread = std::thread([this]() { + model().acquire_compute_resources(); + model().restore_weights_to_gpu(); + }); + } + if (!resolve_reference_prompt_locked(params, ref_audio, ref_codes, T_prompt, effective_prompt_text, ref_encode_ms)) { + if (pre_restore_started) pre_restore_thread.join(); sink.on_error("Failed to resolve reference prompt"); return false; } + if (pre_restore_started) { + pre_restore_thread.join(); + safe_print_ln("[Pipeline] Hot-Swap: Background pre-fetch complete."); + } + PipelineParams effective_params = params; effective_params.prompt_text = std::move(effective_prompt_text); @@ -886,6 +1032,22 @@ bool Pipeline::synthesize_streaming_prompt_codes_locked(const PipelineParams & p return false; } + 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_ && !codec().is_weights_on_gpu()) { + safe_print_ln("[Pipeline] Streaming: Restoring Audio Codec to VRAM..."); + codec().restore_weights_to_gpu(); + } + } + CodecDecodeCacheScope codec_decode_cache_scope(codec()); model().clear_kv_cache(); @@ -1099,6 +1261,31 @@ bool Pipeline::synthesize_streaming_prompt_codes_locked(const PipelineParams & p " ms/frame, ar_avg=" + std::to_string(ar_ms_per_frame) + " ms/frame, total_rtf=" + std::to_string(total_rtf) + ", max_rss=" + std::to_string(get_max_rss_mb()) + " MB"); + + 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_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_weights_on_gpu()) codec().free_gpu_weights(); + } + } else { + safe_print_ln("[Pipeline] Single-shot mode: Skipping post-stream VRAM restore."); + } + } return true; } From 8bca5de6ffc90cccd6b9e22400876f89c48217ce Mon Sep 17 00:00:00 2001 From: Skyrion9 Date: Wed, 22 Jul 2026 15:15:59 +0300 Subject: [PATCH 3/5] perf: use sequential access hints for weight loading 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 --- src/s2_mapped_file.cpp | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/src/s2_mapped_file.cpp b/src/s2_mapped_file.cpp index e0c3021..7275ff2 100644 --- a/src/s2_mapped_file.cpp +++ b/src/s2_mapped_file.cpp @@ -16,8 +16,11 @@ bool MappedFile::open(const std::string& path) { close(); #ifdef _WIN32 + // GGUF is mostly sequential HANDLE fh = CreateFileA(path.c_str(), GENERIC_READ, FILE_SHARE_READ, - nullptr, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, nullptr); + nullptr, OPEN_EXISTING, + FILE_ATTRIBUTE_NORMAL | FILE_FLAG_SEQUENTIAL_SCAN, + nullptr); if (fh == INVALID_HANDLE_VALUE) { return false; } @@ -75,7 +78,15 @@ bool MappedFile::open(const std::string& path) { return false; } - ::madvise(addr, size_, MADV_RANDOM); + // GGUF is mostly sequential + ::madvise(addr, size_, MADV_SEQUENTIAL); + +#ifdef __linux__ + // Huge pages reduce TLB pressure during multi-GB weight loads. + ::madvise(addr, size_, MADV_HUGEPAGE); + // Exclude from core dumps so we don't write entire model to disk in a crash. + ::madvise(addr, size_, MADV_DONTDUMP); +#endif data_ = addr; fd_ = fd; From 92e48f0d9f3bda648bc336ec78616f712eb12065 Mon Sep 17 00:00:00 2001 From: Skyrion9 Date: Thu, 30 Jul 2026 19:32:44 +0300 Subject: [PATCH 4/5] feat: more VRAM split (codec), deferred init, and overlapped decode pipeline 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. --- include/s2_codec.h | 9 + include/s2_mapped_file.h | 1 + include/s2_model.h | 8 +- include/s2_pipeline.h | 3 + src/main.cpp | 4 + src/s2_codec.cpp | 185 +++++++++++++-- src/s2_mapped_file.cpp | 160 ++++++++----- src/s2_model.cpp | 26 ++- src/s2_pipeline.cpp | 481 ++++++++++++++++++++++++++++----------- src/s2_server.cpp | 1 + 10 files changed, 670 insertions(+), 208 deletions(-) diff --git a/include/s2_codec.h b/include/s2_codec.h index 8a7e791..96a9e56 100644 --- a/include/s2_codec.h +++ b/include/s2_codec.h @@ -43,6 +43,15 @@ class AudioCodec { bool free_gpu_weights(); + bool free_encoder_weights(); + bool restore_encoder_weights(); + bool free_decoder_weights(); + bool restore_decoder_weights(); + bool is_encoder_on_gpu() const; + bool is_decoder_on_gpu() const; + size_t get_encoder_gpu_bytes() const; + size_t get_decoder_gpu_bytes() const; + bool is_weights_on_gpu() const; bool refresh_host_caches_from_mmap(); diff --git a/include/s2_mapped_file.h b/include/s2_mapped_file.h index c355a6d..b4619e6 100644 --- a/include/s2_mapped_file.h +++ b/include/s2_mapped_file.h @@ -20,6 +20,7 @@ class MappedFile { bool open(const std::string& path); void close(); void drop_page_cache(); + void warm_page_cache(); bool is_open() const { return data_ != nullptr; } const uint8_t* data() const { return static_cast(data_); } diff --git a/include/s2_model.h b/include/s2_model.h index 123f73d..338b014 100644 --- a/include/s2_model.h +++ b/include/s2_model.h @@ -96,9 +96,9 @@ class SlowARModel { SlowARModel(); ~SlowARModel(); - bool load(const std::string & gguf_path, int32_t gpu_device = -1, BackendType backend_type = BackendType::CPU, int32_t n_gpu_layers = -1); + bool load(const std::string & gguf_path, int32_t gpu_device = -1, BackendType backend_type = BackendType::CPU, int32_t n_gpu_layers = -1, bool fast_decoder_cpu = false, bool codebook_embeddings_cpu = false); - bool load_shared(gguf_context * gguf_ctx, const std::string & gguf_path, int32_t gpu_device = -1, BackendType backend_type = BackendType::CPU, int32_t n_gpu_layers = -1); + bool load_shared(gguf_context * gguf_ctx, const std::string & gguf_path, int32_t gpu_device = -1, BackendType backend_type = BackendType::CPU, int32_t n_gpu_layers = -1, bool fast_decoder_cpu = false, bool codebook_embeddings_cpu = false); ggml_context * weights_ctx() { return weights_.ctx_w; } const std::unordered_set & weight_tensor_set() const { return weight_tensor_set_; } @@ -125,6 +125,8 @@ class SlowARModel { bool is_weights_on_gpu() const { return weights_on_gpu_; } + bool prefers_gpu() const { return !original_gpu_weights_.empty(); } + private: bool eval_cached(const std::vector & flat_tokens, int32_t n_tokens, int32_t n_threads, @@ -175,6 +177,8 @@ class SlowARModel { std::vector original_cpu_weights_; bool weights_on_gpu_ = false; bool weights_allocated_ = false; + bool fast_decoder_cpu_ = false; + bool codebook_embeddings_cpu_ = false; MappedFile mapped_gguf_; diff --git a/include/s2_pipeline.h b/include/s2_pipeline.h index de276e3..d930230 100644 --- a/include/s2_pipeline.h +++ b/include/s2_pipeline.h @@ -38,6 +38,8 @@ struct PipelineParams { int32_t n_gpu_layers = -1; bool codec_auto_backend = true; bool codec_follow_backend = true; + bool fast_decoder_cpu = false; + bool codebook_embeddings_cpu = false; int32_t stream_decode_stride_frames = 0; int32_t stream_holdback_frames = -1; int32_t codec_decode_context_frames = -1; @@ -50,6 +52,7 @@ struct PipelineParams { bool enable_vram_swap = true; bool enable_hot_swap = false; bool is_persistent = false; + bool more_segments_pending = false; }; class Pipeline { diff --git a/src/main.cpp b/src/main.cpp index eff3531..eddcf4a 100644 --- a/src/main.cpp +++ b/src/main.cpp @@ -90,6 +90,8 @@ void print_uso() { safe_print(" --codec-auto Benchmark codec backends and keep the fastest (default)\n"); safe_print(" --codec-follow-backend Force codec to follow the selected GPU backend\n"); safe_print(" --codec-cpu Force codec on CPU even when model uses GPU\n"); + safe_print(" --fast-decoder-cpu Force the 4-layer fast decoder on CPU (enables cached graph batching, saves ~200-400 MB VRAM at Q4)\n"); + safe_print(" --codebook-cpu Force codebook_embeddings on CPU even at full GPU offload (saves ~56 MB VRAM at Q4, costs ~4-5 ms/frame)\n"); safe_print(" --no-vram-swap Disable phase-gated VRAM swapping between Slow-AR and Codec (keeps both in VRAM simultaneously instead)\n"); safe_print(" --hot-swap Aggressively evict weights and OS page cache after each server request (~100 MB RAM, ~25 MB VRAM idle)\n"); safe_print(" --stream-file Write output WAV through the streaming path\n"); @@ -197,6 +199,8 @@ int main(int argc, char** argv) { else if (arg == "--codec-auto") { params.codec_auto_backend = true; params.codec_follow_backend = true; } else if (arg == "--codec-follow-backend") { params.codec_auto_backend = false; params.codec_follow_backend = true; } else if (arg == "--codec-cpu") { params.codec_auto_backend = false; params.codec_follow_backend = false; } + else if (arg == "--fast-decoder-cpu") { params.fast_decoder_cpu = true; } + else if (arg == "--codebook-cpu") { params.codebook_embeddings_cpu = true; } else if (arg == "--no-vram-swap") { params.enable_vram_swap = false; } else if (arg == "--hot-swap") { params.enable_hot_swap = true; } else if (arg == "--stream-file") { use_stream_file = true; } diff --git a/src/s2_codec.cpp b/src/s2_codec.cpp index fae1ced..03aac86 100755 --- a/src/s2_codec.cpp +++ b/src/s2_codec.cpp @@ -65,7 +65,6 @@ struct AudioCodec::Impl { ggml_backend_t backend = nullptr; ggml_backend_t backend_cpu = nullptr; ggml_context * ctx_w = nullptr; - ggml_backend_buffer_t model_buf = nullptr; std::string tprefix; int32_t sample_rate = 0; @@ -107,13 +106,23 @@ struct AudioCodec::Impl { codec_decode_cache decode_cache; std::string gguf_path; size_t gguf_data_offset = 0; - std::unordered_map tensor_offsets; std::vector original_gpu_weights; std::vector original_cpu_weights; - std::vector all_codec_weights; - bool weights_on_gpu = false; MappedFile mapped_gguf_; + + std::unordered_map tensor_offsets; bool weights_allocated_ = false; + std::vector all_codec_weights; + std::vector encoder_weights; + std::vector decoder_weights; + + ggml_backend_buffer_t model_buf = nullptr; + ggml_backend_buffer_t encoder_buf = nullptr; + ggml_backend_buffer_t decoder_buf = nullptr; + + bool weights_on_gpu = false; + bool encoder_on_gpu = false; + bool decoder_on_gpu = false; }; static const char * backend_type_name(BackendType backend_type) { @@ -192,6 +201,14 @@ static void reset_decode_cache(codec_decode_cache & cache, bool preserve_failed_ static void reset_codec_impl(AudioCodec::Impl & impl) { reset_decode_cache(impl.decode_cache, false); + if (impl.encoder_buf) { + ggml_backend_buffer_free(impl.encoder_buf); + impl.encoder_buf = nullptr; + } + if (impl.decoder_buf) { + ggml_backend_buffer_free(impl.decoder_buf); + impl.decoder_buf = nullptr; + } if (impl.model_buf) { ggml_backend_buffer_free(impl.model_buf); impl.model_buf = nullptr; @@ -748,6 +765,13 @@ void AudioCodec::clear_decode_cache() { } } +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; + return false; +} + bool AudioCodec::load_shared(SlowARModel* Model, gguf_context * shared_gguf_ctx, const std::string & gguf_path, int32_t gpu_device, BackendType backend_type) { if (!impl_) { impl_ = new Impl(); @@ -942,6 +966,8 @@ bool AudioCodec::load_shared(SlowARModel* Model, gguf_context * shared_gguf_ctx, const auto & model_weights = Model ? Model->weight_tensor_set() : std::unordered_set(); impl_->all_codec_weights.clear(); + impl_->encoder_weights.clear(); + impl_->decoder_weights.clear(); impl_->tensor_offsets.clear(); impl_->original_gpu_weights.clear(); impl_->original_cpu_weights.clear(); @@ -950,13 +976,18 @@ bool AudioCodec::load_shared(SlowARModel* Model, gguf_context * shared_gguf_ctx, const char * tname = gguf_get_tensor_name(shared_gguf_ctx, ti); ggml_tensor * t = ggml_get_tensor(impl_->ctx_w, tname); if (!t) continue; - - // Skip tensors that belong to the SlowAR Model if (Model && model_weights.find(t) != model_weights.end()) continue; impl_->all_codec_weights.push_back(t); impl_->tensor_offsets[t] = gguf_get_tensor_offset(shared_gguf_ctx, ti); + std::string name_str(tname); + if (is_encoder_tensor(name_str, impl_->tprefix)) { + impl_->encoder_weights.push_back(t); + } else { + impl_->decoder_weights.push_back(t); + } + if (!ggml_backend_is_cpu(impl_->backend)) { impl_->original_gpu_weights.push_back(t); } else { @@ -965,6 +996,8 @@ bool AudioCodec::load_shared(SlowARModel* Model, gguf_context * shared_gguf_ctx, } impl_->weights_on_gpu = false; + impl_->encoder_on_gpu = false; + impl_->decoder_on_gpu = false; impl_->mapped_gguf_.open(gguf_path); if (!impl_->mapped_gguf_.is_open()) { @@ -1483,25 +1516,24 @@ bool AudioCodec::is_weights_on_gpu() const { } bool AudioCodec::free_gpu_weights() { - if (!impl_ || !impl_->weights_on_gpu) return true; - + if (!impl_) return true; + if (impl_->encoder_on_gpu) free_encoder_weights(); + if (impl_->decoder_on_gpu) free_decoder_weights(); + if (!impl_->weights_on_gpu && !impl_->model_buf) return true; S2_LOG_INFO_STREAM("[Codec] >>> FREEING Audio Codec GPU weights..." << std::endl); const auto t0 = std::chrono::steady_clock::now(); - ggml_backend_synchronize(impl_->backend); - if (impl_->model_buf) { ggml_backend_buffer_free(impl_->model_buf); impl_->model_buf = nullptr; } - for (ggml_tensor * t : impl_->all_codec_weights) { if (t) { t->data = nullptr; t->buffer = nullptr; } } - impl_->weights_allocated_ = false; impl_->weights_on_gpu = false; - + impl_->encoder_on_gpu = false; + impl_->decoder_on_gpu = false; const auto t1 = std::chrono::steady_clock::now(); const double free_ms = std::chrono::duration(t1 - t0).count(); S2_LOG_INFO_STREAM("[Codec] <<< Audio Codec GPU weights FREED in " << free_ms << " ms" << std::endl); @@ -1520,9 +1552,126 @@ bool AudioCodec::restore_weights_to_gpu() { return true; } +bool AudioCodec::free_encoder_weights() { + if (!impl_ || !impl_->encoder_on_gpu) return true; + S2_LOG_INFO_STREAM("[Codec] >>> FREEING encoder GPU weights..." << std::endl); + const auto t0 = std::chrono::steady_clock::now(); + ggml_backend_synchronize(impl_->backend); + if (impl_->encoder_buf) { + ggml_backend_buffer_free(impl_->encoder_buf); + impl_->encoder_buf = nullptr; + } + for (ggml_tensor * t : impl_->encoder_weights) { + if (t) { t->data = nullptr; t->buffer = nullptr; } + } + impl_->encoder_on_gpu = false; + impl_->weights_allocated_ = false; + const auto t1 = std::chrono::steady_clock::now(); + const double enc_free_ms = std::chrono::duration(t1 - t0).count(); + S2_LOG_INFO_STREAM("[Codec] <<< Encoder FREED in " << enc_free_ms << " ms" << std::endl); + return true; +} + +bool AudioCodec::restore_encoder_weights() { + if (!impl_ || impl_->encoder_on_gpu || impl_->encoder_weights.empty()) return true; + if (ggml_backend_is_cpu(impl_->backend)) return true; + if (!impl_->mapped_gguf_.is_open()) return false; + S2_LOG_INFO_STREAM("[Codec] >>> RESTORING encoder weights to GPU..." << std::endl); + const auto t0 = std::chrono::steady_clock::now(); + + size_t b = 0; std::string e; + if (!allocate_codec_buffers(impl_->backend, impl_->encoder_weights, impl_->encoder_buf, b, e)) { + std::cerr << "[Codec] Encoder alloc failed: " << e << std::endl; + return false; + } + const uint8_t * base = impl_->mapped_gguf_.data(); + for (ggml_tensor * t : impl_->encoder_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_->encoder_on_gpu = true; + const auto t1 = std::chrono::steady_clock::now(); + const double enc_restore_ms = std::chrono::duration(t1 - t0).count(); + S2_LOG_INFO_STREAM("[Codec] <<< Encoder RESTORED in " << enc_restore_ms << " ms" << std::endl); + return true; +} + +bool AudioCodec::free_decoder_weights() { + if (!impl_ || !impl_->decoder_on_gpu) return true; + S2_LOG_INFO_STREAM("[Codec] >>> FREEING decoder GPU weights..." << std::endl); + const auto t0 = std::chrono::steady_clock::now(); + ggml_backend_synchronize(impl_->backend); + if (impl_->decoder_buf) { + ggml_backend_buffer_free(impl_->decoder_buf); + impl_->decoder_buf = nullptr; + } + for (ggml_tensor * t : impl_->decoder_weights) { + if (t) { t->data = nullptr; t->buffer = nullptr; } + } + impl_->decoder_on_gpu = false; + impl_->weights_allocated_ = false; + + reset_decode_cache(impl_->decode_cache, false); + const auto t1 = std::chrono::steady_clock::now(); + const double dec_free_ms = std::chrono::duration(t1 - t0).count(); + S2_LOG_INFO_STREAM("[Codec] <<< Decoder FREED in " << dec_free_ms << " ms" << std::endl); + return true; +} + +bool AudioCodec::restore_decoder_weights() { + if (!impl_ || impl_->decoder_on_gpu || impl_->decoder_weights.empty()) return true; + if (ggml_backend_is_cpu(impl_->backend)) return true; + if (!impl_->mapped_gguf_.is_open()) return false; + S2_LOG_INFO_STREAM("[Codec] >>> RESTORING decoder weights to GPU..." << std::endl); + const auto t0 = std::chrono::steady_clock::now(); + + size_t b = 0; std::string e; + if (!allocate_codec_buffers(impl_->backend, impl_->decoder_weights, impl_->decoder_buf, b, e)) { + std::cerr << "[Codec] Decoder alloc failed: " << e << std::endl; + return false; + } + const uint8_t * base = impl_->mapped_gguf_.data(); + for (ggml_tensor * t : impl_->decoder_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_->decoder_on_gpu = true; + const auto t1 = std::chrono::steady_clock::now(); + const double dec_restore_ms = std::chrono::duration(t1 - t0).count(); + S2_LOG_INFO_STREAM("[Codec] <<< Decoder RESTORED in " << dec_restore_ms << " ms" << std::endl); + return true; +} + +bool AudioCodec::is_encoder_on_gpu() const { + return impl_ ? impl_->encoder_on_gpu : false; +} + +bool AudioCodec::is_decoder_on_gpu() const { + return impl_ ? impl_->decoder_on_gpu : false; +} + +size_t AudioCodec::get_encoder_gpu_bytes() const { + if (!impl_ || !impl_->encoder_buf || !impl_->encoder_on_gpu) return 0; + return ggml_backend_buffer_get_size(impl_->encoder_buf); +} + +size_t AudioCodec::get_decoder_gpu_bytes() const { + if (!impl_ || !impl_->decoder_buf || !impl_->decoder_on_gpu) return 0; + return ggml_backend_buffer_get_size(impl_->decoder_buf); +} + size_t AudioCodec::get_gpu_memory_usage_bytes() const { - if (!impl_ || !impl_->model_buf || !impl_->weights_on_gpu) return 0; - return ggml_backend_buffer_get_size(impl_->model_buf); + if (!impl_) return 0; + size_t total = 0; + if (impl_->model_buf && impl_->weights_on_gpu) + total += ggml_backend_buffer_get_size(impl_->model_buf); + if (impl_->encoder_buf && impl_->encoder_on_gpu) + total += ggml_backend_buffer_get_size(impl_->encoder_buf); + if (impl_->decoder_buf && impl_->decoder_on_gpu) + total += ggml_backend_buffer_get_size(impl_->decoder_buf); + return total; } bool AudioCodec::refresh_host_caches_from_mmap() { @@ -1567,6 +1716,10 @@ bool AudioCodec::refresh_host_caches_from_mmap() { 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; @@ -1581,6 +1734,8 @@ bool AudioCodec::ensure_weights_loaded() { } 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; } diff --git a/src/s2_mapped_file.cpp b/src/s2_mapped_file.cpp index 7275ff2..7a54afc 100644 --- a/src/s2_mapped_file.cpp +++ b/src/s2_mapped_file.cpp @@ -2,12 +2,15 @@ #ifdef _WIN32 #include +#include + #else #include #include #include #include #include + #endif namespace s2 { @@ -16,8 +19,13 @@ bool MappedFile::open(const std::string& path) { close(); #ifdef _WIN32 - // GGUF is mostly sequential - HANDLE fh = CreateFileA(path.c_str(), GENERIC_READ, FILE_SHARE_READ, + int wlen = MultiByteToWideChar(CP_UTF8, 0, path.c_str(), -1, nullptr, 0); + if (wlen <= 0) return false; + std::wstring wpath(static_cast(wlen), L'\0'); + MultiByteToWideChar(CP_UTF8, 0, path.c_str(), -1, &wpath[0], wlen); + wpath.resize(wcslen(wpath.c_str())); + + HANDLE fh = CreateFileW(wpath.c_str(), GENERIC_READ, FILE_SHARE_READ, nullptr, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL | FILE_FLAG_SEQUENTIAL_SCAN, nullptr); @@ -26,32 +34,32 @@ bool MappedFile::open(const std::string& path) { } LARGE_INTEGER file_size; - if (!GetFileSizeEx(fh, &file_size)) { - CloseHandle(fh); - return false; + if (!GetFileSizeEx(fh, &file_size)) { + CloseHandle(fh); + return false; } size_ = static_cast(file_size.QuadPart); - if (size_ == 0) { - CloseHandle(fh); + if (size_ == 0) { + CloseHandle(fh); return true; } - HANDLE mh = CreateFileMappingA(fh, nullptr, PAGE_READONLY, 0, 0, nullptr); - if (!mh) { - CloseHandle(fh); - return false; + 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; + if (!addr) { + CloseHandle(mh); + CloseHandle(fh); + return false; } - data_ = addr; - file_handle_ = static_cast(fh); + data_ = addr; + file_handle_ = static_cast(fh); mapping_handle_ = static_cast(mh); #else @@ -61,34 +69,31 @@ bool MappedFile::open(const std::string& path) { } struct stat st; - if (fstat(fd, &st) < 0) { - ::close(fd); - return false; + if (fstat(fd, &st) < 0) { + ::close(fd); + return false; } size_ = static_cast(st.st_size); - if (size_ == 0) { - ::close(fd); + if (size_ == 0) { + ::close(fd); return true; } void* addr = ::mmap(nullptr, size_, PROT_READ, MAP_PRIVATE, fd, 0); - if (addr == MAP_FAILED) { - ::close(fd); - return false; + if (addr == MAP_FAILED) { + ::close(fd); + return false; } - // GGUF is mostly sequential ::madvise(addr, size_, MADV_SEQUENTIAL); #ifdef __linux__ - // Huge pages reduce TLB pressure during multi-GB weight loads. ::madvise(addr, size_, MADV_HUGEPAGE); - // Exclude from core dumps so we don't write entire model to disk in a crash. ::madvise(addr, size_, MADV_DONTDUMP); #endif - data_ = addr; + data_ = addr; fd_ = fd; #endif @@ -124,42 +129,42 @@ void MappedFile::close() { MappedFile::MappedFile(MappedFile&& other) noexcept { #ifdef _WIN32 - data_ = other.data_; - other.data_ = nullptr; - size_ = other.size_; + data_ = other.data_; + other.data_ = nullptr; + size_ = other.size_; other.size_ = 0; - file_handle_ = other.file_handle_; + file_handle_ = other.file_handle_; other.file_handle_ = nullptr; - mapping_handle_ = other.mapping_handle_; + mapping_handle_ = other.mapping_handle_; other.mapping_handle_ = nullptr; #else - data_ = other.data_; - other.data_ = nullptr; - size_ = other.size_; - other.size_ = 0; - fd_ = other.fd_; + data_ = other.data_; + other.data_ = nullptr; + size_ = other.size_; + other.size_ = 0; + fd_ = other.fd_; other.fd_ = -1; #endif } MappedFile& MappedFile::operator=(MappedFile&& other) noexcept { - if (this != &other) { + if (this != &other) { close(); #ifdef _WIN32 - data_ = other.data_; - other.data_ = nullptr; - size_ = other.size_; + data_ = other.data_; + other.data_ = nullptr; + size_ = other.size_; other.size_ = 0; - file_handle_ = other.file_handle_; + file_handle_ = other.file_handle_; other.file_handle_ = nullptr; - mapping_handle_ = other.mapping_handle_; + mapping_handle_ = other.mapping_handle_; other.mapping_handle_ = nullptr; #else - data_ = other.data_; - other.data_ = nullptr; - size_ = other.size_; - other.size_ = 0; - fd_ = other.fd_; + data_ = other.data_; + other.data_ = nullptr; + size_ = other.size_; + other.size_ = 0; + fd_ = other.fd_; other.fd_ = -1; #endif } @@ -167,16 +172,65 @@ MappedFile& MappedFile::operator=(MappedFile&& other) noexcept { } void MappedFile::drop_page_cache() { -#if defined(__linux__) || defined(__APPLE__) +#ifdef __linux__ if (data_ && data_ != MAP_FAILED && size_ > 0) { ::madvise(data_, size_, MADV_DONTNEED); } if (fd_ >= 0) { ::posix_fadvise(fd_, 0, 0, POSIX_FADV_DONTNEED); } +#elif defined(__APPLE__) + if (data_ && data_ != MAP_FAILED && size_ > 0) { + ::madvise(data_, size_, MADV_DONTNEED); + } + if (fd_ >= 0) { + ::fcntl(fd_, F_NOCACHE, 1); + ::fcntl(fd_, F_RDAHEAD, 0); + } +#elif defined(_WIN32) + (void)data_; + (void)size_; + SetProcessWorkingSetSizeEx( + GetCurrentProcess(), + static_cast(-1), + static_cast(-1), + QUOTA_LIMITS_HARDWS_MIN_DISABLE); +#endif +} + +void MappedFile::warm_page_cache() { + if (!data_ || size_ == 0) return; + +#ifdef __linux__ + ::madvise(data_, size_, MADV_WILLNEED); + ::madvise(data_, size_, MADV_SEQUENTIAL); + +#ifdef MADV_COLD + ::madvise(data_, size_, MADV_COLD); +#endif + +#elif defined(__APPLE__) + ::madvise(data_, size_, MADV_WILLNEED); + ::madvise(data_, size_, MADV_SEQUENTIAL); + + if (fd_ >= 0) { + ::fcntl(fd_, F_RDAHEAD, 1); + ::fcntl(fd_, F_NOCACHE, 0); + } + #elif defined(_WIN32) - if (data_ && size_ > 0) { - ::VirtualUnlock(data_, size_); + WIN32_MEMORY_RANGE_ENTRY entry; + entry.VirtualAddress = data_; + entry.NumberOfBytes = size_; + if (!PrefetchVirtualMemory(GetCurrentProcess(), 1, &entry, 0)) { + SYSTEM_INFO si; + GetSystemInfo(&si); + const size_t page_size = si.dwPageSize; + volatile uint8_t sink = 0; + const uint8_t * base = static_cast(data_); + for (size_t off = 0; off < size_; off += page_size) + sink = base[off]; + (void)sink; } #endif } diff --git a/src/s2_model.cpp b/src/s2_model.cpp index 992b398..cc6a185 100755 --- a/src/s2_model.cpp +++ b/src/s2_model.cpp @@ -106,6 +106,14 @@ static bool allocate_weight_buffers(ggml_backend_t backend, size_t & max_buffer_bytes, std::string & error_message) { free_backend_buffers(out_buffers); + + for (ggml_tensor * tensor : tensors) { + if (tensor) { + tensor->data = nullptr; + tensor->buffer = nullptr; + } + } + total_bytes = 0; max_buffer_bytes = 0; error_message.clear(); @@ -202,7 +210,7 @@ SlowARModel::~SlowARModel() { weights_.ctx_w = nullptr; } -bool SlowARModel::load_shared(gguf_context * ctx_gguf, const std::string & gguf_path, int32_t gpu_device, BackendType backend_type, int32_t n_gpu_layers) { +bool SlowARModel::load_shared(gguf_context * ctx_gguf, const std::string & gguf_path, int32_t gpu_device, BackendType backend_type, int32_t n_gpu_layers, bool fast_decoder_cpu, bool codebook_embeddings_cpu) { backend_cpu_ = ggml_backend_cpu_init(); if (!backend_cpu_) { @@ -313,6 +321,9 @@ bool SlowARModel::load_shared(gguf_context * ctx_gguf, const std::string & gguf_ n_gpu_layers = hparams_.block_count; } n_gpu_layers_ = n_gpu_layers; + fast_decoder_cpu_ = fast_decoder_cpu; + codebook_embeddings_cpu_ = codebook_embeddings_cpu; + S2_LOG_INFO_STREAM("[Model] GPU layers: " << n_gpu_layers_ << " / " << hparams_.block_count << std::endl); if (n_gpu_layers_ > 0 && wants_gpu_backend) { @@ -496,6 +507,15 @@ bool SlowARModel::load_shared(gguf_context * ctx_gguf, const std::string & gguf_ n_gpu_layers_ == hparams_.block_count; auto get_weight_backend = [&](const std::string & name) -> ggml_backend_t { + if (fast_decoder_cpu_ && + (name.rfind("fast_layers.", 0) == 0 || + name.rfind("fast_", 0) == 0)) { + return backend_cpu_; + } + + if (codebook_embeddings_cpu_ && name == "codebook_embeddings.weight") { + return backend_cpu_; + } if (full_model_offload) { return backend_gpu_; @@ -549,14 +569,14 @@ bool SlowARModel::load_shared(gguf_context * ctx_gguf, const std::string & gguf_ return true; } -bool SlowARModel::load(const std::string & gguf_path, int32_t gpu_device, BackendType backend_type, int32_t n_gpu_layers) { +bool SlowARModel::load(const std::string & gguf_path, int32_t gpu_device, BackendType backend_type, int32_t n_gpu_layers, bool fast_decoder_cpu, bool codebook_embeddings_cpu) { struct gguf_init_params params = { true, nullptr }; gguf_context * ctx_gguf = gguf_init_from_file(gguf_path.c_str(), params); if (!ctx_gguf) { std::cerr << "[Model] Failed to load GGUF from " << gguf_path << std::endl; return false; } - if (!load_shared(ctx_gguf, gguf_path, gpu_device, backend_type, n_gpu_layers)) { + if (!load_shared(ctx_gguf, gguf_path, gpu_device, backend_type, n_gpu_layers, fast_decoder_cpu, codebook_embeddings_cpu)) { gguf_free(ctx_gguf); return false; } diff --git a/src/s2_pipeline.cpp b/src/s2_pipeline.cpp index d194c63..1447f2f 100644 --- a/src/s2_pipeline.cpp +++ b/src/s2_pipeline.cpp @@ -8,6 +8,8 @@ #include #include #include +#include +#include #ifdef __linux__ #include @@ -331,7 +333,7 @@ bool Pipeline::init(const PipelineParams & params) { } const auto model_t0 = std::chrono::steady_clock::now(); - if (!model().load_shared(shared_gguf, params.model_path, params.gpu_device, params.backend_type, params.n_gpu_layers)) { + if (!model().load_shared(shared_gguf, params.model_path, params.gpu_device, params.backend_type, params.n_gpu_layers, params.fast_decoder_cpu, params.codebook_embeddings_cpu)) { safe_print_error_ln("Pipeline error: could not load model from " + params.model_path); gguf_free(shared_gguf); return false; @@ -431,9 +433,16 @@ bool Pipeline::init(const PipelineParams & params) { const auto codec_t1 = std::chrono::steady_clock::now(); const auto model_weights_t0 = std::chrono::steady_clock::now(); - if (!model().allocate_and_load_weights()) { - safe_print_error_ln("Pipeline error: failed to allocate and load Slow-AR weights"); - return false; + const bool defer_weight_loading = params.enable_vram_swap && model().prefers_gpu(); + + if (!defer_weight_loading) { + if (!model().allocate_and_load_weights()) { + safe_print_error_ln("Pipeline error: failed to allocate and load Slow-AR weights"); + return false; + } + } else { + safe_print_ln("[Pipeline] Deferring Slow-AR weight loading to first request (VRAM swap active)."); + model().mapped_file().warm_page_cache(); } const auto model_weights_t1 = std::chrono::steady_clock::now(); @@ -441,7 +450,7 @@ bool Pipeline::init(const PipelineParams & params) { initialized_ = true; - model_prefers_gpu_ = model().is_weights_on_gpu(); + model_prefers_gpu_ = model().prefers_gpu(); codec_prefers_gpu_ = use_gpu_codec; if (model_prefers_gpu_ && codec_prefers_gpu_) { @@ -464,7 +473,8 @@ bool Pipeline::init(const PipelineParams & params) { std::to_string(std::chrono::duration(codec_t1 - codec_t0).count()) + " ms (" + codec().backend_name() + "), model_weights=" + std::to_string(std::chrono::duration(model_weights_t1 - model_weights_t0).count()) + - " ms, total=" + + (defer_weight_loading ? " ms (deferred)" : " ms") + + ", total=" + std::to_string(std::chrono::duration(init_t1 - init_t0).count()) + " ms, max_rss=" + std::to_string(get_max_rss_mb()) + " MB"); @@ -533,6 +543,14 @@ bool Pipeline::resolve_reference_prompt_locked(const PipelineParams & params, Au voice_mgr_.set_storage_dir(params.voice_storage_dir); if (!ref_audio.samples.empty()) { + const bool need_encoder_vram = params.enable_vram_swap && codec_prefers_gpu_; + if (need_encoder_vram) { + if (codec().is_decoder_on_gpu()) { + codec().free_decoder_weights(); + } + codec().restore_encoder_weights(); + } + const auto ref_t0 = std::chrono::steady_clock::now(); if (!codec().encode(ref_audio.samples.data(), static_cast(ref_audio.samples.size()), params.gen.n_threads, ref_codes, T_prompt)) { @@ -543,6 +561,10 @@ bool Pipeline::resolve_reference_prompt_locked(const PipelineParams & params, Au const auto ref_t1 = std::chrono::steady_clock::now(); ref_encode_ms = std::chrono::duration(ref_t1 - ref_t0).count(); + if (need_encoder_vram) { + codec().free_encoder_weights(); + } + if (!ref_codes.empty() && params.save_voice && !params.voice_id.empty()) { save_voice_profile_locked(params.voice_id, ref_codes, T_prompt, effective_prompt_text, params); @@ -719,6 +741,7 @@ bool Pipeline::encode_prompt_audio_data(const AudioData & ref_audio, int32_t n_t bool Pipeline::synthesize_raw(const PipelineParams & params, AudioData & ref_audio, std::vector& audio_out) { std::lock_guard lock(synthesize_mutex_); + std::vector ref_codes; int32_t T_prompt = 0; double ref_encode_ms = 0.0; @@ -729,30 +752,17 @@ bool Pipeline::synthesize_raw(const PipelineParams & params, AudioData & ref_aud return false; } - std::thread pre_restore_thread; - bool pre_restore_started = false; - - if (params.enable_vram_swap && params.is_persistent && - model_prefers_gpu_ && !model().is_weights_on_gpu()) { - safe_print_ln("[Pipeline] Hot-Swap: Pre-fetching Slow-AR to VRAM in background..."); - pre_restore_started = true; - pre_restore_thread = std::thread([this]() { - model().acquire_compute_resources(); - model().restore_weights_to_gpu(); - }); + if (params.enable_vram_swap) { + if (pending_offload_thread_.joinable()) { + pending_offload_thread_.join(); + } } if (!resolve_reference_prompt_locked(params, ref_audio, ref_codes, T_prompt, effective_prompt_text, ref_encode_ms)) { - if (pre_restore_started) pre_restore_thread.join(); return false; } - if (pre_restore_started) { - pre_restore_thread.join(); - safe_print_ln("[Pipeline] Hot-Swap: Background pre-fetch complete."); - } - PipelineParams effective_params = params; effective_params.prompt_text = std::move(effective_prompt_text); @@ -799,133 +809,335 @@ bool Pipeline::synthesize_prompt_codes_locked(const PipelineParams & params, con } CodecDecodeCacheScope codec_decode_cache_scope(codec()); - model().clear_kv_cache(); safe_print_ln("--- Pipeline Synthesize ---"); safe_print_ln("Text: " + params.text); + std::thread vram_phase1_thread; + bool vram_phase1_ok = true; + if (params.enable_vram_swap) { - if (model_prefers_gpu_ && !model().is_weights_on_gpu()) { - safe_print_ln("[Pipeline] Restoring Slow-AR to VRAM for generation..."); - model().acquire_compute_resources(); - model().restore_weights_to_gpu(); - 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"); - } + 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_weights_on_gpu()) { - safe_print_ln("[Pipeline] Pre-loading Audio Codec to VRAM (hiding behind CPU gen)..."); - codec().restore_weights_to_gpu(); - safe_print_ln("[VRAM Diag] Post-Codec 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_ && 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"); + 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; + + 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; + + kv_init_thread = std::thread([&]() { + kv_init_ok = model().init_kv_cache(max_seq_len); + }); + + 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 bool can_overlap_decode = + model_prefers_gpu_ && !codec_prefers_gpu_; - if (res.n_frames == 0) { - safe_print_error_ln("Pipeline error: generation produced no frames."); - return false; - } + const int32_t offline_decode_stride_frames = + params.stream_decode_stride_frames > 0 ? params.stream_decode_stride_frames : 16; - if (params.enable_vram_swap) { - if (model_prefers_gpu_ && model().is_weights_on_gpu()) { - if (params.is_persistent) { - if (codec_prefers_gpu_) { - safe_print_ln("[Pipeline] Freeing Slow-AR from VRAM to make room for GPU Audio Codec..."); + GenerateResult res; + double gen_ms = 0.0; + double decode_ms = 0.0; + int32_t decode_batches = 0; + double decode_wall_ms = 0.0; + + GenerateParams gen_params = params.gen; + + if (can_overlap_decode) { + const int32_t codec_context_frames = + params.codec_decode_context_frames >= 0 + ? params.codec_decode_context_frames + : offline_decode_stride_frames; + const size_t samples_per_frame = + static_cast(std::max(1, codec().samples_per_code_frame())); + + std::vector> accum(num_codebooks); + for (auto & row : accum) + row.reserve(static_cast(params.gen.max_new_tokens)); + + std::mutex decode_mtx; + std::condition_variable decode_cv; + std::atomic frames_available{0}; + std::atomic gen_done{false}; + bool decode_failed = false; + + std::vector audio_accum; + audio_accum.reserve( + static_cast(params.gen.max_new_tokens) * samples_per_frame); + int32_t committed_frames = 0; + + auto decode_window = [&](int32_t total_frames, bool finalize) -> bool { + if (total_frames <= 0 || total_frames <= committed_frames) + return true; + const int32_t stable_frames = total_frames; + if (stable_frames <= committed_frames && !finalize) + return true; + const int32_t window_start = + std::max(0, committed_frames - codec_context_frames); + const int32_t window_frames = total_frames - window_start; + if (window_frames <= 0) + return true; + std::vector codes( + static_cast(num_codebooks) * window_frames); + { + std::lock_guard lock(decode_mtx); + for (int32_t cb = 0; cb < num_codebooks; ++cb) { + std::copy( + accum[cb].begin() + window_start, + accum[cb].begin() + total_frames, + codes.begin() + static_cast(cb) * window_frames); + } + } + std::vector pcm; + const auto t0 = std::chrono::steady_clock::now(); + if (!codec().decode(codes.data(), window_frames, + params.gen.n_threads, pcm)) { + return false; + } + const auto t1 = std::chrono::steady_clock::now(); + decode_ms += std::chrono::duration(t1 - t0).count(); + decode_batches++; + const size_t emit_begin = + static_cast(std::max(0, committed_frames - window_start)) + * samples_per_frame; + const size_t emit_end = finalize + ? pcm.size() + : std::min(pcm.size(), + static_cast( + std::max(0, stable_frames - window_start)) + * samples_per_frame); + if (emit_end > emit_begin) { + audio_accum.insert(audio_accum.end(), + pcm.begin() + emit_begin, + pcm.begin() + emit_end); + } + committed_frames = finalize ? total_frames : stable_frames; + return true; + }; + + const auto decode_thread_t0 = std::chrono::steady_clock::now(); + std::thread decode_thread([&]() { + int32_t last_committed = 0; + while (true) { + std::unique_lock 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 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); + const auto gen_t1 = std::chrono::steady_clock::now(); + gen_ms = std::chrono::duration(gen_t1 - gen_t0).count(); + + { + std::lock_guard lock(decode_mtx); + gen_done.store(true); + } + decode_cv.notify_one(); + decode_thread.join(); + const auto decode_thread_t1 = std::chrono::steady_clock::now(); + decode_wall_ms = std::chrono::duration( + decode_thread_t1 - decode_thread_t0).count(); + + if (res.n_frames == 0) { + safe_print_error_ln("Pipeline error: generation produced no frames."); + return false; + } + if (decode_failed) { + safe_print_error_ln("Pipeline error: overlapped decode failed."); + return false; + } + + if (params.enable_vram_swap) { + if (model_prefers_gpu_ && model().is_weights_on_gpu()) { + if (params.is_persistent) { + if (!params.more_segments_pending) { + safe_print_ln("[Pipeline] Freeing Slow-AR from VRAM (request complete)..."); + model().free_gpu_weights(); + model().free_compute_buffers(); + safe_print_ln("[VRAM Diag] Post-SlowAR 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"); + } + } else { + safe_print_ln("[Pipeline] Single-shot: Freeing Slow-AR from VRAM..."); model().free_gpu_weights(); - model().clear_kv_cache(); + model().free_compute_buffers(); safe_print_ln("[VRAM Diag] Post-SlowAR 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"); } - } else { - safe_print_ln("[Pipeline] Single-shot: Freeing Slow-AR from VRAM..."); - model().free_gpu_weights(); - model().free_compute_buffers(); - safe_print_ln("[VRAM Diag] Post-SlowAR 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"); } } - - if (codec_prefers_gpu_ && !codec().is_weights_on_gpu()) { - safe_print_ln("[Pipeline] Restoring Audio Codec to VRAM for decode..."); - codec().restore_weights_to_gpu(); - safe_print_ln("[VRAM Diag] Post-Codec 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 (params.enable_vram_swap && params.is_persistent && + params.enable_hot_swap && !params.more_segments_pending) { + 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); } - } - const int32_t offline_decode_stride_frames = - params.stream_decode_stride_frames > 0 ? params.stream_decode_stride_frames : 16; - double decode_ms = 0.0; - int32_t decode_batches = 0; - const auto decode_t0 = std::chrono::steady_clock::now(); + audio_out = std::move(audio_accum); - bool decode_ok = decode_codes_windowed(codec(), res.codes.data(), res.n_frames, num_codebooks, - params.gen.n_threads, offline_decode_stride_frames, - params.codec_decode_context_frames, - audio_out, &decode_ms, &decode_batches); - - const auto decode_t1 = std::chrono::steady_clock::now(); + } else { + const auto gen_t0 = std::chrono::steady_clock::now(); + res = generate(model(), tokenizer().config(), prompt, gen_params); + const auto gen_t1 = std::chrono::steady_clock::now(); + gen_ms = std::chrono::duration(gen_t1 - gen_t0).count(); - 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_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); + if (res.n_frames == 0) { + safe_print_error_ln("Pipeline error: generation produced no frames."); + return false; + } + + if (params.enable_vram_swap) { + if (codec_prefers_gpu_ && !codec().is_decoder_on_gpu()) { + safe_print_ln("[Pipeline] Restoring Audio Codec DECODER to VRAM (alongside Slow-AR)..."); + 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"); + } + } + + const auto decode_t0 = std::chrono::steady_clock::now(); + bool decode_ok = decode_codes_windowed(codec(), res.codes.data(), res.n_frames, num_codebooks, + params.gen.n_threads, offline_decode_stride_frames, + params.codec_decode_context_frames, + audio_out, &decode_ms, &decode_batches); + const auto decode_t1 = std::chrono::steady_clock::now(); + decode_wall_ms = std::chrono::duration(decode_t1 - decode_t0).count(); + + if (params.enable_vram_swap) { + if (params.is_persistent) { + if (params.enable_hot_swap && !params.more_segments_pending) { + 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(); + + if (!params.more_segments_pending) { + model().free_gpu_weights(); + model().free_compute_buffers(); + } + } } else { - if (codec().is_weights_on_gpu()) codec().free_gpu_weights(); + 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-decode VRAM restore."); } - } - if (!decode_ok) { - safe_print_error_ln("Pipeline error: decode failed."); - return false; + if (!decode_ok) { + safe_print_error_ln("Pipeline error: decode failed."); + return false; + } } model().clear_kv_cache(); + const auto synth_t1 = std::chrono::steady_clock::now(); const double kv_ms = std::chrono::duration(kv_t1 - kv_t0).count(); - const double gen_ms = std::chrono::duration(gen_t1 - gen_t0).count(); - const double decode_wall_ms = std::chrono::duration(decode_t1 - decode_t0).count(); const double total_ms = std::chrono::duration(synth_t1 - synth_t0).count(); const double audio_seconds = codec().sample_rate() > 0 ? (static_cast(audio_out.size()) / codec().sample_rate()) @@ -945,7 +1157,10 @@ bool Pipeline::synthesize_prompt_codes_locked(const PipelineParams & params, con " 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") + + (params.more_segments_pending ? ", vram=held" : "") + + ", total=" + std::to_string(total_ms) + " ms, gen_avg=" + std::to_string(gen_ms_per_frame) + " ms/frame, total_avg=" + std::to_string(total_ms_per_frame) + " ms/frame, gen_rtf=" + std::to_string(gen_rtf) + @@ -953,7 +1168,7 @@ bool Pipeline::synthesize_prompt_codes_locked(const PipelineParams & params, con ", max_rss=" + std::to_string(get_max_rss_mb()) + " MB"); if (params.enable_vram_swap) { - safe_print_ln("[VRAM Diag] Post-Phase3: Slow-AR=" + + safe_print_ln("[VRAM Diag] Post-Phase3: 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"); } @@ -964,6 +1179,7 @@ bool Pipeline::synthesize_prompt_codes_locked(const PipelineParams & params, con bool Pipeline::synthesize_streaming_raw(const PipelineParams & params, AudioData & ref_audio, StreamingSink & sink) { std::lock_guard lock(synthesize_mutex_); + std::vector ref_codes; int32_t T_prompt = 0; double ref_encode_ms = 0.0; @@ -975,31 +1191,18 @@ bool Pipeline::synthesize_streaming_raw(const PipelineParams & params, AudioData return false; } - std::thread pre_restore_thread; - bool pre_restore_started = false; - - if (params.enable_vram_swap && params.is_persistent && - model_prefers_gpu_ && !model().is_weights_on_gpu()) { - safe_print_ln("[Pipeline] Hot-Swap: Pre-fetching Slow-AR to VRAM in background..."); - pre_restore_started = true; - pre_restore_thread = std::thread([this]() { - model().acquire_compute_resources(); - model().restore_weights_to_gpu(); - }); + if (params.enable_vram_swap) { + if (pending_offload_thread_.joinable()) { + pending_offload_thread_.join(); + } } if (!resolve_reference_prompt_locked(params, ref_audio, ref_codes, T_prompt, effective_prompt_text, ref_encode_ms)) { - if (pre_restore_started) pre_restore_thread.join(); sink.on_error("Failed to resolve reference prompt"); return false; } - if (pre_restore_started) { - pre_restore_thread.join(); - safe_print_ln("[Pipeline] Hot-Swap: Background pre-fetch complete."); - } - PipelineParams effective_params = params; effective_params.prompt_text = std::move(effective_prompt_text); @@ -1042,9 +1245,15 @@ bool Pipeline::synthesize_streaming_prompt_codes_locked(const PipelineParams & p model().acquire_compute_resources(); model().restore_weights_to_gpu(); } - if (codec_prefers_gpu_ && !codec().is_weights_on_gpu()) { - safe_print_ln("[Pipeline] Streaming: Restoring Audio Codec to VRAM..."); - codec().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(); + } } } @@ -1271,16 +1480,18 @@ bool Pipeline::synthesize_streaming_prompt_codes_locked(const PipelineParams & p 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_weights_on_gpu()) codec().free_gpu_weights(); + 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."); diff --git a/src/s2_server.cpp b/src/s2_server.cpp index 08fea2d..f23aa45 100644 --- a/src/s2_server.cpp +++ b/src/s2_server.cpp @@ -234,6 +234,7 @@ static bool synthesize_segmented_to_sink(s2::Pipeline & pipeline, s2::PipelineParams segment_params = base_params; segment_params.text = segments[i]; segment_params.prompt_text = effective_prompt_text; + segment_params.more_segments_pending = (i + 1 < segments.size()); std::vector audio_out; const bool ok = pipeline.synthesize_with_prompt_codes( From 0dc1aeec486b829d63feafbe76409a03b8ae647f Mon Sep 17 00:00:00 2001 From: Skyrion9 Date: Mon, 3 Aug 2026 17:53:37 +0300 Subject: [PATCH 5/5] feat: KV cache reuse and prompt prefill caching 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. --- include/s2_generate.h | 3 +- include/s2_model.h | 10 +++ include/s2_pipeline.h | 18 ++++++ src/main.cpp | 4 ++ src/s2_generate.cpp | 56 ++++++++++------- src/s2_model.cpp | 74 ++++++++++++++++++++++ src/s2_pipeline.cpp | 141 +++++++++++++++++++++++++++++++++++++++--- 7 files changed, 273 insertions(+), 33 deletions(-) diff --git a/include/s2_generate.h b/include/s2_generate.h index c6b7dd4..8c9c606 100644 --- a/include/s2_generate.h +++ b/include/s2_generate.h @@ -41,7 +41,8 @@ GenerateResult generate( SlowARModel & model, const TokenizerConfig & config, const PromptTensor & prompt, - const GenerateParams & params + const GenerateParams & params, + const StepResult * initial_state = nullptr ); } diff --git a/include/s2_model.h b/include/s2_model.h index 338b014..d165772 100644 --- a/include/s2_model.h +++ b/include/s2_model.h @@ -109,6 +109,16 @@ class SlowARModel { void clear_kv_cache(); + void reset_kv_cache(); + int32_t kv_max_seq_len() const { return max_seq_len_; } + int32_t n_past() const { return n_past_; } + void set_n_past(int32_t n) { n_past_ = n; } + bool save_kv_state(std::vector & k_out, std::vector & v_out, + int32_t n_positions); + bool restore_kv_state(const std::vector & k_data, + const std::vector & v_data, + int32_t n_past); + MappedFile& mapped_file() { return mapped_gguf_; } bool allocate_and_load_weights(); diff --git a/include/s2_pipeline.h b/include/s2_pipeline.h index d930230..8a60deb 100644 --- a/include/s2_pipeline.h +++ b/include/s2_pipeline.h @@ -53,6 +53,8 @@ struct PipelineParams { bool enable_hot_swap = false; bool is_persistent = false; bool more_segments_pending = false; + bool enable_kv_reuse = true; + bool kv_cache_vram = false; }; class Pipeline { @@ -120,6 +122,22 @@ class Pipeline { bool initialized_ = false; bool model_prefers_gpu_ = false; bool codec_prefers_gpu_ = false; + + struct PrefillCacheEntry { + std::string cache_key; + int32_t n_past = 0; + int32_t max_seq_len = 0; + StepResult state; + std::vector k_data; + std::vector v_data; + bool vram_resident = false; + bool valid = false; + }; + PrefillCacheEntry prefill_cache_; + + static std::string compute_prefill_cache_key(const PipelineParams & params, + const int32_t * ref_codes, + int32_t T_prompt); }; } diff --git a/src/main.cpp b/src/main.cpp index eddcf4a..3abd60d 100644 --- a/src/main.cpp +++ b/src/main.cpp @@ -94,6 +94,8 @@ void print_uso() { safe_print(" --codebook-cpu Force codebook_embeddings on CPU even at full GPU offload (saves ~56 MB VRAM at Q4, costs ~4-5 ms/frame)\n"); safe_print(" --no-vram-swap Disable phase-gated VRAM swapping between Slow-AR and Codec (keeps both in VRAM simultaneously instead)\n"); safe_print(" --hot-swap Aggressively evict weights and OS page cache after each server request (~100 MB RAM, ~25 MB VRAM idle)\n"); + safe_print(" --no-kv-reuse Disable KV cache reuse and prompt prefill caching (old per-request free/realloc behavior)\n"); + safe_print(" --kv-cache-vram Pin prefill cache in VRAM between requests (default: system RAM, lower VRAM usage)\n"); safe_print(" --stream-file Write output WAV through the streaming path\n"); safe_print(" --stream-decode-stride Decode cadence in frames (0 = auto: server 4, file/offline 16)\n"); safe_print(" --codec-context-frames Override codec decode history (lower uses less VRAM, default: auto)\n"); @@ -203,6 +205,8 @@ int main(int argc, char** argv) { else if (arg == "--codebook-cpu") { params.codebook_embeddings_cpu = true; } else if (arg == "--no-vram-swap") { params.enable_vram_swap = false; } else if (arg == "--hot-swap") { params.enable_hot_swap = true; } + else if (arg == "--no-kv-reuse") { params.enable_kv_reuse = false; } + else if (arg == "--kv-cache-vram") { params.kv_cache_vram = true; } else if (arg == "--stream-file") { use_stream_file = true; } else if (arg == "--stream-decode-stride") { if (i+1 < argc) { diff --git a/src/s2_generate.cpp b/src/s2_generate.cpp index 6ee4c83..dcef695 100755 --- a/src/s2_generate.cpp +++ b/src/s2_generate.cpp @@ -12,8 +12,9 @@ GenerateResult generate( SlowARModel & model, const TokenizerConfig & config, const PromptTensor & prompt, - const GenerateParams & params -) { + const GenerateParams & params, + const StepResult * initial_state) +{ const auto generate_t0 = std::chrono::steady_clock::now(); GenerateResult out; out.num_codebooks = model.hparams().num_codebooks; @@ -34,25 +35,31 @@ GenerateResult generate( sem_mask[im_end_id] = 0.0f; } - const int32_t rows = prompt.rows; - const int32_t cols = prompt.cols; - std::vector prompt_tm(static_cast(rows) * cols); - for (int32_t r = 0; r < rows; ++r) { - for (int32_t c = 0; c < cols; ++c) { - prompt_tm[static_cast(c) * rows + r] = prompt.data[static_cast(r) * cols + c]; - } - } - 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; + } else { + const int32_t rows = prompt.rows; + const int32_t cols = prompt.cols; + std::vector prompt_tm(static_cast(rows) * cols); + for (int32_t r = 0; r < rows; ++r) + for (int32_t c = 0; c < cols; ++c) + prompt_tm[static_cast(c) * rows + r] = + prompt.data[static_cast(r) * cols + c]; + + 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; + } + const auto prefill_t1 = std::chrono::steady_clock::now(); + prefill_ms = std::chrono::duration(prefill_t1 - prefill_t0).count(); } - const auto prefill_t1 = std::chrono::steady_clock::now(); auto apply_mask_and_sample = [&](const std::vector & logits, bool block_im_end) -> int32_t { @@ -183,14 +190,17 @@ GenerateResult generate( } if (params.verbose && log_enabled(LogLevel::Info)) { - const auto loop_t1 = std::chrono::steady_clock::now(); + const auto loop_t1 = std::chrono::steady_clock::now(); const auto generate_t1 = std::chrono::steady_clock::now(); - const double prefill_ms = std::chrono::duration(prefill_t1 - prefill_t0).count(); - const double loop_ms = std::chrono::duration(loop_t1 - loop_t0).count(); + + const double loop_ms = std::chrono::duration( + loop_t1 - loop_t0).count(); const double total_ms = std::chrono::duration(generate_t1 - generate_t0).count(); const double ms_per_frame = out.n_frames > 0 ? (loop_ms / out.n_frames) : 0.0; std::cout << std::endl; - std::cout << "[Generate] Done: " << out.n_frames << " frames generated." << std::endl; + std::cout << "[Generate] Done: " << out.n_frames + << " frames generated." + << (initial_state ? " (prefill cached)" : "") << std::endl; std::cout << "[Metrics] Generate: prefill=" << prefill_ms << " ms, loop=" << loop_ms << " ms, total=" << total_ms diff --git a/src/s2_model.cpp b/src/s2_model.cpp index cc6a185..a3912a2 100755 --- a/src/s2_model.cpp +++ b/src/s2_model.cpp @@ -660,6 +660,80 @@ void SlowARModel::clear_kv_cache() { max_seq_len_ = 0; n_past_ = 0; } +void SlowARModel::reset_kv_cache() { + if (memory_k_ && kv_buf_) { + ggml_backend_tensor_memset(memory_k_, 0, 0, ggml_nbytes(memory_k_)); + ggml_backend_tensor_memset(memory_v_, 0, 0, ggml_nbytes(memory_v_)); + } + n_past_ = 0; +} + +bool SlowARModel::save_kv_state(std::vector & k_out, + std::vector & v_out, + int32_t n_positions) { + if (!memory_k_ || !memory_v_ || n_positions <= 0) return false; + if (n_positions > max_seq_len_) return false; + + const int32_t n_layer = hparams_.block_count; + const int32_t n_head_kv = hparams_.head_count_kv; + + const size_t pos_bytes = static_cast(n_positions) * memory_k_->nb[1]; + const size_t out_bytes = static_cast(n_layer) * n_head_kv * pos_bytes; + + const size_t full_bytes = ggml_nbytes(memory_k_); + std::vector k_full(full_bytes); + std::vector v_full(full_bytes); + ggml_backend_tensor_get(memory_k_, k_full.data(), 0, full_bytes); + ggml_backend_tensor_get(memory_v_, v_full.data(), 0, full_bytes); + + k_out.resize(out_bytes); + v_out.resize(out_bytes); + size_t out_offset = 0; + for (int32_t l = 0; l < n_layer; ++l) { + for (int32_t h = 0; h < n_head_kv; ++h) { + const size_t src = static_cast(l) * memory_k_->nb[3] + + static_cast(h) * memory_k_->nb[2]; + std::memcpy(k_out.data() + out_offset, k_full.data() + src, pos_bytes); + std::memcpy(v_out.data() + out_offset, v_full.data() + src, pos_bytes); + out_offset += pos_bytes; + } + } + return true; +} + +bool SlowARModel::restore_kv_state(const std::vector & k_data, + const std::vector & v_data, + int32_t n_past) { + if (!memory_k_ || !memory_v_) return false; + + const int32_t n_layer = hparams_.block_count; + const int32_t n_head_kv = hparams_.head_count_kv; + const size_t pos_bytes = static_cast(n_past) * memory_k_->nb[1]; + const size_t expected = static_cast(n_layer) * n_head_kv * pos_bytes; + if (k_data.size() != expected || v_data.size() != expected) return false; + + const size_t full_bytes = ggml_nbytes(memory_k_); + std::vector k_full(full_bytes, 0); + std::vector v_full(full_bytes, 0); + + size_t in_offset = 0; + for (int32_t l = 0; l < n_layer; ++l) { + for (int32_t h = 0; h < n_head_kv; ++h) { + const size_t dst = static_cast(l) * memory_k_->nb[3] + + static_cast(h) * memory_k_->nb[2]; + std::memcpy(k_full.data() + dst, k_data.data() + in_offset, pos_bytes); + std::memcpy(v_full.data() + dst, v_data.data() + in_offset, pos_bytes); + in_offset += pos_bytes; + } + } + + ggml_backend_tensor_set(memory_k_, k_full.data(), 0, full_bytes); + ggml_backend_tensor_set(memory_v_, v_full.data(), 0, full_bytes); + + n_past_ = n_past; + return true; +} + bool SlowARModel::prefill_fast(const std::vector & flat_tokens, int32_t n_tokens, int32_t n_threads, StepResult & result) { return eval_cached(flat_tokens, n_tokens, n_threads, result); diff --git a/src/s2_pipeline.cpp b/src/s2_pipeline.cpp index 1447f2f..0acfab8 100644 --- a/src/s2_pipeline.cpp +++ b/src/s2_pipeline.cpp @@ -302,6 +302,25 @@ static void sync_tokenizer_config_from_model(Tokenizer& tokenizer, const SlowARM if (hp.vocab_size > 0) tc.vocab_size = hp.vocab_size; } +std::string Pipeline::compute_prefill_cache_key(const PipelineParams & params, + const int32_t * ref_codes, + int32_t T_prompt) { + std::string key; + if (!params.voice_id.empty()) { + key = "voice:" + params.voice_id; + } else if (ref_codes && T_prompt > 0) { + key = "prompt:" + params.prompt_text; + const int32_t n = std::min(T_prompt, 16); + for (int32_t t = 0; t < n; ++t) + key += "," + std::to_string(ref_codes[t]); + } else { + key = "noprompt"; + } + + key += "|" + params.text; + return key; +} + Pipeline::Pipeline() {} Pipeline::~Pipeline() { if (pending_offload_thread_.joinable()) { @@ -863,15 +882,30 @@ bool Pipeline::synthesize_prompt_codes_locked(const PipelineParams & params, con ref_codes, num_codebooks, T_prompt); int32_t max_seq_len = prompt.cols + params.gen.max_new_tokens; - model().clear_kv_cache(); + 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(); std::thread kv_init_thread; bool kv_init_ok = true; - kv_init_thread = std::thread([&]() { - kv_init_ok = model().init_kv_cache(max_seq_len); - }); + 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(); @@ -889,6 +923,77 @@ bool Pipeline::synthesize_prompt_codes_locked(const PipelineParams & params, con const auto kv_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 (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 + ")"); + } + } + } + + StepResult prefill_state; + double prefill_ms = 0.0; + std::thread kv_save_thread; + + if (!prefill_hit) { + const int32_t rows = prompt.rows; + const int32_t cols = prompt.cols; + std::vector prompt_tm(static_cast(rows) * cols); + for (int32_t r = 0; r < rows; ++r) + for (int32_t c = 0; c < cols; ++c) + prompt_tm[static_cast(c) * rows + r] = + prompt.data[static_cast(r) * cols + c]; + + safe_print_ln("[Generate] Prefilling " + std::to_string(prompt.cols) + " tokens..."); + const auto pf_t0 = std::chrono::steady_clock::now(); + if (!model().prefill_fast(prompt_tm, prompt.cols, params.gen.n_threads, prefill_state)) { + safe_print_error_ln("Pipeline error: prefill failed."); + return false; + } + const auto pf_t1 = std::chrono::steady_clock::now(); + prefill_ms = std::chrono::duration(pf_t1 - pf_t0).count(); + + if (kv_reuse) { + prefill_cache_.cache_key = cache_key; + prefill_cache_.n_past = model().n_past(); + prefill_cache_.max_seq_len = model().kv_max_seq_len(); + prefill_cache_.state = prefill_state; + prefill_cache_.vram_resident = false; + prefill_cache_.k_data.clear(); + prefill_cache_.v_data.clear(); + + 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)")); + } + } + const bool can_overlap_decode = model_prefers_gpu_ && !codec_prefers_gpu_; @@ -1009,10 +1114,15 @@ bool Pipeline::synthesize_prompt_codes_locked(const PipelineParams & params, con }; const auto gen_t0 = std::chrono::steady_clock::now(); - res = generate(model(), tokenizer().config(), prompt, gen_params); + 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(gen_t1 - gen_t0).count(); + if (kv_save_thread.joinable()) { + kv_save_thread.join(); + } + { std::lock_guard lock(decode_mtx); gen_done.store(true); @@ -1038,7 +1148,7 @@ bool Pipeline::synthesize_prompt_codes_locked(const PipelineParams & params, con if (!params.more_segments_pending) { safe_print_ln("[Pipeline] Freeing Slow-AR from VRAM (request complete)..."); model().free_gpu_weights(); - model().free_compute_buffers(); + if (!keep_kv_on_gpu) model().free_compute_buffers(); safe_print_ln("[VRAM Diag] Post-SlowAR 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"); } } else { @@ -1071,10 +1181,15 @@ bool Pipeline::synthesize_prompt_codes_locked(const PipelineParams & params, con } else { const auto gen_t0 = std::chrono::steady_clock::now(); - res = generate(model(), tokenizer().config(), prompt, gen_params); + 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(gen_t1 - gen_t0).count(); + if (kv_save_thread.joinable()) { + kv_save_thread.join(); + } + if (res.n_frames == 0) { safe_print_error_ln("Pipeline error: generation produced no frames."); return false; @@ -1117,7 +1232,7 @@ bool Pipeline::synthesize_prompt_codes_locked(const PipelineParams & params, con if (!params.more_segments_pending) { model().free_gpu_weights(); - model().free_compute_buffers(); + if (!keep_kv_on_gpu) model().free_compute_buffers(); } } } else { @@ -1133,7 +1248,13 @@ bool Pipeline::synthesize_prompt_codes_locked(const PipelineParams & params, con } } - model().clear_kv_cache(); + if (!params.more_segments_pending) { + if (!kv_reuse) { + model().clear_kv_cache(); + } else if (!keep_kv_on_gpu) { + model().set_n_past(0); + } + } const auto synth_t1 = std::chrono::steady_clock::now(); @@ -1152,6 +1273,7 @@ bool Pipeline::synthesize_prompt_codes_locked(const PipelineParams & params, con ", audio_s=" + std::to_string(audio_seconds) + ", ref_encode=" + std::to_string(ref_encode_ms) + " ms, kv_init=" + std::to_string(kv_ms) + + " 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) + @@ -1159,6 +1281,7 @@ bool Pipeline::synthesize_prompt_codes_locked(const PipelineParams & params, con ", decode_stride=" + std::to_string(offline_decode_stride_frames) + " frames" + (can_overlap_decode ? ", decode_mode=overlapped" : ", decode_mode=sequential") + + (prefill_hit ? ", prefill=cached" : ", prefill=computed") + (params.more_segments_pending ? ", vram=held" : "") + ", total=" + std::to_string(total_ms) + " ms, gen_avg=" + std::to_string(gen_ms_per_frame) +