Enhance FireRedTTS3 with streaming, batching, and server integration

CMakeLists.txt

@@ -1531,6 +1531,7 @@ audiocpp_add_model(fireredtts3
     SOURCES
         src/models/fireredtts3/assets.cpp
         src/models/fireredtts3/ar.cpp
+        src/models/fireredtts3/batch_scheduler.cpp
         src/models/fireredtts3/flow.cpp
         src/models/fireredtts3/pipeline.cpp
         src/models/fireredtts3/redae.cpp
@@ -1539,6 +1540,7 @@ audiocpp_add_model(fireredtts3
     INCLUDES
         engine/models/fireredtts3/assets.h
         engine/models/fireredtts3/ar.h
+        engine/models/fireredtts3/batch_scheduler.h
         engine/models/fireredtts3/flow.h
         engine/models/fireredtts3/pipeline.h
         engine/models/fireredtts3/redae.h

app/server/config.cpp

@@ -306,6 +306,14 @@ ServerConfig load_server_config(const std::filesystem::path & path) {
         model.load_options = options_from_object(item.find("load_options"));
         model.session_options = options_from_object(item.find("session_options"));
         model.default_request_options = options_from_object(item.find("default_request_options"));
+        if (const auto * value = item.find("instance_count")) {
+            const int n = value->as_i64();
+            if (n < 1 || n > 64) {
+                throw std::runtime_error(
+                    "instance_count for model " + model.id + " must be in [1, 64]");
+            }
+            model.instance_count = n;
+        }
         if (const auto * voice_presets = item.find("voice_presets")) {
             if (!voice_presets->is_object()) {
                 throw std::runtime_error("voice_presets for model " + model.id + " must be an object");

app/server/config.h

@@ -50,6 +50,10 @@ struct ServerModelConfig {
     // magnitude (a short TTS clip vs. minutes of music generation), so one fleet-wide
     // bound is either too tight for the slow models or useless for the fast ones.
     std::optional<int> busy_timeout_ms;
+    // Number of concurrent session instances for this model (a runtime pool).
+    // Each instance has its own graph arena + reference cache, enabling true
+    // multi-request concurrency within one loaded model. Default 1 (serialized).
+    int instance_count = 1;
     // Only meaningful for a streaming model reachable over the live-ingest route;
     // ignored otherwise, since no other route delivers its body incrementally.
     LiveIngestOverrides live_ingest;

app/server/firered_server.json

@@ -0,0 +1,27 @@
+{
+  "host": "0.0.0.0",
+  "port": 8007,
+  "backend": "cuda",
+  "device": 0,
+  "threads": 4,
+  "lazy_load": true,
+  "busy_timeout_ms": 0,
+  "models": [
+    {
+      "id": "firered-base",
+      "family": "fireredtts3",
+      "path": "/data/megastore/Projects/DuJing/models/FireRedTTS3-Base-GGUF/fireredtts3-base-q8_0.gguf",
+      "model_spec_override": "/data/megastore/Projects/DuJing/code/audio.cpp/model_specs/fireredtts3.json",
+      "task": "clon",
+      "mode": "streaming",
+      "instance_count": 3,
+      "session_options": {
+        "fireredtts3.reference_cache_slots": "8",
+        "fireredtts3.chunk_sizes": "3,12,12,12,12,12,12,12,12,12,12,12,12,12,12,12"
+      },
+      "default_request_options": {
+        "num_inference_steps": "5"
+      }
+    }
+  ]
+}

app/server/runtime.h

@@ -13,6 +13,10 @@
 #include "engine/framework/runtime/model.h"
 #include "engine/framework/runtime/session.h"
 
+#include <deque>
+#include <mutex>
+#include <thread>
+
 #include <atomic>
 #include <cstdint>
 #include <filesystem>
@@ -44,6 +48,30 @@ class ServerState final : public IHttpHandler {
     LiveIngestLimits live_ingest_limits(const HttpRequest & request) const override;
 
 private:
+    struct LoadedModel;
+
+    // session 池借用锁（RAII）：借一个空闲 session 实例，析构归还。
+    class SessionPoolLock {
+    public:
+        SessionPoolLock() = default;
+        SessionPoolLock(LoadedModel & model, size_t index);
+        SessionPoolLock(SessionPoolLock && other) noexcept;
+        SessionPoolLock & operator=(SessionPoolLock && other) noexcept;
+        SessionPoolLock(const SessionPoolLock &) = delete;
+        SessionPoolLock & operator=(const SessionPoolLock &) = delete;
+        ~SessionPoolLock();
+
+        // 借到的 session 下标。public 且是唯一存储：构造函数/move/赋值都写这里，
+        // release() 也读这里归还。曾有个 private index_ 与 public index 并存，
+        // 构造函数只写 index_ 而调用处全读 public index → 恒为 0 → 所有并发请求
+        // 都绑 session 0（跨请求串音/截断/double free 根因）。
+        size_t index = 0;
+
+    private:
+        void release();
+        LoadedModel * model_ = nullptr;
+    };
+
     struct LoadedModel {
         struct RuntimeVoicePreset {
             std::optional<std::string> voice_id;
@@ -54,10 +82,14 @@ class ServerState final : public IHttpHandler {
         ServerModelConfig config;
         engine::runtime::TaskSpec task;
         std::unique_ptr<engine::runtime::ILoadedVoiceModel> model;
-        std::unique_ptr<engine::runtime::IVoiceTaskSession> session;
+        // 并发 session 池：每个实例独立 graph arena + reference cache。
+        std::vector<std::unique_ptr<engine::runtime::IVoiceTaskSession>> sessions;
         engine::runtime::IOfflineVoiceTaskSession * offline = nullptr;
         engine::runtime::IStreamingVoiceTaskSession * streaming = nullptr;
         std::atomic<bool> loaded{false};
+        // 空闲 session 索引队列（受 pool_mutex 保护）
+        std::mutex pool_mutex;
+        std::deque<size_t> free_sessions;
         // Steady-clock ms of the most recent load or run of this model. Orders
         // eviction when max_loaded_models forces an unload: the least recently
         // used idle model goes first.
@@ -85,6 +117,9 @@ class ServerState final : public IHttpHandler {
     // (-> HTTP 503) once the effective timeout has elapsed.
     BusyGuard::Lock acquire_model_run(LoadedModel & model, std::optional<int> request_timeout_ms);
 
+    // 从 session 池借一个空闲实例（真并发）；析构自动归还。池满时阻塞/超时。
+    SessionPoolLock borrow_session(LoadedModel & model, std::optional<int> request_timeout_ms);
+
     // Server policy for this model: its own busy_timeout_ms if set, else the
     // top-level config value.
     engine::runtime::RunMode model_run_mode(const LoadedModel & model) const;

app/streaming/streaming.cpp

@@ -138,6 +138,9 @@ engine::runtime::TaskResult run_stream(
         return result;
     } catch (...) {
         session.set_stream_event_sink(nullptr);
+        // 中断/异常请求：reset() 清掉可能遗留的流式状态（含 scheduler 中仍 Active 的 slot），
+        // 避免该 session 带活 slot 归还池中、被下一请求复用而串音/崩溃。
+        session.reset();
         throw;
     }
 }

include/engine/framework/audio/istft_graph.h

@@ -67,6 +67,16 @@ class HostLogMagnitudePhaseISTFT {
         const std::vector<float> & log_magnitude_phase,
         const std::vector<float> & window);
 
+    // --- 增量 overlap-add iSTFT（用于流式）---
+    // 按块喂入 log-magnitude+phase 帧，内部累积 overlap-add，
+    // 返回"已能被窗口包络完整覆盖"的音频样本（块间平滑衔接）。
+    // 首次 append 自动初始化；finish 收尾 flush 尾部并复位。
+    std::vector<float> append_incremental(
+        const std::vector<float> & log_magnitude_phase,
+        int64_t frames,
+        const std::vector<float> & window);
+    std::vector<float> finish_incremental();
+
 private:
     class Impl;
     std::unique_ptr<Impl> impl_;

include/engine/framework/codecs/redae_codec_runtime.h

@@ -1,12 +1,15 @@
 #pragma once
 
 #include "engine/framework/assets/tensor_source.h"
+#include "engine/framework/audio/istft_graph.h"
 #include "engine/framework/core/execution_context.h"
+#include "engine/framework/runtime/kv_cache.h"
 #include "engine/framework/runtime/session.h"
 
 #include <cstddef>
 #include <cstdint>
 #include <memory>
+#include <optional>
 #include <string>
 #include <vector>
 
@@ -76,6 +79,24 @@ class RedAeCodecRuntime {
 
     std::vector<float> encode(const std::vector<float> & audio_24k);
     runtime::AudioBuffer decode(const std::vector<float> & latents);
+
+    // --- 增量解码（流式）---
+    // 每个并发 slot 独立持有解码状态（decoder KV + 增量 iSTFT），
+    // 使得多 slot 交错的增量解码互不干扰。
+    struct DecodeState {
+        std::optional<engine::runtime::TransformerKVState> dec_state;
+        int64_t dec_qwen_frames = 0;
+        std::unique_ptr<audio::HostLogMagnitudePhaseISTFT> inc_istft;
+        int64_t inc_istft_frames = 0;
+    };
+    // 重置解码器 KV 状态（每次新请求开始时调用）。
+    void decode_reset(DecodeState & state);
+    // 解码一批 latent（chunk），返回该块对应的音频（float32 24k mono）。
+    // 内部用 decoder Qwen 的 KV 缓存跨块保持上下文，并用增量 iSTFT 逐块输出。
+    runtime::AudioBuffer decode_incremental(DecodeState & state, const std::vector<float> & latents);
+    // flush 增量 iSTFT 尾部样本（生成结束时调用）。
+    runtime::AudioBuffer flush_incremental(DecodeState & state);
+
     void release_runtime_graphs();
 
 private:

include/engine/framework/modules/transformers/qwen_causal_decode_runtime.h

@@ -68,6 +68,14 @@ class QwenCausalDecodeRuntime {
 
     QwenCausalPrefillResult prefill_tokens(const std::vector<int32_t> & token_ids);
     QwenCausalPrefillResult prefill_embeddings(const std::vector<float> & embeddings, int64_t steps);
+    // 固定 graph 的 padded prefill：graph 按 padded_steps 建一次并复用（避免因
+    // 不同 steps 重建 prefill graph 破坏 CUDA pool 逆序约束）。embeddings 必须
+    // 是 padded_steps × hidden（padding 零），只有前 valid_steps 参与位置/mask。
+    // 返回 state 只含 valid_steps（截断）。若 padded_steps < 当前已建 steps 则复用。
+    QwenCausalPrefillResult prefill_embeddings_padded(
+        const std::vector<float> & embeddings,
+        int64_t padded_steps,
+        int64_t valid_steps);
 
     QwenCausalBatchedPrefillResult prefill_tokens_batched(
         const std::vector<int32_t> & token_ids,
@@ -90,9 +98,16 @@ class QwenCausalDecodeRuntime {
         const runtime::TransformerBatchedKVState & state,
         int64_t required_cache_steps);
     QwenCausalDecodeStepResult decode_tokens_batched(const std::vector<int32_t> & tokens);
+    // 每步 batched decode。active_mask（可选，长度==batch_size）：只有置 1 的行才
+    // 真正前进一步、mask 才暴露其前缀；置 0 的行全 -inf（不读自身 stale KV）、
+    // 不 advance —— 非活跃行彻底 inert，杜绝"冻结行携带上一请求 stale KV 参与
+    // decode"导致的跨请求串音。
     QwenCausalDecodeStepResult decode_embeddings_batched(
         const std::vector<float> & embeddings,
-        int64_t batch_size);
+        int64_t batch_size,
+        const std::vector<uint8_t> & active_mask = {});
+    // 冻结/重置某 batch 行的解码位置（非活跃行 end=0，mask 全 -inf）。
+    void set_batched_member_end(int64_t batch, int64_t end);
 
     // Snapshot of the batched decode KV cache (host vectors), suitable for
     // replication and re-import via start_decode_*_batched with a different
@@ -104,6 +119,9 @@ class QwenCausalDecodeRuntime {
     int64_t decode_valid_steps() const noexcept;
     void release_runtime_graphs();
 
+    // [DIAG] 每 batch 行的当前解码结束位置（member_ends_ 拷贝；未启动则空）。
+    std::vector<int64_t> batched_member_ends() const;
+
 private:
     class Impl;
     std::unique_ptr<Impl> impl_;

include/engine/framework/modules/transformers/qwen_causal_decoder.h

@@ -182,12 +182,15 @@ void write_qwen_cached_step_mask(
     int64_t visible_prefix_steps,
     int64_t current_slot);
 
+// active_mask（可选，长度==batch_size）：置 0 的行整行 -inf（即使其 cache 段残留
+// stale KV 也不 attend）；nullptr = 全部活跃（原行为）。
 void write_qwen_batched_cached_step_mask(
     ggml_tensor * tensor,
     std::vector<ggml_fp16_t> & scratch,
     int64_t batch_size,
     int64_t mask_steps,
-    int64_t visible_prefix_steps,
-    int64_t current_slot);
+    const std::vector<int64_t> & member_ends,
+    const std::vector<int32_t> & cache_slots,
+    const std::vector<uint8_t> * active_mask = nullptr);
 
 }  // namespace engine::modules

include/engine/framework/runtime/kv_cache.h

@@ -79,6 +79,8 @@ struct BatchedKVLayerState {
 struct TransformerBatchedKVState {
     int64_t batch_size = 0;
     int64_t current_end = 0;
+    // 可选的 per-member 结束位置（大小 == batch_size）。空 = 均匀（current_end 生效）。
+    std::vector<int64_t> current_ends;
     std::vector<BatchedKVLayerState> layers;
 };
 
@@ -109,6 +111,13 @@ class TransformerBatchedKVCache {
     int64_t current_end() const noexcept;
     int64_t cache_steps() const noexcept;
 
+    // --- per-member 结束位置（不同序列可处于不同位置）---
+    int64_t member_end(int64_t batch) const noexcept;
+    void set_member_end(int64_t batch, int64_t end) noexcept;
+    void advance_member(int64_t batch, int64_t steps) noexcept;
+    // 返回 per-member ends（空=均匀，调用方回退到 cache_slots）
+    const std::vector<int64_t> & member_ends_for_mask() const noexcept { return member_ends_; }
+
 private:
     struct LayerCache {
         core::TensorValue key_tensor;
@@ -122,6 +131,8 @@ class TransformerBatchedKVCache {
     int64_t row_elems_ = 0;
     int64_t valid_steps_ = 0;
     int64_t current_end_ = 0;
+    // per-member 结束位置；空 = 均匀（用 current_end_）
+    std::vector<int64_t> member_ends_;
     TransformerKVCacheOptions options_;
     std::vector<LayerCache> layers_;
 };

include/engine/framework/runtime/session_base.h

@@ -9,6 +9,7 @@
 #include "engine/framework/runtime/session.h"
 #include "engine/framework/runtime/workspace.h"
 
+#include <memory>
 #include <stdexcept>
 #include <string>
 #include <string_view>
@@ -18,6 +19,13 @@ namespace engine::runtime {
 class RuntimeSessionBase {
 public:
     explicit RuntimeSessionBase(const SessionOptions & options);
+    // 共享 backend 模式（llama.cpp 单 context 多 slot）：`external_context` 非空时
+    // 本 session 不自建 ExecutionContext/backend，而是借用外部持有的那个（所有权在
+    // 调用方，通常是 model/scheduler 级，生命周期须长于本 session 及其 runtime）。
+    // nullptr = 原行为（每个 session 自建自己的 context）。
+    RuntimeSessionBase(
+        const SessionOptions & options,
+        std::shared_ptr<engine::core::ExecutionContext> external_context);
     virtual ~RuntimeSessionBase() = default;
 
 protected:
@@ -39,7 +47,9 @@ class RuntimeSessionBase {
 
 private:
     SessionOptions options_;
-    engine::core::ExecutionContext execution_context_;
+    // 本 session 的 backend context。默认自建（shared_ptr 持有）；共享模式外部传入。
+    // 借用的外部 context 同样以 shared_ptr 持有，保证它在本 session 存活期间不析构。
+    std::shared_ptr<engine::core::ExecutionContext> context_;
     ArtifactStore artifacts_;
     RuntimeCache cache_;
     RuntimeWorkspace workspace_;

include/engine/models/fireredtts3/ar.h

@@ -6,6 +6,7 @@
 #include "engine/models/fireredtts3/assets.h"
 
 #include <cstddef>
+#include <cstdint>
 #include <memory>
 #include <vector>
 
@@ -35,9 +36,24 @@ class FireRedArRuntime {
     std::vector<float> text_logits(const std::vector<float> & hidden);
 
     engine::modules::QwenCausalPrefillResult prefill_embeddings(const std::vector<float> & embeddings, int64_t steps);
+    engine::modules::QwenCausalPrefillResult prefill_embeddings_padded(
+        const std::vector<float> & embeddings, int64_t padded_steps, int64_t valid_steps);
     void start_decode_embeddings(const engine::runtime::TransformerKVState & state, int64_t required_cache_steps);
     engine::modules::QwenCausalDecodeStepResult decode_embedding(const std::vector<float> & embedding);
 
+    // --- batch（多 slot 并发推理）passthroughs ---
+    void start_decode_embeddings_batched(
+        const engine::runtime::TransformerBatchedKVState & state, int64_t required_cache_steps);
+    engine::modules::QwenCausalDecodeStepResult decode_embeddings_batched(
+        const std::vector<float> & embeddings, int64_t batch_size,
+        const std::vector<uint8_t> & active_mask = {});
+    engine::runtime::TransformerBatchedKVState export_batched_decode_state() const;
+    // 冻结/重置某 batch 行的解码位置：非活跃行应保持 end=0（mask 全 -inf，不参与 attention），
+    // 避免 run_batched_decode_step 对空行 advance_member 导致其位置递增、mask 污染活跃行。
+    void set_batched_member_end(int64_t batch, int64_t end);
+    // [DIAG] 当前 batched decode 各行的解码结束位置（未启动则空）。
+    std::vector<int64_t> batched_member_ends() const;
+
     void release_graphs();
     void release_backbone_graphs();