feat: mmap lazy loading and phase-gated VRAM swap

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)

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 <cstdint>
 #include <string>
 #include <vector>
+#include <unordered_set>
+#include <unordered_map>
 
 #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_; }

include/s2_mapped_file.h

@@ -0,0 +1,41 @@
+#pragma once
+
+#include <string>
+#include <cstddef>
+#include <cstdint>
+
+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<const uint8_t*>(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
+};
+
+}

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 <cstdint>
 #include <string>
 #include <unordered_set>
+#include <unordered_map>
 #include <vector>
 
 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<ggml_tensor *> & 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<int32_t> & flat_tokens,
         int32_t n_tokens, int32_t n_threads,
@@ -152,6 +170,18 @@ class SlowARModel {
 
     std::unordered_set<ggml_tensor *> weight_tensor_set_;
 
+    std::string gguf_path_;
+    size_t gguf_data_offset_ = 0;
+    std::unordered_map<ggml_tensor*, size_t> tensor_offsets_;
+    std::vector<ggml_tensor*> original_gpu_weights_;
+    std::vector<ggml_tensor*> 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_;
+
 };
 
 }

include/s2_pipeline.h

@@ -11,6 +11,7 @@
 #include <cstdint>
 #include <string>
 #include <mutex>
+#include <thread>
 
 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;
 };
 
 }

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 <n>  Decode cadence in frames (0 = auto: server 4, file/offline 16)\n");
     safe_print("  --codec-context-frames <n>  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");