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..96a9e56 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,29 @@ class AudioCodec { void clear_decode_cache(); + MappedFile& mapped_file(); + + bool restore_weights_to_gpu(); + + 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(); + + 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..b4619e6 --- /dev/null +++ b/include/s2_mapped_file.h @@ -0,0 +1,41 @@ +#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(); + void warm_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..338b014 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 { @@ -94,11 +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 read_tensor_data(const std::string & gguf_path, gguf_context * gguf_ctx); + 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_; } @@ -109,6 +109,24 @@ 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_; } + + 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, @@ -152,6 +170,18 @@ 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; + 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 19b3c3d..d930230 100644 --- a/include/s2_pipeline.h +++ b/include/s2_pipeline.h @@ -11,6 +11,7 @@ #include #include #include +#include namespace s2 { @@ -37,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; @@ -46,6 +49,10 @@ 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; + bool more_segments_pending = false; }; class Pipeline { @@ -108,8 +115,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..eddcf4a 100644 --- a/src/main.cpp +++ b/src/main.cpp @@ -90,6 +90,10 @@ 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"); 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 +190,20 @@ 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 == "--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; } else if (arg == "--stream-decode-stride") { if (i+1 < argc) { try { params.stream_decode_stride_frames = std::stoi(argv[++i]); } catch(...) {} @@ -309,6 +317,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_codec.cpp b/src/s2_codec.cpp index e4dc3bc..03aac86 100755 --- a/src/s2_codec.cpp +++ b/src/s2_codec.cpp @@ -20,6 +20,7 @@ #include #include #include +#include namespace s2 { @@ -61,9 +62,9 @@ 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; std::string tprefix; int32_t sample_rate = 0; @@ -103,6 +104,25 @@ struct AudioCodec::Impl { std::vector residual_vq; codec_decode_cache decode_cache; + std::string gguf_path; + size_t gguf_data_offset = 0; + std::vector original_gpu_weights; + std::vector original_cpu_weights; + 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) { @@ -115,6 +135,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); @@ -134,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; @@ -690,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(); @@ -873,103 +955,85 @@ 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_->encoder_weights.clear(); + impl_->decoder_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; + 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 { + 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; + impl_->encoder_on_gpu = false; + impl_->decoder_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 +1044,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 +1364,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 +1511,234 @@ 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_) 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); + 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; +} + +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_) 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() { + 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_->decoder_on_gpu || impl_->encoder_on_gpu) { + impl_->weights_allocated_ = true; + return true; + } + if (!impl_->mapped_gguf_.is_open()) return false; + S2_LOG_INFO_STREAM("[Codec] >>> Allocating and loading Audio Codec weights on demand..." << std::endl); + size_t b = 0; std::string e; + if (!allocate_codec_buffers(impl_->backend, impl_->all_codec_weights, impl_->model_buf, b, e)) { + std::cerr << "[Codec] Alloc failed: " << e << std::endl; return false; + } + const uint8_t* base = impl_->mapped_gguf_.data(); + for (ggml_tensor * t : impl_->all_codec_weights) { + auto it = impl_->tensor_offsets.find(t); + if (it != impl_->tensor_offsets.end()) + ggml_backend_tensor_set(t, base + impl_->gguf_data_offset + it->second, 0, ggml_nbytes(t)); + } + impl_->weights_allocated_ = true; + impl_->weights_on_gpu = !ggml_backend_is_cpu(impl_->backend); + impl_->encoder_on_gpu = impl_->weights_on_gpu; + impl_->decoder_on_gpu = impl_->weights_on_gpu; + return true; +} + +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..7a54afc --- /dev/null +++ b/src/s2_mapped_file.cpp @@ -0,0 +1,238 @@ +#include "../include/s2_mapped_file.h" + +#ifdef _WIN32 +#include +#include + +#else +#include +#include +#include +#include +#include + +#endif + +namespace s2 { + +bool MappedFile::open(const std::string& path) { + close(); + +#ifdef _WIN32 + 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); + 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 = 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; + } + + 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_SEQUENTIAL); + +#ifdef __linux__ + ::madvise(addr, size_, MADV_HUGEPAGE); + ::madvise(addr, size_, MADV_DONTDUMP); +#endif + + 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() { +#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) + 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 9ced623..cc6a185 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 @@ -105,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(); @@ -201,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_) { @@ -217,6 +226,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); @@ -310,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) { @@ -478,12 +492,30 @@ 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 && 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_; @@ -521,196 +553,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; - - 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); + weights_allocated_ = false; + weights_on_gpu_ = false; + - 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) { - +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; } - - 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 +1129,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..1447f2f 100644 --- a/src/s2_pipeline.cpp +++ b/src/s2_pipeline.cpp @@ -8,6 +8,8 @@ #include #include #include +#include +#include #ifdef __linux__ #include @@ -301,91 +303,10 @@ 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; +Pipeline::~Pipeline() { + if (pending_offload_thread_.joinable()) { + pending_offload_thread_.join(); } - - 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) { @@ -412,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; @@ -477,45 +398,71 @@ 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); + 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(); - gguf_free(shared_gguf); + const auto model_weights_t0 = std::chrono::steady_clock::now(); + const bool defer_weight_loading = params.enable_vram_swap && model().prefers_gpu(); - const auto codec_t1 = std::chrono::steady_clock::now(); + 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(); sync_tokenizer_config_from_model(tokenizer(), model()); initialized_ = true; + + model_prefers_gpu_ = model().prefers_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( "[Metrics] Init: tokenizer=" + @@ -524,7 +471,10 @@ 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()) + + (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"); @@ -593,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)) { @@ -603,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); @@ -779,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; @@ -789,6 +752,12 @@ bool Pipeline::synthesize_raw(const PipelineParams & params, AudioData & ref_aud return false; } + 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)) { return false; @@ -833,56 +802,342 @@ 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); - const int32_t num_codebooks = model().hparams().num_codebooks; + std::thread vram_phase1_thread; + bool vram_phase1_ok = true; + + if (params.enable_vram_swap) { + vram_phase1_thread = std::thread([this, ¶ms, &vram_phase1_ok]() { + if (model_prefers_gpu_ && !model().is_weights_on_gpu()) { + safe_print_ln("[Pipeline] Restoring Slow-AR to VRAM for generation..."); + model().acquire_compute_resources(); + if (!model().restore_weights_to_gpu()) { + safe_print_error_ln("Pipeline error: Slow-AR weight restore failed."); + vram_phase1_ok = false; + return; + } + safe_print_ln("[VRAM Diag] Post-SlowAR restore: Slow-AR=" + std::to_string(model().get_gpu_memory_usage_bytes() / 1024 / 1024) + " MB, Codec=" + std::to_string(codec().get_gpu_memory_usage_bytes() / 1024 / 1024) + " MB"); + } + if (!model_prefers_gpu_ && codec_prefers_gpu_ && !codec().is_decoder_on_gpu()) { + safe_print_ln("[Pipeline] Pre-loading Audio Codec decoder to VRAM (hiding behind CPU gen)..."); + codec().restore_decoder_weights(); + safe_print_ln("[VRAM Diag] Post-Decoder restore: Slow-AR=" + std::to_string(model().get_gpu_memory_usage_bytes() / 1024 / 1024) + " MB, Codec=" + std::to_string(codec().get_gpu_memory_usage_bytes() / 1024 / 1024) + " MB"); + } + + if (model_prefers_gpu_ && codec_prefers_gpu_) { + if (codec().is_encoder_on_gpu()) { + safe_print_ln("[Pipeline] Freeing codec encoder from VRAM (not needed during generation)..."); + codec().free_encoder_weights(); + } + if (codec().is_decoder_on_gpu()) { + safe_print_ln("[Pipeline] Freeing codec decoder from VRAM (not needed during generation)..."); + codec().free_decoder_weights(); + } + if (codec().is_weights_on_gpu()) { + safe_print_ln("[Pipeline] Freeing Audio Codec from VRAM for Slow-AR generation..."); + codec().free_gpu_weights(); + } + safe_print_ln("[VRAM Diag] Post-Codec free: Slow-AR=" + std::to_string(model().get_gpu_memory_usage_bytes() / 1024 / 1024) + " MB, Codec=" + std::to_string(codec().get_gpu_memory_usage_bytes() / 1024 / 1024) + " MB"); + } + + safe_print_ln("[VRAM Diag] End-Phase1: Slow-AR=" + + std::to_string(model().get_gpu_memory_usage_bytes() / 1024 / 1024) + " MB, Codec=" + + std::to_string(codec().get_gpu_memory_usage_bytes() / 1024 / 1024) + " MB"); + }); + } + + const int32_t num_codebooks = model().hparams().num_codebooks; PromptTensor prompt = build_prompt( tokenizer(), params.text, params.prompt_text, - ref_codes, - num_codebooks, T_prompt); - + ref_codes, num_codebooks, T_prompt); int32_t max_seq_len = prompt.cols + params.gen.max_new_tokens; + + model().clear_kv_cache(); + const auto kv_t0 = std::chrono::steady_clock::now(); - if (!model().init_kv_cache(max_seq_len)) { - safe_print_error_ln("Pipeline error: init_kv_cache failed."); + 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; } - 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(); - if (res.n_frames == 0) { - safe_print_error_ln("Pipeline error: generation produced no frames."); + 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 bool can_overlap_decode = + model_prefers_gpu_ && !codec_prefers_gpu_; + 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, - params.gen.n_threads, offline_decode_stride_frames, - params.codec_decode_context_frames, - audio_out, &decode_ms, &decode_batches)) { - safe_print_error_ln("Pipeline error: decode failed."); - return false; + + 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().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 (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); + } + + audio_out = std::move(audio_accum); + + } 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 (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_decoder_on_gpu()) codec().free_decoder_weights(); + model().free_gpu_weights(); + model().free_compute_buffers(); + } + } + + 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(); 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()) @@ -902,18 +1157,29 @@ 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) + ", 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; } 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; @@ -925,6 +1191,12 @@ bool Pipeline::synthesize_streaming_raw(const PipelineParams & params, AudioData return false; } + 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)) { sink.on_error("Failed to resolve reference prompt"); @@ -963,6 +1235,28 @@ 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_) { + 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(); + } + } + } + CodecDecodeCacheScope codec_decode_cache_scope(codec()); model().clear_kv_cache(); @@ -1176,6 +1470,33 @@ 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_decoder_on_gpu()) codec().free_decoder_weights(); + if (codec().is_encoder_on_gpu()) codec().free_encoder_weights(); + if (codec().is_weights_on_gpu()) codec().free_gpu_weights(); + model().mapped_file().drop_page_cache(); + codec().mapped_file().drop_page_cache(); + safe_print_ln("[Pipeline] Hot-Swap: Background VRAM & RAM free complete."); + }); + pending_offload_thread_ = std::move(offload_thread); + } else { + if (codec().is_decoder_on_gpu()) codec().free_decoder_weights(); + model().free_gpu_weights(); + model().free_compute_buffers(); + } + } else { + safe_print_ln("[Pipeline] Single-shot mode: Skipping post-stream VRAM restore."); + } + } return true; } 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(