From 7e4bcbeda13abb0b62e2399ef84605bb3136d50d Mon Sep 17 00:00:00 2001 From: dujing Date: Mon, 31 Aug 2026 21:22:31 +0800 Subject: [PATCH 01/12] Add FireRedTTS3 audiocpp_server config (firered-base, offline clone) - app/server/firered_server.json: model entry for FireRedTTS3-Base q8_0 GGUF, task=clon, mode=offline, CUDA device 0, reference_cache_slots=8. Co-Authored-By: Claude --- app/server/firered_server.json | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) create mode 100644 app/server/firered_server.json diff --git a/app/server/firered_server.json b/app/server/firered_server.json new file mode 100644 index 000000000..b92f183ad --- /dev/null +++ b/app/server/firered_server.json @@ -0,0 +1,21 @@ +{ + "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", + "task": "clon", + "mode": "offline", + "session_options": { + "fireredtts3.reference_cache_slots": "8" + } + } + ] +} From 254e1ee2b48a75d8d9c3eecb9d1c924663e9b5c3 Mon Sep 17 00:00:00 2001 From: dujing Date: Tue, 1 Sep 2026 10:43:29 +0800 Subject: [PATCH 02/12] Add incremental streaming to FireRedTTS3 Base (offline + streaming dual mode) - istft_graph: incremental overlap-add iSTFT (append_incremental/finish_incremental) for streaming audio output. - redae_codec_runtime: decode_incremental (chunked decoder Qwen with KV chaining via start_decode_embeddings/decode_embedding) + flush_incremental. - fireredtts3 redae wrapper: expose decode_reset/decode_incremental/flush_incremental. - fireredtts3 pipeline: begin_streaming -> FireRedTTS3StreamSession pull-based generator (AR loop emits one chunk per next_chunk(); chunk_patches config e.g. 3,12,12 -> first chunk ~0.5s, then ~2s). - fireredtts3 session: implement IStreamingVoiceTaskSession (PullEvents); parse fireredtts3.chunk_sizes session option; allow RunMode::Streaming. - model_specs/fireredtts3.json: add streaming mode + chunk_sizes session option. - app/server/firered_server.json: mode=streaming, chunk_sizes, model_spec_override. Verified: stream=True first chunk ~0.3s latency, then 2s chunks at ~0.62s intervals (RTF ~0.3); stream=False still returns whole utterance. Co-Authored-By: Claude --- app/server/firered_server.json | 6 +- include/engine/framework/audio/istft_graph.h | 10 ++ .../framework/codecs/redae_codec_runtime.h | 10 ++ include/engine/models/fireredtts3/pipeline.h | 21 +++ include/engine/models/fireredtts3/redae.h | 6 + include/engine/models/fireredtts3/session.h | 25 ++- model_specs/fireredtts3.json | 10 +- src/framework/audio/istft_graph.cpp | 154 ++++++++++++++++++ src/framework/codecs/redae_codec_runtime.cpp | 88 ++++++++++ src/models/fireredtts3/pipeline.cpp | 146 +++++++++++++++++ src/models/fireredtts3/redae.cpp | 24 +++ src/models/fireredtts3/session.cpp | 112 ++++++++++++- 12 files changed, 605 insertions(+), 7 deletions(-) diff --git a/app/server/firered_server.json b/app/server/firered_server.json index b92f183ad..cf17ef0b3 100644 --- a/app/server/firered_server.json +++ b/app/server/firered_server.json @@ -11,10 +11,12 @@ "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": "offline", + "mode": "streaming", "session_options": { - "fireredtts3.reference_cache_slots": "8" + "fireredtts3.reference_cache_slots": "8", + "fireredtts3.chunk_sizes": "3,12,12,12,12,12,12,12,12,12,12,12,12,12,12,12" } } ] diff --git a/include/engine/framework/audio/istft_graph.h b/include/engine/framework/audio/istft_graph.h index 6330ea212..d273b852e 100644 --- a/include/engine/framework/audio/istft_graph.h +++ b/include/engine/framework/audio/istft_graph.h @@ -67,6 +67,16 @@ class HostLogMagnitudePhaseISTFT { const std::vector & log_magnitude_phase, const std::vector & window); + // --- 增量 overlap-add iSTFT（用于流式）--- + // 按块喂入 log-magnitude+phase 帧，内部累积 overlap-add， + // 返回"已能被窗口包络完整覆盖"的音频样本（块间平滑衔接）。 + // 首次 append 自动初始化；finish 收尾 flush 尾部并复位。 + std::vector append_incremental( + const std::vector & log_magnitude_phase, + int64_t frames, + const std::vector & window); + std::vector finish_incremental(); + private: class Impl; std::unique_ptr impl_; diff --git a/include/engine/framework/codecs/redae_codec_runtime.h b/include/engine/framework/codecs/redae_codec_runtime.h index 1b8d34293..ba9ae36e6 100644 --- a/include/engine/framework/codecs/redae_codec_runtime.h +++ b/include/engine/framework/codecs/redae_codec_runtime.h @@ -76,6 +76,16 @@ class RedAeCodecRuntime { std::vector encode(const std::vector & audio_24k); runtime::AudioBuffer decode(const std::vector & latents); + + // --- 增量解码（流式）--- + // 重置解码器 KV 状态（每次新请求开始时调用）。 + void decode_reset(); + // 解码一批 latent（chunk），返回该块对应的音频（float32 24k mono）。 + // 内部用 decoder Qwen 的 KV 缓存跨块保持上下文，并用增量 iSTFT 逐块输出。 + runtime::AudioBuffer decode_incremental(const std::vector & latents); + // flush 增量 iSTFT 尾部样本（生成结束时调用）。 + runtime::AudioBuffer flush_incremental(); + void release_runtime_graphs(); private: diff --git a/include/engine/models/fireredtts3/pipeline.h b/include/engine/models/fireredtts3/pipeline.h index 9dc2484e3..248c09861 100644 --- a/include/engine/models/fireredtts3/pipeline.h +++ b/include/engine/models/fireredtts3/pipeline.h @@ -6,6 +6,7 @@ #include "engine/models/fireredtts3/assets.h" #include +#include #include #include #include @@ -53,6 +54,18 @@ struct FireRedTTS3InstructResult { std::string generated_text; }; +// 流式块回调：每生成一块音频调用一次（chunk_index 从 0 起）。 +using FireRedStreamChunkCallback = std::function; + +// 增量流式会话：持有 AR 循环状态，逐块产出音频。 +// 由 FireRedTTS3BaseRuntime 创建；调用方每取一块调 next_chunk()。 +class FireRedTTS3StreamSession { +public: + virtual ~FireRedTTS3StreamSession() = default; + // 返回下一块音频；流结束返回 empty（samples 为空）。 + virtual engine::runtime::AudioBuffer next_chunk() = 0; +}; + class FireRedTTS3BaseRuntime { public: FireRedTTS3BaseRuntime( @@ -70,6 +83,14 @@ class FireRedTTS3BaseRuntime { FireRedTTS3BaseRuntime & operator=(const FireRedTTS3BaseRuntime &) = delete; engine::runtime::AudioBuffer generate(const FireRedTTS3BaseRequest & request); + + // 增量流式：启动一个流式会话（AR 逐 patch 生成，按 chunk 边界增量解码）。 + // chunk_patches[i] 为第 i 块的 patch 数（如 {3,12,12,...} 首块 0.5s 后续 2s）。 + // 返回的会话每次 next_chunk() 产出一块音频；结束后 next_chunk() 返回空 audio。 + std::unique_ptr begin_streaming( + const FireRedTTS3BaseRequest & request, + const std::vector & chunk_patches); + void release_graphs(); private: diff --git a/include/engine/models/fireredtts3/redae.h b/include/engine/models/fireredtts3/redae.h index f82391a6f..aed3df16b 100644 --- a/include/engine/models/fireredtts3/redae.h +++ b/include/engine/models/fireredtts3/redae.h @@ -31,6 +31,12 @@ class FireRedRedAeRuntime { std::vector encode(const std::vector & audio_24k); engine::runtime::AudioBuffer decode(const std::vector & latents); + + // 增量解码（流式）：重置 + 逐块解码 + flush 尾部 + void decode_reset(); + engine::runtime::AudioBuffer decode_incremental(const std::vector & latents); + engine::runtime::AudioBuffer flush_incremental(); + void release_graphs(); private: diff --git a/include/engine/models/fireredtts3/session.h b/include/engine/models/fireredtts3/session.h index 95cb9d9c8..b83997630 100644 --- a/include/engine/models/fireredtts3/session.h +++ b/include/engine/models/fireredtts3/session.h @@ -12,12 +12,14 @@ namespace engine::models::fireredtts3 { class FireRedTTS3BaseRuntime; class FireRedTTS3InstructRuntime; +class FireRedTTS3StreamSession; std::shared_ptr make_fireredtts3_loader(); class FireRedTTS3Session final : public engine::runtime::RuntimeSessionBase - , public engine::runtime::IOfflineVoiceTaskSession { + , public engine::runtime::IOfflineVoiceTaskSession + , public engine::runtime::IStreamingVoiceTaskSession { public: FireRedTTS3Session( engine::runtime::TaskSpec task, @@ -32,7 +34,19 @@ class FireRedTTS3Session final void prepare(const engine::runtime::SessionPreparationRequest & request) override; engine::runtime::TaskResult run(const engine::runtime::TaskRequest & request) override; + // 流式（增量）接口 + engine::runtime::StreamingPolicy streaming_policy() const override; + void start_stream(const engine::runtime::TaskRequest & request) override; + std::optional next_stream_event() override; + void set_stream_event_sink(engine::runtime::StreamEventCallback sink) override; + engine::runtime::TaskResult finish_stream() override; + void reset() override; + engine::runtime::StreamEvent process_audio_chunk(const engine::runtime::AudioChunk & chunk) override; + engine::runtime::TaskResult finalize() override; + private: + void initialize_stream_request(const engine::runtime::TaskRequest & request); + engine::runtime::TaskSpec task_; std::shared_ptr assets_; std::shared_ptr contract_; @@ -40,6 +54,15 @@ class FireRedTTS3Session final std::unique_ptr runtime_; std::unique_ptr instruct_runtime_; bool mem_saver_ = false; + + // 流式状态 + bool stream_started_ = false; + engine::runtime::TaskRequest stream_request_; + std::vector stream_chunk_patches_; + size_t stream_chunk_index_ = 0; + engine::runtime::AudioBuffer stream_merged_audio_; + std::string stream_generated_text_; + std::unique_ptr stream_session_; }; } // namespace engine::models::fireredtts3 diff --git a/model_specs/fireredtts3.json b/model_specs/fireredtts3.json index f34f27686..5bf902d56 100644 --- a/model_specs/fireredtts3.json +++ b/model_specs/fireredtts3.json @@ -11,7 +11,8 @@ "design" ], "modes": [ - "offline" + "offline", + "streaming" ], "languages": [ "Chinese", @@ -215,6 +216,13 @@ "description": "Release cached runtime graphs after each request to reduce peak VRAM; default false.", "required": false, "default": false + }, + { + "name": "chunk_sizes", + "type": "string", + "description": "Streaming chunk sizes in AR latent patches (comma-separated). e.g. 3,12,12 -> first chunk ~0.5s, then ~2s each.", + "required": false, + "default": "3,12,12,12,12,12,12,12,12,12,12,12,12,12,12,12" } ], "load": [] diff --git a/src/framework/audio/istft_graph.cpp b/src/framework/audio/istft_graph.cpp index 700854291..ec457fc4d 100644 --- a/src/framework/audio/istft_graph.cpp +++ b/src/framework/audio/istft_graph.cpp @@ -211,6 +211,149 @@ class HostLogMagnitudePhaseISTFT::Impl { HostLogMagnitudePhaseISTFTConfig config_; int64_t freq_bins_ = 0; Workspace workspace_; + + // --- 增量状态 --- + bool inc_initialized_ = false; + std::vector inc_window_; + std::vector inc_folded_; + std::vector inc_envelope_; + int64_t inc_frames_emitted_ = 0; // 已折叠的 frame 绝对计数 + int64_t inc_emitted_samples_ = 0; // 已输出的样本数（累计） + + // 初始化滚动累加器：buffer 长度 = (本块 frame 数 - 1) * hop + n_fft（可增长）， + // envelope 按已加入的 frame 计算。 + void inc_ensure(int64_t total_frames, const std::vector & window) { + const int64_t need = (total_frames - 1) * config_.hop_length + config_.n_fft; + if (static_cast(inc_folded_.size()) < need) { + inc_folded_.resize(static_cast(need), 0.0F); + inc_envelope_.resize(static_cast(need), 0.0F); + } + inc_window_ = window; + } + +public: + std::vector append_incremental( + const std::vector & log_magnitude_phase, + int64_t frames, + const std::vector & window) { + require(static_cast(window.size()) == config_.n_fft, "host incremental ISTFT window size mismatch"); + require(static_cast(log_magnitude_phase.size()) == frames * config_.out_dim, + "host incremental ISTFT input size mismatch"); + require(frames > 0, "host incremental ISTFT requires positive frames"); + if (!inc_initialized_) { + inc_folded_.clear(); + inc_envelope_.clear(); + inc_frames_emitted_ = 0; + inc_emitted_samples_ = 0; + inc_initialized_ = true; + } + // 目标总帧数（含本块） + const int64_t new_total = inc_frames_emitted_ + frames; + inc_ensure(new_total, window); + + // 折叠本块 spectrum（frame 绝对索引从 inc_frames_emitted_ 起） + fold_frames(log_magnitude_phase, frames, window); + + // 重新计算 envelope（基于已加入的帧数） + std::fill(inc_envelope_.begin(), inc_envelope_.end(), 0.0F); + for (int64_t f = 0; f < new_total; ++f) { + const int64_t start = f * config_.hop_length; + for (int64_t i = 0; i < config_.n_fft; ++i) { + const int64_t pos = start + i; + if (pos >= static_cast(inc_envelope_.size())) { + break; + } + const float w = window[static_cast(i)]; + inc_envelope_[static_cast(pos)] += w * w; + } + } + + // 可输出样本：从 pad 到 (buffer 末尾 - pad)，且只输出"自上次以来新增"的部分。 + // 首块从 pad 起；后续从上次输出的绝对位置 inc_emitted_samples_ 起。 + const int64_t pad = (config_.n_fft - config_.hop_length) / 2; + const int64_t total_samples = (new_total - 1) * config_.hop_length + config_.n_fft; + const int64_t out_end = total_samples - pad; + const int64_t out_start = inc_emitted_samples_ == 0 ? pad : inc_emitted_samples_; + std::vector out; + if (out_end > out_start) { + out.resize(static_cast(out_end - out_start)); + for (int64_t i = out_start; i < out_end; ++i) { + const float denom = inc_envelope_[static_cast(i)]; + out[static_cast(i - out_start)] = denom <= 1.0e-11F ? 0.0F + : inc_folded_[static_cast(i)] / denom; + } + } + // 记录已输出的绝对样本数（不含 pad 前缀，保持与输出对齐） + inc_emitted_samples_ = out_end; + return out; + } + + std::vector finish_incremental() { + if (!inc_initialized_) { + return {}; + } + // flush 剩余全部样本（含尾部） + std::vector out(inc_folded_.size()); + for (size_t i = 0; i < inc_folded_.size(); ++i) { + const float denom = inc_envelope_[i]; + out[i] = denom <= 1.0e-11F ? 0.0F : inc_folded_[i] / denom; + } + inc_folded_.clear(); + inc_envelope_.clear(); + inc_frames_emitted_ = 0; + inc_emitted_samples_ = 0; + inc_initialized_ = false; + return out; + } + + // 折叠一批 log-magnitude+phase 帧到滚动累加器（frame 绝对索引从 inc_frames_emitted_ 起）。 + void fold_frames(const std::vector & log_magnitude_phase, int64_t frames, const std::vector & window) { + // 1) 逐帧把 log-magnitude+phase 转成复数 spectrum + std::vector> spectrum(static_cast(frames * freq_bins_)); + for (int64_t frame = 0; frame < frames; ++frame) { + const float * row = log_magnitude_phase.data() + static_cast(frame * config_.out_dim); + auto * spectrum_row = spectrum.data() + static_cast(frame * freq_bins_); + for (int64_t freq = 0; freq < freq_bins_; ++freq) { + const float mag = std::min(std::exp(row[freq]), 100.0F); + const float phase = row[freq_bins_ + freq]; + spectrum_row[static_cast(freq)] = { + mag * std::cos(phase), + mag * std::sin(phase), + }; + } + } + // 2) inverse FFT + std::vector framed(static_cast(frames * config_.n_fft), 0.0F); + real_fft_inverse( + {static_cast(frames), static_cast(config_.n_fft)}, + { + static_cast(freq_bins_ * static_cast(sizeof(std::complex))), + static_cast(sizeof(std::complex)), + }, + { + static_cast(config_.n_fft * static_cast(sizeof(float))), + static_cast(sizeof(float)), + }, + 1, + spectrum.data(), + framed.data(), + 1.0F / static_cast(config_.n_fft), + config_.threads); + // 3) windowed overlap-add 到滚动累加器 + for (int64_t frame = 0; frame < frames; ++frame) { + const int64_t abs_start = (inc_frames_emitted_ + frame) * config_.hop_length; + const float * src = framed.data() + static_cast(frame * config_.n_fft); + for (int64_t i = 0; i < config_.n_fft; ++i) { + const int64_t pos = abs_start + i; + if (pos >= static_cast(inc_folded_.size())) { + continue; + } + const float w = window[static_cast(i)]; + inc_folded_[static_cast(pos)] += src[i] * w; + } + } + inc_frames_emitted_ += frames; + } }; class CudaLogMagnitudePhaseISTFT::Impl { @@ -285,6 +428,17 @@ HostLogMagnitudePhaseISTFTResult HostLogMagnitudePhaseISTFT::compute( return impl_->compute(log_magnitude_phase, window); } +std::vector HostLogMagnitudePhaseISTFT::append_incremental( + const std::vector & log_magnitude_phase, + int64_t frames, + const std::vector & window) { + return impl_->append_incremental(log_magnitude_phase, frames, window); +} + +std::vector HostLogMagnitudePhaseISTFT::finish_incremental() { + return impl_->finish_incremental(); +} + CudaLogMagnitudePhaseISTFT::CudaLogMagnitudePhaseISTFT( const CudaLogMagnitudePhaseISTFTConfig & config) : impl_(std::make_unique(config)) {} diff --git a/src/framework/codecs/redae_codec_runtime.cpp b/src/framework/codecs/redae_codec_runtime.cpp index bb9cce972..494844ccd 100644 --- a/src/framework/codecs/redae_codec_runtime.cpp +++ b/src/framework/codecs/redae_codec_runtime.cpp @@ -643,6 +643,76 @@ class RedAeRuntime { return audio; } + // --- 增量解码（流式）--- + void decode_reset() { + dec_state_.reset(); + dec_qwen_frames_ = 0; + inc_host_istft_.reset(); + inc_istft_frames_ = 0; + } + + runtime::AudioBuffer decode_incremental(const std::vector & latents) { + if (latents.empty() || static_cast(latents.size()) % config_.bottleneck_dim != 0) { + throw std::runtime_error("RedAE codec incremental decode latent size mismatch"); + } + const int64_t latent_frames = static_cast(latents.size()) / config_.bottleneck_dim; + auto decoder_in = encoder_out_.decode_in(latents, latent_frames); + const int64_t qwen_frames = latent_frames * config_.enc_extra_downsample_rate; + std::vector decoder_hidden; + + if (!dec_state_.has_value()) { + // 首块：一次性 prefill（返回 KV state 用于后续 chaining） + auto decoder = decoder_qwen_.prefill_embeddings(decoder_in, qwen_frames); + decoder_hidden = std::move(decoder.hidden); + dec_state_ = std::move(decoder.state); + } else { + // 后续块：用 decode_embedding 逐 qwen frame 解码，复用前序 KV + const int64_t required = dec_qwen_frames_ + qwen_frames; + decoder_qwen_.start_decode_embeddings(*dec_state_, required); + decoder_hidden.reserve(static_cast(qwen_frames * config_.dec_hidden_size)); + for (int64_t f = 0; f < qwen_frames; ++f) { + std::vector row( + decoder_in.begin() + static_cast(f * config_.dec_hidden_size), + decoder_in.begin() + static_cast((f + 1) * config_.dec_hidden_size)); + auto step = decoder_qwen_.decode_embedding(row); + decoder_hidden.insert( + decoder_hidden.end(), step.hidden.begin(), step.hidden.end()); + } + } + dec_qwen_frames_ += qwen_frames; + + auto spec = encoder_out_.istft_head(decoder_hidden, qwen_frames); + + // 增量 iSTFT + if (!inc_host_istft_) { + audio::HostLogMagnitudePhaseISTFTConfig cfg; + cfg.frames = qwen_frames; + cfg.n_fft = config_.audio_patch_size * 4; + cfg.hop_length = config_.audio_patch_size; + cfg.out_dim = config_.audio_patch_size * 4 + 2; + cfg.threads = static_cast(std::max(1, execution_.config().threads)); + inc_host_istft_ = std::make_unique(cfg); + inc_istft_frames_ = qwen_frames; + } + auto chunk_audio = inc_host_istft_->append_incremental(spec, qwen_frames, istft_window_); + + runtime::AudioBuffer audio; + audio.sample_rate = static_cast(config_.sample_rate); + audio.channels = 1; + audio.samples = std::move(chunk_audio); + return audio; + } + + runtime::AudioBuffer flush_incremental() { + runtime::AudioBuffer audio; + audio.sample_rate = static_cast(config_.sample_rate); + audio.channels = 1; + if (inc_host_istft_) { + audio.samples = inc_host_istft_->finish_incremental(); + } + return audio; + } + void release_graphs() { encoder_in_.release_graph(); encoder_out_.release_graph(); @@ -722,6 +792,12 @@ class RedAeRuntime { std::vector istft_window_; std::unique_ptr host_istft_; int64_t host_istft_frames_ = 0; + + // 增量解码状态 + std::optional dec_state_; + int64_t dec_qwen_frames_ = 0; + std::unique_ptr inc_host_istft_; + int64_t inc_istft_frames_ = 0; }; @@ -766,6 +842,18 @@ runtime::AudioBuffer RedAeCodecRuntime::decode(const std::vector & latent return impl_->runtime.decode(latents); } +void RedAeCodecRuntime::decode_reset() { + impl_->runtime.decode_reset(); +} + +runtime::AudioBuffer RedAeCodecRuntime::decode_incremental(const std::vector & latents) { + return impl_->runtime.decode_incremental(latents); +} + +runtime::AudioBuffer RedAeCodecRuntime::flush_incremental() { + return impl_->runtime.flush_incremental(); +} + void RedAeCodecRuntime::release_runtime_graphs() { impl_->runtime.release_graphs(); } diff --git a/src/models/fireredtts3/pipeline.cpp b/src/models/fireredtts3/pipeline.cpp index d914c4e73..222fc4682 100644 --- a/src/models/fireredtts3/pipeline.cpp +++ b/src/models/fireredtts3/pipeline.cpp @@ -256,6 +256,146 @@ class FireRedTTS3BaseRuntime::Impl { return decoded; } + // 启动流式会话（增量 AR + 逐块解码） + std::unique_ptr begin_streaming( + const FireRedTTS3BaseRequest & request, + const std::vector & chunk_patches) { + struct StreamSession final : public FireRedTTS3StreamSession { + FireRedTTS3BaseRuntime::Impl & owner; + FireRedTTS3BaseRequest req; + std::vector chunks; + // AR 循环状态 + std::vector latents_gen; + std::vector backbone_cond; + std::vector schedule; + std::vector next_input; + std::vector prefill_hidden; + int64_t prefill_steps = 0; + std::vector spk_dit; + int64_t prompt_latent_frames = 0; + bool started = false; + bool finished = false; + int64_t step = 0; + int64_t generated_patches = 0; + size_t chunk_index = 0; + int64_t chunk_target = 0; + std::vector chunk_latents; + + StreamSession(FireRedTTS3BaseRuntime::Impl & o, + FireRedTTS3BaseRequest r, + std::vector c) + : owner(o), req(std::move(r)), chunks(std::move(c)) {} + + runtime::AudioBuffer next_chunk() override { + if (finished) { + return {}; + } + if (!started) { + start(); + } + // 逐 patch 生成直到达到当前块边界 + const int64_t max_steps = 400; + while (step < max_steps) { + std::vector hidden; + int64_t hidden_rows = 0; + if (step == 0) { + hidden = prefill_hidden; + hidden_rows = prefill_steps; + } else { + auto decoded = owner.ar_->decode_embedding(next_input); + hidden = decoded.hidden; + hidden_rows = 1; + } + const auto last = last_rows(hidden, hidden_rows, owner.assets_->base.hidden_size, 1); + const float stop = owner.ar_->stop(last); + if (stop >= req.stop_threshold && step >= 6) { + break; + } + std::vector one_backbone; + if (step == 0) { + const int64_t prompt_patches = prompt_latent_frames / owner.assets_->base.patch_size; + one_backbone = last_rows(hidden, hidden_rows, owner.assets_->base.hidden_size, prompt_patches); + } else { + one_backbone = last; + } + backbone_cond.insert(backbone_cond.end(), one_backbone.begin(), one_backbone.end()); + const int64_t cond_rows = static_cast(backbone_cond.size()) / owner.assets_->base.hidden_size; + auto cond3 = last_rows( + backbone_cond, cond_rows, owner.assets_->base.hidden_size, + owner.assets_->base.history_patches + 1); + auto dit_cond3 = owner.ar_->dit_head(cond3, owner.assets_->base.history_patches + 1); + const auto next_latent = owner.flow_one_patch( + latents_gen, dit_cond3, spk_dit, + req.guidance_scale, schedule, req.seed, + static_cast(step)); + latents_gen.insert(latents_gen.end(), next_latent.begin(), next_latent.end()); + next_input = owner.ar_->patch_encode(next_latent); + chunk_latents.insert(chunk_latents.end(), next_latent.begin(), next_latent.end()); + ++generated_patches; + ++step; + + if (generated_patches >= chunk_target && chunk_index < chunks.size()) { + auto audio = owner.redae_->decode_incremental(chunk_latents); + chunk_latents.clear(); + chunk_index++; + chunk_target = (chunk_index < chunks.size()) + ? chunk_target + chunks[chunk_index] + : chunk_target; + return std::move(audio); + } + } + // AR 结束：flush 剩余 latents + iSTFT 尾部 + if (!chunk_latents.empty() && generated_patches > 0) { + auto tail = owner.redae_->decode_incremental(chunk_latents); + chunk_latents.clear(); + finished = true; + return std::move(tail); + } + auto flush = owner.redae_->flush_incremental(); + finished = true; + return std::move(flush); + } + + void start() { + if (chunks.empty()) { + throw std::runtime_error("FireRedTTS3 streaming requires at least one chunk"); + } + const auto & reference = owner.prepare_reference_voice(req.prompt_audio); + const auto & prompt_latents = reference.prompt_latents; + prompt_latent_frames = static_cast(prompt_latents.size()) / owner.assets_->base.redae_dim; + auto spk_llm = owner.ar_->speaker_llm(reference.speaker_embedding); + spk_dit = owner.ar_->speaker_dit(reference.speaker_embedding); + auto text_embeds = owner.ar_->token_embedding(req.token_ids); + if (owner.mem_saver_) { + owner.ar_->release_graphs(); + } + auto patch_prompt = owner.ar_->patch_encode(prompt_latents); + + std::vector input_embeddings; + input_embeddings.reserve(spk_llm.size() + text_embeds.size() + patch_prompt.size()); + input_embeddings.insert(input_embeddings.end(), spk_llm.begin(), spk_llm.end()); + input_embeddings.insert(input_embeddings.end(), text_embeds.begin(), text_embeds.end()); + input_embeddings.insert(input_embeddings.end(), patch_prompt.begin(), patch_prompt.end()); + prefill_steps = 1 + static_cast(req.token_ids.size()) + prompt_latent_frames / owner.assets_->base.patch_size; + auto prefill = owner.ar_->prefill_embeddings(input_embeddings, prefill_steps); + owner.ar_->start_decode_embeddings(prefill.state, prefill_steps + 400 + 1); + prefill_hidden = std::move(prefill.hidden); + + latents_gen.assign(static_cast(owner.assets_->base.history_patches * owner.assets_->base.patch_size * owner.assets_->base.redae_dim), 0.0F); + latents_gen.insert(latents_gen.end(), prompt_latents.begin(), prompt_latents.end()); + backbone_cond.assign(static_cast(owner.assets_->base.history_patches * owner.assets_->base.hidden_size), 0.0F); + schedule = firered_cosine_time_schedule(req.num_inference_steps); + next_input = {}; + chunk_latents.reserve(static_cast(owner.assets_->base.patch_size * owner.assets_->base.redae_dim)); + chunk_target = chunks[0]; + owner.redae_->decode_reset(); + started = true; + } + }; + auto session = std::make_unique(*this, request, chunk_patches); + return session; + } + void release_graphs() { if (redae_) { redae_->release_graphs(); @@ -855,6 +995,12 @@ engine::runtime::AudioBuffer FireRedTTS3BaseRuntime::generate(const FireRedTTS3B return impl_->generate(request); } +std::unique_ptr FireRedTTS3BaseRuntime::begin_streaming( + const FireRedTTS3BaseRequest & request, + const std::vector & chunk_patches) { + return impl_->begin_streaming(request, chunk_patches); +} + void FireRedTTS3BaseRuntime::release_graphs() { impl_->release_graphs(); } diff --git a/src/models/fireredtts3/redae.cpp b/src/models/fireredtts3/redae.cpp index 83a4688aa..914e9a48b 100644 --- a/src/models/fireredtts3/redae.cpp +++ b/src/models/fireredtts3/redae.cpp @@ -130,6 +130,18 @@ class FireRedRedAeRuntime::Impl { return runtime_.decode(latents); } + void decode_reset() { + runtime_.decode_reset(); + } + + runtime::AudioBuffer decode_incremental(const std::vector & latents) { + return runtime_.decode_incremental(latents); + } + + runtime::AudioBuffer flush_incremental() { + return runtime_.flush_incremental(); + } + void release_graphs() { runtime_.release_runtime_graphs(); } @@ -157,6 +169,18 @@ engine::runtime::AudioBuffer FireRedRedAeRuntime::decode(const std::vectordecode(latents); } +void FireRedRedAeRuntime::decode_reset() { + impl_->decode_reset(); +} + +engine::runtime::AudioBuffer FireRedRedAeRuntime::decode_incremental(const std::vector & latents) { + return impl_->decode_incremental(latents); +} + +engine::runtime::AudioBuffer FireRedRedAeRuntime::flush_incremental() { + return impl_->flush_incremental(); +} + void FireRedRedAeRuntime::release_graphs() { impl_->release_graphs(); } diff --git a/src/models/fireredtts3/session.cpp b/src/models/fireredtts3/session.cpp index 79313eb3c..7f3d3ba09 100644 --- a/src/models/fireredtts3/session.cpp +++ b/src/models/fireredtts3/session.cpp @@ -9,6 +9,7 @@ #include "engine/framework/text/text_normalization.h" #include "engine/models/fireredtts3/pipeline.h" +#include #include #include @@ -233,9 +234,7 @@ FireRedTTS3Session::FireRedTTS3Session( } else if (task_.task != runtime::VoiceTaskKind::VoiceCloning) { throw std::runtime_error("FireRedTTS3 Base supports the voice clone task"); } - if (task_.mode != runtime::RunMode::Offline) { - throw std::runtime_error("FireRedTTS3 supports offline sessions"); - } + // 双模式：Offline 与 Streaming 都允许（Base clone 支持离线 + 增量流式） using T = engine::assets::TensorStorageType; const auto storage_type = runtime::parse_tensor_storage_option( options.options, @@ -365,6 +364,113 @@ runtime::TaskResult FireRedTTS3Session::run(const runtime::TaskRequest & request return result; } +// ------------------------------------------------------------------ // +// 流式（增量）接口 +// ------------------------------------------------------------------ // +runtime::StreamingPolicy FireRedTTS3Session::streaming_policy() const { + runtime::StreamingPolicy policy; + policy.input = runtime::StreamingInputKind::None; + policy.output = runtime::StreamingOutputKind::PullEvents; + return policy; +} + +void FireRedTTS3Session::initialize_stream_request(const runtime::TaskRequest & request) { + if (is_instruct_variant(assets_->variant)) { + throw std::runtime_error("FireRedTTS3 Instruct streaming is not implemented yet; use offline mode"); + } + stream_request_ = request; + // chunk_patches: 每个块的 patch 数（首块 ~0.5s=3 patch, 后续 ~2s=12 patch） + std::string sizes_str = runtime::find_option(request.options, {"fireredtts3.chunk_sizes"}).value_or("3,12,12,12,12,12,12,12,12,12,12,12,12,12,12,12"); + stream_chunk_patches_.clear(); + std::stringstream ss(sizes_str); + std::string item; + while (std::getline(ss, item, ',')) { + if (!item.empty()) { + stream_chunk_patches_.push_back(std::stoll(item)); + } + } + if (stream_chunk_patches_.empty()) { + stream_chunk_patches_.push_back(3); + } + stream_chunk_index_ = 0; + stream_merged_audio_ = runtime::AudioBuffer{}; + stream_generated_text_.clear(); + stream_session_.reset(); +} + +void FireRedTTS3Session::start_stream(const runtime::TaskRequest & request) { + require_prepared("FireRedTTS3 streaming"); + if (task_.mode != runtime::RunMode::Streaming) { + throw std::runtime_error("FireRedTTS3 start_stream requires a streaming session"); + } + reset(); + initialize_stream_request(request); + auto base = make_base_request(*tokenizer_, request); + stream_session_ = runtime_->begin_streaming(base, stream_chunk_patches_); + stream_started_ = true; +} + +std::optional FireRedTTS3Session::next_stream_event() { + if (!stream_started_) { + throw std::runtime_error("FireRedTTS3 streaming has not been started"); + } + if (!stream_session_) { + return std::nullopt; + } + auto audio = stream_session_->next_chunk(); + if (audio.samples.empty()) { + stream_session_.reset(); + return std::nullopt; + } + runtime::append_audio_buffer(stream_merged_audio_, audio); + runtime::StreamEvent event; + event.named_audio_outputs.push_back({ + "chunk_" + std::to_string(stream_chunk_index_++), + std::move(audio), + {}, + }); + return event; +} + +void FireRedTTS3Session::set_stream_event_sink(runtime::StreamEventCallback sink) { + (void)sink; +} + +runtime::TaskResult FireRedTTS3Session::finish_stream() { + if (!stream_started_) { + throw std::runtime_error("FireRedTTS3 streaming has not been started"); + } + // 排空剩余块 + while (next_stream_event().has_value()) { + } + runtime::TaskResult result; + result.audio_output = std::move(stream_merged_audio_); + if (!stream_generated_text_.empty()) { + result.text_output = runtime::Transcript{std::move(stream_generated_text_), request_language(stream_request_)}; + } + reset(); + return result; +} + +void FireRedTTS3Session::reset() { + stream_request_ = {}; + stream_chunk_patches_.clear(); + stream_chunk_index_ = 0; + stream_merged_audio_ = runtime::AudioBuffer{}; + stream_generated_text_.clear(); + stream_session_.reset(); + stream_started_ = false; +} + +runtime::StreamEvent FireRedTTS3Session::process_audio_chunk(const runtime::AudioChunk & chunk) { + (void)chunk; + throw std::runtime_error("FireRedTTS3 streaming does not consume audio chunks"); +} + +runtime::TaskResult FireRedTTS3Session::finalize() { + return finish_stream(); +} + std::shared_ptr make_fireredtts3_loader() { runtime::SpecBackedVoiceModelConfig config; config.family = kFamily; From 8c8d608194ade88e966fa7dd9b49ab815156452b Mon Sep 17 00:00:00 2001 From: dujing Date: Tue, 1 Sep 2026 14:32:58 +0800 Subject: [PATCH 03/12] Add per-member position/mask/cache-slot to batched decode (P2a-1) - kv_cache: TransformerBatchedKVState gains current_ends; TransformerBatchedKVCache supports per-member ends (member_end/set_member_end/advance_member). - qwen_causal_decode_runtime: run_batched_decode_step uses per-member positions + cache slots; positions tensor is {batch_size}; fixed max_batch graph reuse. - qwen_causal_decoder: batched cached-step mask writer takes per-member ends+slots. - qwen_decoder: batched RoPE applies per-member positions via permute (batch <-> token axis) so different-prefix-length sequences batch correctly (llama.cpp-style). Single-path regression passes (rms 0.048). Enables FireRedTTS3 batch scheduler. Co-Authored-By: Claude --- .../transformers/qwen_causal_decoder.h | 4 +- include/engine/framework/runtime/kv_cache.h | 11 +++ .../qwen_causal_decode_runtime.cpp | 48 ++++++------ .../transformers/qwen_causal_decoder.cpp | 13 ++-- .../modules/transformers/qwen_decoder.cpp | 32 +++++++- src/framework/runtime/kv_cache.cpp | 76 +++++++++++++++---- 6 files changed, 136 insertions(+), 48 deletions(-) diff --git a/include/engine/framework/modules/transformers/qwen_causal_decoder.h b/include/engine/framework/modules/transformers/qwen_causal_decoder.h index 4231fcaeb..9e95de8dd 100644 --- a/include/engine/framework/modules/transformers/qwen_causal_decoder.h +++ b/include/engine/framework/modules/transformers/qwen_causal_decoder.h @@ -187,7 +187,7 @@ void write_qwen_batched_cached_step_mask( std::vector & scratch, int64_t batch_size, int64_t mask_steps, - int64_t visible_prefix_steps, - int64_t current_slot); + const std::vector & member_ends, + const std::vector & cache_slots); } // namespace engine::modules diff --git a/include/engine/framework/runtime/kv_cache.h b/include/engine/framework/runtime/kv_cache.h index db55f1176..8a2f7a73f 100644 --- a/include/engine/framework/runtime/kv_cache.h +++ b/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 current_ends; std::vector 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 & 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 member_ends_; TransformerKVCacheOptions options_; std::vector layers_; }; diff --git a/src/framework/modules/transformers/qwen_causal_decode_runtime.cpp b/src/framework/modules/transformers/qwen_causal_decode_runtime.cpp index 8d0e65ad9..e89548511 100644 --- a/src/framework/modules/transformers/qwen_causal_decode_runtime.cpp +++ b/src/framework/modules/transformers/qwen_causal_decode_runtime.cpp @@ -203,11 +203,10 @@ void write_batched_cached_step_mask( std::vector & scratch, int64_t batch_size, int64_t mask_steps, - int64_t visible_prefix_steps, - int64_t current_slot, - int64_t position) { + const std::vector & member_ends, + const std::vector & cache_slots) { if (config.sliding_window <= 0) { - write_qwen_batched_cached_step_mask(tensor, scratch, batch_size, mask_steps, visible_prefix_steps, current_slot); + write_qwen_batched_cached_step_mask(tensor, scratch, batch_size, mask_steps, member_ends, cache_slots); return; } if (tensor == nullptr) { @@ -219,12 +218,6 @@ void write_batched_cached_step_mask( if (mask_steps <= 0) { throw std::runtime_error("QwenCausalDecodeRuntime sliding batched cached mask requires positive steps"); } - if (visible_prefix_steps < 0 || visible_prefix_steps > mask_steps) { - throw std::runtime_error("QwenCausalDecodeRuntime sliding batched cached mask visible prefix is out of range"); - } - if (current_slot < 0 || current_slot >= mask_steps) { - throw std::runtime_error("QwenCausalDecodeRuntime sliding batched cached mask current slot is out of range"); - } const auto masked = ggml_fp32_to_fp16(-std::numeric_limits::infinity()); const auto visible = ggml_fp32_to_fp16(0.0F); const size_t row_size = static_cast(mask_steps); @@ -232,14 +225,16 @@ void write_batched_cached_step_mask( if (scratch.size() != total_size) { scratch.resize(total_size); } - const int64_t begin = std::max(0, position - config.sliding_window + 1); for (int64_t batch = 0; batch < batch_size; ++batch) { + const int64_t end = member_ends.empty() ? cache_slots[static_cast(batch)] : member_ends[static_cast(batch)]; + const int64_t current_slot = cache_slots[static_cast(batch)]; + const int64_t begin = std::max(0, end - config.sliding_window + 1); const size_t offset = static_cast(batch) * row_size; std::fill( scratch.begin() + static_cast(offset), scratch.begin() + static_cast(offset + row_size), masked); - for (int64_t i = begin; i < visible_prefix_steps; ++i) { + for (int64_t i = begin; i < end; ++i) { scratch[offset + static_cast(i)] = visible; } scratch[offset + static_cast(current_slot)] = visible; @@ -1178,11 +1173,12 @@ class QwenCausalDecodeRuntime::Impl { batched_decode_input_ = input.tensor; x = input; } - batched_decode_positions_ = ggml_new_tensor_1d(batched_decode_ctx_.get(), GGML_TYPE_I32, 1); + batched_decode_positions_ = ggml_new_tensor_1d(batched_decode_ctx_.get(), GGML_TYPE_I32, batch_size); auto positions = core::wrap_tensor( batched_decode_positions_, - core::TensorShape::from_dims({1}), + core::TensorShape::from_dims({batch_size}), GGML_TYPE_I32); + batched_decode_positions_host_.resize(static_cast(batch_size), 0); auto slot = core::make_tensor(ctx, GGML_TYPE_I32, core::TensorShape::from_dims({batch_size})); batched_decode_cache_slot_ = slot.tensor; auto attention = core::make_tensor( @@ -1303,13 +1299,18 @@ class QwenCausalDecodeRuntime::Impl { if (batched_decode_cache_.valid_steps() >= batched_decode_cache_steps_) { throw std::runtime_error("QwenCausalDecodeRuntime batched decode cache exhausted"); } - const int32_t position = static_cast(batched_decode_cache_.current_end()); - ggml_backend_tensor_set(batched_decode_positions_, &position, 0, sizeof(int32_t)); - const int32_t cache_slot = static_cast(batched_decode_cache_.valid_steps()); + // per-member positions + cache slots（不同序列可处于不同位置） for (int64_t batch = 0; batch < batched_decode_batch_size_; ++batch) { + const int32_t pos = static_cast(batched_decode_cache_.member_end(batch)); + batched_decode_positions_host_[static_cast(batch)] = pos; batched_decode_cache_slots_[static_cast(batch)] = - static_cast(batch * batched_decode_cache_steps_ + cache_slot); + static_cast(batch * batched_decode_cache_steps_ + pos); } + ggml_backend_tensor_set( + batched_decode_positions_, + batched_decode_positions_host_.data(), + 0, + batched_decode_positions_host_.size() * sizeof(int32_t)); ggml_backend_tensor_set( batched_decode_cache_slot_, batched_decode_cache_slots_.data(), @@ -1321,9 +1322,8 @@ class QwenCausalDecodeRuntime::Impl { batched_decode_attention_mask_values_, batched_decode_batch_size_, batched_decode_cache_steps_, - batched_decode_cache_.valid_steps(), - cache_slot, - position); + batched_decode_cache_.member_ends_for_mask(), + batched_decode_cache_slots_); core::set_backend_threads(backend_, threads_); const ggml_status status = core::compute_backend_graph(backend_, batched_decode_graph_); ggml_backend_synchronize(backend_); @@ -1348,7 +1348,10 @@ class QwenCausalDecodeRuntime::Impl { out.hidden.size() * sizeof(float)); round_readback(out.hidden, config_); } - batched_decode_cache_.advance_after_direct_append(1); + // 每个 member 前进 1 步 + for (int64_t batch = 0; batch < batched_decode_batch_size_; ++batch) { + batched_decode_cache_.advance_member(batch, 1); + } return out; } @@ -1523,6 +1526,7 @@ class QwenCausalDecodeRuntime::Impl { ggml_backend_buffer_t batched_decode_buffer_ = nullptr; std::vector batched_decode_attention_mask_values_; std::vector batched_decode_cache_slots_; + std::vector batched_decode_positions_host_; runtime::TransformerBatchedKVCache batched_decode_cache_; int64_t batched_decode_batch_size_ = 0; int64_t batched_decode_cache_steps_ = 0; diff --git a/src/framework/modules/transformers/qwen_causal_decoder.cpp b/src/framework/modules/transformers/qwen_causal_decoder.cpp index 94ca0050c..8be371389 100644 --- a/src/framework/modules/transformers/qwen_causal_decoder.cpp +++ b/src/framework/modules/transformers/qwen_causal_decoder.cpp @@ -502,8 +502,8 @@ void write_qwen_batched_cached_step_mask( std::vector & scratch, int64_t batch_size, int64_t mask_steps, - int64_t visible_prefix_steps, - int64_t current_slot) { + const std::vector & member_ends, + const std::vector & cache_slots) { if (tensor == nullptr) { throw std::runtime_error("write_qwen_batched_cached_step_mask requires a tensor"); } @@ -511,12 +511,6 @@ void write_qwen_batched_cached_step_mask( throw std::runtime_error("write_qwen_batched_cached_step_mask requires positive batch size"); } validate_steps(mask_steps, "write_qwen_batched_cached_step_mask"); - if (visible_prefix_steps < 0 || visible_prefix_steps > mask_steps) { - throw std::runtime_error("write_qwen_batched_cached_step_mask visible prefix is out of range"); - } - if (current_slot < 0 || current_slot >= mask_steps) { - throw std::runtime_error("write_qwen_batched_cached_step_mask current slot is out of range"); - } const auto masked = ggml_fp32_to_fp16(-INFINITY); const auto visible = ggml_fp32_to_fp16(0.0F); const size_t row_size = static_cast(mask_steps); @@ -525,6 +519,9 @@ void write_qwen_batched_cached_step_mask( scratch.resize(total_size); } for (int64_t batch = 0; batch < batch_size; ++batch) { + const int64_t end = member_ends.empty() ? cache_slots[static_cast(batch)] : member_ends[static_cast(batch)]; + const int64_t visible_prefix_steps = end; + const int64_t current_slot = cache_slots[static_cast(batch)]; const size_t offset = static_cast(batch) * row_size; std::fill( scratch.begin() + static_cast(offset), diff --git a/src/framework/modules/transformers/qwen_decoder.cpp b/src/framework/modules/transformers/qwen_decoder.cpp index a50b7191c..5a0fd5b54 100644 --- a/src/framework/modules/transformers/qwen_decoder.cpp +++ b/src/framework/modules/transformers/qwen_decoder.cpp @@ -861,8 +861,36 @@ QwenDecoderLayerOutputs QwenDecoderLayerModule::build_with_static_cache_tail_bat const core::TensorValue * rope_factors = weights.rope_frequency_factors.has_value() ? &*weights.rope_frequency_factors : nullptr; - q = RoPEModule({dim, config_.rope_type, config_.rope_theta}).build(ctx, q, positions, rope_factors); - k = RoPEModule({dim, config_.rope_type, config_.rope_theta}).build(ctx, k, positions, rope_factors); + // per-member positions: positions len == batch (each sequence its own position). + // Permute q/k from [batch, heads, 1, dim] to [1, batch, heads, dim] so batch sits + // on ggml ne[2] (token axis); RoPE takes per-member pos; permute back. + if (positions.shape.dims[0] > 1 && positions.shape.dims[0] == q.shape.dims[0]) { + const int64_t batch = q.shape.dims[0]; + const int64_t heads = q.shape.dims[1]; + auto perm_batch_to_token = [&](core::TensorValue x) { + const auto contiguous = core::ensure_backend_addressable_layout(ctx, x); + return core::wrap_tensor( + ggml_permute(ctx.ggml, contiguous.tensor, 1, 0, 2, 3), + core::TensorShape::from_dims({1, batch, heads, x.shape.last_dim()}), + x.type); + }; + auto perm_back = [&](core::TensorValue x) { + const auto contiguous = core::ensure_backend_addressable_layout(ctx, x); + return core::wrap_tensor( + ggml_permute(ctx.ggml, contiguous.tensor, 1, 0, 2, 3), + core::TensorShape::from_dims({batch, heads, 1, x.shape.last_dim()}), + x.type); + }; + q = perm_batch_to_token(q); + k = perm_batch_to_token(k); + q = RoPEModule({dim, config_.rope_type, config_.rope_theta}).build(ctx, q, positions, rope_factors); + k = RoPEModule({dim, config_.rope_type, config_.rope_theta}).build(ctx, k, positions, rope_factors); + q = perm_back(q); + k = perm_back(k); + } else { + q = RoPEModule({dim, config_.rope_type, config_.rope_theta}).build(ctx, q, positions, rope_factors); + k = RoPEModule({dim, config_.rope_type, config_.rope_theta}).build(ctx, k, positions, rope_factors); + } if (config_.activation_cast.enabled && config_.activation_cast.after_rope) { q = activation_cast(ctx, q, config_.activation_cast); k = activation_cast(ctx, k, config_.activation_cast); diff --git a/src/framework/runtime/kv_cache.cpp b/src/framework/runtime/kv_cache.cpp index 567aac4de..2bf307ec3 100644 --- a/src/framework/runtime/kv_cache.cpp +++ b/src/framework/runtime/kv_cache.cpp @@ -260,6 +260,14 @@ void TransformerBatchedKVCache::import_state(const TransformerBatchedKVState & s throw std::runtime_error("TransformerBatchedKVCache state batch size does not match cache batch size"); } current_end_ = state.current_end; + // per-member ends + member_ends_.clear(); + if (!state.current_ends.empty()) { + if (static_cast(state.current_ends.size()) != batch_size_) { + throw std::runtime_error("TransformerBatchedKVCache state current_ends size does not match batch size"); + } + member_ends_ = state.current_ends; + } if (layers_.empty()) { valid_steps_ = 0; return; @@ -267,27 +275,32 @@ void TransformerBatchedKVCache::import_state(const TransformerBatchedKVState & s if (state.layers.size() != layers_.size()) { throw std::runtime_error("TransformerBatchedKVCache state layer count does not match cache layer count"); } - const int64_t state_steps = state.layers.empty() ? 0 : state.layers.front().valid_steps; - if (state_steps > cache_steps_) { + // 每 member 的有效步数（per-member 或均匀） + std::vector member_steps(static_cast(batch_size_)); + int64_t max_steps = 0; + for (int64_t b = 0; b < batch_size_; ++b) { + member_steps[static_cast(b)] = member_ends_.empty() ? current_end_ : member_ends_[static_cast(b)]; + max_steps = std::max(max_steps, member_steps[static_cast(b)]); + } + if (max_steps > cache_steps_) { throw std::runtime_error("TransformerBatchedKVCache state valid_steps exceeds cache capacity"); } - valid_steps_ = state_steps; - const size_t copy_elems = static_cast(state_steps * row_elems_); + valid_steps_ = max_steps; for (size_t layer = 0; layer < layers_.size(); ++layer) { auto & cache = layers_[layer]; const auto & source = state.layers[layer]; - if (source.valid_steps != state_steps) { - throw std::runtime_error("TransformerBatchedKVCache requires consistent valid_steps across all layers"); - } - const size_t state_elems = static_cast(batch_size_) * copy_elems; + // source 按 batch * max_steps * row 布局（导出时统一 max_steps） + const size_t state_elems = static_cast(batch_size_) * static_cast(max_steps * row_elems_); if (source.key.size() != source.value.size() || source.key.size() != state_elems) { - throw std::runtime_error("TransformerBatchedKVCache source tensors do not match batch * valid_steps * row_elems"); + throw std::runtime_error("TransformerBatchedKVCache source tensors do not match batch * max_steps * row_elems"); } std::fill(cache.import_key_scratch.begin(), cache.import_key_scratch.end(), 0.0F); std::fill(cache.import_value_scratch.begin(), cache.import_value_scratch.end(), 0.0F); - for (int64_t batch = 0; batch < batch_size_; ++batch) { - const size_t src_offset = static_cast(batch) * copy_elems; - const size_t dst_offset = static_cast(batch * cache_steps_ * row_elems_); + for (int64_t b = 0; b < batch_size_; ++b) { + const int64_t steps = member_steps[static_cast(b)]; + const size_t src_offset = static_cast(b) * static_cast(max_steps * row_elems_); + const size_t dst_offset = static_cast(b * cache_steps_ * row_elems_); + const size_t copy_elems = static_cast(steps * row_elems_); std::copy( source.key.begin() + static_cast(src_offset), source.key.begin() + static_cast(src_offset + copy_elems), @@ -306,12 +319,16 @@ TransformerBatchedKVState TransformerBatchedKVCache::export_state() const { TransformerBatchedKVState state; state.batch_size = batch_size_; state.current_end = current_end_; + if (!member_ends_.empty()) { + state.current_ends = member_ends_; + } state.layers.resize(layers_.size()); - const size_t copy_elems = static_cast(valid_steps_ * row_elems_); + const int64_t max_steps = valid_steps_; + const size_t copy_elems = static_cast(max_steps * row_elems_); const size_t state_elems = static_cast(batch_size_) * copy_elems; for (size_t layer = 0; layer < layers_.size(); ++layer) { auto & out = state.layers[layer]; - out.valid_steps = valid_steps_; + out.valid_steps = max_steps; out.key.resize(state_elems); out.value.resize(state_elems); if (copy_elems == 0) { @@ -344,6 +361,11 @@ void TransformerBatchedKVCache::advance_after_direct_append(int64_t steps) { } valid_steps_ += steps; current_end_ += steps; + if (!member_ends_.empty()) { + for (auto & end : member_ends_) { + end += steps; + } + } } int64_t TransformerBatchedKVCache::batch_size() const noexcept { @@ -362,6 +384,32 @@ int64_t TransformerBatchedKVCache::cache_steps() const noexcept { return cache_steps_; } +int64_t TransformerBatchedKVCache::member_end(int64_t batch) const noexcept { + if (member_ends_.empty()) { + return current_end_; + } + return member_ends_[static_cast(batch)]; +} + +void TransformerBatchedKVCache::set_member_end(int64_t batch, int64_t end) noexcept { + if (member_ends_.empty()) { + member_ends_.assign(static_cast(batch_size_), current_end_); + } + member_ends_[static_cast(batch)] = end; + current_end_ = std::max(current_end_, end); + valid_steps_ = std::max(valid_steps_, end); +} + +void TransformerBatchedKVCache::advance_member(int64_t batch, int64_t steps) noexcept { + const int64_t new_end = member_end(batch) + steps; + if (member_ends_.empty()) { + member_ends_.assign(static_cast(batch_size_), current_end_); + } + member_ends_[static_cast(batch)] = new_end; + current_end_ = std::max(current_end_, new_end); + valid_steps_ = std::max(valid_steps_, new_end); +} + core::TensorValue view_transformer_kv_cache_steps( core::ModuleBuildContext & ctx, const core::TensorValue & cache, From bb001c473efb4991b6edea87d0c69bd581e9b980 Mon Sep 17 00:00:00 2001 From: dujing Date: Tue, 1 Sep 2026 14:39:05 +0800 Subject: [PATCH 04/12] Add batch passthroughs + per-slot RedAE decode state (P2a-2, P2a-3) - FireRedArRuntime: expose start_decode_embeddings_batched / decode_embeddings_batched / export_batched_decode_state (thin wrappers over the Qwen backbone) so the scheduler can batch the AR decode. - RedAE codec: incremental decode state (decoder KV + ISTFT) hoisted into RedAeCodecRuntime::DecodeState, passed explicitly to decode_reset / decode_incremental / flush_incremental, enabling interleaved multi-slot streaming. FireRedRedAeRuntime threads the state; StreamSession owns a per-slot DecodeState. Co-Authored-By: Claude --- .../framework/codecs/redae_codec_runtime.h | 17 +++++- include/engine/models/fireredtts3/ar.h | 7 +++ include/engine/models/fireredtts3/redae.h | 10 ++-- src/framework/codecs/redae_codec_runtime.cpp | 56 +++++++++---------- src/models/fireredtts3/ar.cpp | 27 +++++++++ src/models/fireredtts3/pipeline.cpp | 9 +-- src/models/fireredtts3/redae.cpp | 24 ++++---- 7 files changed, 96 insertions(+), 54 deletions(-) diff --git a/include/engine/framework/codecs/redae_codec_runtime.h b/include/engine/framework/codecs/redae_codec_runtime.h index ba9ae36e6..c97ace51f 100644 --- a/include/engine/framework/codecs/redae_codec_runtime.h +++ b/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 #include #include +#include #include #include @@ -78,13 +81,21 @@ class RedAeCodecRuntime { runtime::AudioBuffer decode(const std::vector & latents); // --- 增量解码（流式）--- + // 每个并发 slot 独立持有解码状态（decoder KV + 增量 iSTFT）， + // 使得多 slot 交错的增量解码互不干扰。 + struct DecodeState { + std::optional dec_state; + int64_t dec_qwen_frames = 0; + std::unique_ptr inc_istft; + int64_t inc_istft_frames = 0; + }; // 重置解码器 KV 状态（每次新请求开始时调用）。 - void decode_reset(); + void decode_reset(DecodeState & state); // 解码一批 latent（chunk），返回该块对应的音频（float32 24k mono）。 // 内部用 decoder Qwen 的 KV 缓存跨块保持上下文，并用增量 iSTFT 逐块输出。 - runtime::AudioBuffer decode_incremental(const std::vector & latents); + runtime::AudioBuffer decode_incremental(DecodeState & state, const std::vector & latents); // flush 增量 iSTFT 尾部样本（生成结束时调用）。 - runtime::AudioBuffer flush_incremental(); + runtime::AudioBuffer flush_incremental(DecodeState & state); void release_runtime_graphs(); diff --git a/include/engine/models/fireredtts3/ar.h b/include/engine/models/fireredtts3/ar.h index 4da49135a..d3e5c5094 100644 --- a/include/engine/models/fireredtts3/ar.h +++ b/include/engine/models/fireredtts3/ar.h @@ -38,6 +38,13 @@ class FireRedArRuntime { void start_decode_embeddings(const engine::runtime::TransformerKVState & state, int64_t required_cache_steps); engine::modules::QwenCausalDecodeStepResult decode_embedding(const std::vector & 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 & embeddings, int64_t batch_size); + engine::runtime::TransformerBatchedKVState export_batched_decode_state() const; + void release_graphs(); void release_backbone_graphs(); diff --git a/include/engine/models/fireredtts3/redae.h b/include/engine/models/fireredtts3/redae.h index aed3df16b..7b79b8ff9 100644 --- a/include/engine/models/fireredtts3/redae.h +++ b/include/engine/models/fireredtts3/redae.h @@ -1,6 +1,7 @@ #pragma once #include "engine/framework/assets/tensor_source.h" +#include "engine/framework/codecs/redae_codec_runtime.h" #include "engine/framework/core/execution_context.h" #include "engine/framework/runtime/session.h" #include "engine/models/fireredtts3/assets.h" @@ -32,10 +33,11 @@ class FireRedRedAeRuntime { std::vector encode(const std::vector & audio_24k); engine::runtime::AudioBuffer decode(const std::vector & latents); - // 增量解码（流式）：重置 + 逐块解码 + flush 尾部 - void decode_reset(); - engine::runtime::AudioBuffer decode_incremental(const std::vector & latents); - engine::runtime::AudioBuffer flush_incremental(); + // 增量解码（流式）：重置 + 逐块解码 + flush 尾部（per-slot 状态） + using DecodeState = codecs::RedAeCodecRuntime::DecodeState; + void decode_reset(DecodeState & state); + engine::runtime::AudioBuffer decode_incremental(DecodeState & state, const std::vector & latents); + engine::runtime::AudioBuffer flush_incremental(DecodeState & state); void release_graphs(); diff --git a/src/framework/codecs/redae_codec_runtime.cpp b/src/framework/codecs/redae_codec_runtime.cpp index 494844ccd..37485e7d4 100644 --- a/src/framework/codecs/redae_codec_runtime.cpp +++ b/src/framework/codecs/redae_codec_runtime.cpp @@ -643,15 +643,15 @@ class RedAeRuntime { return audio; } - // --- 增量解码（流式）--- - void decode_reset() { - dec_state_.reset(); - dec_qwen_frames_ = 0; - inc_host_istft_.reset(); - inc_istft_frames_ = 0; + // --- 增量解码（流式）--- per-slot state + void decode_reset(RedAeCodecRuntime::DecodeState & st) { + st.dec_state.reset(); + st.dec_qwen_frames = 0; + st.inc_istft.reset(); + st.inc_istft_frames = 0; } - runtime::AudioBuffer decode_incremental(const std::vector & latents) { + runtime::AudioBuffer decode_incremental(RedAeCodecRuntime::DecodeState & st, const std::vector & latents) { if (latents.empty() || static_cast(latents.size()) % config_.bottleneck_dim != 0) { throw std::runtime_error("RedAE codec incremental decode latent size mismatch"); } @@ -660,15 +660,15 @@ class RedAeRuntime { const int64_t qwen_frames = latent_frames * config_.enc_extra_downsample_rate; std::vector decoder_hidden; - if (!dec_state_.has_value()) { + if (!st.dec_state.has_value()) { // 首块：一次性 prefill（返回 KV state 用于后续 chaining） auto decoder = decoder_qwen_.prefill_embeddings(decoder_in, qwen_frames); decoder_hidden = std::move(decoder.hidden); - dec_state_ = std::move(decoder.state); + st.dec_state = std::move(decoder.state); } else { // 后续块：用 decode_embedding 逐 qwen frame 解码，复用前序 KV - const int64_t required = dec_qwen_frames_ + qwen_frames; - decoder_qwen_.start_decode_embeddings(*dec_state_, required); + const int64_t required = st.dec_qwen_frames + qwen_frames; + decoder_qwen_.start_decode_embeddings(*st.dec_state, required); decoder_hidden.reserve(static_cast(qwen_frames * config_.dec_hidden_size)); for (int64_t f = 0; f < qwen_frames; ++f) { std::vector row( @@ -679,22 +679,22 @@ class RedAeRuntime { decoder_hidden.end(), step.hidden.begin(), step.hidden.end()); } } - dec_qwen_frames_ += qwen_frames; + st.dec_qwen_frames += qwen_frames; auto spec = encoder_out_.istft_head(decoder_hidden, qwen_frames); // 增量 iSTFT - if (!inc_host_istft_) { + if (!st.inc_istft) { audio::HostLogMagnitudePhaseISTFTConfig cfg; cfg.frames = qwen_frames; cfg.n_fft = config_.audio_patch_size * 4; cfg.hop_length = config_.audio_patch_size; cfg.out_dim = config_.audio_patch_size * 4 + 2; cfg.threads = static_cast(std::max(1, execution_.config().threads)); - inc_host_istft_ = std::make_unique(cfg); - inc_istft_frames_ = qwen_frames; + st.inc_istft = std::make_unique(cfg); + st.inc_istft_frames = qwen_frames; } - auto chunk_audio = inc_host_istft_->append_incremental(spec, qwen_frames, istft_window_); + auto chunk_audio = st.inc_istft->append_incremental(spec, qwen_frames, istft_window_); runtime::AudioBuffer audio; audio.sample_rate = static_cast(config_.sample_rate); @@ -703,12 +703,12 @@ class RedAeRuntime { return audio; } - runtime::AudioBuffer flush_incremental() { + runtime::AudioBuffer flush_incremental(RedAeCodecRuntime::DecodeState & st) { runtime::AudioBuffer audio; audio.sample_rate = static_cast(config_.sample_rate); audio.channels = 1; - if (inc_host_istft_) { - audio.samples = inc_host_istft_->finish_incremental(); + if (st.inc_istft) { + audio.samples = st.inc_istft->finish_incremental(); } return audio; } @@ -792,12 +792,6 @@ class RedAeRuntime { std::vector istft_window_; std::unique_ptr host_istft_; int64_t host_istft_frames_ = 0; - - // 增量解码状态 - std::optional dec_state_; - int64_t dec_qwen_frames_ = 0; - std::unique_ptr inc_host_istft_; - int64_t inc_istft_frames_ = 0; }; @@ -842,16 +836,16 @@ runtime::AudioBuffer RedAeCodecRuntime::decode(const std::vector & latent return impl_->runtime.decode(latents); } -void RedAeCodecRuntime::decode_reset() { - impl_->runtime.decode_reset(); +void RedAeCodecRuntime::decode_reset(RedAeCodecRuntime::DecodeState & state) { + impl_->runtime.decode_reset(state); } -runtime::AudioBuffer RedAeCodecRuntime::decode_incremental(const std::vector & latents) { - return impl_->runtime.decode_incremental(latents); +runtime::AudioBuffer RedAeCodecRuntime::decode_incremental(RedAeCodecRuntime::DecodeState & state, const std::vector & latents) { + return impl_->runtime.decode_incremental(state, latents); } -runtime::AudioBuffer RedAeCodecRuntime::flush_incremental() { - return impl_->runtime.flush_incremental(); +runtime::AudioBuffer RedAeCodecRuntime::flush_incremental(RedAeCodecRuntime::DecodeState & state) { + return impl_->runtime.flush_incremental(state); } void RedAeCodecRuntime::release_runtime_graphs() { diff --git a/src/models/fireredtts3/ar.cpp b/src/models/fireredtts3/ar.cpp index f340c59ea..27d956372 100644 --- a/src/models/fireredtts3/ar.cpp +++ b/src/models/fireredtts3/ar.cpp @@ -1011,6 +1011,18 @@ class FireRedArRuntime::Impl { return backbone_->decode_embedding(embedding); } + void start_decode_embeddings_batched(const runtime::TransformerBatchedKVState & state, int64_t required_cache_steps) { + backbone_->start_decode_embeddings_batched(state, required_cache_steps); + } + + modules::QwenCausalDecodeStepResult decode_embeddings_batched(const std::vector & embeddings, int64_t batch_size) { + return backbone_->decode_embeddings_batched(embeddings, batch_size); + } + + runtime::TransformerBatchedKVState export_batched_decode_state() const { + return backbone_->export_batched_decode_state(); + } + void release_graphs() { if (token_embedding_) { token_embedding_->release_graph(); @@ -1116,6 +1128,21 @@ engine::modules::QwenCausalDecodeStepResult FireRedArRuntime::decode_embedding(c return impl_->decode_embedding(embedding); } +void FireRedArRuntime::start_decode_embeddings_batched( + const engine::runtime::TransformerBatchedKVState & state, + int64_t required_cache_steps) { + impl_->start_decode_embeddings_batched(state, required_cache_steps); +} + +engine::modules::QwenCausalDecodeStepResult FireRedArRuntime::decode_embeddings_batched( + const std::vector & embeddings, int64_t batch_size) { + return impl_->decode_embeddings_batched(embeddings, batch_size); +} + +engine::runtime::TransformerBatchedKVState FireRedArRuntime::export_batched_decode_state() const { + return impl_->export_batched_decode_state(); +} + void FireRedArRuntime::release_graphs() { impl_->release_graphs(); } diff --git a/src/models/fireredtts3/pipeline.cpp b/src/models/fireredtts3/pipeline.cpp index 222fc4682..9b7161a09 100644 --- a/src/models/fireredtts3/pipeline.cpp +++ b/src/models/fireredtts3/pipeline.cpp @@ -280,6 +280,7 @@ class FireRedTTS3BaseRuntime::Impl { size_t chunk_index = 0; int64_t chunk_target = 0; std::vector chunk_latents; + FireRedRedAeRuntime::DecodeState redae_state; StreamSession(FireRedTTS3BaseRuntime::Impl & o, FireRedTTS3BaseRequest r, @@ -335,7 +336,7 @@ class FireRedTTS3BaseRuntime::Impl { ++step; if (generated_patches >= chunk_target && chunk_index < chunks.size()) { - auto audio = owner.redae_->decode_incremental(chunk_latents); + auto audio = owner.redae_->decode_incremental(redae_state, chunk_latents); chunk_latents.clear(); chunk_index++; chunk_target = (chunk_index < chunks.size()) @@ -346,12 +347,12 @@ class FireRedTTS3BaseRuntime::Impl { } // AR 结束：flush 剩余 latents + iSTFT 尾部 if (!chunk_latents.empty() && generated_patches > 0) { - auto tail = owner.redae_->decode_incremental(chunk_latents); + auto tail = owner.redae_->decode_incremental(redae_state, chunk_latents); chunk_latents.clear(); finished = true; return std::move(tail); } - auto flush = owner.redae_->flush_incremental(); + auto flush = owner.redae_->flush_incremental(redae_state); finished = true; return std::move(flush); } @@ -388,7 +389,7 @@ class FireRedTTS3BaseRuntime::Impl { next_input = {}; chunk_latents.reserve(static_cast(owner.assets_->base.patch_size * owner.assets_->base.redae_dim)); chunk_target = chunks[0]; - owner.redae_->decode_reset(); + owner.redae_->decode_reset(redae_state); started = true; } }; diff --git a/src/models/fireredtts3/redae.cpp b/src/models/fireredtts3/redae.cpp index 914e9a48b..63e15afcd 100644 --- a/src/models/fireredtts3/redae.cpp +++ b/src/models/fireredtts3/redae.cpp @@ -130,16 +130,16 @@ class FireRedRedAeRuntime::Impl { return runtime_.decode(latents); } - void decode_reset() { - runtime_.decode_reset(); + void decode_reset(FireRedRedAeRuntime::DecodeState & st) { + runtime_.decode_reset(st); } - runtime::AudioBuffer decode_incremental(const std::vector & latents) { - return runtime_.decode_incremental(latents); + runtime::AudioBuffer decode_incremental(FireRedRedAeRuntime::DecodeState & st, const std::vector & latents) { + return runtime_.decode_incremental(st, latents); } - runtime::AudioBuffer flush_incremental() { - return runtime_.flush_incremental(); + runtime::AudioBuffer flush_incremental(FireRedRedAeRuntime::DecodeState & st) { + return runtime_.flush_incremental(st); } void release_graphs() { @@ -169,16 +169,16 @@ engine::runtime::AudioBuffer FireRedRedAeRuntime::decode(const std::vectordecode(latents); } -void FireRedRedAeRuntime::decode_reset() { - impl_->decode_reset(); +void FireRedRedAeRuntime::decode_reset(DecodeState & st) { + impl_->decode_reset(st); } -engine::runtime::AudioBuffer FireRedRedAeRuntime::decode_incremental(const std::vector & latents) { - return impl_->decode_incremental(latents); +engine::runtime::AudioBuffer FireRedRedAeRuntime::decode_incremental(DecodeState & st, const std::vector & latents) { + return impl_->decode_incremental(st, latents); } -engine::runtime::AudioBuffer FireRedRedAeRuntime::flush_incremental() { - return impl_->flush_incremental(); +engine::runtime::AudioBuffer FireRedRedAeRuntime::flush_incremental(DecodeState & st) { + return impl_->flush_incremental(st); } void FireRedRedAeRuntime::release_graphs() { From 30d9a3794102df12b3fed7f7a471557c7fc1e751 Mon Sep 17 00:00:00 2001 From: dujing Date: Tue, 1 Sep 2026 17:25:29 +0800 Subject: [PATCH 05/12] Add shared batch scheduler + server integration for FireRedTTS3 (P2a-4, P2a-5) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 真·batch 并发推理：多 session 共享一个 FireRedTTS3BatchScheduler， 并发请求各占一个 slot，单调度线程每轮把活跃 slot 的 AR decode 合并成 一个 batch（per-member position/mask/cache-slot），实现共享 GPU batch。 - batch_scheduler: llama.cpp update_slots 的模拟。固定 max_batch graph （kMaxDecodeCacheSteps=700 不重建），wave batching prefill（规避 CUDA pool 逆序 free 约束），step-0 AR 不 consume decode，stop/flush/RedAE 增量与 pipeline.cpp 一致。ggml CUDA pool 逆序约束 -> 唯一 GPU decode 线程（调度线程）。 - padded prefill: prefill_embeddings_padded 固定 kMaxPrefillSteps graph （grow-only），避免 prefill graph 重建破坏 pool；padding 行保持因果 mask（非 -inf，避免 softmax NaN 污染 KV）。 - per-member RoPE 修复: q 布局 [batch,steps,heads,dim]，permute(0,1,3,2) 把 batch 移到 token 轴（原 permute(1,0,2,3) heads/dim 交换导致 flash 输出 shape mismatch）。 - write_qwen_batched_cached_step_mask 修复: current_slot 是绝对位置， mask 是 batch-major 布局，需 offset + pos（原 offset + current_slot 越界）。 - server 集成: instance_count 注入 fireredtts3.max_batch，session 池共享 scheduler（session.cpp 注册表）。runtime.cpp 多实例池保留，每 session 变 slot shim。 - session 流式加锁 + free_slots_ 防重复归还，修复并发竞态崩溃。 验证: 单请求/连续请求/2/3 并发全部 HTTP 200，并发请求与单请求音频 逐位一致（MD5 相同）。正确性门通过。 Co-Authored-By: Claude --- CMakeLists.txt | 2 + app/server/config.cpp | 8 + app/server/config.h | 4 + app/server/firered_server.json | 1 + app/server/runtime.cpp | 173 +++- app/server/runtime.h | 34 +- .../transformers/qwen_causal_decode_runtime.h | 8 + include/engine/models/fireredtts3/ar.h | 2 + .../models/fireredtts3/batch_scheduler.h | 87 ++ include/engine/models/fireredtts3/session.h | 10 + model_specs/fireredtts3.json | 8 + src/framework/core/module.cpp | 1 + .../scaled_dot_product_attention.cpp | 1 + .../qwen_causal_decode_runtime.cpp | 76 +- .../transformers/qwen_causal_decoder.cpp | 8 +- .../modules/transformers/qwen_decoder.cpp | 16 +- src/framework/runtime/kv_cache.cpp | 1 + src/models/fireredtts3/ar.cpp | 12 + src/models/fireredtts3/batch_scheduler.cpp | 896 ++++++++++++++++++ src/models/fireredtts3/flow.cpp | 2 + src/models/fireredtts3/session.cpp | 124 ++- 21 files changed, 1417 insertions(+), 57 deletions(-) create mode 100644 include/engine/models/fireredtts3/batch_scheduler.h create mode 100644 src/models/fireredtts3/batch_scheduler.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 63195ff21..27353b80f 100644 --- a/CMakeLists.txt +++ b/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 diff --git a/app/server/config.cpp b/app/server/config.cpp index 1199dd15d..d23b97ff4 100644 --- a/app/server/config.cpp +++ b/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"); diff --git a/app/server/config.h b/app/server/config.h index 43f83f020..183ba2d84 100644 --- a/app/server/config.h +++ b/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 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; diff --git a/app/server/firered_server.json b/app/server/firered_server.json index cf17ef0b3..aa570993c 100644 --- a/app/server/firered_server.json +++ b/app/server/firered_server.json @@ -14,6 +14,7 @@ "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" diff --git a/app/server/runtime.cpp b/app/server/runtime.cpp index 74d1e3807..2e2812df0 100644 --- a/app/server/runtime.cpp +++ b/app/server/runtime.cpp @@ -1,3 +1,4 @@ +#include #include "runtime.h" #include "base64.h" @@ -202,6 +203,13 @@ ServerModelConfig model_config_from_json( } model.load_options = options_from_object(body.find("load_options")); model.session_options = options_from_object(body.find("session_options")); + if (const auto * value = body.find("instance_count")) { + const int n = value->as_i64(); + if (n < 1 || n > 64) { + throw std::runtime_error("instance_count must be in [1, 64]"); + } + model.instance_count = n; + } return model; } @@ -1710,7 +1718,7 @@ void ServerState::evict_for_model_limit(const LoadedModel & loading) { { std::lock_guard state_lock(models_mutex_); for (const auto & model : models_) { - if (model.get() != &loading && model->session != nullptr) { + if (model.get() != &loading && !model->sessions.empty()) { resident.push_back(model.get()); } } @@ -1748,7 +1756,7 @@ void ServerState::evict_for_model_limit(const LoadedModel & loading) { void ServerState::ensure_model_loaded_locked(LoadedModel & model) { last_activity_ms_.store(steady_now_ms(), std::memory_order_relaxed); model.last_used_ms.store(steady_now_ms(), std::memory_order_relaxed); - if (model.session != nullptr) { + if (!model.sessions.empty()) { return; } // Serialize the whole "evict -> memory check -> load" sequence only when a guard @@ -1785,6 +1793,13 @@ void ServerState::ensure_model_loaded_locked(LoadedModel & model) { session_options.backend.device = config_.device; session_options.backend.threads = config_.threads; session_options.options = model.config.session_options; + // 并发：把 instance_count 注入 session option，供 session 层共享 scheduler 使用。 + // instance_count 个 session 共享一个 scheduler（max_batch = instance_count）， + // 并发请求经各 session 进入 scheduler 的不同 slot，实现共享 GPU batch decode。 + if (model.config.instance_count > 1) { + session_options.options["fireredtts3.max_batch"] = + std::to_string(model.config.instance_count); + } engine::debug::trace_log_scalar("server.model.id", model.config.id); engine::debug::trace_log_scalar("server.model.path", model.config.path.string()); @@ -1806,17 +1821,29 @@ void ServerState::ensure_model_loaded_locked(LoadedModel & model) { } auto loaded_model = registry.load(load_request); - auto session = loaded_model->create_task_session(model.task, session_options); - auto * offline = dynamic_cast(session.get()); - auto * streaming = dynamic_cast(session.get()); - if (model.task.mode == engine::runtime::RunMode::Offline && offline == nullptr) { - throw std::runtime_error("configured model does not provide offline execution: " + model.config.id); - } - if (model.task.mode == engine::runtime::RunMode::Streaming && streaming == nullptr) { - throw std::runtime_error("configured model does not provide streaming execution: " + model.config.id); + // 创建 session 池（instance_count 个实例） + const int instance_count = std::max(1, model.config.instance_count); + model.sessions.clear(); + model.sessions.reserve(static_cast(instance_count)); + model.free_sessions.clear(); + engine::runtime::IOfflineVoiceTaskSession * offline = nullptr; + engine::runtime::IStreamingVoiceTaskSession * streaming = nullptr; + for (int i = 0; i < instance_count; ++i) { + auto session = loaded_model->create_task_session(model.task, session_options); + if (i == 0) { + offline = dynamic_cast(session.get()); + streaming = dynamic_cast(session.get()); + if (model.task.mode == engine::runtime::RunMode::Offline && offline == nullptr) { + throw std::runtime_error("configured model does not provide offline execution: " + model.config.id); + } + if (model.task.mode == engine::runtime::RunMode::Streaming && streaming == nullptr) { + throw std::runtime_error("configured model does not provide streaming execution: " + model.config.id); + } + } + model.sessions.push_back(std::move(session)); + model.free_sessions.push_back(static_cast(i)); } model.model = std::move(loaded_model); - model.session = std::move(session); model.offline = offline; model.streaming = streaming; model.loaded.store(true); @@ -2010,6 +2037,61 @@ struct ServerState::TimedTaskResult { std::optional ttft_ms; }; +ServerState::SessionPoolLock::SessionPoolLock(ServerState::LoadedModel & model, size_t index) + : model_(&model), index_(index) {} + +ServerState::SessionPoolLock::SessionPoolLock(SessionPoolLock && other) noexcept + : model_(other.model_), index_(other.index_) { other.model_ = nullptr; } + +ServerState::SessionPoolLock & ServerState::SessionPoolLock::operator=(SessionPoolLock && other) noexcept { + if (this != &other) { + release(); + model_ = other.model_; + index_ = other.index_; + other.model_ = nullptr; + } + return *this; +} + +ServerState::SessionPoolLock::~SessionPoolLock() { release(); } + +void ServerState::SessionPoolLock::release() { + if (model_ != nullptr) { + std::lock_guard lock(model_->pool_mutex); + model_->free_sessions.push_back(index_); + model_ = nullptr; + } +} + +ServerState::SessionPoolLock ServerState::borrow_session( + LoadedModel & model, + std::optional request_timeout_ms) { + // 等待一个空闲 session（阻塞，超时抛 503-like 异常） + const int timeout_ms = request_timeout_ms.value_or(0); + const auto deadline = timeout_ms > 0 + ? std::chrono::steady_clock::now() + std::chrono::milliseconds(timeout_ms) + : std::chrono::steady_clock::time_point::max(); + while (true) { + { + std::unique_lock lock(model.pool_mutex); + if (!model.free_sessions.empty()) { + const size_t idx = model.free_sessions.front(); + model.free_sessions.pop_front(); + return SessionPoolLock(model, idx); + } + if (timeout_ms > 0 && std::chrono::steady_clock::now() >= deadline) { + throw ServerBusyError( + "model '" + model.config.id + "' is busy: all session instances are in use"); + } + } + if (timeout_ms > 0) { + std::this_thread::sleep_for(std::chrono::milliseconds(5)); + } else { + std::this_thread::sleep_for(std::chrono::milliseconds(5)); + } + } +} + engine::runtime::RunMode ServerState::model_run_mode(const LoadedModel & model) const { std::shared_lock metadata_lock(model.metadata_mutex); return model.task.mode; @@ -2034,14 +2116,19 @@ ServerState::TimedTaskResult ServerState::run_model( LoadedModel & model, const engine::runtime::TaskRequest & request, std::optional busy_timeout_ms) { - BusyGuard::Lock lock = acquire_model_run(model, busy_timeout_ms); ensure_model_loaded_locked(model); + SessionPoolLock pool = borrow_session(model, busy_timeout_ms); if (model.offline == nullptr) { throw std::runtime_error("configured model does not provide offline execution: " + model.config.id); } + auto * session = model.sessions[pool.index].get(); + auto * offline = dynamic_cast(session); + if (offline == nullptr) { + throw std::runtime_error("configured model does not provide offline execution: " + model.config.id); + } const auto started = Clock::now(); - model.session->prepare(engine::runtime::build_preparation_request(request)); - auto result = model.offline->run(request); + session->prepare(engine::runtime::build_preparation_request(request)); + auto result = offline->run(request); // Mark activity at completion too: idle unload must measure from when the // request finished, not when it started, or a long inference would look idle // (and be unloaded) the moment it returns. @@ -2059,13 +2146,18 @@ ServerState::TimedTaskResult ServerState::run_streaming_model_impl( const minitts::app::AudioChunkStream * audio, const std::function & event_sink, std::optional busy_timeout_ms) { - BusyGuard::Lock lock = acquire_model_run(model, busy_timeout_ms); ensure_model_loaded_locked(model); + SessionPoolLock pool = borrow_session(model, busy_timeout_ms); if (model.streaming == nullptr) { throw std::runtime_error("configured model does not provide streaming execution: " + model.config.id); } + auto * session = model.sessions[pool.index].get(); + auto * streaming = dynamic_cast(session); + if (streaming == nullptr) { + throw std::runtime_error("configured model does not provide streaming execution: " + model.config.id); + } const auto started = Clock::now(); - model.session->prepare(engine::runtime::build_preparation_request(request)); + session->prepare(engine::runtime::build_preparation_request(request)); TimedTaskResult timed_result; const auto sink = [&](const engine::runtime::StreamEvent & event) { if (!timed_result.ttft_ms.has_value() && stream_event_has_output(event)) { @@ -2076,8 +2168,8 @@ ServerState::TimedTaskResult ServerState::run_streaming_model_impl( } }; auto result = audio != nullptr - ? minitts::app::run_streaming_task(*model.streaming, request, sink, *audio) - : minitts::app::run_streaming_task(*model.streaming, request, sink); + ? minitts::app::run_streaming_task(*streaming, request, sink, *audio) + : minitts::app::run_streaming_task(*streaming, request, sink); timed_result.result = std::move(result); timed_result.wall_ms = elapsed_ms(started); if (!timed_result.ttft_ms.has_value() && task_result_has_output(timed_result.result)) { @@ -2884,7 +2976,11 @@ std::string ServerState::get_allowed_origin(const HttpRequest & request) const { void ServerState::LoadedModel::unload() { offline = nullptr; streaming = nullptr; - session.reset(); + { + std::lock_guard lock(pool_mutex); + sessions.clear(); + free_sessions.clear(); + } model.reset(); loaded.store(false); } @@ -2916,18 +3012,20 @@ void ServerState::unload_idle_models() { { std::lock_guard state_lock(models_mutex_); for (const auto & model : models_) { - if (model->session != nullptr) { + if (!model->sessions.empty()) { resident.push_back(model.get()); } } } int unloaded = 0; for (LoadedModel * model : resident) { - // Never unload a model mid-inference; a busy model keeps its slot and the - // next idle pass retries it. - const auto lock = model->busy.try_acquire(); - if (!lock.has_value()) { - continue; + // Never unload a model mid-inference; a borrowed session keeps it loaded and + // the next idle pass retries it. + { + std::lock_guard pool_lock(model->pool_mutex); + if (model->free_sessions.size() != model->sessions.size()) { + continue; + } } model->unload(); ++unloaded; @@ -3014,13 +3112,15 @@ HttpResponse ServerState::handle_unload_models(const std::string & body_text) { continue; } LoadedModel & model = *models_.at(it->second); - // Only unload if the model is currently loaded in memory. Acquire the busy - // lock for the duration of the unload so no inference starts mid-operation. - if (model.session != nullptr) { - [[maybe_unused]] BusyGuard::Lock lock = model.busy.acquire(0, model.config.id); - model.unload(); - unloaded.push_back(id); + // Only unload if loaded AND all sessions idle (none borrowed). + { + std::lock_guard pool_lock(model.pool_mutex); + if (model.sessions.empty() || model.free_sessions.size() != model.sessions.size()) { + continue; + } } + model.unload(); + unloaded.push_back(id); } std::ostringstream out; @@ -3043,11 +3143,14 @@ HttpResponse ServerState::handle_unload_all_models() { std::vector unloaded; for (auto & model : models_) { - if (model->session != nullptr) { - [[maybe_unused]] BusyGuard::Lock lock = model->busy.acquire(0, model->config.id); - model->unload(); - unloaded.push_back(model->config.id); + { + std::lock_guard pool_lock(model->pool_mutex); + if (model->sessions.empty() || model->free_sessions.size() != model->sessions.size()) { + continue; + } } + model->unload(); + unloaded.push_back(model->config.id); } std::ostringstream out; diff --git a/app/server/runtime.h b/app/server/runtime.h index 07c5a52b9..4e66cbfdc 100644 --- a/app/server/runtime.h +++ b/app/server/runtime.h @@ -13,6 +13,10 @@ #include "engine/framework/runtime/model.h" #include "engine/framework/runtime/session.h" +#include +#include +#include + #include #include #include @@ -44,6 +48,27 @@ 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(); + + size_t index = 0; + + private: + void release(); + LoadedModel * model_ = nullptr; + size_t index_ = 0; + }; + struct LoadedModel { struct RuntimeVoicePreset { std::optional voice_id; @@ -54,10 +79,14 @@ class ServerState final : public IHttpHandler { ServerModelConfig config; engine::runtime::TaskSpec task; std::unique_ptr model; - std::unique_ptr session; + // 并发 session 池：每个实例独立 graph arena + reference cache。 + std::vector> sessions; engine::runtime::IOfflineVoiceTaskSession * offline = nullptr; engine::runtime::IStreamingVoiceTaskSession * streaming = nullptr; std::atomic loaded{false}; + // 空闲 session 索引队列（受 pool_mutex 保护） + std::mutex pool_mutex; + std::deque 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 +114,9 @@ class ServerState final : public IHttpHandler { // (-> HTTP 503) once the effective timeout has elapsed. BusyGuard::Lock acquire_model_run(LoadedModel & model, std::optional request_timeout_ms); + // 从 session 池借一个空闲实例（真并发）；析构自动归还。池满时阻塞/超时。 + SessionPoolLock borrow_session(LoadedModel & model, std::optional 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; diff --git a/include/engine/framework/modules/transformers/qwen_causal_decode_runtime.h b/include/engine/framework/modules/transformers/qwen_causal_decode_runtime.h index 15063d25a..1e813db39 100644 --- a/include/engine/framework/modules/transformers/qwen_causal_decode_runtime.h +++ b/include/engine/framework/modules/transformers/qwen_causal_decode_runtime.h @@ -68,6 +68,14 @@ class QwenCausalDecodeRuntime { QwenCausalPrefillResult prefill_tokens(const std::vector & token_ids); QwenCausalPrefillResult prefill_embeddings(const std::vector & 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 & embeddings, + int64_t padded_steps, + int64_t valid_steps); QwenCausalBatchedPrefillResult prefill_tokens_batched( const std::vector & token_ids, diff --git a/include/engine/models/fireredtts3/ar.h b/include/engine/models/fireredtts3/ar.h index d3e5c5094..82e99aaa5 100644 --- a/include/engine/models/fireredtts3/ar.h +++ b/include/engine/models/fireredtts3/ar.h @@ -35,6 +35,8 @@ class FireRedArRuntime { std::vector text_logits(const std::vector & hidden); engine::modules::QwenCausalPrefillResult prefill_embeddings(const std::vector & embeddings, int64_t steps); + engine::modules::QwenCausalPrefillResult prefill_embeddings_padded( + const std::vector & 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 & embedding); diff --git a/include/engine/models/fireredtts3/batch_scheduler.h b/include/engine/models/fireredtts3/batch_scheduler.h new file mode 100644 index 000000000..38c7a4e0f --- /dev/null +++ b/include/engine/models/fireredtts3/batch_scheduler.h @@ -0,0 +1,87 @@ +#pragma once + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/core/execution_context.h" +#include "engine/framework/runtime/kv_cache.h" +#include "engine/framework/runtime/session.h" +#include "engine/models/fireredtts3/assets.h" +#include "engine/models/fireredtts3/pipeline.h" +#include "engine/models/fireredtts3/redae.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::fireredtts3 { + + + +// 真·batch 并发调度器（llama.cpp update_slots 的模拟）。 +// 一个共享 FireRedArRuntime + FireRedFlowRuntime + FireRedRedAeRuntime， +// 多个 slot（请求）每轮把活跃 slot 的 AR decode 合并成一个 batch +// （decode_embeddings_batched，per-member position）。固定 max_batch graph， +// 运行期不重建；slot 中途 stop 掉队（dead 行零填充），新 slot 通过 +// export -> splice -> import 加入。 +class FireRedTTS3BatchScheduler { +public: + // 一个并发请求 = 一个 slot（持有其 AR + RedAE 增量解码状态 + chunk 队列） + struct Slot { + enum class State { Idle, Active, Dead, Done, Failed }; + State state = State::Idle; + FireRedTTS3BaseRequest request; + std::vector chunks; // 块 patch 数（流式） + // AR 状态 + std::vector latents_gen, backbone_cond, schedule, next_input, prefill_hidden; + std::vector spk_dit; + int64_t prefill_steps = 0, prompt_latent_frames = 0, step = 0, generated_patches = 0; + int64_t chunk_index = 0, chunk_target = 0; + std::vector chunk_latents; + // 参考音色 prep 产物（wave 形成时拷进 slot，不持有 cache 引用跨轮） + std::vector prompt_latents; + // prefill 单路径 KV state（CPU，用于拼进 batched KV） + engine::runtime::TransformerKVState prefill_state; + bool prefill_done = false; + // RedAE 增量解码状态（per-slot） + FireRedRedAeRuntime::DecodeState redae_state; + // 流式输出 + std::deque chunk_queue; + bool finished = false; + std::exception_ptr error; + }; + + FireRedTTS3BatchScheduler( + std::shared_ptr assets, + engine::core::ExecutionContext & execution, + size_t graph_arena_bytes, + size_t helper_graph_arena_bytes, + size_t weight_context_bytes, + engine::assets::TensorStorageType storage_type, + size_t reference_cache_slots, + bool mem_saver, + int64_t max_batch); + ~FireRedTTS3BatchScheduler(); + + FireRedTTS3BatchScheduler(const FireRedTTS3BatchScheduler &) = delete; + FireRedTTS3BatchScheduler & operator=(const FireRedTTS3BatchScheduler &) = delete; + + // 启动一个请求（分配 slot），返回 slot id。返回 -1 表示槽池耗尽。 + int64_t launch(const FireRedTTS3BaseRequest & request, const std::vector & chunk_patches); + // 取下一个音频块（驱动调度轮次直到该 slot 产出块或结束）。 + engine::runtime::AudioBuffer next_chunk(int64_t slot_id); + // 离线整句：launch + drain + 拼接（trim prompt）。 + engine::runtime::AudioBuffer generate(const FireRedTTS3BaseRequest & request); + void release_graphs(); + + // 槽池容量（max_batch）。 + int64_t max_batch() const noexcept; + +private: + class Impl; + std::unique_ptr impl_; +}; + +} // namespace engine::models::fireredtts3 diff --git a/include/engine/models/fireredtts3/session.h b/include/engine/models/fireredtts3/session.h index b83997630..b530dea9a 100644 --- a/include/engine/models/fireredtts3/session.h +++ b/include/engine/models/fireredtts3/session.h @@ -4,9 +4,11 @@ #include "engine/framework/runtime/session_base.h" #include "engine/framework/model_spec/metadata.h" #include "engine/models/fireredtts3/assets.h" +#include "engine/models/fireredtts3/batch_scheduler.h" #include "engine/models/fireredtts3/tokenizer_text.h" #include +#include namespace engine::models::fireredtts3 { @@ -54,6 +56,11 @@ class FireRedTTS3Session final std::unique_ptr runtime_; std::unique_ptr instruct_runtime_; bool mem_saver_ = false; + // 共享 batch scheduler（同一 assets 的所有 session 共享）；Base clone 用。 + std::shared_ptr scheduler_; + // 当前流式请求在 scheduler 中占用的 slot（-1 = 无）。 + int64_t scheduler_slot_ = -1; + bool scheduler_enabled_ = false; // 流式状态 bool stream_started_ = false; @@ -63,6 +70,9 @@ class FireRedTTS3Session final engine::runtime::AudioBuffer stream_merged_audio_; std::string stream_generated_text_; std::unique_ptr stream_session_; + // 流式方法互斥（start_stream / next_stream_event / finish_stream）： + // 防止同一 session 被并发驱动（防御性，borrow 池应保证独占）。 + std::mutex stream_mutex_; }; } // namespace engine::models::fireredtts3 diff --git a/model_specs/fireredtts3.json b/model_specs/fireredtts3.json index 5bf902d56..d24ec36e5 100644 --- a/model_specs/fireredtts3.json +++ b/model_specs/fireredtts3.json @@ -223,6 +223,14 @@ "description": "Streaming chunk sizes in AR latent patches (comma-separated). e.g. 3,12,12 -> first chunk ~0.5s, then ~2s each.", "required": false, "default": "3,12,12,12,12,12,12,12,12,12,12,12,12,12,12,12" + }, + { + "name": "max_batch", + "type": "int", + "description": "Shared batch-scheduler concurrency slots. Set >1 to enable true batch decode across concurrent requests sharing one GPU model; default 1 (per-session serial).", + "required": false, + "min": 1, + "default": 1 } ], "load": [] diff --git a/src/framework/core/module.cpp b/src/framework/core/module.cpp index d724d3e5f..2fa1b2020 100644 --- a/src/framework/core/module.cpp +++ b/src/framework/core/module.cpp @@ -1,3 +1,4 @@ +#include #include "engine/framework/core/module.h" #include diff --git a/src/framework/modules/attention/scaled_dot_product_attention.cpp b/src/framework/modules/attention/scaled_dot_product_attention.cpp index 6514f1d58..c18818e2c 100644 --- a/src/framework/modules/attention/scaled_dot_product_attention.cpp +++ b/src/framework/modules/attention/scaled_dot_product_attention.cpp @@ -1,3 +1,4 @@ +#include #include "engine/framework/modules/attention/scaled_dot_product_attention.h" #include "attention_internal.h" diff --git a/src/framework/modules/transformers/qwen_causal_decode_runtime.cpp b/src/framework/modules/transformers/qwen_causal_decode_runtime.cpp index e89548511..2e1168a82 100644 --- a/src/framework/modules/transformers/qwen_causal_decode_runtime.cpp +++ b/src/framework/modules/transformers/qwen_causal_decode_runtime.cpp @@ -477,6 +477,54 @@ class QwenCausalDecodeRuntime::Impl { return run_prefill(); } + QwenCausalPrefillResult prefill_embeddings_padded( + const std::vector & embeddings, + int64_t padded_steps, + int64_t valid_steps) { + if (padded_steps <= 0 || valid_steps <= 0 || valid_steps > padded_steps) { + throw std::runtime_error("QwenCausalDecodeRuntime padded prefill requires 0 < valid_steps <= padded_steps"); + } + const size_t expected = static_cast(padded_steps * config_.decoder.stack.hidden_size); + if (embeddings.size() != expected) { + throw std::runtime_error("QwenCausalDecodeRuntime padded prefill embedding size mismatch"); + } + // grow-only graph：padded_steps 建一次，之后复用（不重建）。 + ensure_prefill_embedding_graph(padded_steps); + ggml_backend_tensor_set( + prefill_input_, + embeddings.data(), + 0, + embeddings.size() * sizeof(float)); + // 位置：有效部分 0..valid_steps-1，padding 部分 valid_steps 起继续递增。 + // 注意：padding 行保持标准因果 mask（不设 -inf），softmax 有界（非 NaN）， + // 避免 NaN 污染 KV cache 与后续 decode 的有效位置 attention。 + if (prefill_positions_values_.size() != static_cast(padded_steps)) { + prefill_positions_values_ = qwen_position_ids(padded_steps); + } + ggml_backend_tensor_set( + prefill_positions_, + prefill_positions_values_.data(), + 0, + prefill_positions_values_.size() * sizeof(int32_t)); + // mask：标准因果（padding 行也 attend 前面的 token，softmax 有界）。 + if (prefill_attention_mask_values_.size() != static_cast(padded_steps * padded_steps)) { + prefill_attention_mask_values_ = prefill_attention_mask_values(config_, 1, padded_steps); + } + ggml_backend_tensor_set( + prefill_attention_mask_, + prefill_attention_mask_values_.data(), + 0, + prefill_attention_mask_values_.size() * sizeof(ggml_fp16_t)); + if (prefill_logits_readback_token_ids_ != nullptr) { + upload_logits_readback_token_ids(prefill_logits_readback_token_ids_, config_); + } + // 导出时截断到 valid_steps + prefill_export_steps_ = valid_steps; + auto result = run_prefill(); + prefill_export_steps_ = -1; + return result; + } + QwenCausalBatchedPrefillResult prefill_tokens_batched( const std::vector & token_ids, int64_t batch_size, @@ -643,7 +691,10 @@ class QwenCausalDecodeRuntime::Impl { } void ensure_prefill_embedding_graph(int64_t steps) { - if (prefill_graph_ != nullptr && prefill_input_kind_ == InputKind::Embedding && prefill_steps_ == steps) { + // grow-only：graph 建为请求过的最大 steps，之后 steps <= 已建则复用 + // （padded prefill 需要；避免因 steps 变化重建 prefill graph 破坏 + // CUDA pool 逆序 free 约束）。 + if (prefill_graph_ != nullptr && prefill_input_kind_ == InputKind::Embedding && prefill_steps_ >= steps) { debug::timing_log_scalar(config_.trace_name + ".prefill.graph.build_ms", 0.0); debug::trace_log_scalar(config_.trace_name + ".prefill.steps", steps); return; @@ -795,13 +846,14 @@ class QwenCausalDecodeRuntime::Impl { ggml_backend_tensor_get(prefill_hidden_, out.hidden.data(), 0, out.hidden.size() * sizeof(float)); round_readback(out.hidden, config_); } - out.state.current_end = prefill_steps_; + const int64_t export_steps = prefill_export_steps_ >= 0 ? prefill_export_steps_ : prefill_steps_; + out.state.current_end = export_steps; out.state.layers.resize(prefill_keys_.size()); const size_t layer_values = static_cast( - prefill_steps_ * config_.decoder.stack.num_key_value_heads * config_.decoder.stack.head_dim); + export_steps * config_.decoder.stack.num_key_value_heads * config_.decoder.stack.head_dim); for (size_t layer = 0; layer < prefill_keys_.size(); ++layer) { auto & state = out.state.layers[layer]; - state.valid_steps = prefill_steps_; + state.valid_steps = export_steps; state.key.resize(layer_values); state.value.resize(layer_values); ggml_backend_tensor_get(prefill_keys_[layer], state.key.data(), 0, state.key.size() * sizeof(float)); @@ -809,6 +861,13 @@ class QwenCausalDecodeRuntime::Impl { round_readback(state.key, config_); round_readback(state.value, config_); } + // padded prefill：hidden 也截断到 valid_steps 行（scheduler 只取有效部分） + if (prefill_export_steps_ >= 0 && !out.hidden.empty()) { + const size_t hidden_elems = static_cast(export_steps * config_.decoder.stack.hidden_size); + if (out.hidden.size() > hidden_elems) { + out.hidden.resize(hidden_elems); + } + } return out; } @@ -1479,6 +1538,8 @@ class QwenCausalDecodeRuntime::Impl { std::vector prefill_positions_values_; std::vector prefill_attention_mask_values_; int64_t prefill_steps_ = 0; + // >=0 时表示 padded prefill 的导出 steps（截断 state/hidden 到 valid_steps）；-1 = 正常。 + int64_t prefill_export_steps_ = -1; InputKind prefill_input_kind_ = InputKind::None; std::unique_ptr batched_prefill_ctx_; @@ -1551,6 +1612,13 @@ QwenCausalPrefillResult QwenCausalDecodeRuntime::prefill_embeddings( return impl_->prefill_embeddings(embeddings, steps); } +QwenCausalPrefillResult QwenCausalDecodeRuntime::prefill_embeddings_padded( + const std::vector & embeddings, + int64_t padded_steps, + int64_t valid_steps) { + return impl_->prefill_embeddings_padded(embeddings, padded_steps, valid_steps); +} + QwenCausalBatchedPrefillResult QwenCausalDecodeRuntime::prefill_tokens_batched( const std::vector & token_ids, int64_t batch_size, diff --git a/src/framework/modules/transformers/qwen_causal_decoder.cpp b/src/framework/modules/transformers/qwen_causal_decoder.cpp index 8be371389..bdc33d3e4 100644 --- a/src/framework/modules/transformers/qwen_causal_decoder.cpp +++ b/src/framework/modules/transformers/qwen_causal_decoder.cpp @@ -1,3 +1,4 @@ +#include #include "engine/framework/modules/transformers/qwen_causal_decoder.h" #include "engine/framework/core/backend.h" @@ -530,7 +531,12 @@ void write_qwen_batched_cached_step_mask( for (int64_t i = 0; i < visible_prefix_steps; ++i) { scratch[offset + static_cast(i)] = visible; } - scratch[offset + static_cast(current_slot)] = visible; + // current_slot 是绝对 cache 位置（batch*cache_steps + pos）。mask 是 batch-major + // 布局（每行 offset = batch*row_size），因此当前 step 的相对位置 = current_slot - batch*row_size。 + const int64_t pos = current_slot - batch * static_cast(row_size); + if (pos >= 0 && pos < static_cast(row_size)) { + scratch[offset + static_cast(pos)] = visible; + } } ggml_backend_tensor_set(tensor, scratch.data(), 0, scratch.size() * sizeof(ggml_fp16_t)); } diff --git a/src/framework/modules/transformers/qwen_decoder.cpp b/src/framework/modules/transformers/qwen_decoder.cpp index 5a0fd5b54..3dd9b4536 100644 --- a/src/framework/modules/transformers/qwen_decoder.cpp +++ b/src/framework/modules/transformers/qwen_decoder.cpp @@ -1,3 +1,4 @@ +#include #include "engine/framework/modules/transformers/qwen_decoder.h" #include "engine/framework/modules/activation_modules.h" @@ -865,20 +866,25 @@ QwenDecoderLayerOutputs QwenDecoderLayerModule::build_with_static_cache_tail_bat // Permute q/k from [batch, heads, 1, dim] to [1, batch, heads, dim] so batch sits // on ggml ne[2] (token axis); RoPE takes per-member pos; permute back. if (positions.shape.dims[0] > 1 && positions.shape.dims[0] == q.shape.dims[0]) { + // q 布局是 [batch, steps, heads, dim]（reshape_qwen_heads 产生）， + // ggml ne = {dim, heads, steps, batch}。RoPE 沿 token 轴（ne[2]）应用位置， + // 而 per-member 的 positions 长度 == batch，因此要把 batch 移到 token 轴： + // permute(0,1,3,2) 交换 ne[2](steps) 与 ne[3](batch)，逻辑 shape 变 [1, batch, heads, dim]。 const int64_t batch = q.shape.dims[0]; - const int64_t heads = q.shape.dims[1]; auto perm_batch_to_token = [&](core::TensorValue x) { const auto contiguous = core::ensure_backend_addressable_layout(ctx, x); + const int64_t x_heads = x.shape.dims[2]; return core::wrap_tensor( - ggml_permute(ctx.ggml, contiguous.tensor, 1, 0, 2, 3), - core::TensorShape::from_dims({1, batch, heads, x.shape.last_dim()}), + ggml_permute(ctx.ggml, contiguous.tensor, 0, 1, 3, 2), + core::TensorShape::from_dims({1, batch, x_heads, x.shape.last_dim()}), x.type); }; auto perm_back = [&](core::TensorValue x) { const auto contiguous = core::ensure_backend_addressable_layout(ctx, x); + const int64_t x_heads = x.shape.dims[2]; return core::wrap_tensor( - ggml_permute(ctx.ggml, contiguous.tensor, 1, 0, 2, 3), - core::TensorShape::from_dims({batch, heads, 1, x.shape.last_dim()}), + ggml_permute(ctx.ggml, contiguous.tensor, 0, 1, 3, 2), + core::TensorShape::from_dims({batch, 1, x_heads, x.shape.last_dim()}), x.type); }; q = perm_batch_to_token(q); diff --git a/src/framework/runtime/kv_cache.cpp b/src/framework/runtime/kv_cache.cpp index 2bf307ec3..03a579268 100644 --- a/src/framework/runtime/kv_cache.cpp +++ b/src/framework/runtime/kv_cache.cpp @@ -1,3 +1,4 @@ +#include #include "engine/framework/runtime/kv_cache.h" #include "engine/framework/core/backend.h" diff --git a/src/models/fireredtts3/ar.cpp b/src/models/fireredtts3/ar.cpp index 27d956372..9cf82ae21 100644 --- a/src/models/fireredtts3/ar.cpp +++ b/src/models/fireredtts3/ar.cpp @@ -1003,6 +1003,11 @@ class FireRedArRuntime::Impl { return backbone_->prefill_embeddings(embeddings, steps); } + modules::QwenCausalPrefillResult prefill_embeddings_padded( + const std::vector & embeddings, int64_t padded_steps, int64_t valid_steps) { + return backbone_->prefill_embeddings_padded(embeddings, padded_steps, valid_steps); + } + void start_decode_embeddings(const runtime::TransformerKVState & state, int64_t required_cache_steps) { backbone_->start_decode_embeddings(state, required_cache_steps); } @@ -1118,6 +1123,13 @@ engine::modules::QwenCausalPrefillResult FireRedArRuntime::prefill_embeddings( return impl_->prefill_embeddings(embeddings, steps); } +engine::modules::QwenCausalPrefillResult FireRedArRuntime::prefill_embeddings_padded( + const std::vector & embeddings, + int64_t padded_steps, + int64_t valid_steps) { + return impl_->prefill_embeddings_padded(embeddings, padded_steps, valid_steps); +} + void FireRedArRuntime::start_decode_embeddings( const engine::runtime::TransformerKVState & state, int64_t required_cache_steps) { diff --git a/src/models/fireredtts3/batch_scheduler.cpp b/src/models/fireredtts3/batch_scheduler.cpp new file mode 100644 index 000000000..0ab9ef489 --- /dev/null +++ b/src/models/fireredtts3/batch_scheduler.cpp @@ -0,0 +1,896 @@ +#include "engine/models/fireredtts3/batch_scheduler.h" + +#include "engine/framework/audio/conversion.h" +#include "engine/framework/audio/kaldi_fbank.h" +#include "engine/framework/audio/resampling.h" +#include "engine/framework/modules/speech_encoders/campplus_encoder.h" +#include "engine/framework/runtime/cache_slots.h" +#include "engine/framework/sampling/torch_random.h" +#include "engine/models/fireredtts3/ar.h" +#include "engine/models/fireredtts3/flow.h" +#include "engine/models/fireredtts3/redae.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::fireredtts3 { +namespace { + +namespace core = engine::core; +namespace modules = engine::modules; +using Clock = std::chrono::steady_clock; + +constexpr int64_t kMaxArSteps = 400; +// prefill graph 固定 steps（grow-only 复用）：所有 wave 的 prefill 都 padding 到该长度， +// 避免因 steps 变化重建 prefill graph 破坏 CUDA pool 逆序 free 约束。 +constexpr int64_t kMaxPrefillSteps = 1024; +// batched decode KV cache 固定上限（首次 build 后永不重建）： +// 最大可能 prefill（参考 + 长文本 ~300） + kMaxArSteps(400) + 尾部余量。 +// 若实际 prefill_steps 超此预算，decode 会 cache 耗尽（罕见，可调大）。 +constexpr int64_t kMaxDecodeCacheSteps = 700; + +// 与 pipeline.cpp 一致的参考音色 key/entry（避免引用失效跨轮，slot 拷贝产物）。 +struct ReferenceVoiceCacheKey { + int sample_rate = 0; + int channels = 0; + uint64_t sample_count = 0; + uint64_t sample_hash = 0; +}; + +struct ReferenceVoiceCacheKeyEqual { + bool operator()(const ReferenceVoiceCacheKey & a, const ReferenceVoiceCacheKey & b) const noexcept { + return a.sample_rate == b.sample_rate && a.channels == b.channels && + a.sample_count == b.sample_count && a.sample_hash == b.sample_hash; + } +}; + +struct ReferenceVoiceCacheEntry { + std::vector prompt_audio_24k; + std::vector prompt_latents; + std::vector speaker_embedding; +}; + +uint64_t hash_audio_samples(const runtime::AudioBuffer & audio) { + uint64_t key = 1469598103934665603ull; + for (const float sample : audio.samples) { + uint32_t bits = 0; + std::memcpy(&bits, &sample, sizeof(bits)); + key ^= bits; + key *= 1099511628211ull; + } + return key; +} + +std::vector last_rows( + const std::vector & values, + int64_t rows, + int64_t width, + int64_t keep) { + if (keep <= 0 || rows < keep || static_cast(values.size()) != rows * width) { + throw std::runtime_error("FireRedTTS3 hidden row slice is out of range"); + } + return std::vector( + values.begin() + static_cast((rows - keep) * width), + values.end()); +} + +std::vector campplus_fbank(const std::vector & prompt_24k) { + auto audio_16k = audio::resample_mono_torchaudio_sinc_hann(prompt_24k, 24000, 16000); + audio::KaldiFbankOptions options; + options.sample_rate = 16000; + options.num_mels = 80; + options.window_type = audio::KaldiFbankWindowType::Povey; + options.lfr_m = 1; + options.lfr_n = 1; + options.apply_cmvn = false; + options.upscale_samples = false; + auto features = audio::extract_kaldi_fbank(audio_16k, options); + if (features.frames <= 0 || features.feature_dim != 80) { + throw std::runtime_error("FireRedTTS3 CAM++ fbank extraction failed"); + } + for (int64_t m = 0; m < features.feature_dim; ++m) { + double sum = 0.0; + for (int64_t t = 0; t < features.frames; ++t) { + sum += features.values[static_cast(t * features.feature_dim + m)]; + } + const float mean = static_cast(sum / static_cast(features.frames)); + for (int64_t t = 0; t < features.frames; ++t) { + features.values[static_cast(t * features.feature_dim + m)] -= mean; + } + } + return std::move(features.values); +} + +} // namespace + +class FireRedTTS3BatchScheduler::Impl { +public: + Impl( + std::shared_ptr assets, + engine::core::ExecutionContext & execution, + size_t graph_arena_bytes, + size_t helper_graph_arena_bytes, + size_t weight_context_bytes, + engine::assets::TensorStorageType storage_type, + size_t reference_cache_slots, + bool mem_saver, + int64_t max_batch) + : assets_(std::move(assets)), + execution_(execution), + mem_saver_(mem_saver), + max_batch_(std::max(1, max_batch)), + sampling_policy_(sampling::resolve_torch_cuda_sampling_policy( + execution.backend_type(), + execution.config().device, + "fireredtts3", + "FireRedTTS3Batch", + sampling::TorchCudaSamplingPolicyFailureMode::FallbackToDefault)), + ar_(std::make_unique( + assets_, execution_, graph_arena_bytes, helper_graph_arena_bytes, + weight_context_bytes, storage_type, false)), + redae_(std::make_unique( + assets_, execution_, graph_arena_bytes, weight_context_bytes, storage_type)), + flow_(std::make_unique( + assets_, execution_, graph_arena_bytes, weight_context_bytes, storage_type, false)), + reference_voice_cache_( + runtime::CacheSlots( + reference_cache_slots)) { + if (assets_ == nullptr) { + throw std::runtime_error("FireRedTTS3BatchScheduler requires assets"); + } + slots_.resize(static_cast(max_batch_)); + for (size_t i = 0; i < slots_.size(); ++i) { + free_slots_.push_back(static_cast(i)); + } + // 专用调度线程：唯一执行 GPU decode graph 的线程。 + // ggml CUDA memory pool 要求分配/释放严格逆序（GGML_ASSERT(ptr == pool_addr + pool_used)）， + // 多请求线程并发 tick 会破坏该不变量导致崩溃。单调度线程 = llama.cpp update_slots 模式。 + scheduler_thread_ = std::thread(&Impl::scheduler_loop, this); + } + + ~Impl() { + { + std::lock_guard lock(mu_); + stop_ = true; + } + cv_wake_.notify_all(); + if (scheduler_thread_.joinable()) { + scheduler_thread_.join(); + } + } + + int64_t launch(const FireRedTTS3BaseRequest & request, const std::vector & chunk_patches) { + std::lock_guard lock(mu_); + if (free_slots_.empty()) { + return -1; + } + const int64_t slot_id = free_slots_.front(); + free_slots_.pop_front(); + auto & slot = slots_[static_cast(slot_id)]; + slot = Slot{}; + slot.state = Slot::State::Active; + slot.request = request; + slot.chunks = chunk_patches.empty() ? std::vector{400} : chunk_patches; + slot.chunk_target = slot.chunks[0]; + pending_.push_back(slot_id); + cv_wake_.notify_all(); + return slot_id; + } + + engine::runtime::AudioBuffer next_chunk(int64_t slot_id) { + std::unique_lock lock(mu_); + auto & slot = slots_[static_cast(slot_id)]; + if (slot.error) { + std::rethrow_exception(slot.error); + } + if (!slot.chunk_queue.empty()) { + auto audio = std::move(slot.chunk_queue.front()); + slot.chunk_queue.pop_front(); + return audio; + } + if (slot.finished) { + return {}; + } + // 纯等待模式：调度线程负责 tick（batch decode），本线程只等自己的 chunk。 + // ggml CUDA pool 逆序约束 -> 唯一 GPU decode 线程必须是调度线程。 + // 用带谓词的 wait（wait 内部重查条件），避免"检查 queue 后、wait 前"的 + // notify 竞态窗口导致死等。 + cv_.wait(lock, [&] { + return slot.finished || !slot.chunk_queue.empty() || slot.state == Slot::State::Failed || slot.error != nullptr; + }); + if (slot.error) { + std::rethrow_exception(slot.error); + } + if (!slot.chunk_queue.empty()) { + auto audio = std::move(slot.chunk_queue.front()); + slot.chunk_queue.pop_front(); + return audio; + } + return {}; + } + + engine::runtime::AudioBuffer generate(const FireRedTTS3BaseRequest & request) { + const int64_t slot_id = launch(request, {}); + if (slot_id < 0) { + throw std::runtime_error("FireRedTTS3BatchScheduler slot pool exhausted"); + } + engine::runtime::AudioBuffer merged; + merged.sample_rate = static_cast(assets_->redae.sample_rate); + merged.channels = 1; + while (true) { + auto chunk = next_chunk(slot_id); + if (chunk.samples.empty()) { + break; + } + runtime::append_audio_buffer(merged, chunk); + } + return merged; + } + + int64_t max_batch() const noexcept { + return max_batch_; + } + + void release_graphs() { + std::lock_guard lock(mu_); + if (redae_) { + redae_->release_graphs(); + } + if (ar_) { + ar_->release_graphs(); + ar_->release_backbone_graphs(); + } + if (flow_) { + flow_->release_graph(); + } + campplus_.release_runtime_graph(); + } + +private: + // ---- 参考音色（锁内串行，产物拷进 slot）---- + const ReferenceVoiceCacheEntry & prepare_reference_voice_locked(const runtime::AudioBuffer & audio) { + ReferenceVoiceCacheKey key; + key.sample_rate = audio.sample_rate; + key.channels = audio.channels; + key.sample_count = static_cast(audio.samples.size()); + key.sample_hash = hash_audio_samples(audio); + const ReferenceVoiceCacheEntry * cached = reference_voice_cache_.find(key); + if (cached != nullptr) { + return *cached; + } + ReferenceVoiceCacheEntry entry; + entry.prompt_audio_24k = prepare_firered_prompt_audio_24k(audio, assets_->redae, assets_->base.patch_size); + entry.prompt_latents = redae_->encode(entry.prompt_audio_24k); + ensure_campplus_locked(); + auto speaker_features = campplus_fbank(entry.prompt_audio_24k); + auto speaker = campplus_.embed_from_features( + speaker_features, static_cast(speaker_features.size()) / 80, 80); + entry.speaker_embedding = std::move(speaker.embedding); + reference_voice_cache_.put(key, std::move(entry)); + const ReferenceVoiceCacheEntry * stored = reference_voice_cache_.find(key); + if (stored == nullptr) { + throw std::runtime_error("FireRedTTS3 reference voice cache insert failed"); + } + return *stored; + } + + void ensure_campplus_locked() { + if (campplus_.weights() != nullptr) { + return; + } + modules::CampplusEncoderConfig cfg; + cfg.feat_dim = 80; + cfg.embedding_size = assets_->base.speaker_dim; + cfg.tensor_prefix = "campplus"; + cfg.weight_storage_type = storage_type_; + campplus_ = modules::CampplusEncoderComponent::load_from_tensor_source( + assets_->campplus_weights, execution_.config(), std::move(cfg)); + } + + // 单个 slot 的 prefill（锁内串行）：参考 prep + 单路径 prefill，KV 存 CPU。 + void prefill_slot_locked(int64_t slot_id) { + auto & slot = slots_[static_cast(slot_id)]; + const auto & ref = prepare_reference_voice_locked(slot.request.prompt_audio); + slot.prompt_latents = ref.prompt_latents; + slot.prompt_latent_frames = static_cast(ref.prompt_latents.size()) / assets_->base.redae_dim; + auto spk_llm = ar_->speaker_llm(ref.speaker_embedding); + slot.spk_dit = ar_->speaker_dit(ref.speaker_embedding); + auto text_embeds = ar_->token_embedding(slot.request.token_ids); + auto patch_prompt = ar_->patch_encode(ref.prompt_latents); + std::vector input_embeddings; + input_embeddings.reserve(spk_llm.size() + text_embeds.size() + patch_prompt.size()); + input_embeddings.insert(input_embeddings.end(), spk_llm.begin(), spk_llm.end()); + input_embeddings.insert(input_embeddings.end(), text_embeds.begin(), text_embeds.end()); + input_embeddings.insert(input_embeddings.end(), patch_prompt.begin(), patch_prompt.end()); + slot.prefill_steps = 1 + static_cast(slot.request.token_ids.size()) + + slot.prompt_latent_frames / assets_->base.patch_size; + if (slot.prefill_steps > kMaxPrefillSteps) { + throw std::runtime_error("FireRedTTS3 prefill steps exceed scheduler max (increase kMaxPrefillSteps)"); + } + // padded prefill：input padding 到 kMaxPrefillSteps（零 embedding）， + // 固定 prefill graph（不重建），返回 state/hidden 截断到 prefill_steps。 + std::vector padded( + static_cast(kMaxPrefillSteps) * static_cast(assets_->base.hidden_size), 0.0F); + std::copy(input_embeddings.begin(), input_embeddings.end(), padded.begin()); + auto prefill = ar_->prefill_embeddings_padded( + padded, kMaxPrefillSteps, slot.prefill_steps); + slot.prefill_hidden = std::move(prefill.hidden); + slot.prefill_state = std::move(prefill.state); + slot.prefill_done = true; + + slot.latents_gen.assign( + static_cast(assets_->base.history_patches * assets_->base.patch_size * assets_->base.redae_dim), 0.0F); + slot.latents_gen.insert(slot.latents_gen.end(), slot.prompt_latents.begin(), slot.prompt_latents.end()); + slot.backbone_cond.assign(static_cast(assets_->base.history_patches * assets_->base.hidden_size), 0.0F); + slot.schedule = firered_cosine_time_schedule(slot.request.num_inference_steps); + redae_->decode_reset(slot.redae_state); + } + + // 是否存在活跃 decode slot（step>=1，已在 batched KV 中）。 + bool has_active_decode_locked() const { + for (int64_t i = 0; i < max_batch_; ++i) { + const auto & slot = slots_[static_cast(i)]; + if (slot.state == Slot::State::Active && slot.prefill_done && slot.step >= 1) { + return true; + } + } + return false; + } + + // ---- 调度主循环（锁内）---- + // 与 pipeline.cpp 的单路径 AR 循环逐位对应： + // step 0: 用 prefill_hidden 的 stop/one_backbone 走 AR（不 consume decode） + // step>=1: decode_embeddings_batched 产出 hidden 后走 AR + void tick_locked() { + // 1. 处理 pending 启动（prefill 新 slot）。仅当无活跃 decode slot 时才 prefill + // （wave batching）：prefill graph 与 batched decode graph 共用同一 CUDA + // device pool，若在 decode 进行中 prefill（重建 prefill graph），会破坏 + // pool 的逆序 free 约束（GGML_ASSERT(ptr == pool_addr + pool_used)）。 + if (!pending_.empty() && !has_active_decode_locked()) { + form_wave_locked(); + } + // 2. 对 step==0 且已 prefill 的 slot 走 step-0 AR（不进 batch decode， + // 否则零行会被 advance_member 污染 KV 位置）。 + for (int64_t i = 0; i < max_batch_; ++i) { + auto & slot = slots_[static_cast(i)]; + if (slot.state == Slot::State::Active && slot.prefill_done && slot.step == 0) { + advance_slot_from_prefill_locked(i); + } + } + // 3. 收集 step>=1 的活跃 slot（参与 batched decode） + std::vector active; + for (int64_t i = 0; i < max_batch_; ++i) { + auto & slot = slots_[static_cast(i)]; + if (slot.state == Slot::State::Active && slot.prefill_done && slot.step >= 1) { + active.push_back(i); + } + } + if (active.empty()) { + // 没有需要 decode 的 slot；若 step-0 AR 让某些 slot 进入 Dead， + // 本轮结束前清理。 + for (int64_t i = 0; i < max_batch_; ++i) { + auto & slot = slots_[static_cast(i)]; + if (slot.state == Slot::State::Dead || slot.state == Slot::State::Done) { + finish_slot_locked(i); + } + } + return; + } + // 4. 若 batched decode 未启动，或需要 splice 新 slot（step-0 AR 后），重建。 + if (!batched_started_ || !imported_slots_.empty()) { + rebuild_batched_state_locked(); + } + // 5. 收集 embeddings；非活跃行零填充 + const int64_t hidden = assets_->base.hidden_size; + std::vector embeddings(static_cast(max_batch_ * hidden), 0.0F); + for (int64_t i = 0; i < max_batch_; ++i) { + auto & slot = slots_[static_cast(i)]; + if (slot.state == Slot::State::Active && slot.prefill_done && slot.step >= 1) { + std::copy(slot.next_input.begin(), slot.next_input.end(), + embeddings.begin() + static_cast(i * hidden)); + } + } + // 6. 一次 batched decode（每成员独立 pos/cache-slot/mask） + auto out = ar_->decode_embeddings_batched(embeddings, max_batch_); + // 7. 分发 hidden 到各 step>=1 活跃 slot + for (int64_t i = 0; i < max_batch_; ++i) { + auto & slot = slots_[static_cast(i)]; + if (slot.state != Slot::State::Active || !slot.prefill_done || slot.step < 1) { + continue; + } + std::vector row( + out.hidden.begin() + static_cast(i * hidden), + out.hidden.begin() + static_cast((i + 1) * hidden)); + advance_slot_locked(i, row); + } + // 8. 处理完成/死亡的 slot（flush RedAE 尾部，标记 finished，归还槽位） + for (int64_t i = 0; i < max_batch_; ++i) { + auto & slot = slots_[static_cast(i)]; + if (slot.state == Slot::State::Dead || slot.state == Slot::State::Done) { + finish_slot_locked(i); + } + } + } + + // step-0 AR：用 prefill_hidden 的 stop/one_backbone 走 AR，不 consume decode。 + void advance_slot_from_prefill_locked(int64_t slot_id) { + auto & slot = slots_[static_cast(slot_id)]; + const int64_t hidden = assets_->base.hidden_size; + const int64_t prompt_patches = slot.prompt_latent_frames / assets_->base.patch_size; + + const std::vector last( + slot.prefill_hidden.end() - hidden, slot.prefill_hidden.end()); + const float stop = ar_->stop(last); + if (stop >= slot.request.stop_threshold && slot.step >= 6) { + slot.state = Slot::State::Dead; + return; + } + std::vector one_backbone( + slot.prefill_hidden.end() - static_cast(prompt_patches * hidden), + slot.prefill_hidden.end()); + slot.backbone_cond.insert(slot.backbone_cond.end(), one_backbone.begin(), one_backbone.end()); + const int64_t cond_rows = static_cast(slot.backbone_cond.size()) / hidden; + auto cond3 = last_rows(slot.backbone_cond, cond_rows, hidden, assets_->base.history_patches + 1); + auto dit_cond3 = ar_->dit_head(cond3, assets_->base.history_patches + 1); + const auto next_latent = flow_one_patch_locked(slot, dit_cond3, 0); + slot.latents_gen.insert(slot.latents_gen.end(), next_latent.begin(), next_latent.end()); + slot.next_input = ar_->patch_encode(next_latent); + slot.chunk_latents.insert(slot.chunk_latents.end(), next_latent.begin(), next_latent.end()); + slot.generated_patches++; + slot.step = 1; + } + + // 单 slot 一步 AR（step>=1：hidden 来自一次 batched decode）。 + void advance_slot_locked(int64_t slot_id, const std::vector & hidden_row) { + auto & slot = slots_[static_cast(slot_id)]; + const int64_t hidden = assets_->base.hidden_size; + + // AR 步数上限（与 pipeline.cpp 的 max_steps=400 一致）：stop 未触发时强制结束。 + if (slot.step >= kMaxArSteps) { + slot.state = Slot::State::Dead; + return; + } + const float stop = ar_->stop(hidden_row); + if (stop >= slot.request.stop_threshold && slot.step >= 6) { + slot.state = Slot::State::Dead; + return; + } + + slot.backbone_cond.insert(slot.backbone_cond.end(), hidden_row.begin(), hidden_row.end()); + const int64_t cond_rows = static_cast(slot.backbone_cond.size()) / hidden; + auto cond3 = last_rows(slot.backbone_cond, cond_rows, hidden, assets_->base.history_patches + 1); + auto dit_cond3 = ar_->dit_head(cond3, assets_->base.history_patches + 1); + const auto next_latent = flow_one_patch_locked(slot, dit_cond3, static_cast(slot.step)); + slot.latents_gen.insert(slot.latents_gen.end(), next_latent.begin(), next_latent.end()); + slot.next_input = ar_->patch_encode(next_latent); + slot.chunk_latents.insert(slot.chunk_latents.end(), next_latent.begin(), next_latent.end()); + slot.generated_patches++; + slot.step++; + + // chunk 边界（与 pipeline.cpp 相同） + if (slot.generated_patches >= slot.chunk_target && slot.chunk_index < slot.chunks.size()) { + auto audio = redae_->decode_incremental(slot.redae_state, slot.chunk_latents); + slot.chunk_latents.clear(); + slot.chunk_index++; + if (slot.chunk_index < slot.chunks.size()) { + slot.chunk_target += slot.chunks[slot.chunk_index]; + } + if (!audio.samples.empty()) { + slot.chunk_queue.push_back(std::move(audio)); + cv_.notify_all(); // 唤醒等待该 slot 的 next_chunk + } + } + } + + // flow + CFG 单 patch（与 pipeline.cpp 的 flow_one_patch 一致，但使用 slot 状态）。 + std::vector flow_one_patch_locked(Slot & slot, const std::vector & dit_cond3, uint64_t step_index) { + const int64_t history_tokens = assets_->base.history_patches * assets_->base.patch_size; + const int64_t tokens = history_tokens + assets_->base.patch_size; + const int64_t input_channels = assets_->base.redae_dim + assets_->base.dit_hidden_size + assets_->base.speaker_dim; + auto history = last_rows( + slot.latents_gen, + static_cast(slot.latents_gen.size()) / assets_->base.redae_dim, + assets_->base.redae_dim, + history_tokens); + const uint64_t noise_offset = + sampling::torch_cuda_tensor_iterator_offset_blocks( + static_cast(assets_->base.patch_size * assets_->base.redae_dim), + sampling_policy_) * + step_index; + auto current = sampling::generate_torch_cuda_tensor_iterator_randn( + static_cast(assets_->base.patch_size * assets_->base.redae_dim), + slot.request.seed, + noise_offset, + sampling_policy_, + sampling::TorchRandnPrecision::Float32); + + const int64_t batch = slot.request.guidance_scale > 0.0F ? 2 : 1; + for (size_t i = 0; i + 1 < slot.schedule.size(); ++i) { + std::vector x_in(static_cast(batch * tokens * input_channels), 0.0F); + for (int64_t b = 0; b < batch; ++b) { + for (int64_t t = 0; t < tokens; ++t) { + float * row = x_in.data() + static_cast((b * tokens + t) * input_channels); + if (t < history_tokens) { + std::copy( + history.begin() + static_cast(t * assets_->base.redae_dim), + history.begin() + static_cast((t + 1) * assets_->base.redae_dim), + row); + } else { + const int64_t local = t - history_tokens; + std::copy( + current.begin() + static_cast(local * assets_->base.redae_dim), + current.begin() + static_cast((local + 1) * assets_->base.redae_dim), + row); + } + if (b == 0) { + const int64_t cond_row = t / assets_->base.patch_size; + std::copy( + dit_cond3.begin() + static_cast(cond_row * assets_->base.dit_hidden_size), + dit_cond3.begin() + static_cast((cond_row + 1) * assets_->base.dit_hidden_size), + row + assets_->base.redae_dim); + std::copy(slot.spk_dit.begin(), slot.spk_dit.end(), + row + assets_->base.redae_dim + assets_->base.dit_hidden_size); + } + } + } + auto te = firered_timestep_embedding(slot.schedule[i]); + std::vector time(static_cast(batch * 256)); + for (int64_t b = 0; b < batch; ++b) { + std::copy(te.begin(), te.end(), time.begin() + static_cast(b * 256)); + } + auto pred = flow_->run(x_in, time, batch); + const float dt = slot.schedule[i + 1] - slot.schedule[i]; + for (int64_t t = 0; t < assets_->base.patch_size; ++t) { + for (int64_t c = 0; c < assets_->base.redae_dim; ++c) { + const size_t idx = static_cast(t * assets_->base.redae_dim + c); + float vt = pred[idx]; + if (batch == 2) { + const float uncond = pred[static_cast((assets_->base.patch_size + t) * assets_->base.redae_dim + c)]; + vt = (1.0F + slot.request.guidance_scale) * vt - slot.request.guidance_scale * uncond; + } + current[idx] += dt * vt; + } + } + } + return current; + } + + // 完成一个 slot：flush RedAE 尾部、标记 finished、归还槽位。 + // 只处理 Dead（Done 已处理过并归还）；Done 状态由 launch 复用前的 reset 清除。 + void finish_slot_locked(int64_t slot_id) { + auto & slot = slots_[static_cast(slot_id)]; + if (slot.state != Slot::State::Dead) { + return; + } + // 还有未解完的 chunk_latents → flush + if (!slot.chunk_latents.empty() && slot.generated_patches > 0) { + auto audio = redae_->decode_incremental(slot.redae_state, slot.chunk_latents); + slot.chunk_latents.clear(); + if (!audio.samples.empty()) { + slot.chunk_queue.push_back(std::move(audio)); + } + } + auto flush = redae_->flush_incremental(slot.redae_state); + if (!flush.samples.empty()) { + slot.chunk_queue.push_back(std::move(flush)); + } + slot.state = Slot::State::Idle; + slot.finished = true; + // 防重复归还：仅当 slot 不在 free_slots_ 时才 push（避免并发下重复导致 + // 同一 slot 被多个 session launch）。 + if (std::find(free_slots_.begin(), free_slots_.end(), slot_id) == free_slots_.end()) { + free_slots_.push_back(slot_id); + } + // 该行不再参与 decode：从 running 集合移除（下次 join 时重建）。 + running_active_.erase( + std::remove(running_active_.begin(), running_active_.end(), slot_id), + running_active_.end()); + imported_slots_.erase(slot_id); + // 该行 KV 在 GPU 上仍残留 decode 产物；标记脏，下次 rebuild 时清零。 + dirty_rows_.insert(slot_id); + cv_.notify_all(); + } + + // 专用调度线程主循环（唯一 GPU decode 线程）。 + // 持锁跑 tick（batch decode 推进所有活跃 slot），无活跃 slot 时等待 launch/stop。 + void scheduler_loop() { + std::unique_lock lock(mu_); + while (!stop_) { + tick_locked(); + // 有工作（Active 未完成 / Dead 待 finish）则继续 tick；否则等新工作。 + bool has_work = false; + for (int64_t i = 0; i < max_batch_; ++i) { + auto & slot = slots_[static_cast(i)]; + if (slot.state == Slot::State::Active || + slot.state == Slot::State::Dead) { + has_work = true; + break; + } + } + if (!has_work) { + cv_wake_.wait(lock); + } + } + } + + // 形成 wave：prefill pending slot（锁内串行）。prefill 后标记需 splice。 + void form_wave_locked() { + // 一次最多 prefill 直到槽池满（但同一轮只 prefill，不挤占已活跃 slot 的 GPU） + std::vector newly; + while (!pending_.empty() && newly.size() < static_cast(max_batch_)) { + const int64_t slot_id = pending_.front(); + pending_.pop_front(); + newly.push_back(slot_id); + } + for (int64_t slot_id : newly) { + try { + prefill_slot_locked(slot_id); + imported_slots_.insert(slot_id); // 标记需要 splice 进 batched KV + } catch (...) { + auto & slot = slots_[static_cast(slot_id)]; + slot.error = std::current_exception(); + slot.state = Slot::State::Failed; + slot.finished = true; + if (std::find(free_slots_.begin(), free_slots_.end(), slot_id) == free_slots_.end()) { + free_slots_.push_back(slot_id); + } + } + } + } + + // 重建 batched KV state：export 当前 GPU state + splice 新 prefill slot 的 CPU KV + 重新 start。 + // 固定 max_batch graph，不重建（batch_size 不变、cache_steps 不减）。 + void rebuild_batched_state_locked() { + // 需要 splice 的新 slot（尚未在 batched 中） + std::vector joining; + for (int64_t i = 0; i < max_batch_; ++i) { + auto & slot = slots_[static_cast(i)]; + if (slot.prefill_done && imported_slots_.count(i) > 0) { + joining.push_back(i); + } + } + if (batched_started_ && !joining.empty()) { + // 从 GPU 导出当前 batched state（active slot 的中间 KV 保留） + runtime::TransformerBatchedKVState gpu = ar_->export_batched_decode_state(); + // 层数/row_elems 校验（与 import_state 的期望布局一致） + if (gpu.layers.size() != slot_prefill_layer_count()) { + throw std::runtime_error("FireRedTTS3BatchScheduler batched state layer mismatch"); + } + // 组装全量 batched state：max_steps = 所有"有数据行"（imported/running）的最大 end。 + // 空行（无 prefill、非 running）end=0，不参与 max_steps（避免 gpu 残留的空行 + // 高 end 撑大 state，导致与 current_ends 的 import max 不一致）。 + runtime::TransformerBatchedKVState state; + state.batch_size = max_batch_; + state.current_ends.assign(static_cast(max_batch_), 0); + for (int64_t i = 0; i < max_batch_; ++i) { + if (imported_slots_.count(i) > 0) { + state.current_ends[static_cast(i)] = slots_[static_cast(i)].prefill_steps; + } else if (std::find(running_active_.begin(), running_active_.end(), i) != running_active_.end()) { + state.current_ends[static_cast(i)] = gpu_end_for(gpu, i); + } + } + int64_t max_steps = 0; + for (int64_t i = 0; i < max_batch_; ++i) { + max_steps = std::max(max_steps, state.current_ends[static_cast(i)]); + } + state.current_end = max_steps; + state.layers.resize(gpu.layers.size()); + for (size_t layer = 0; layer < state.layers.size(); ++layer) { + auto & out_layer = state.layers[layer]; + out_layer.valid_steps = max_steps; + const size_t row_elems = row_elems_for_layer(layer); + out_layer.key.assign(static_cast(max_batch_) * static_cast(max_steps) * row_elems, 0.0F); + out_layer.value.assign(static_cast(max_batch_) * static_cast(max_steps) * row_elems, 0.0F); + } + // 填充每行：已运行 slot（running）从 gpu 拷贝其 ends；新 slot 从 prefill_state 拷贝； + // 其余（dead/脏/空行）保持全零。 + for (int64_t b = 0; b < max_batch_; ++b) { + if (imported_slots_.count(b) == 0 && std::find(running_active_.begin(), running_active_.end(), b) == running_active_.end()) { + continue; // 非活跃行：全零（不拷 gpu 残留） + } + const int64_t end = (imported_slots_.count(b) > 0) + ? slots_[static_cast(b)].prefill_steps + : gpu_end_for(gpu, b); + if (end <= 0) { + continue; + } + const int64_t current_end = end; + for (size_t layer = 0; layer < state.layers.size(); ++layer) { + const size_t row_elems = row_elems_for_layer(layer); + const size_t copy_elems = static_cast(current_end) * row_elems; + float * dst_key = state.layers[layer].key.data() + + static_cast(b) * static_cast(max_steps) * row_elems; + float * dst_value = state.layers[layer].value.data() + + static_cast(b) * static_cast(max_steps) * row_elems; + if (imported_slots_.count(b) > 0) { + const auto & src = slots_[static_cast(b)].prefill_state.layers[layer]; + std::copy(src.key.begin(), src.key.end(), dst_key); + std::copy(src.value.begin(), src.value.end(), dst_value); + } else { + const auto & src = gpu.layers[layer]; + const size_t src_row_elems = row_elems; + const int64_t gpu_valid_steps = gpu.layers.empty() ? 0 : gpu.layers.front().valid_steps; + const size_t src_offset = static_cast(b) * static_cast(gpu_valid_steps) * src_row_elems; + std::copy( + src.key.begin() + static_cast(src_offset), + src.key.begin() + static_cast(src_offset + copy_elems), + dst_key); + std::copy( + src.value.begin() + static_cast(src_offset), + src.value.begin() + static_cast(src_offset + copy_elems), + dst_value); + } + } + } + const int64_t required_steps = compute_required_steps(max_steps); + ar_->start_decode_embeddings_batched(state, required_steps); + for (int64_t i : joining) { + imported_slots_.erase(i); + } + } else if (!joining.empty()) { + // 首次启动（无 GPU 导出，全部新 slot） + int64_t max_steps = 0; + for (int64_t i = 0; i < max_batch_; ++i) { + auto & slot = slots_[static_cast(i)]; + if (slot.prefill_done) { + max_steps = std::max(max_steps, slot.prefill_steps); + } + } + runtime::TransformerBatchedKVState state; + state.batch_size = max_batch_; + state.current_end = max_steps; + state.current_ends.assign(static_cast(max_batch_), 0); + state.layers.resize(slot_prefill_layer_count()); + for (size_t layer = 0; layer < state.layers.size(); ++layer) { + auto & out_layer = state.layers[layer]; + out_layer.valid_steps = max_steps; + const size_t row_elems = row_elems_for_layer(layer); + out_layer.key.assign(static_cast(max_batch_) * static_cast(max_steps) * row_elems, 0.0F); + out_layer.value.assign(static_cast(max_batch_) * static_cast(max_steps) * row_elems, 0.0F); + } + for (int64_t b = 0; b < max_batch_; ++b) { + auto & slot = slots_[static_cast(b)]; + if (!slot.prefill_done) { + continue; + } + state.current_ends[static_cast(b)] = slot.prefill_steps; + for (size_t layer = 0; layer < state.layers.size(); ++layer) { + const size_t row_elems = row_elems_for_layer(layer); + const size_t copy_elems = static_cast(slot.prefill_steps) * row_elems; + float * dst_key = state.layers[layer].key.data() + + static_cast(b) * static_cast(max_steps) * row_elems; + float * dst_value = state.layers[layer].value.data() + + static_cast(b) * static_cast(max_steps) * row_elems; + const auto & src = slot.prefill_state.layers[layer]; + std::copy(src.key.begin(), src.key.end(), dst_key); + std::copy(src.value.begin(), src.value.end(), dst_value); + } + } + const int64_t required_steps = compute_required_steps(max_steps); + ar_->start_decode_embeddings_batched(state, required_steps); + // 全量分支：所有 prefill slot 已 splice 进 batched KV，清空 imported 标记， + // 否则每次 tick 都会重复 rebuild（用 prefill_steps 覆盖 GPU 增量 KV）。 + imported_slots_.clear(); + } + batched_started_ = true; + dirty_rows_.clear(); + running_active_.clear(); + for (int64_t i = 0; i < max_batch_; ++i) { + auto & slot = slots_[static_cast(i)]; + if (slot.state == Slot::State::Active && slot.prefill_done) { + running_active_.push_back(i); + } + } + } + + // 计算 batched decode graph 所需 cache steps。为避免 graph 重建（昂贵）， + // 首次 build 用保守上限：max_steps + 完整 AR 预算。后续 join 若 max_steps + // 增长超预算才重建（罕见），通常复用首次 graph。 + int64_t compute_required_steps(int64_t /*max_prefill_steps*/) const { + // 固定 cache 上限：首次 build 后永不重建（避免 export 空 state / pool 顺序破坏）。 + // 预算 = 最大可能 prefill（参考+文本） + kMaxArSteps + 尾部。 + // KV 显存 = max_batch × cache_steps × layers × row_elems ≈ 可控。 + return kMaxDecodeCacheSteps; + } + + // GPU batched state 中某行的当前 end（运行中 slot 用它作为拷贝长度） + int64_t gpu_end_for(const runtime::TransformerBatchedKVState & gpu, int64_t batch) const { + if (!gpu.current_ends.empty()) { + if (static_cast(gpu.current_ends.size()) == gpu.batch_size) { + return gpu.current_ends[static_cast(batch)]; + } + } + return gpu.current_end; + } + + size_t slot_prefill_layer_count() const { + if (max_batch_ == 0) { + return 0; + } + for (int64_t i = 0; i < max_batch_; ++i) { + auto & slot = slots_[static_cast(i)]; + if (slot.prefill_done && !slot.prefill_state.layers.empty()) { + return slot.prefill_state.layers.size(); + } + } + return 0; + } + + size_t row_elems_for_layer(size_t /*layer*/) const { + // FireRed TTS3 AR: kv_heads × head_dim（同 pipeline prefill state 的层元素数）。 + return static_cast(assets_->base.kv_heads * assets_->base.head_dim); + } + + std::shared_ptr assets_; + engine::core::ExecutionContext & execution_; + bool mem_saver_ = false; + int64_t max_batch_ = 1; + engine::assets::TensorStorageType storage_type_ = engine::assets::TensorStorageType::Native; + sampling::TorchCudaSamplingPolicy sampling_policy_; + std::unique_ptr ar_; + std::unique_ptr redae_; + std::unique_ptr flow_; + modules::CampplusEncoderComponent campplus_; + runtime::CacheSlots reference_voice_cache_; + + std::mutex mu_; + std::condition_variable cv_; // chunk 产出 / slot 结束通知（next_chunk 等待） + std::condition_variable cv_wake_; // 调度线程唤醒（launch / stop） + std::vector slots_; + std::deque free_slots_; + std::deque pending_; + bool batched_started_ = false; + bool stop_ = false; + std::thread scheduler_thread_; + // 已 spliced 进 batched KV 的 slot（运行中）；新 prefill 的 slot 在此集合外，需 join。 + std::vector running_active_; + // 已 prefill 但尚未 splice 进 batched KV 的 slot id + std::set imported_slots_; + // 刚结束、GPU 行残留 KV 的 slot id（下次 rebuild 时该行清零） + std::set dirty_rows_; +}; + +FireRedTTS3BatchScheduler::FireRedTTS3BatchScheduler( + std::shared_ptr assets, + engine::core::ExecutionContext & execution, + size_t graph_arena_bytes, + size_t helper_graph_arena_bytes, + size_t weight_context_bytes, + engine::assets::TensorStorageType storage_type, + size_t reference_cache_slots, + bool mem_saver, + int64_t max_batch) + : impl_(std::make_unique( + std::move(assets), execution, graph_arena_bytes, helper_graph_arena_bytes, + weight_context_bytes, storage_type, reference_cache_slots, mem_saver, max_batch)) {} + +FireRedTTS3BatchScheduler::~FireRedTTS3BatchScheduler() = default; + +int64_t FireRedTTS3BatchScheduler::launch(const FireRedTTS3BaseRequest & request, const std::vector & chunk_patches) { + return impl_->launch(request, chunk_patches); +} + +engine::runtime::AudioBuffer FireRedTTS3BatchScheduler::next_chunk(int64_t slot_id) { + return impl_->next_chunk(slot_id); +} + +engine::runtime::AudioBuffer FireRedTTS3BatchScheduler::generate(const FireRedTTS3BaseRequest & request) { + return impl_->generate(request); +} + +int64_t FireRedTTS3BatchScheduler::max_batch() const noexcept { + return impl_->max_batch(); +} + +void FireRedTTS3BatchScheduler::release_graphs() { + impl_->release_graphs(); +} + +} // namespace engine::models::fireredtts3 diff --git a/src/models/fireredtts3/flow.cpp b/src/models/fireredtts3/flow.cpp index dd90ccebb..85bc187f5 100644 --- a/src/models/fireredtts3/flow.cpp +++ b/src/models/fireredtts3/flow.cpp @@ -16,7 +16,9 @@ #include #include +#include #include +#include #include #include #include diff --git a/src/models/fireredtts3/session.cpp b/src/models/fireredtts3/session.cpp index 7f3d3ba09..4cc3c849b 100644 --- a/src/models/fireredtts3/session.cpp +++ b/src/models/fireredtts3/session.cpp @@ -8,9 +8,13 @@ #include "engine/framework/text/chinese_normalization.h" #include "engine/framework/text/text_normalization.h" #include "engine/models/fireredtts3/pipeline.h" +#include "engine/models/fireredtts3/batch_scheduler.h" +#include +#include #include #include +#include #include namespace engine::models::fireredtts3 { @@ -22,6 +26,48 @@ constexpr const char * kInstructName = "FireRedTTS3 Instruct"; constexpr int64_t kDefaultTextChunkSize = 600; constexpr size_t kDefaultReferenceCacheSlots = 4; +// ---- 共享 batch scheduler 注册表（assets-keyed）---- +// 同一 assets 的所有 session（server session 池）共享一个 scheduler： +// 并发请求经各 session 进入 scheduler 的不同 slot，实现共享 GPU batch decode。 +// model unload（全部 session 析构）时 scheduler 随最后一个 weak_ptr 消亡。 +struct SchedulerRegistry { + std::mutex mu; + std::unordered_map< + const FireRedTTS3Assets *, + std::weak_ptr> by_assets; + + std::shared_ptr get_or_create( + std::shared_ptr assets, + engine::core::ExecutionContext & execution, + size_t graph_arena_bytes, + size_t helper_graph_arena_bytes, + size_t weight_context_bytes, + engine::assets::TensorStorageType storage_type, + size_t reference_cache_slots, + bool mem_saver, + int64_t max_batch) { + std::lock_guard lock(mu); + auto it = by_assets.find(assets.get()); + if (it != by_assets.end()) { + if (auto existing = it->second.lock()) { + return existing; + } + by_assets.erase(it); + } + const FireRedTTS3Assets * key = assets.get(); + auto scheduler = std::make_shared( + std::move(assets), execution, graph_arena_bytes, helper_graph_arena_bytes, + weight_context_bytes, storage_type, reference_cache_slots, mem_saver, max_batch); + by_assets.emplace(key, scheduler); + return scheduler; + } +}; + +SchedulerRegistry & scheduler_registry() { + static SchedulerRegistry registry; + return registry; +} + const char * variant_name(FireRedTTS3Variant variant) { return variant == FireRedTTS3Variant::Instruct ? kInstructName : kBaseName; } @@ -272,6 +318,25 @@ FireRedTTS3Session::FireRedTTS3Session( weight_context_bytes, storage_type, mem_saver_); + return; + } + // Base clone：优先走共享 batch scheduler（多 session 共享 GPU batch decode）。 + const int64_t max_batch = runtime::parse_i64_option( + options.options, + {"fireredtts3.max_batch"}) + .value_or(1); + if (max_batch > 1) { + scheduler_ = scheduler_registry().get_or_create( + assets_, + execution_context(), + graph_arena_bytes, + helper_graph_arena_bytes, + weight_context_bytes, + storage_type, + static_cast(reference_cache_slots), + mem_saver_, + max_batch); + scheduler_enabled_ = true; } else { runtime_ = std::make_unique( assets_, @@ -336,15 +401,22 @@ runtime::TaskResult FireRedTTS3Session::run(const runtime::TaskRequest & request if (i > 0) { parsed.seed += static_cast(i); } - runtime::append_audio_buffer(merged_audio, runtime_->generate(parsed)); + if (scheduler_enabled_ && scheduler_) { + // 共享 scheduler 整句生成（并发请求共享 GPU batch decode） + runtime::append_audio_buffer(merged_audio, scheduler_->generate(parsed)); + } else { + runtime::append_audio_buffer(merged_audio, runtime_->generate(parsed)); + } } } } catch (...) { if (mem_saver_) { if (is_instruct) { instruct_runtime_->release_graphs(); - } else { + } else if (runtime_) { runtime_->release_graphs(); + } else if (scheduler_) { + scheduler_->release_graphs(); } } throw; @@ -352,8 +424,10 @@ runtime::TaskResult FireRedTTS3Session::run(const runtime::TaskRequest & request if (mem_saver_) { if (is_instruct) { instruct_runtime_->release_graphs(); - } else { + } else if (runtime_) { runtime_->release_graphs(); + } else if (scheduler_) { + scheduler_->release_graphs(); } } result.audio_output = std::move(merged_audio); @@ -399,6 +473,7 @@ void FireRedTTS3Session::initialize_stream_request(const runtime::TaskRequest & } void FireRedTTS3Session::start_stream(const runtime::TaskRequest & request) { + std::lock_guard lock(stream_mutex_); require_prepared("FireRedTTS3 streaming"); if (task_.mode != runtime::RunMode::Streaming) { throw std::runtime_error("FireRedTTS3 start_stream requires a streaming session"); @@ -406,21 +481,42 @@ void FireRedTTS3Session::start_stream(const runtime::TaskRequest & request) { reset(); initialize_stream_request(request); auto base = make_base_request(*tokenizer_, request); - stream_session_ = runtime_->begin_streaming(base, stream_chunk_patches_); + if (scheduler_enabled_ && scheduler_) { + // 共享 scheduler：launch 一个 slot（并发请求共享 GPU batch decode） + scheduler_slot_ = scheduler_->launch(base, stream_chunk_patches_); + if (scheduler_slot_ < 0) { + throw std::runtime_error("FireRedTTS3 scheduler slot pool exhausted"); + } + } else { + stream_session_ = runtime_->begin_streaming(base, stream_chunk_patches_); + } stream_started_ = true; } std::optional FireRedTTS3Session::next_stream_event() { + std::lock_guard lock(stream_mutex_); if (!stream_started_) { throw std::runtime_error("FireRedTTS3 streaming has not been started"); } - if (!stream_session_) { - return std::nullopt; - } - auto audio = stream_session_->next_chunk(); - if (audio.samples.empty()) { - stream_session_.reset(); - return std::nullopt; + runtime::AudioBuffer audio; + if (scheduler_enabled_ && scheduler_) { + if (scheduler_slot_ < 0) { + return std::nullopt; + } + audio = scheduler_->next_chunk(scheduler_slot_); + if (audio.samples.empty()) { + scheduler_slot_ = -1; + return std::nullopt; + } + } else { + if (!stream_session_) { + return std::nullopt; + } + audio = stream_session_->next_chunk(); + if (audio.samples.empty()) { + stream_session_.reset(); + return std::nullopt; + } } runtime::append_audio_buffer(stream_merged_audio_, audio); runtime::StreamEvent event; @@ -458,6 +554,12 @@ void FireRedTTS3Session::reset() { stream_chunk_index_ = 0; stream_merged_audio_ = runtime::AudioBuffer{}; stream_generated_text_.clear(); + if (scheduler_enabled_ && scheduler_ && scheduler_slot_ >= 0) { + // 排空剩余块，确保 slot 正常结束并归还（finish_slot 内部释放槽位）。 + while (scheduler_->next_chunk(scheduler_slot_).samples.empty() == false) { + } + scheduler_slot_ = -1; + } stream_session_.reset(); stream_started_ = false; } From 6354b391fb7b938ae9715e3d6256410376c335d7 Mon Sep 17 00:00:00 2001 From: dujing Date: Tue, 1 Sep 2026 18:20:19 +0800 Subject: [PATCH 06/12] Fix streaming regression: yield scheduler lock after chunk output for true incremental streaming MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 调度线程在 next_chunk 的 cv_.wait 让锁期间持锁持续 tick，一次性生成完整个 slot， 导致所有 chunk 攒齐才返回，SSE 流式退化成离线（chunk 同时到达）。 修复：调度线程每 tick 后检测是否产出了新 chunk，有则 cv_wake_.wait_for(1ms) 短暂 让锁，让 next_chunk 抢到锁取走 chunk，实现真增量流式。 验证：44s 长文本流式 16 个 chunk 逐块到达（间隔 ~0.6s），每 chunk RTF≈0.32； 2 并发下两路 chunk 交错增量到达，整体 RTF 1.12。 Co-Authored-By: Claude --- src/models/fireredtts3/batch_scheduler.cpp | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/src/models/fireredtts3/batch_scheduler.cpp b/src/models/fireredtts3/batch_scheduler.cpp index 0ab9ef489..b9b23d876 100644 --- a/src/models/fireredtts3/batch_scheduler.cpp +++ b/src/models/fireredtts3/batch_scheduler.cpp @@ -603,7 +603,22 @@ class FireRedTTS3BatchScheduler::Impl { void scheduler_loop() { std::unique_lock lock(mu_); while (!stop_) { + // 记录 tick 前各 slot 的已产出 chunk 数 + size_t chunks_before = 0; + for (int64_t i = 0; i < max_batch_; ++i) { + chunks_before += slots_[static_cast(i)].chunk_queue.size(); + } tick_locked(); + // 若本 tick 产出了新 chunk，短暂让出锁：让 next_chunk（cv_.wait）能抢到锁 + // 取走 chunk，实现真正的增量流式（否则调度线程持锁持续 tick，一次性 + // 生成完整个 slot，流式退化成离线）。 + size_t chunks_after = 0; + for (int64_t i = 0; i < max_batch_; ++i) { + chunks_after += slots_[static_cast(i)].chunk_queue.size(); + } + if (chunks_after > chunks_before) { + cv_wake_.wait_for(lock, std::chrono::milliseconds(1)); + } // 有工作（Active 未完成 / Dead 待 finish）则继续 tick；否则等新工作。 bool has_work = false; for (int64_t i = 0; i < max_batch_; ++i) { From 1ced5f04c212928130f24984d8cdbb35ed9a478f Mon Sep 17 00:00:00 2001 From: dujing Date: Wed, 2 Sep 2026 19:10:14 +0800 Subject: [PATCH 07/12] =?UTF-8?q?fix(firered):=20=E6=B5=81=E5=BC=8F?= =?UTF-8?q?=E6=9C=AB=E5=8F=A5=202=C3=97=20=E6=97=B6=E9=95=BF=20+=20?= =?UTF-8?q?=E6=89=B9=E5=A4=84=E7=90=86=E7=A9=BA=E8=A1=8C=E5=86=BB=E7=BB=93?= =?UTF-8?q?/=E5=8F=8C=E8=BE=93=E5=87=BA=E4=BF=AE=E5=A4=8D?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - istft_graph: finish_incremental 只输出未输出尾部，不再从 0 重放整段 （append 已输出 [pad,total-pad)，此前重放导致流式 2× 时长/末句重复） - batch_scheduler: finish_slot 的 tail decode 与 flush 改为 if/else 二选一， 避免 decode+flush 双输出；冻结非活跃 batch 行解码位置（set_batched_member_end）， 避免空行 mask 污染活跃行致 stop 判定错乱/内容重复 - qwen_decode_runtime/ar: 新增 set_batched_member_end 贯穿实现 - redae_codec_runtime: encoder/downsample 关 bf16_autocast（配合 per-slot 状态） - runtime/server: [REQ]/[LAUNCH]/[FINISH]/[REBUILD]/[SPLICE] 诊断日志； build_speech_request 透传 reference_text - pipeline.cpp: 清理调试残留（空块/孤立空行） 验证: 单路流式末句 20.35s（原 26.62s 2×），ASR 末句仅一遍；离线/流式对齐 16.32s Co-Authored-By: Claude --- app/server/runtime.cpp | 6 +++ .../transformers/qwen_causal_decode_runtime.h | 2 + include/engine/models/fireredtts3/ar.h | 3 ++ src/framework/audio/istft_graph.cpp | 20 +++++-- src/framework/codecs/redae_codec_runtime.cpp | 5 +- .../qwen_causal_decode_runtime.cpp | 8 +++ src/models/fireredtts3/ar.cpp | 8 +++ src/models/fireredtts3/batch_scheduler.cpp | 52 +++++++++++++++++-- 8 files changed, 92 insertions(+), 12 deletions(-) diff --git a/app/server/runtime.cpp b/app/server/runtime.cpp index 2e2812df0..cc0cee19c 100644 --- a/app/server/runtime.cpp +++ b/app/server/runtime.cpp @@ -1908,6 +1908,7 @@ engine::runtime::TaskRequest ServerState::build_speech_request(const LoadedModel add_option_from_json(request.options, body, "repetition_penalty", "repetition_penalty"); add_option_from_json(request.options, body, "guidance_scale", "guidance_scale"); add_option_from_json(request.options, body, "num_inference_steps", "num_inference_steps"); + add_option_from_json(request.options, body, "reference_text", "reference_text"); if (const auto * value = body.find("instructions")) { request.options["instruction"] = value->as_string(); } @@ -2157,6 +2158,11 @@ ServerState::TimedTaskResult ServerState::run_streaming_model_impl( throw std::runtime_error("configured model does not provide streaming execution: " + model.config.id); } const auto started = Clock::now(); + if (request.text_input.has_value()) { + const std::string txt = request.text_input->text.substr(0, 12); + fprintf(stderr, "[REQ] session_idx=%zu text=%s\n", pool.index, txt.c_str()); + fflush(stderr); + } session->prepare(engine::runtime::build_preparation_request(request)); TimedTaskResult timed_result; const auto sink = [&](const engine::runtime::StreamEvent & event) { diff --git a/include/engine/framework/modules/transformers/qwen_causal_decode_runtime.h b/include/engine/framework/modules/transformers/qwen_causal_decode_runtime.h index 1e813db39..52a42ac53 100644 --- a/include/engine/framework/modules/transformers/qwen_causal_decode_runtime.h +++ b/include/engine/framework/modules/transformers/qwen_causal_decode_runtime.h @@ -101,6 +101,8 @@ class QwenCausalDecodeRuntime { QwenCausalDecodeStepResult decode_embeddings_batched( const std::vector & embeddings, int64_t batch_size); + // 冻结/重置某 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 diff --git a/include/engine/models/fireredtts3/ar.h b/include/engine/models/fireredtts3/ar.h index 82e99aaa5..cdd45c06d 100644 --- a/include/engine/models/fireredtts3/ar.h +++ b/include/engine/models/fireredtts3/ar.h @@ -46,6 +46,9 @@ class FireRedArRuntime { engine::modules::QwenCausalDecodeStepResult decode_embeddings_batched( const std::vector & embeddings, int64_t batch_size); 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); void release_graphs(); void release_backbone_graphs(); diff --git a/src/framework/audio/istft_graph.cpp b/src/framework/audio/istft_graph.cpp index ec457fc4d..4bcc22e63 100644 --- a/src/framework/audio/istft_graph.cpp +++ b/src/framework/audio/istft_graph.cpp @@ -292,11 +292,21 @@ class HostLogMagnitudePhaseISTFT::Impl { if (!inc_initialized_) { return {}; } - // flush 剩余全部样本（含尾部） - std::vector out(inc_folded_.size()); - for (size_t i = 0; i < inc_folded_.size(); ++i) { - const float denom = inc_envelope_[i]; - out[i] = denom <= 1.0e-11F ? 0.0F : inc_folded_[i] / denom; + // append_incremental 已输出 [pad, total_samples-pad) 的完整内部覆盖区， + // 与离线 compute() 的裁剪语义（输出 = output_size - 2*pad）一致。 + // 这里只输出"尚未输出"的尾部（若有），避免把整段折叠缓冲区重复倒出 + // （此前从索引 0 全量重放会导致流式输出 2 倍时长）。 + const int64_t start = inc_emitted_samples_; + const int64_t end = static_cast(inc_folded_.size()); + std::vector out; + if (end > start) { + out.resize(static_cast(end - start)); + for (int64_t i = start; i < end; ++i) { + const float denom = inc_envelope_[static_cast(i)]; + out[static_cast(i - start)] = + denom <= 1.0e-11F ? 0.0F + : inc_folded_[static_cast(i)] / denom; + } } inc_folded_.clear(); inc_envelope_.clear(); diff --git a/src/framework/codecs/redae_codec_runtime.cpp b/src/framework/codecs/redae_codec_runtime.cpp index 37485e7d4..5acacc922 100644 --- a/src/framework/codecs/redae_codec_runtime.cpp +++ b/src/framework/codecs/redae_codec_runtime.cpp @@ -1,3 +1,4 @@ +#include #include "engine/framework/codecs/redae_codec_runtime.h" #include "engine/framework/audio/istft_graph.h" @@ -264,7 +265,7 @@ std::shared_ptr load_redae_weights( options.graph_arena_bytes, options.graph_arena_bytes, execution.backend_type(), - true, + false, c.enc_sliding_window); auto encoder_load_config = encoder_config; encoder_load_config.decoder.stack.runtime.attention.prefill_mode = modules::QwenDecoderAttentionMode::ManualRepeat; @@ -284,7 +285,7 @@ std::shared_ptr load_redae_weights( options.graph_arena_bytes, options.graph_arena_bytes, execution.backend_type(), - true, + false, 0); auto downsample_load_config = downsample_config; downsample_load_config.decoder.stack.runtime.attention.prefill_mode = modules::QwenDecoderAttentionMode::ManualRepeat; diff --git a/src/framework/modules/transformers/qwen_causal_decode_runtime.cpp b/src/framework/modules/transformers/qwen_causal_decode_runtime.cpp index 2e1168a82..0ad543290 100644 --- a/src/framework/modules/transformers/qwen_causal_decode_runtime.cpp +++ b/src/framework/modules/transformers/qwen_causal_decode_runtime.cpp @@ -615,6 +615,10 @@ class QwenCausalDecodeRuntime::Impl { return batched_decode_cache_.export_state(); } + void set_batched_member_end(int64_t batch, int64_t end) { + batched_decode_cache_.set_member_end(batch, end); + } + void start_decode_embeddings_batched( const runtime::TransformerBatchedKVState & state, int64_t required_cache_steps) { @@ -1665,6 +1669,10 @@ void QwenCausalDecodeRuntime::start_decode_embeddings_batched( impl_->start_decode_embeddings_batched(state, required_cache_steps); } +void QwenCausalDecodeRuntime::set_batched_member_end(int64_t batch, int64_t end) { + impl_->set_batched_member_end(batch, end); +} + QwenCausalDecodeStepResult QwenCausalDecodeRuntime::decode_tokens_batched(const std::vector & tokens) { return impl_->decode_tokens_batched(tokens); } diff --git a/src/models/fireredtts3/ar.cpp b/src/models/fireredtts3/ar.cpp index 9cf82ae21..24c7b983b 100644 --- a/src/models/fireredtts3/ar.cpp +++ b/src/models/fireredtts3/ar.cpp @@ -1024,6 +1024,10 @@ class FireRedArRuntime::Impl { return backbone_->decode_embeddings_batched(embeddings, batch_size); } + void set_batched_member_end(int64_t batch, int64_t end) { + backbone_->set_batched_member_end(batch, end); + } + runtime::TransformerBatchedKVState export_batched_decode_state() const { return backbone_->export_batched_decode_state(); } @@ -1151,6 +1155,10 @@ engine::modules::QwenCausalDecodeStepResult FireRedArRuntime::decode_embeddings_ return impl_->decode_embeddings_batched(embeddings, batch_size); } +void FireRedArRuntime::set_batched_member_end(int64_t batch, int64_t end) { + impl_->set_batched_member_end(batch, end); +} + engine::runtime::TransformerBatchedKVState FireRedArRuntime::export_batched_decode_state() const { return impl_->export_batched_decode_state(); } diff --git a/src/models/fireredtts3/batch_scheduler.cpp b/src/models/fireredtts3/batch_scheduler.cpp index b9b23d876..73c0178ff 100644 --- a/src/models/fireredtts3/batch_scheduler.cpp +++ b/src/models/fireredtts3/batch_scheduler.cpp @@ -180,6 +180,8 @@ class FireRedTTS3BatchScheduler::Impl { slot.chunks = chunk_patches.empty() ? std::vector{400} : chunk_patches; slot.chunk_target = slot.chunks[0]; pending_.push_back(slot_id); + fprintf(stderr, "[LAUNCH] slot=%ld tokens=%zu\n", slot_id, request.token_ids.size()); + fflush(stderr); cv_wake_.notify_all(); return slot_id; } @@ -397,6 +399,14 @@ class FireRedTTS3BatchScheduler::Impl { embeddings.begin() + static_cast(i * hidden)); } } + // 5.5 冻结非活跃行的解码位置（end=0，mask 全 -inf），避免空行被 + // advance_member 推高位置后 mask 污染活跃行（导致 stop 判定错乱/内容重复）。 + for (int64_t i = 0; i < max_batch_; ++i) { + auto & slot = slots_[static_cast(i)]; + if (slot.state != Slot::State::Active || !slot.prefill_done || slot.step < 1) { + ar_->set_batched_member_end(i, 0); + } + } // 6. 一次 batched decode（每成员独立 pos/cache-slot/mask） auto out = ar_->decode_embeddings_batched(embeddings, max_batch_); // 7. 分发 hidden 到各 step>=1 活跃 slot @@ -566,20 +576,25 @@ class FireRedTTS3BatchScheduler::Impl { // 只处理 Dead（Done 已处理过并归还）；Done 状态由 launch 复用前的 reset 清除。 void finish_slot_locked(int64_t slot_id) { auto & slot = slots_[static_cast(slot_id)]; + fprintf(stderr, "[FINISH] slot=%ld tokens=%zu gen_patches=%ld\n", + slot_id, slot.request.token_ids.size(), slot.generated_patches); + fflush(stderr); if (slot.state != Slot::State::Dead) { return; } - // 还有未解完的 chunk_latents → flush + // 有剩余 chunk_latents → decode_incremental 输出 tail（含 istft 尾部）； + // 无剩余 → flush_incremental 输出尾部。二选一，避免 decode + flush 双输出导致重复。 if (!slot.chunk_latents.empty() && slot.generated_patches > 0) { auto audio = redae_->decode_incremental(slot.redae_state, slot.chunk_latents); slot.chunk_latents.clear(); if (!audio.samples.empty()) { slot.chunk_queue.push_back(std::move(audio)); } - } - auto flush = redae_->flush_incremental(slot.redae_state); - if (!flush.samples.empty()) { - slot.chunk_queue.push_back(std::move(flush)); + } else { + auto flush = redae_->flush_incremental(slot.redae_state); + if (!flush.samples.empty()) { + slot.chunk_queue.push_back(std::move(flush)); + } } slot.state = Slot::State::Idle; slot.finished = true; @@ -691,6 +706,26 @@ class FireRedTTS3BatchScheduler::Impl { state.current_ends[static_cast(i)] = gpu_end_for(gpu, i); } } + static int rebuild_dbg = 0; + if (rebuild_dbg++ < 10) { + std::string ends; + for (int64_t i = 0; i < max_batch_; ++i) { + ends += std::to_string(state.current_ends[static_cast(i)]) + ","; + } + std::string gpu_ends; + if (!gpu.current_ends.empty()) { + for (size_t i = 0; i < gpu.current_ends.size(); ++i) { + gpu_ends += std::to_string(gpu.current_ends[i]) + ","; + } + } + fprintf(stderr, "[REBUILD] batch=%ld state_ends=[%s] gpu_ends=[%s] running=[", + (long)max_batch_, ends.c_str(), gpu_ends.c_str()); + for (size_t i = 0; i < running_active_.size(); ++i) { + fprintf(stderr, "%ld,", (long)running_active_[i]); + } + fprintf(stderr, "]\n"); + fflush(stderr); + } int64_t max_steps = 0; for (int64_t i = 0; i < max_batch_; ++i) { max_steps = std::max(max_steps, state.current_ends[static_cast(i)]); @@ -716,6 +751,13 @@ class FireRedTTS3BatchScheduler::Impl { if (end <= 0) { continue; } + static int splice_dbg = 0; + if (splice_dbg++ < 20) { + fprintf(stderr, "[SPLICE] row=%ld src=%s tokens=%zu end=%ld\n", b, + imported_slots_.count(b) > 0 ? "new-prefill" : "gpu-running", + slots_[static_cast(b)].request.token_ids.size(), end); + fflush(stderr); + } const int64_t current_end = end; for (size_t layer = 0; layer < state.layers.size(); ++layer) { const size_t row_elems = row_elems_for_layer(layer); From acc3b67750ef6fc81069064049b1626b92ae04e2 Mon Sep 17 00:00:00 2001 From: dujing Date: Wed, 2 Sep 2026 20:06:07 +0800 Subject: [PATCH 08/12] =?UTF-8?q?fix(firered):=20=E5=85=B1=E4=BA=AB=20sche?= =?UTF-8?q?duler=20=E5=B9=B6=E5=8F=91=20double-free=20=E5=B4=A9=E6=BA=83?= =?UTF-8?q?=EF=BC=88epoch=20=E5=8F=A5=E6=9F=84=20+=20owner=20=E5=BB=B6?= =?UTF-8?q?=E8=BF=9F=E5=BD=92=E8=BF=98=EF=BC=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 根因：finish_slot_locked 在 slot 变 Dead 时立即 push 回 free_slots_，而所属 session 尚未 drain 完尾块；新请求 launch 随即复用同 slot（slot=Slot{} 重置）， 旧 session 的 next_chunk 读到新请求状态（串音），Slot{} 析构与 drain 交错 导致 free(): double free。 修复： - Slot 增加 epoch；launch 返回 SlotHandle{id, epoch}；next_chunk 首行校验 epoch，stale 句柄立即返回空（不读新请求状态，防串音） - finish_slot 不再 push free_slots_（slot 留保留态等 owner），owner 在 next_chunk 见空/reset/generate drain 完时显式 release_slot 归还 - 新增 abort（Active→Dead，reset/异常清理用） - session scheduler_slot_ 句柄化；next_stream_event 见空即 release - streaming run_stream catch 补 session.reset()（中断请求遗留 slot 立即回收） 验证：3 路并发无 double-free，slot 生命周期严格 LAUNCH(epoch)→FINISH→DONE_WAIT_RELEASE→RELEASED→LAUNCH。 注：并发串音（GPU decode 层）独立于本崩溃，另定位。 Co-Authored-By: Claude --- app/streaming/streaming.cpp | 3 + .../models/fireredtts3/batch_scheduler.h | 25 +++- include/engine/models/fireredtts3/session.h | 4 +- src/models/fireredtts3/batch_scheduler.cpp | 137 +++++++++++++++--- src/models/fireredtts3/session.cpp | 24 ++- 5 files changed, 155 insertions(+), 38 deletions(-) diff --git a/app/streaming/streaming.cpp b/app/streaming/streaming.cpp index 55bedf3ac..53eea8775 100644 --- a/app/streaming/streaming.cpp +++ b/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; } } diff --git a/include/engine/models/fireredtts3/batch_scheduler.h b/include/engine/models/fireredtts3/batch_scheduler.h index 38c7a4e0f..b9700b492 100644 --- a/include/engine/models/fireredtts3/batch_scheduler.h +++ b/include/engine/models/fireredtts3/batch_scheduler.h @@ -51,6 +51,16 @@ class FireRedTTS3BatchScheduler { std::deque chunk_queue; bool finished = false; std::exception_ptr error; + // 代次：launch 时 +1（重置后回写），release_slot 时再 +1。句柄失效/幂等判定用。 + uint64_t epoch = 0; + }; + + // 句柄：{slot id, 代次}。持有句柄才能读块/归还/终止。 + // slot 被释放并复用时 epoch +1，旧句柄立即失效（stale 访问视作流结束，不读新请求状态）。 + struct SlotHandle { + int64_t id = -1; + uint64_t epoch = 0; + bool valid() const noexcept { return id >= 0; } }; FireRedTTS3BatchScheduler( @@ -68,11 +78,16 @@ class FireRedTTS3BatchScheduler { FireRedTTS3BatchScheduler(const FireRedTTS3BatchScheduler &) = delete; FireRedTTS3BatchScheduler & operator=(const FireRedTTS3BatchScheduler &) = delete; - // 启动一个请求（分配 slot），返回 slot id。返回 -1 表示槽池耗尽。 - int64_t launch(const FireRedTTS3BaseRequest & request, const std::vector & chunk_patches); - // 取下一个音频块（驱动调度轮次直到该 slot 产出块或结束）。 - engine::runtime::AudioBuffer next_chunk(int64_t slot_id); - // 离线整句：launch + drain + 拼接（trim prompt）。 + // 启动一个请求（分配 slot），返回句柄。槽池耗尽返回无效句柄（id == -1）。 + SlotHandle launch(const FireRedTTS3BaseRequest & request, const std::vector & chunk_patches); + // 取下一音频块（驱动调度轮次直到该 slot 产出块或结束）。 + // 句柄失效（stale / 已 release）→ 立即返回空块，不读新请求状态（防串音）。 + engine::runtime::AudioBuffer next_chunk(const SlotHandle & handle); + // 归还已完成并排空的 slot 到空闲池（owner 在确认结束时显式调用；幂等）。 + void release_slot(const SlotHandle & handle); + // 快速终止仍 Active 的 slot（供 reset/异常清理），终止后需排空 + release_slot。 + void abort(const SlotHandle & handle); + // 离线整句：launch + drain + 内部 release + 拼接（trim prompt）。 engine::runtime::AudioBuffer generate(const FireRedTTS3BaseRequest & request); void release_graphs(); diff --git a/include/engine/models/fireredtts3/session.h b/include/engine/models/fireredtts3/session.h index b530dea9a..59836ac66 100644 --- a/include/engine/models/fireredtts3/session.h +++ b/include/engine/models/fireredtts3/session.h @@ -58,8 +58,8 @@ class FireRedTTS3Session final bool mem_saver_ = false; // 共享 batch scheduler（同一 assets 的所有 session 共享）；Base clone 用。 std::shared_ptr scheduler_; - // 当前流式请求在 scheduler 中占用的 slot（-1 = 无）。 - int64_t scheduler_slot_ = -1; + // 当前流式请求在 scheduler 中占用的 slot 句柄（无效 = 无）。 + FireRedTTS3BatchScheduler::SlotHandle scheduler_slot_; bool scheduler_enabled_ = false; // 流式状态 diff --git a/src/models/fireredtts3/batch_scheduler.cpp b/src/models/fireredtts3/batch_scheduler.cpp index 73c0178ff..922215330 100644 --- a/src/models/fireredtts3/batch_scheduler.cpp +++ b/src/models/fireredtts3/batch_scheduler.cpp @@ -166,29 +166,42 @@ class FireRedTTS3BatchScheduler::Impl { } } - int64_t launch(const FireRedTTS3BaseRequest & request, const std::vector & chunk_patches) { + SlotHandle launch(const FireRedTTS3BaseRequest & request, const std::vector & chunk_patches) { std::lock_guard lock(mu_); if (free_slots_.empty()) { - return -1; + return SlotHandle{}; } const int64_t slot_id = free_slots_.front(); free_slots_.pop_front(); auto & slot = slots_[static_cast(slot_id)]; + // epoch 先取后回写：slot = Slot{} 会把 epoch 清零，故先 +1 再重置、回写单调代次。 + const uint64_t epoch = slot.epoch + 1; slot = Slot{}; + slot.epoch = epoch; slot.state = Slot::State::Active; slot.request = request; slot.chunks = chunk_patches.empty() ? std::vector{400} : chunk_patches; slot.chunk_target = slot.chunks[0]; pending_.push_back(slot_id); - fprintf(stderr, "[LAUNCH] slot=%ld tokens=%zu\n", slot_id, request.token_ids.size()); + fprintf(stderr, "[LAUNCH] slot=%ld epoch=%llu tokens=%zu\n", + slot_id, (unsigned long long)epoch, request.token_ids.size()); fflush(stderr); cv_wake_.notify_all(); - return slot_id; + return SlotHandle{slot_id, epoch}; } - engine::runtime::AudioBuffer next_chunk(int64_t slot_id) { + engine::runtime::AudioBuffer next_chunk(const SlotHandle & handle) { std::unique_lock lock(mu_); - auto & slot = slots_[static_cast(slot_id)]; + // 首行所有权校验（在读 slot.error/queue/finished 之前）： + // 句柄无效 / 越界 / 代次不匹配（stale，slot 已被 release 或复用）→ 立即返回空 = 流结束， + // 绝不读新请求的 slot 状态（防串音），也避免旧 owner 在已复用 slot 上继续 move/析构。 + if (!handle.valid() || handle.id >= max_batch_) { + return {}; + } + auto & slot = slots_[static_cast(handle.id)]; + if (slot.epoch != handle.epoch) { + return {}; + } if (slot.error) { std::rethrow_exception(slot.error); } @@ -219,20 +232,34 @@ class FireRedTTS3BatchScheduler::Impl { } engine::runtime::AudioBuffer generate(const FireRedTTS3BaseRequest & request) { - const int64_t slot_id = launch(request, {}); - if (slot_id < 0) { + const SlotHandle handle = launch(request, {}); + if (!handle.valid()) { throw std::runtime_error("FireRedTTS3BatchScheduler slot pool exhausted"); } engine::runtime::AudioBuffer merged; merged.sample_rate = static_cast(assets_->redae.sample_rate); merged.channels = 1; - while (true) { - auto chunk = next_chunk(slot_id); - if (chunk.samples.empty()) { - break; + try { + while (true) { + auto chunk = next_chunk(handle); + if (chunk.samples.empty()) { + break; + } + runtime::append_audio_buffer(merged, chunk); + } + } catch (...) { + // 异常：快速终结仍 Active 的 slot，排空到 finished，归还，再抛。 + abort(handle); + try { + while (!next_chunk(handle).samples.empty()) { + } + } catch (...) { } - runtime::append_audio_buffer(merged, chunk); + release_slot(handle); + throw; } + // drain 结束（slot 必已 finished）：显式归还，slot 才可被复用。 + release_slot(handle); return merged; } @@ -255,6 +282,58 @@ class FireRedTTS3BatchScheduler::Impl { campplus_.release_runtime_graph(); } + // owner 在排空结束（next_chunk 见空 / reset / generate drain 完）时显式归还 slot。 + // 只回收 terminal（finished 且非 Active/Dead）的 slot；过早调用 no-op。 + // 幂等：epoch 不匹配（已 release/复用）直接返回。 + void release_slot(const SlotHandle & handle) { + std::lock_guard lock(mu_); + if (!handle.valid() || handle.id >= max_batch_) { + return; + } + auto & slot = slots_[static_cast(handle.id)]; + if (slot.epoch != handle.epoch) { + return; // 已 release 或复用 → 幂等 + } + if (slot.state == Slot::State::Active || slot.state == Slot::State::Dead) { + return; // 过早：调用方应先 abort+drain，不回收活 slot + } + if (!slot.finished) { + return; + } + slot.epoch++; // 残留句柄立即失效 + if (std::find(free_slots_.begin(), free_slots_.end(), handle.id) == free_slots_.end()) { + free_slots_.push_back(handle.id); + } + fprintf(stderr, "[RELEASED] slot=%ld epoch=%llu\n", + handle.id, (unsigned long long)slot.epoch); + fflush(stderr); + cv_.notify_all(); + } + + // 快速终止仍 Active 的 slot（reset/异常清理用）： + // Active → Dead（若尚未被 form_wave 领走则从 pending 摘除），调度线程下一轮 finish 收尾。 + void abort(const SlotHandle & handle) { + std::lock_guard lock(mu_); + if (!handle.valid() || handle.id >= max_batch_) { + return; + } + auto & slot = slots_[static_cast(handle.id)]; + if (slot.epoch != handle.epoch) { + return; + } + if (slot.state != Slot::State::Active) { + return; // 已 Dead/Idle/Failed，无需干预 + } + if (!slot.prefill_done) { + // 尚未被 form_wave 领走：直接从 pending 摘除，避免 Dead slot 被 prefill。 + pending_.erase( + std::remove(pending_.begin(), pending_.end(), handle.id), + pending_.end()); + } + slot.state = Slot::State::Dead; + cv_wake_.notify_all(); // 唤醒可能休眠的调度线程收尾 + } + private: // ---- 参考音色（锁内串行，产物拷进 slot）---- const ReferenceVoiceCacheEntry & prepare_reference_voice_locked(const runtime::AudioBuffer & audio) { @@ -598,11 +677,12 @@ class FireRedTTS3BatchScheduler::Impl { } slot.state = Slot::State::Idle; slot.finished = true; - // 防重复归还：仅当 slot 不在 free_slots_ 时才 push（避免并发下重复导致 - // 同一 slot 被多个 session launch）。 - if (std::find(free_slots_.begin(), free_slots_.end(), slot_id) == free_slots_.end()) { - free_slots_.push_back(slot_id); - } + // 不再 push free_slots_：slot 暂归 owner 所有（保留态），等 owner 在 + // next_chunk 排空见空后显式 release_slot 归还，杜绝"未 drain 完的 slot 被复用" + // （旧 owner 串读新请求 chunk / Slot{} 析构与 drain 交错 double-free）。 + fprintf(stderr, "[DONE_WAIT_RELEASE] slot=%ld gen_patches=%ld\n", + slot_id, slot.generated_patches); + fflush(stderr); // 该行不再参与 decode：从 running 集合移除（下次 join 时重建）。 running_active_.erase( std::remove(running_active_.begin(), running_active_.end(), slot_id), @@ -668,9 +748,8 @@ class FireRedTTS3BatchScheduler::Impl { slot.error = std::current_exception(); slot.state = Slot::State::Failed; slot.finished = true; - if (std::find(free_slots_.begin(), free_slots_.end(), slot_id) == free_slots_.end()) { - free_slots_.push_back(slot_id); - } + // 不 push free_slots_：失败 slot 留待 owner next_chunk 抛错 → reset/异常路径 + // 显式 release_slot 回收，避免失败槽在 owner 处理错误前被复用。 } } } @@ -930,12 +1009,22 @@ FireRedTTS3BatchScheduler::FireRedTTS3BatchScheduler( FireRedTTS3BatchScheduler::~FireRedTTS3BatchScheduler() = default; -int64_t FireRedTTS3BatchScheduler::launch(const FireRedTTS3BaseRequest & request, const std::vector & chunk_patches) { +FireRedTTS3BatchScheduler::SlotHandle FireRedTTS3BatchScheduler::launch( + const FireRedTTS3BaseRequest & request, + const std::vector & chunk_patches) { return impl_->launch(request, chunk_patches); } -engine::runtime::AudioBuffer FireRedTTS3BatchScheduler::next_chunk(int64_t slot_id) { - return impl_->next_chunk(slot_id); +engine::runtime::AudioBuffer FireRedTTS3BatchScheduler::next_chunk(const SlotHandle & handle) { + return impl_->next_chunk(handle); +} + +void FireRedTTS3BatchScheduler::release_slot(const SlotHandle & handle) { + impl_->release_slot(handle); +} + +void FireRedTTS3BatchScheduler::abort(const SlotHandle & handle) { + impl_->abort(handle); } engine::runtime::AudioBuffer FireRedTTS3BatchScheduler::generate(const FireRedTTS3BaseRequest & request) { diff --git a/src/models/fireredtts3/session.cpp b/src/models/fireredtts3/session.cpp index 4cc3c849b..0dda462e6 100644 --- a/src/models/fireredtts3/session.cpp +++ b/src/models/fireredtts3/session.cpp @@ -484,7 +484,7 @@ void FireRedTTS3Session::start_stream(const runtime::TaskRequest & request) { if (scheduler_enabled_ && scheduler_) { // 共享 scheduler：launch 一个 slot（并发请求共享 GPU batch decode） scheduler_slot_ = scheduler_->launch(base, stream_chunk_patches_); - if (scheduler_slot_ < 0) { + if (!scheduler_slot_.valid()) { throw std::runtime_error("FireRedTTS3 scheduler slot pool exhausted"); } } else { @@ -500,12 +500,15 @@ std::optional FireRedTTS3Session::next_stream_event() { } runtime::AudioBuffer audio; if (scheduler_enabled_ && scheduler_) { - if (scheduler_slot_ < 0) { + if (!scheduler_slot_.valid()) { return std::nullopt; } audio = scheduler_->next_chunk(scheduler_slot_); if (audio.samples.empty()) { - scheduler_slot_ = -1; + // 流结束：显式归还 slot（此后该 id 才可能被复用），并失效本句柄， + // 杜绝"本 session 尚未完全 drain / 仍在持句柄时 slot 被其他请求复用"。 + scheduler_->release_slot(scheduler_slot_); + scheduler_slot_ = FireRedTTS3BatchScheduler::SlotHandle{}; return std::nullopt; } } else { @@ -554,11 +557,18 @@ void FireRedTTS3Session::reset() { stream_chunk_index_ = 0; stream_merged_audio_ = runtime::AudioBuffer{}; stream_generated_text_.clear(); - if (scheduler_enabled_ && scheduler_ && scheduler_slot_ >= 0) { - // 排空剩余块，确保 slot 正常结束并归还（finish_slot 内部释放槽位）。 - while (scheduler_->next_chunk(scheduler_slot_).samples.empty() == false) { + if (scheduler_enabled_ && scheduler_ && scheduler_slot_.valid()) { + // 若上一次流式请求因异常/中断遗留仍 Active 的 slot：先 abort 快速收尾， + // 再排空（next_chunk 等待至 finished），最后显式 release_slot 归还。 + scheduler_->abort(scheduler_slot_); + try { + while (!scheduler_->next_chunk(scheduler_slot_).samples.empty()) { + } + } catch (...) { + // 旧请求的 prefill/decode 错误只在此吞掉清理，不得泄漏到下一个请求。 } - scheduler_slot_ = -1; + scheduler_->release_slot(scheduler_slot_); + scheduler_slot_ = FireRedTTS3BatchScheduler::SlotHandle{}; } stream_session_.reset(); stream_started_ = false; From 0490201a499aef55b53ed8d82ce2346a88f8fe9b Mon Sep 17 00:00:00 2001 From: dujing Date: Thu, 3 Sep 2026 11:36:26 +0800 Subject: [PATCH 09/12] =?UTF-8?q?fix(firered):=20=E5=85=B1=E4=BA=AB=20sche?= =?UTF-8?q?duler=20=E5=8A=A0=E8=BD=BD=E7=AB=9E=E6=80=81=20+=20=E6=AF=8F=20?= =?UTF-8?q?model=20=E5=85=B1=E4=BA=AB=E5=8D=95=E4=B8=80=20CUDA=20context?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 两个独立修复，消除并发冷启动的串音与崩溃根因： 1. server session 池加载竞态（ensure_model_loaded_locked）： - 旧逻辑仅当 max_loaded_models/min_free_memory_mb 配置时才持 model_load_mutex_ - 两者都没设时并发首请求同时进入，各自 clear()+push 同一 session 池 → free_sessions 重复索引，所有请求借到同一 session（串音/截断的根因） - 改为无条件持锁 + 以 model.loaded 为门，只建一个干净 session 池 2. 每 model 共享单一 ExecutionContext/backend（对齐 llama.cpp 单 context 多 slot）： - 旧实现每 session 在 RuntimeSessionBase 自建 ExecutionContext（自建 CUDA backend），instance_count=3 → 3 个 device-0 context；scheduler 借用第一个 session 的引用（session 析构即悬垂）→ 冷并发 "CUDA error: invalid device context" - RuntimeSessionBase 增加"借外部 context"构造（nullptr=自建原行为，其余 ~46 个 session 不受影响）；共享注册表 assets-keyed 持 context，所有权随 session - scheduler-mode（Base + max_batch>1）session 借用共享 context；Instruct/单 Base 保持各自独立 context（非单调度线程，共享会撞后端线程安全） 验证：单路回归 WER 0.14 不回归；并发冷启动不再 crash invalid device context。 注：并发请求在共享 scheduler 串行轮转下的内容残留/截断（同输入串行两请求 44 vs 3 patches）是独立 pre-existing bug，另定位。 Co-Authored-By: Claude --- app/server/runtime.cpp | 22 ++--- .../engine/framework/runtime/session_base.h | 12 ++- src/framework/runtime/session_base.cpp | 19 ++++- src/models/fireredtts3/session.cpp | 84 +++++++++++++++---- 4 files changed, 107 insertions(+), 30 deletions(-) diff --git a/app/server/runtime.cpp b/app/server/runtime.cpp index cc0cee19c..2c6d82adc 100644 --- a/app/server/runtime.cpp +++ b/app/server/runtime.cpp @@ -1756,19 +1756,16 @@ void ServerState::evict_for_model_limit(const LoadedModel & loading) { void ServerState::ensure_model_loaded_locked(LoadedModel & model) { last_activity_ms_.store(steady_now_ms(), std::memory_order_relaxed); model.last_used_ms.store(steady_now_ms(), std::memory_order_relaxed); - if (!model.sessions.empty()) { + if (model.loaded.load(std::memory_order_acquire)) { return; } - // Serialize the whole "evict -> memory check -> load" sequence only when a guard - // actually needs it: eviction (max_loaded_models > 0) or the memory pre-check - // (min_free_memory_mb > 0). With both off there is nothing for concurrent loads - // to race over, so unrelated first-load requests keep their original concurrency. - const bool serialize_load = - config_.max_loaded_models > 0 || config_.min_free_memory_mb > 0; - std::unique_lock load_lock(model_load_mutex_, std::defer_lock); - if (serialize_load) { - load_lock.lock(); - } + // 无条件串行化整个 "check -> evict -> load -> create session pool" 序列。 + // 若不持锁（旧逻辑只在 max_loaded_models / min_free_memory_mb 配置时才 serialize）， + // 多个并发首请求会在 model.sessions 仍空时同时进入，各自 clear()+push 同一份 + // LoadedModel.sessions / free_sessions（数据竞争），导致池被重复填充、 + // free_sessions 出现重复索引 —— 并发请求全部借到同一 session 下标，共享一个 + // session 的 scheduler slot 与 chunk 队列，造成跨请求串音/截断。 + std::unique_lock load_lock(model_load_mutex_); if (config_.max_loaded_models > 0) { evict_for_model_limit(model); // Note: if ensure_model_fits_memory() below then refuses the load, the @@ -1843,6 +1840,9 @@ void ServerState::ensure_model_loaded_locked(LoadedModel & model) { model.sessions.push_back(std::move(session)); model.free_sessions.push_back(static_cast(i)); } + fprintf(stderr, "[DIAG] model '%s' created %d sessions, free=%zu\n", + model.config.id.c_str(), instance_count, model.free_sessions.size()); + fflush(stderr); model.model = std::move(loaded_model); model.offline = offline; model.streaming = streaming; diff --git a/include/engine/framework/runtime/session_base.h b/include/engine/framework/runtime/session_base.h index 027e7914f..ce4dd7bf6 100644 --- a/include/engine/framework/runtime/session_base.h +++ b/include/engine/framework/runtime/session_base.h @@ -9,6 +9,7 @@ #include "engine/framework/runtime/session.h" #include "engine/framework/runtime/workspace.h" +#include #include #include #include @@ -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 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 context_; ArtifactStore artifacts_; RuntimeCache cache_; RuntimeWorkspace workspace_; diff --git a/src/framework/runtime/session_base.cpp b/src/framework/runtime/session_base.cpp index ea62fc108..84bba37cd 100644 --- a/src/framework/runtime/session_base.cpp +++ b/src/framework/runtime/session_base.cpp @@ -3,16 +3,27 @@ namespace engine::runtime { RuntimeSessionBase::RuntimeSessionBase(const SessionOptions & options) + : RuntimeSessionBase(options, nullptr) {} + +RuntimeSessionBase::RuntimeSessionBase( + const SessionOptions & options, + std::shared_ptr external_context) : options_(options), - execution_context_(options.backend), - graph_executor_(execution_context_) {} + context_(external_context != nullptr + ? std::move(external_context) + : std::make_shared(options.backend)), + graph_executor_(*context_) { + if (context_ == nullptr) { + throw std::runtime_error("RuntimeSessionBase requires a non-null execution context"); + } +} engine::core::ExecutionContext & RuntimeSessionBase::execution_context() noexcept { - return execution_context_; + return *context_; } const engine::core::ExecutionContext & RuntimeSessionBase::execution_context() const noexcept { - return execution_context_; + return *context_; } ArtifactStore & RuntimeSessionBase::artifacts() noexcept { diff --git a/src/models/fireredtts3/session.cpp b/src/models/fireredtts3/session.cpp index 0dda462e6..328111b35 100644 --- a/src/models/fireredtts3/session.cpp +++ b/src/models/fireredtts3/session.cpp @@ -26,19 +26,52 @@ constexpr const char * kInstructName = "FireRedTTS3 Instruct"; constexpr int64_t kDefaultTextChunkSize = 600; constexpr size_t kDefaultReferenceCacheSlots = 4; -// ---- 共享 batch scheduler 注册表（assets-keyed）---- -// 同一 assets 的所有 session（server session 池）共享一个 scheduler： -// 并发请求经各 session 进入 scheduler 的不同 slot，实现共享 GPU batch decode。 -// model unload（全部 session 析构）时 scheduler 随最后一个 weak_ptr 消亡。 +// ---- 共享 batch scheduler + 共享 ExecutionContext 注册表（assets-keyed）---- +// 同一 assets 的所有 session（server session 池）共享一个 scheduler（llama.cpp +// update_slots 的模拟）和一个 ExecutionContext/backend（llama.cpp 的单 context）。 +// +// 为什么 context 也必须共享：旧实现每个 session 在 RuntimeSessionBase 自建一个 +// ExecutionContext（自建 CUDA backend），instance_count=3 → 3 个 device-0 context， +// 而 scheduler 只借用第一个 session 的（引用绑定其成员，session 析构即悬垂）。 +// 冷并发首请求多 context 并存 → "CUDA error: invalid device context"。 +// 对齐 llama.cpp：一个 context 随 model 建一次、所有 session 借它。 +// +// 所有权/生命周期：context 与 scheduler 都由每个 session 强持有（session 基类持 +// context 的 shared_ptr，session 持 scheduler 的 shared_ptr）。注册表只存 weak_ptr， +// 不延长存活——最后一个 session 析构时 context+scheduler 随之释放。析构顺序安全： +// FireRedTTS3Session 的派生成员（含 scheduler_）先于基类析构，故 scheduler（内部持 +// context 引用）先于基类持有的 context 释放。 struct SchedulerRegistry { std::mutex mu; std::unordered_map< const FireRedTTS3Assets *, - std::weak_ptr> by_assets; + std::weak_ptr> context_by_assets; + std::unordered_map< + const FireRedTTS3Assets *, + std::weak_ptr> scheduler_by_assets; + + // 取 assets 对应的共享 backend context（首次创建）。返回强副本，由调用方 + // （session 基类）持有，保证该 context 在本 session 存活期不析构。 + std::shared_ptr get_or_create_context( + const std::shared_ptr & assets, + const engine::core::BackendConfig & backend_config) { + std::lock_guard lock(mu); + auto it = context_by_assets.find(assets.get()); + if (it != context_by_assets.end()) { + if (auto existing = it->second.lock()) { + return existing; + } + context_by_assets.erase(it); + } + auto context = std::make_shared(backend_config); + context_by_assets.emplace(assets.get(), context); + return context; + } + // 取 assets 对应的共享 scheduler（首次创建时用共享 context 构造）。返回强副本。 std::shared_ptr get_or_create( std::shared_ptr assets, - engine::core::ExecutionContext & execution, + std::shared_ptr context, size_t graph_arena_bytes, size_t helper_graph_arena_bytes, size_t weight_context_bytes, @@ -47,18 +80,18 @@ struct SchedulerRegistry { bool mem_saver, int64_t max_batch) { std::lock_guard lock(mu); - auto it = by_assets.find(assets.get()); - if (it != by_assets.end()) { + auto it = scheduler_by_assets.find(assets.get()); + if (it != scheduler_by_assets.end()) { if (auto existing = it->second.lock()) { return existing; } - by_assets.erase(it); + scheduler_by_assets.erase(it); } const FireRedTTS3Assets * key = assets.get(); auto scheduler = std::make_shared( - std::move(assets), execution, graph_arena_bytes, helper_graph_arena_bytes, + std::move(assets), *context, graph_arena_bytes, helper_graph_arena_bytes, weight_context_bytes, storage_type, reference_cache_slots, mem_saver, max_batch); - by_assets.emplace(key, scheduler); + scheduler_by_assets.emplace(key, scheduler); return scheduler; } }; @@ -91,6 +124,28 @@ std::shared_ptr require_contract( return contract; } +// scheduler-mode（Base clone + fireredtts3.max_batch>1）时，返回 assets 对应的 +// 共享 ExecutionContext（注册表持有，所有 session 借用 → llama.cpp 单 context）。 +// 其它情况（Instruct、非 scheduler 的单 Base）返回 nullptr → session 自建 context： +// 这些路径不集中到单调度线程，每个 session 独立并发跑自己的 graph，共享 backend +// 会撞 ggml 后端线程安全，故必须各持一个。 +std::shared_ptr firered_scheduler_shared_context( + const runtime::SessionOptions & options, + const std::shared_ptr & assets) { + if (assets == nullptr) { + return nullptr; // require_assets 会给出明确报错；此处不提前解引用 + } + if (is_instruct_variant(assets->variant)) { + return nullptr; + } + const int64_t max_batch = runtime::parse_i64_option( + options.options, {"fireredtts3.max_batch"}).value_or(1); + if (max_batch <= 1) { + return nullptr; + } + return scheduler_registry().get_or_create_context(assets, options.backend); +} + std::string request_language(const runtime::TaskRequest & request) { if (const auto language = runtime::find_option(request.options, {"language"})) { return *language; @@ -263,9 +318,9 @@ FireRedTTS3Session::FireRedTTS3Session( runtime::SessionOptions options, std::shared_ptr assets, std::shared_ptr contract) - : runtime::RuntimeSessionBase(options), + : runtime::RuntimeSessionBase(options, firered_scheduler_shared_context(options, assets)), task_(task), - assets_(require_assets(std::move(assets))), + assets_(require_assets(assets)), contract_(require_contract(std::move(contract))), tokenizer_(std::make_unique(assets_)) { const bool is_instruct = is_instruct_variant(assets_->variant); @@ -326,9 +381,10 @@ FireRedTTS3Session::FireRedTTS3Session( {"fireredtts3.max_batch"}) .value_or(1); if (max_batch > 1) { + // scheduler 用共享 context 构造（与基类 borrow 的是同一个）。 scheduler_ = scheduler_registry().get_or_create( assets_, - execution_context(), + scheduler_registry().get_or_create_context(assets_, options.backend), graph_arena_bytes, helper_graph_arena_bytes, weight_context_bytes, From e3a406ff7177b33f5366fc6b69aa0f8e53f32430 Mon Sep 17 00:00:00 2001 From: dujing Date: Fri, 4 Sep 2026 10:11:21 +0800 Subject: [PATCH 10/12] =?UTF-8?q?fix(firered):=20=E9=87=8D=E5=86=99?= =?UTF-8?q?=E5=85=B1=E4=BA=AB=20scheduler=20=E4=B8=BA=20round=20=E5=88=B6?= =?UTF-8?q?=E7=9C=9F=E5=B9=B6=E5=8F=91=20+=20runtime=20active=5Fmask=20?= =?UTF-8?q?=E8=A1=8C=E9=9A=94=E7=A6=BB?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 并发串音调查：scheduler KV 层确有 stale 行泄漏（slot 完成后 GPU 行带完整 K/V 不被 清零，Idle 行 freeze→advance 振荡且 mask 暴露 stale pos0），但非串音根因 —— 彻底重写后单路逐位一致、多路 decode 行互不污染，串音仍复现于 server session 并发 层（double-free + 并发 [REQ] 同 session_idx，另查）。 重写（round 制，删除 export/splice/rebuild/dirty/freeze 整机）： - 一轮并发请求 = 一次全量干净 import（只写本轮 prefill KV，其余行全零） - 每 tick 一次 batched decode 真并发推进所有活跃行 - runtime batched decode 新增 active_mask：非活跃行整行 -inf 且不 advance， 即使 cache 段残留上一请求内容也不 attend（inert 冻结） - 保留 AR/flow/RedAE 单路径语义（与 pipeline.cpp 逐位一致） - 保留 [ROUND]/[NAN]/[FINISH] 诊断日志 + batched_member_ends 只读 accessor 验证：单路流式 audio=7.04s / WER 0.1429，与重写前逐位一致（不回归）。 Co-Authored-By: Claude --- .../transformers/qwen_causal_decode_runtime.h | 10 +- .../transformers/qwen_causal_decoder.h | 5 +- include/engine/models/fireredtts3/ar.h | 6 +- .../models/fireredtts3/batch_scheduler.h | 33 +- .../qwen_causal_decode_runtime.cpp | 59 +- .../transformers/qwen_causal_decoder.cpp | 18 +- src/models/fireredtts3/ar.cpp | 19 +- src/models/fireredtts3/batch_scheduler.cpp | 559 +++++++----------- 8 files changed, 327 insertions(+), 382 deletions(-) diff --git a/include/engine/framework/modules/transformers/qwen_causal_decode_runtime.h b/include/engine/framework/modules/transformers/qwen_causal_decode_runtime.h index 52a42ac53..d18cf757b 100644 --- a/include/engine/framework/modules/transformers/qwen_causal_decode_runtime.h +++ b/include/engine/framework/modules/transformers/qwen_causal_decode_runtime.h @@ -98,9 +98,14 @@ class QwenCausalDecodeRuntime { const runtime::TransformerBatchedKVState & state, int64_t required_cache_steps); QwenCausalDecodeStepResult decode_tokens_batched(const std::vector & 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 & embeddings, - int64_t batch_size); + int64_t batch_size, + const std::vector & active_mask = {}); // 冻结/重置某 batch 行的解码位置（非活跃行 end=0，mask 全 -inf）。 void set_batched_member_end(int64_t batch, int64_t end); @@ -114,6 +119,9 @@ class QwenCausalDecodeRuntime { int64_t decode_valid_steps() const noexcept; void release_runtime_graphs(); + // [DIAG] 每 batch 行的当前解码结束位置（member_ends_ 拷贝；未启动则空）。 + std::vector batched_member_ends() const; + private: class Impl; std::unique_ptr impl_; diff --git a/include/engine/framework/modules/transformers/qwen_causal_decoder.h b/include/engine/framework/modules/transformers/qwen_causal_decoder.h index 9e95de8dd..867665669 100644 --- a/include/engine/framework/modules/transformers/qwen_causal_decoder.h +++ b/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 & scratch, int64_t batch_size, int64_t mask_steps, const std::vector & member_ends, - const std::vector & cache_slots); + const std::vector & cache_slots, + const std::vector * active_mask = nullptr); } // namespace engine::modules diff --git a/include/engine/models/fireredtts3/ar.h b/include/engine/models/fireredtts3/ar.h index cdd45c06d..9e5091bc7 100644 --- a/include/engine/models/fireredtts3/ar.h +++ b/include/engine/models/fireredtts3/ar.h @@ -6,6 +6,7 @@ #include "engine/models/fireredtts3/assets.h" #include +#include #include #include @@ -44,11 +45,14 @@ class FireRedArRuntime { void start_decode_embeddings_batched( const engine::runtime::TransformerBatchedKVState & state, int64_t required_cache_steps); engine::modules::QwenCausalDecodeStepResult decode_embeddings_batched( - const std::vector & embeddings, int64_t batch_size); + const std::vector & embeddings, int64_t batch_size, + const std::vector & 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 batched_member_ends() const; void release_graphs(); void release_backbone_graphs(); diff --git a/include/engine/models/fireredtts3/batch_scheduler.h b/include/engine/models/fireredtts3/batch_scheduler.h index b9700b492..7793b05b8 100644 --- a/include/engine/models/fireredtts3/batch_scheduler.h +++ b/include/engine/models/fireredtts3/batch_scheduler.h @@ -18,19 +18,26 @@ namespace engine::models::fireredtts3 { - - -// 真·batch 并发调度器（llama.cpp update_slots 的模拟）。 -// 一个共享 FireRedArRuntime + FireRedFlowRuntime + FireRedRedAeRuntime， -// 多个 slot（请求）每轮把活跃 slot 的 AR decode 合并成一个 batch -// （decode_embeddings_batched，per-member position）。固定 max_batch graph， -// 运行期不重建；slot 中途 stop 掉队（dead 行零填充），新 slot 通过 -// export -> splice -> import 加入。 +// 真·并发调度器（llama.cpp update_slots 的精神：model/graph 对请求无状态，slot 是 +// 轻量状态机，一个请求 = 一个 slot）。 +// +// 结构（比历史版本大幅简化，删除了 export/splice/rebuild/dirty/freeze 整机）： +// - 共享一个 FireRedArRuntime + FireRedFlowRuntime + FireRedRedAeRuntime， +// 固定 max_batch batched-decode graph（运行期永不重建）。 +// - "round"：一轮并发请求（≤ max_batch）。轮开始时对该轮所有 slot 做一次全量 +// 干净 import（只写本轮的 prefill KV，其余行全零），之后每 tick 一次 batched +// decode 同时推进该轮所有活跃行。轮与轮之间串行 —— 新请求等当前轮 drain。 +// - 行隔离由 decode 层的 active_mask 保证：非活跃行整行 -inf 且不 advance， +// 即使其 cache 段有内容也不参与 attention，绝不残留跨轮内容（每轮 import 全量清零）。 +// +// 不变量： +// - 任何一次 batched decode 运行时，图上只有"当前轮活跃行"带真实内容。 +// - slot release（owner 排空结束）即归还空闲；epoch 单调，stale 句柄立即失效。 class FireRedTTS3BatchScheduler { public: // 一个并发请求 = 一个 slot（持有其 AR + RedAE 增量解码状态 + chunk 队列） struct Slot { - enum class State { Idle, Active, Dead, Done, Failed }; + enum class State { Idle, Active, Dead, Failed }; State state = State::Idle; FireRedTTS3BaseRequest request; std::vector chunks; // 块 patch 数（流式） @@ -42,7 +49,7 @@ class FireRedTTS3BatchScheduler { std::vector chunk_latents; // 参考音色 prep 产物（wave 形成时拷进 slot，不持有 cache 引用跨轮） std::vector prompt_latents; - // prefill 单路径 KV state（CPU，用于拼进 batched KV） + // prefill 单路径 KV state（CPU，拼进 batched KV 用） engine::runtime::TransformerKVState prefill_state; bool prefill_done = false; // RedAE 增量解码状态（per-slot） @@ -51,12 +58,11 @@ class FireRedTTS3BatchScheduler { std::deque chunk_queue; bool finished = false; std::exception_ptr error; - // 代次：launch 时 +1（重置后回写），release_slot 时再 +1。句柄失效/幂等判定用。 + // 代次：launch 时 +1，release_slot 时再 +1。句柄失效/幂等判定用。 uint64_t epoch = 0; }; // 句柄：{slot id, 代次}。持有句柄才能读块/归还/终止。 - // slot 被释放并复用时 epoch +1，旧句柄立即失效（stale 访问视作流结束，不读新请求状态）。 struct SlotHandle { int64_t id = -1; uint64_t epoch = 0; @@ -81,9 +87,8 @@ class FireRedTTS3BatchScheduler { // 启动一个请求（分配 slot），返回句柄。槽池耗尽返回无效句柄（id == -1）。 SlotHandle launch(const FireRedTTS3BaseRequest & request, const std::vector & chunk_patches); // 取下一音频块（驱动调度轮次直到该 slot 产出块或结束）。 - // 句柄失效（stale / 已 release）→ 立即返回空块，不读新请求状态（防串音）。 engine::runtime::AudioBuffer next_chunk(const SlotHandle & handle); - // 归还已完成并排空的 slot 到空闲池（owner 在确认结束时显式调用；幂等）。 + // 归还已完成并排空的 slot 到空闲池（owner 确认结束时显式调用；幂等）。 void release_slot(const SlotHandle & handle); // 快速终止仍 Active 的 slot（供 reset/异常清理），终止后需排空 + release_slot。 void abort(const SlotHandle & handle); diff --git a/src/framework/modules/transformers/qwen_causal_decode_runtime.cpp b/src/framework/modules/transformers/qwen_causal_decode_runtime.cpp index 0ad543290..428205a0e 100644 --- a/src/framework/modules/transformers/qwen_causal_decode_runtime.cpp +++ b/src/framework/modules/transformers/qwen_causal_decode_runtime.cpp @@ -11,6 +11,7 @@ #include #include +#include #include #include #include @@ -204,9 +205,10 @@ void write_batched_cached_step_mask( int64_t batch_size, int64_t mask_steps, const std::vector & member_ends, - const std::vector & cache_slots) { + const std::vector & cache_slots, + const std::vector * active_mask) { if (config.sliding_window <= 0) { - write_qwen_batched_cached_step_mask(tensor, scratch, batch_size, mask_steps, member_ends, cache_slots); + write_qwen_batched_cached_step_mask(tensor, scratch, batch_size, mask_steps, member_ends, cache_slots, active_mask); return; } if (tensor == nullptr) { @@ -218,6 +220,9 @@ void write_batched_cached_step_mask( if (mask_steps <= 0) { throw std::runtime_error("QwenCausalDecodeRuntime sliding batched cached mask requires positive steps"); } + if (active_mask != nullptr && static_cast(active_mask->size()) != batch_size) { + throw std::runtime_error("QwenCausalDecodeRuntime sliding batched cached mask active mask size mismatch"); + } const auto masked = ggml_fp32_to_fp16(-std::numeric_limits::infinity()); const auto visible = ggml_fp32_to_fp16(0.0F); const size_t row_size = static_cast(mask_steps); @@ -226,10 +231,18 @@ void write_batched_cached_step_mask( scratch.resize(total_size); } for (int64_t batch = 0; batch < batch_size; ++batch) { + const size_t offset = static_cast(batch) * row_size; + if (active_mask != nullptr && (*active_mask)[static_cast(batch)] == 0) { + // 非活跃行：整行 -inf（同 write_qwen_batched_cached_step_mask）。 + std::fill( + scratch.begin() + static_cast(offset), + scratch.begin() + static_cast(offset + row_size), + masked); + continue; + } const int64_t end = member_ends.empty() ? cache_slots[static_cast(batch)] : member_ends[static_cast(batch)]; const int64_t current_slot = cache_slots[static_cast(batch)]; const int64_t begin = std::max(0, end - config.sliding_window + 1); - const size_t offset = static_cast(batch) * row_size; std::fill( scratch.begin() + static_cast(offset), scratch.begin() + static_cast(offset + row_size), @@ -638,12 +651,13 @@ class QwenCausalDecodeRuntime::Impl { throw std::runtime_error("QwenCausalDecodeRuntime batched decode token size mismatch"); } ggml_backend_tensor_set(batched_decode_input_, tokens.data(), 0, tokens.size() * sizeof(int32_t)); - return run_batched_decode_step(); + return run_batched_decode_step({}); } QwenCausalDecodeStepResult decode_embeddings_batched( const std::vector & embeddings, - int64_t batch_size) { + int64_t batch_size, + const std::vector & active_mask) { ensure_batched_decode_started(); if (batched_decode_input_kind_ != InputKind::Embedding) { throw std::runtime_error("QwenCausalDecodeRuntime batched decode graph expects embeddings"); @@ -654,8 +668,11 @@ class QwenCausalDecodeRuntime::Impl { if (embeddings.size() != static_cast(batch_size * config_.decoder.stack.hidden_size)) { throw std::runtime_error("QwenCausalDecodeRuntime batched decode embedding size mismatch"); } + if (!active_mask.empty() && static_cast(active_mask.size()) != batch_size) { + throw std::runtime_error("QwenCausalDecodeRuntime batched decode active mask size mismatch"); + } ggml_backend_tensor_set(batched_decode_input_, embeddings.data(), 0, embeddings.size() * sizeof(float)); - return run_batched_decode_step(); + return run_batched_decode_step(active_mask); } int64_t decode_cache_steps() const noexcept { @@ -670,6 +687,13 @@ class QwenCausalDecodeRuntime::Impl { return decode_cache_.valid_steps(); } + std::vector batched_member_ends() const { + if (batched_decode_graph_ == nullptr) { + return {}; + } + return batched_decode_cache_.member_ends_for_mask(); + } + void release_runtime_graphs() { release_prefill_graph(); release_decode_graph(); @@ -1358,10 +1382,11 @@ class QwenCausalDecodeRuntime::Impl { return out; } - QwenCausalDecodeStepResult run_batched_decode_step() { + QwenCausalDecodeStepResult run_batched_decode_step(const std::vector & active_mask) { if (batched_decode_cache_.valid_steps() >= batched_decode_cache_steps_) { throw std::runtime_error("QwenCausalDecodeRuntime batched decode cache exhausted"); } + const bool have_mask = !active_mask.empty(); // per-member positions + cache slots（不同序列可处于不同位置） for (int64_t batch = 0; batch < batched_decode_batch_size_; ++batch) { const int32_t pos = static_cast(batched_decode_cache_.member_end(batch)); @@ -1379,6 +1404,8 @@ class QwenCausalDecodeRuntime::Impl { batched_decode_cache_slots_.data(), 0, batched_decode_cache_slots_.size() * sizeof(int32_t)); + // mask：活跃行暴露 [0, member_end)；非活跃行（active_mask 置 0）整行 -inf， + // 即使其 cache 段残留上一请求的 stale KV 也绝不 attend —— 非活跃行彻底 inert。 write_batched_cached_step_mask( config_, batched_decode_attention_mask_, @@ -1386,7 +1413,8 @@ class QwenCausalDecodeRuntime::Impl { batched_decode_batch_size_, batched_decode_cache_steps_, batched_decode_cache_.member_ends_for_mask(), - batched_decode_cache_slots_); + batched_decode_cache_slots_, + have_mask ? &active_mask : nullptr); core::set_backend_threads(backend_, threads_); const ggml_status status = core::compute_backend_graph(backend_, batched_decode_graph_); ggml_backend_synchronize(backend_); @@ -1411,8 +1439,12 @@ class QwenCausalDecodeRuntime::Impl { out.hidden.size() * sizeof(float)); round_readback(out.hidden, config_); } - // 每个 member 前进 1 步 + // 只有真正活跃的行前进 1 步；非活跃行不 advance（member_end 恒为 0/前值， + // 无 freeze→advance 振荡，位置/cache-slot 稳定）。 for (int64_t batch = 0; batch < batched_decode_batch_size_; ++batch) { + if (have_mask && active_mask[static_cast(batch)] == 0) { + continue; + } batched_decode_cache_.advance_member(batch, 1); } return out; @@ -1679,8 +1711,9 @@ QwenCausalDecodeStepResult QwenCausalDecodeRuntime::decode_tokens_batched(const QwenCausalDecodeStepResult QwenCausalDecodeRuntime::decode_embeddings_batched( const std::vector & embeddings, - int64_t batch_size) { - return impl_->decode_embeddings_batched(embeddings, batch_size); + int64_t batch_size, + const std::vector & active_mask) { + return impl_->decode_embeddings_batched(embeddings, batch_size, active_mask); } runtime::TransformerBatchedKVState QwenCausalDecodeRuntime::export_batched_decode_state() const { @@ -1699,6 +1732,10 @@ int64_t QwenCausalDecodeRuntime::decode_valid_steps() const noexcept { return impl_->decode_valid_steps(); } +std::vector QwenCausalDecodeRuntime::batched_member_ends() const { + return impl_->batched_member_ends(); +} + void QwenCausalDecodeRuntime::release_runtime_graphs() { impl_->release_runtime_graphs(); } diff --git a/src/framework/modules/transformers/qwen_causal_decoder.cpp b/src/framework/modules/transformers/qwen_causal_decoder.cpp index bdc33d3e4..d31ea4278 100644 --- a/src/framework/modules/transformers/qwen_causal_decoder.cpp +++ b/src/framework/modules/transformers/qwen_causal_decoder.cpp @@ -8,6 +8,7 @@ #include #include +#include #include #include #include @@ -504,7 +505,8 @@ void write_qwen_batched_cached_step_mask( int64_t batch_size, int64_t mask_steps, const std::vector & member_ends, - const std::vector & cache_slots) { + const std::vector & cache_slots, + const std::vector * active_mask) { if (tensor == nullptr) { throw std::runtime_error("write_qwen_batched_cached_step_mask requires a tensor"); } @@ -512,6 +514,9 @@ void write_qwen_batched_cached_step_mask( throw std::runtime_error("write_qwen_batched_cached_step_mask requires positive batch size"); } validate_steps(mask_steps, "write_qwen_batched_cached_step_mask"); + if (active_mask != nullptr && static_cast(active_mask->size()) != batch_size) { + throw std::runtime_error("write_qwen_batched_cached_step_mask active mask size mismatch"); + } const auto masked = ggml_fp32_to_fp16(-INFINITY); const auto visible = ggml_fp32_to_fp16(0.0F); const size_t row_size = static_cast(mask_steps); @@ -520,10 +525,19 @@ void write_qwen_batched_cached_step_mask( scratch.resize(total_size); } for (int64_t batch = 0; batch < batch_size; ++batch) { + const size_t offset = static_cast(batch) * row_size; + if (active_mask != nullptr && (*active_mask)[static_cast(batch)] == 0) { + // 非活跃行：整行 -inf。即使该 cache 段残留上一请求的 stale KV，也绝不 + // attend —— 杜绝"冻结行携带旧内容参与 decode"的跨请求串音。 + std::fill( + scratch.begin() + static_cast(offset), + scratch.begin() + static_cast(offset + row_size), + masked); + continue; + } const int64_t end = member_ends.empty() ? cache_slots[static_cast(batch)] : member_ends[static_cast(batch)]; const int64_t visible_prefix_steps = end; const int64_t current_slot = cache_slots[static_cast(batch)]; - const size_t offset = static_cast(batch) * row_size; std::fill( scratch.begin() + static_cast(offset), scratch.begin() + static_cast(offset + row_size), diff --git a/src/models/fireredtts3/ar.cpp b/src/models/fireredtts3/ar.cpp index 24c7b983b..f9263a4d9 100644 --- a/src/models/fireredtts3/ar.cpp +++ b/src/models/fireredtts3/ar.cpp @@ -1020,14 +1020,20 @@ class FireRedArRuntime::Impl { backbone_->start_decode_embeddings_batched(state, required_cache_steps); } - modules::QwenCausalDecodeStepResult decode_embeddings_batched(const std::vector & embeddings, int64_t batch_size) { - return backbone_->decode_embeddings_batched(embeddings, batch_size); + modules::QwenCausalDecodeStepResult decode_embeddings_batched( + const std::vector & embeddings, int64_t batch_size, + const std::vector & active_mask) { + return backbone_->decode_embeddings_batched(embeddings, batch_size, active_mask); } void set_batched_member_end(int64_t batch, int64_t end) { backbone_->set_batched_member_end(batch, end); } + std::vector batched_member_ends() const { + return backbone_->batched_member_ends(); + } + runtime::TransformerBatchedKVState export_batched_decode_state() const { return backbone_->export_batched_decode_state(); } @@ -1151,14 +1157,19 @@ void FireRedArRuntime::start_decode_embeddings_batched( } engine::modules::QwenCausalDecodeStepResult FireRedArRuntime::decode_embeddings_batched( - const std::vector & embeddings, int64_t batch_size) { - return impl_->decode_embeddings_batched(embeddings, batch_size); + const std::vector & embeddings, int64_t batch_size, + const std::vector & active_mask) { + return impl_->decode_embeddings_batched(embeddings, batch_size, active_mask); } void FireRedArRuntime::set_batched_member_end(int64_t batch, int64_t end) { impl_->set_batched_member_end(batch, end); } +std::vector FireRedArRuntime::batched_member_ends() const { + return impl_->batched_member_ends(); +} + engine::runtime::TransformerBatchedKVState FireRedArRuntime::export_batched_decode_state() const { return impl_->export_batched_decode_state(); } diff --git a/src/models/fireredtts3/batch_scheduler.cpp b/src/models/fireredtts3/batch_scheduler.cpp index 922215330..0a158493f 100644 --- a/src/models/fireredtts3/batch_scheduler.cpp +++ b/src/models/fireredtts3/batch_scheduler.cpp @@ -16,7 +16,6 @@ #include #include #include -#include #include #include @@ -28,12 +27,11 @@ namespace modules = engine::modules; using Clock = std::chrono::steady_clock; constexpr int64_t kMaxArSteps = 400; -// prefill graph 固定 steps（grow-only 复用）：所有 wave 的 prefill 都 padding 到该长度， +// prefill graph 固定 steps（grow-only 复用）：所有请求的 prefill 都 padding 到该长度， // 避免因 steps 变化重建 prefill graph 破坏 CUDA pool 逆序 free 约束。 constexpr int64_t kMaxPrefillSteps = 1024; // batched decode KV cache 固定上限（首次 build 后永不重建）： // 最大可能 prefill（参考 + 长文本 ~300） + kMaxArSteps(400) + 尾部余量。 -// 若实际 prefill_steps 超此预算，decode 会 cache 耗尽（罕见，可调大）。 constexpr int64_t kMaxDecodeCacheSteps = 700; // 与 pipeline.cpp 一致的参考音色 key/entry（避免引用失效跨轮，slot 拷贝产物）。 @@ -214,9 +212,6 @@ class FireRedTTS3BatchScheduler::Impl { return {}; } // 纯等待模式：调度线程负责 tick（batch decode），本线程只等自己的 chunk。 - // ggml CUDA pool 逆序约束 -> 唯一 GPU decode 线程必须是调度线程。 - // 用带谓词的 wait（wait 内部重查条件），避免"检查 queue 后、wait 前"的 - // notify 竞态窗口导致死等。 cv_.wait(lock, [&] { return slot.finished || !slot.chunk_queue.empty() || slot.state == Slot::State::Failed || slot.error != nullptr; }); @@ -258,7 +253,6 @@ class FireRedTTS3BatchScheduler::Impl { release_slot(handle); throw; } - // drain 结束（slot 必已 finished）：显式归还，slot 才可被复用。 release_slot(handle); return merged; } @@ -283,8 +277,8 @@ class FireRedTTS3BatchScheduler::Impl { } // owner 在排空结束（next_chunk 见空 / reset / generate drain 完）时显式归还 slot。 - // 只回收 terminal（finished 且非 Active/Dead）的 slot；过早调用 no-op。 - // 幂等：epoch 不匹配（已 release/复用）直接返回。 + // 只回收 terminal（Idle 且 finished）的 slot；过早调用 no-op。幂等。 + // 归还即"slot 回到干净空闲态"，下一次 launch 复用它时行是零（每轮 import 全量清零）。 void release_slot(const SlotHandle & handle) { std::lock_guard lock(mu_); if (!handle.valid() || handle.id >= max_batch_) { @@ -294,11 +288,8 @@ class FireRedTTS3BatchScheduler::Impl { if (slot.epoch != handle.epoch) { return; // 已 release 或复用 → 幂等 } - if (slot.state == Slot::State::Active || slot.state == Slot::State::Dead) { - return; // 过早：调用方应先 abort+drain，不回收活 slot - } - if (!slot.finished) { - return; + if (slot.state != Slot::State::Idle || !slot.finished) { + return; // 过早：调用方应先排空（Dead 在 finish_slot 转 Idle），不回收活 slot } slot.epoch++; // 残留句柄立即失效 if (std::find(free_slots_.begin(), free_slots_.end(), handle.id) == free_slots_.end()) { @@ -307,11 +298,11 @@ class FireRedTTS3BatchScheduler::Impl { fprintf(stderr, "[RELEASED] slot=%ld epoch=%llu\n", handle.id, (unsigned long long)slot.epoch); fflush(stderr); + cv_wake_.notify_all(); // 空闲槽可能让新一轮开跑 cv_.notify_all(); } - // 快速终止仍 Active 的 slot（reset/异常清理用）： - // Active → Dead（若尚未被 form_wave 领走则从 pending 摘除），调度线程下一轮 finish 收尾。 + // 快速终止仍 Active 的 slot（reset/异常清理用）：Active → Dead，调度线程收尾。 void abort(const SlotHandle & handle) { std::lock_guard lock(mu_); if (!handle.valid() || handle.id >= max_batch_) { @@ -325,13 +316,13 @@ class FireRedTTS3BatchScheduler::Impl { return; // 已 Dead/Idle/Failed，无需干预 } if (!slot.prefill_done) { - // 尚未被 form_wave 领走：直接从 pending 摘除，避免 Dead slot 被 prefill。 + // 尚未被领进 round：直接从 pending 摘除，避免 Dead slot 被 prefill。 pending_.erase( std::remove(pending_.begin(), pending_.end(), handle.id), pending_.end()); } slot.state = Slot::State::Dead; - cv_wake_.notify_all(); // 唤醒可能休眠的调度线程收尾 + cv_wake_.notify_all(); } private: @@ -375,7 +366,7 @@ class FireRedTTS3BatchScheduler::Impl { assets_->campplus_weights, execution_.config(), std::move(cfg)); } - // 单个 slot 的 prefill（锁内串行）：参考 prep + 单路径 prefill，KV 存 CPU。 + // 单个 slot 的 prefill（锁内串行）：参考 prep + 单路径 padded prefill，KV 存 CPU。 void prefill_slot_locked(int64_t slot_id) { auto & slot = slots_[static_cast(slot_id)]; const auto & ref = prepare_reference_voice_locked(slot.request.prompt_audio); @@ -414,82 +405,188 @@ class FireRedTTS3BatchScheduler::Impl { redae_->decode_reset(slot.redae_state); } - // 是否存在活跃 decode slot（step>=1，已在 batched KV 中）。 - bool has_active_decode_locked() const { + // 是否存在某 slot 已 prefill（即已占用 decode graph 的一行）。非空 → 当前有 round 在跑。 + bool any_round_participant_locked() const { for (int64_t i = 0; i < max_batch_; ++i) { const auto & slot = slots_[static_cast(i)]; - if (slot.state == Slot::State::Active && slot.prefill_done && slot.step >= 1) { + if (slot.state == Slot::State::Active && slot.prefill_done) { return true; } } return false; } + // 当前 round 的活跃 decode 行（step>=1，真正在 decode）。 + std::vector active_decode_rows_locked() const { + std::vector out; + for (int64_t i = 0; i < max_batch_; ++i) { + const auto & slot = slots_[static_cast(i)]; + if (slot.state == Slot::State::Active && slot.prefill_done && slot.step >= 1) { + out.push_back(i); + } + } + return out; + } + // ---- 调度主循环（锁内）---- - // 与 pipeline.cpp 的单路径 AR 循环逐位对应： - // step 0: 用 prefill_hidden 的 stop/one_backbone 走 AR（不 consume decode） - // step>=1: decode_embeddings_batched 产出 hidden 后走 AR + // round 制： + // begin_round: prefill 所有 pending 请求 → 一次全量干净 import（只写本轮行， + // 其余行零）→ 对本轮每个 slot 走 step-0 AR（AR0，不 consume decode）。 + // 之后每 tick: 一次 batched decode 同时推进本轮所有 step>=1 活跃行。 + // slot 结束（stop/上限）→ Dead → finish_slot 收尾（flush RedAE）→ Idle+finished， + // 等 owner 排空后 release_slot 归还空闲。 + // round 全部 drain 后，新 pending 才开新一轮（全量 import 天然清零旧内容）。 void tick_locked() { - // 1. 处理 pending 启动（prefill 新 slot）。仅当无活跃 decode slot 时才 prefill - // （wave batching）：prefill graph 与 batched decode graph 共用同一 CUDA - // device pool，若在 decode 进行中 prefill（重建 prefill graph），会破坏 - // pool 的逆序 free 约束（GGML_ASSERT(ptr == pool_addr + pool_used)）。 - if (!pending_.empty() && !has_active_decode_locked()) { - form_wave_locked(); + tick_count_++; + // 1. 开新一轮：有待 prefill 的请求，且当前无任何已 prefill 的 round 参与者在跑。 + if (!pending_.empty() && !any_round_participant_locked()) { + begin_round_locked(); + } + // 2. 一次 batched decode 推进当前 round 所有活跃行。 + const auto active = active_decode_rows_locked(); + if (!active.empty()) { + decode_round_step_locked(active); } - // 2. 对 step==0 且已 prefill 的 slot 走 step-0 AR（不进 batch decode， - // 否则零行会被 advance_member 污染 KV 位置）。 + // 3. 收尾 Dead/Failed slot（flush RedAE 尾部、标记 finished、唤醒 owner）。 for (int64_t i = 0; i < max_batch_; ++i) { auto & slot = slots_[static_cast(i)]; - if (slot.state == Slot::State::Active && slot.prefill_done && slot.step == 0) { - advance_slot_from_prefill_locked(i); + if (slot.state == Slot::State::Dead || slot.state == Slot::State::Failed) { + finish_slot_locked(i); } } - // 3. 收集 step>=1 的活跃 slot（参与 batched decode） - std::vector active; + } + + // 开一轮：prefill 全部 pending，全量干净 import，AR0。 + void begin_round_locked() { + std::vector round_slots; + while (!pending_.empty()) { + const int64_t slot_id = pending_.front(); + pending_.pop_front(); + round_slots.push_back(slot_id); + } + for (int64_t slot_id : round_slots) { + try { + prefill_slot_locked(slot_id); + } catch (...) { + auto & slot = slots_[static_cast(slot_id)]; + slot.error = std::current_exception(); + slot.state = Slot::State::Failed; + slot.finished = true; + // 失败 slot 留待 owner next_chunk 抛错后 release_slot 回收。 + } + } + // 成功的 round 成员（Active && prefill_done） + std::vector members; for (int64_t i = 0; i < max_batch_; ++i) { - auto & slot = slots_[static_cast(i)]; - if (slot.state == Slot::State::Active && slot.prefill_done && slot.step >= 1) { - active.push_back(i); + const auto & slot = slots_[static_cast(i)]; + if (slot.state == Slot::State::Active && slot.prefill_done) { + members.push_back(i); } } - if (active.empty()) { - // 没有需要 decode 的 slot；若 step-0 AR 让某些 slot 进入 Dead， - // 本轮结束前清理。 - for (int64_t i = 0; i < max_batch_; ++i) { - auto & slot = slots_[static_cast(i)]; - if (slot.state == Slot::State::Dead || slot.state == Slot::State::Done) { - finish_slot_locked(i); + if (members.empty()) { + return; // 全失败；本 tick 末尾 finish_slot 收尾 + } + if (diag_enabled()) { + std::string m; + for (size_t k = 0; k < members.size(); ++k) { + m += std::to_string(members[k]); + if (k + 1 < members.size()) { + m += ","; } } - return; + fprintf(stderr, "[ROUND] t=%lld members=[%s] prefill_ends=[", + (long long)tick_count_, m.c_str()); + for (int64_t i : members) { + fprintf(stderr, "%lld,", (long long)slots_[static_cast(i)].prefill_steps); + } + fprintf(stderr, "]\n"); + fflush(stderr); } - // 4. 若 batched decode 未启动，或需要 splice 新 slot（step-0 AR 后），重建。 - if (!batched_started_ || !imported_slots_.empty()) { - rebuild_batched_state_locked(); + // 全量干净 import：只写本轮成员行的 prefill KV，其余行全零。 + import_round_state_locked(members); + // AR0：对本轮每个成员走 step-0 AR（用 CPU prefill_hidden，不 consume decode）。 + for (int64_t slot_id : members) { + advance_slot_from_prefill_locked(slot_id); } - // 5. 收集 embeddings；非活跃行零填充 - const int64_t hidden = assets_->base.hidden_size; - std::vector embeddings(static_cast(max_batch_ * hidden), 0.0F); - for (int64_t i = 0; i < max_batch_; ++i) { - auto & slot = slots_[static_cast(i)]; - if (slot.state == Slot::State::Active && slot.prefill_done && slot.step >= 1) { - std::copy(slot.next_input.begin(), slot.next_input.end(), - embeddings.begin() + static_cast(i * hidden)); + } + + // 组装全量 batched state 并 import（写满所有行：成员行=prefill，非成员行=零）。 + void import_round_state_locked(const std::vector & members) { + int64_t max_steps = 0; + for (int64_t i : members) { + max_steps = std::max(max_steps, slots_[static_cast(i)].prefill_steps); + } + const size_t layer_count = slot_prefill_layer_count(); + if (layer_count == 0) { + throw std::runtime_error("FireRedTTS3 round import requires prefill layer state"); + } + const size_t row_elems = static_cast(assets_->base.kv_heads * assets_->base.head_dim); + runtime::TransformerBatchedKVState state; + state.batch_size = max_batch_; + state.current_end = max_steps; + state.current_ends.assign(static_cast(max_batch_), 0); + state.layers.resize(layer_count); + for (size_t layer = 0; layer < layer_count; ++layer) { + auto & out_layer = state.layers[layer]; + out_layer.valid_steps = max_steps; + out_layer.key.assign( + static_cast(max_batch_) * static_cast(max_steps) * row_elems, 0.0F); + out_layer.value.assign( + static_cast(max_batch_) * static_cast(max_steps) * row_elems, 0.0F); + } + for (int64_t b : members) { + auto & slot = slots_[static_cast(b)]; + state.current_ends[static_cast(b)] = slot.prefill_steps; + for (size_t layer = 0; layer < layer_count; ++layer) { + const size_t copy_elems = static_cast(slot.prefill_steps) * row_elems; + float * dst_key = state.layers[layer].key.data() + + static_cast(b) * static_cast(max_steps) * row_elems; + float * dst_value = state.layers[layer].value.data() + + static_cast(b) * static_cast(max_steps) * row_elems; + const auto & src = slot.prefill_state.layers[layer]; + if (src.key.size() != copy_elems || src.value.size() != copy_elems) { + throw std::runtime_error("FireRedTTS3 prefill state size mismatch during round import"); + } + std::copy(src.key.begin(), src.key.end(), dst_key); + std::copy(src.value.begin(), src.value.end(), dst_value); } } - // 5.5 冻结非活跃行的解码位置（end=0，mask 全 -inf），避免空行被 - // advance_member 推高位置后 mask 污染活跃行（导致 stop 判定错乱/内容重复）。 - for (int64_t i = 0; i < max_batch_; ++i) { + // start_decode_embeddings_batched：graph 未建则建（固定 cache），已建则仅 import + // （重写整张 cache tensor —— 本轮成员行写入，非成员行清零）。 + ar_->start_decode_embeddings_batched(state, kMaxDecodeCacheSteps); + } + + // 一次 batched decode 推进本轮所有活跃行（非活跃行 inert）。 + void decode_round_step_locked(const std::vector & active) { + const int64_t hidden = assets_->base.hidden_size; + std::vector embeddings(static_cast(max_batch_ * hidden), 0.0F); + std::vector active_mask(static_cast(max_batch_), 0); + for (int64_t i : active) { auto & slot = slots_[static_cast(i)]; - if (slot.state != Slot::State::Active || !slot.prefill_done || slot.step < 1) { - ar_->set_batched_member_end(i, 0); + std::copy(slot.next_input.begin(), slot.next_input.end(), + embeddings.begin() + static_cast(i * hidden)); + active_mask[static_cast(i)] = 1; + } + auto out = ar_->decode_embeddings_batched(embeddings, max_batch_, active_mask); + if (diag_enabled()) { + for (int64_t i : active) { + const size_t base = static_cast(i) * static_cast(hidden); + bool nan = false; + for (int64_t e = 0; e < hidden; ++e) { + if (std::isnan(out.hidden[base + static_cast(e)])) { + nan = true; + break; + } + } + if (nan) { + fprintf(stderr, "[NAN] t=%lld row=%lld step=%lld\n", + (long long)tick_count_, (long long)i, + (long long)slots_[static_cast(i)].step); + fflush(stderr); + } } } - // 6. 一次 batched decode（每成员独立 pos/cache-slot/mask） - auto out = ar_->decode_embeddings_batched(embeddings, max_batch_); - // 7. 分发 hidden 到各 step>=1 活跃 slot - for (int64_t i = 0; i < max_batch_; ++i) { + for (int64_t i : active) { auto & slot = slots_[static_cast(i)]; if (slot.state != Slot::State::Active || !slot.prefill_done || slot.step < 1) { continue; @@ -499,16 +596,10 @@ class FireRedTTS3BatchScheduler::Impl { out.hidden.begin() + static_cast((i + 1) * hidden)); advance_slot_locked(i, row); } - // 8. 处理完成/死亡的 slot（flush RedAE 尾部，标记 finished，归还槽位） - for (int64_t i = 0; i < max_batch_; ++i) { - auto & slot = slots_[static_cast(i)]; - if (slot.state == Slot::State::Dead || slot.state == Slot::State::Done) { - finish_slot_locked(i); - } - } } // step-0 AR：用 prefill_hidden 的 stop/one_backbone 走 AR，不 consume decode。 + // 生成的第一个 latent 作为 next_input，随后进入 batched decode（step>=1）。 void advance_slot_from_prefill_locked(int64_t slot_id) { auto & slot = slots_[static_cast(slot_id)]; const int64_t hidden = assets_->base.hidden_size; @@ -651,50 +742,57 @@ class FireRedTTS3BatchScheduler::Impl { return current; } - // 完成一个 slot：flush RedAE 尾部、标记 finished、归还槽位。 - // 只处理 Dead（Done 已处理过并归还）；Done 状态由 launch 复用前的 reset 清除。 + // 完成一个 slot：flush RedAE 尾部、标记 finished（Idle），等 owner 排空 + release。 + // 只处理 Dead / Failed；finished 后槽位暂归 owner，绝不在此 push free_slots_。 void finish_slot_locked(int64_t slot_id) { auto & slot = slots_[static_cast(slot_id)]; + if (slot.state != Slot::State::Dead && slot.state != Slot::State::Failed) { + return; + } fprintf(stderr, "[FINISH] slot=%ld tokens=%zu gen_patches=%ld\n", slot_id, slot.request.token_ids.size(), slot.generated_patches); fflush(stderr); - if (slot.state != Slot::State::Dead) { - return; - } - // 有剩余 chunk_latents → decode_incremental 输出 tail（含 istft 尾部）； - // 无剩余 → flush_incremental 输出尾部。二选一，避免 decode + flush 双输出导致重复。 - if (!slot.chunk_latents.empty() && slot.generated_patches > 0) { - auto audio = redae_->decode_incremental(slot.redae_state, slot.chunk_latents); - slot.chunk_latents.clear(); - if (!audio.samples.empty()) { - slot.chunk_queue.push_back(std::move(audio)); - } - } else { - auto flush = redae_->flush_incremental(slot.redae_state); - if (!flush.samples.empty()) { - slot.chunk_queue.push_back(std::move(flush)); + if (slot.state == Slot::State::Dead) { + // 有剩余 chunk_latents → decode_incremental 输出 tail（含 istft 尾部）； + // 无剩余 → flush_incremental 输出尾部。二选一，避免 decode + flush 双输出。 + if (!slot.chunk_latents.empty() && slot.generated_patches > 0) { + auto audio = redae_->decode_incremental(slot.redae_state, slot.chunk_latents); + slot.chunk_latents.clear(); + if (!audio.samples.empty()) { + slot.chunk_queue.push_back(std::move(audio)); + } + } else { + auto flush = redae_->flush_incremental(slot.redae_state); + if (!flush.samples.empty()) { + slot.chunk_queue.push_back(std::move(flush)); + } } } + // Failed slot 的 error 由 owner next_chunk rethrow；此处仅置 finished。 slot.state = Slot::State::Idle; slot.finished = true; - // 不再 push free_slots_：slot 暂归 owner 所有（保留态），等 owner 在 - // next_chunk 排空见空后显式 release_slot 归还，杜绝"未 drain 完的 slot 被复用" - // （旧 owner 串读新请求 chunk / Slot{} 析构与 drain 交错 double-free）。 - fprintf(stderr, "[DONE_WAIT_RELEASE] slot=%ld gen_patches=%ld\n", - slot_id, slot.generated_patches); - fflush(stderr); - // 该行不再参与 decode：从 running 集合移除（下次 join 时重建）。 - running_active_.erase( - std::remove(running_active_.begin(), running_active_.end(), slot_id), - running_active_.end()); - imported_slots_.erase(slot_id); - // 该行 KV 在 GPU 上仍残留 decode 产物；标记脏，下次 rebuild 时清零。 - dirty_rows_.insert(slot_id); + // 该 slot 的 GPU 行内容由下一轮全量 import 清零；期间它不在 active（active_mask + // 置 0 → 整行 -inf 且不 advance），绝不参与 decode、绝不污染其他行。 cv_.notify_all(); } + // 是否有待启动请求或待收尾 slot（决定 scheduler_loop 是否休眠）。 + bool has_work_locked() const { + if (!pending_.empty()) { + return true; + } + for (int64_t i = 0; i < max_batch_; ++i) { + const auto & slot = slots_[static_cast(i)]; + if (slot.state == Slot::State::Active || slot.state == Slot::State::Dead || + slot.state == Slot::State::Failed) { + return true; + } + } + return false; + } + // 专用调度线程主循环（唯一 GPU decode 线程）。 - // 持锁跑 tick（batch decode 推进所有活跃 slot），无活跃 slot 时等待 launch/stop。 + // 持锁跑 tick（一次 batched decode 推进所有活跃行），无工作/无 pending 时等待。 void scheduler_loop() { std::unique_lock lock(mu_); while (!stop_) { @@ -706,7 +804,7 @@ class FireRedTTS3BatchScheduler::Impl { tick_locked(); // 若本 tick 产出了新 chunk，短暂让出锁：让 next_chunk（cv_.wait）能抢到锁 // 取走 chunk，实现真正的增量流式（否则调度线程持锁持续 tick，一次性 - // 生成完整个 slot，流式退化成离线）。 + // 生成完整个 round，流式退化成离线）。 size_t chunks_after = 0; for (int64_t i = 0; i < max_batch_; ++i) { chunks_after += slots_[static_cast(i)].chunk_queue.size(); @@ -714,242 +812,18 @@ class FireRedTTS3BatchScheduler::Impl { if (chunks_after > chunks_before) { cv_wake_.wait_for(lock, std::chrono::milliseconds(1)); } - // 有工作（Active 未完成 / Dead 待 finish）则继续 tick；否则等新工作。 - bool has_work = false; - for (int64_t i = 0; i < max_batch_; ++i) { - auto & slot = slots_[static_cast(i)]; - if (slot.state == Slot::State::Active || - slot.state == Slot::State::Dead) { - has_work = true; - break; - } - } - if (!has_work) { + if (!has_work_locked()) { cv_wake_.wait(lock); } } } - // 形成 wave：prefill pending slot（锁内串行）。prefill 后标记需 splice。 - void form_wave_locked() { - // 一次最多 prefill 直到槽池满（但同一轮只 prefill，不挤占已活跃 slot 的 GPU） - std::vector newly; - while (!pending_.empty() && newly.size() < static_cast(max_batch_)) { - const int64_t slot_id = pending_.front(); - pending_.pop_front(); - newly.push_back(slot_id); - } - for (int64_t slot_id : newly) { - try { - prefill_slot_locked(slot_id); - imported_slots_.insert(slot_id); // 标记需要 splice 进 batched KV - } catch (...) { - auto & slot = slots_[static_cast(slot_id)]; - slot.error = std::current_exception(); - slot.state = Slot::State::Failed; - slot.finished = true; - // 不 push free_slots_：失败 slot 留待 owner next_chunk 抛错 → reset/异常路径 - // 显式 release_slot 回收，避免失败槽在 owner 处理错误前被复用。 - } - } - } - - // 重建 batched KV state：export 当前 GPU state + splice 新 prefill slot 的 CPU KV + 重新 start。 - // 固定 max_batch graph，不重建（batch_size 不变、cache_steps 不减）。 - void rebuild_batched_state_locked() { - // 需要 splice 的新 slot（尚未在 batched 中） - std::vector joining; - for (int64_t i = 0; i < max_batch_; ++i) { - auto & slot = slots_[static_cast(i)]; - if (slot.prefill_done && imported_slots_.count(i) > 0) { - joining.push_back(i); - } - } - if (batched_started_ && !joining.empty()) { - // 从 GPU 导出当前 batched state（active slot 的中间 KV 保留） - runtime::TransformerBatchedKVState gpu = ar_->export_batched_decode_state(); - // 层数/row_elems 校验（与 import_state 的期望布局一致） - if (gpu.layers.size() != slot_prefill_layer_count()) { - throw std::runtime_error("FireRedTTS3BatchScheduler batched state layer mismatch"); - } - // 组装全量 batched state：max_steps = 所有"有数据行"（imported/running）的最大 end。 - // 空行（无 prefill、非 running）end=0，不参与 max_steps（避免 gpu 残留的空行 - // 高 end 撑大 state，导致与 current_ends 的 import max 不一致）。 - runtime::TransformerBatchedKVState state; - state.batch_size = max_batch_; - state.current_ends.assign(static_cast(max_batch_), 0); - for (int64_t i = 0; i < max_batch_; ++i) { - if (imported_slots_.count(i) > 0) { - state.current_ends[static_cast(i)] = slots_[static_cast(i)].prefill_steps; - } else if (std::find(running_active_.begin(), running_active_.end(), i) != running_active_.end()) { - state.current_ends[static_cast(i)] = gpu_end_for(gpu, i); - } - } - static int rebuild_dbg = 0; - if (rebuild_dbg++ < 10) { - std::string ends; - for (int64_t i = 0; i < max_batch_; ++i) { - ends += std::to_string(state.current_ends[static_cast(i)]) + ","; - } - std::string gpu_ends; - if (!gpu.current_ends.empty()) { - for (size_t i = 0; i < gpu.current_ends.size(); ++i) { - gpu_ends += std::to_string(gpu.current_ends[i]) + ","; - } - } - fprintf(stderr, "[REBUILD] batch=%ld state_ends=[%s] gpu_ends=[%s] running=[", - (long)max_batch_, ends.c_str(), gpu_ends.c_str()); - for (size_t i = 0; i < running_active_.size(); ++i) { - fprintf(stderr, "%ld,", (long)running_active_[i]); - } - fprintf(stderr, "]\n"); - fflush(stderr); - } - int64_t max_steps = 0; - for (int64_t i = 0; i < max_batch_; ++i) { - max_steps = std::max(max_steps, state.current_ends[static_cast(i)]); - } - state.current_end = max_steps; - state.layers.resize(gpu.layers.size()); - for (size_t layer = 0; layer < state.layers.size(); ++layer) { - auto & out_layer = state.layers[layer]; - out_layer.valid_steps = max_steps; - const size_t row_elems = row_elems_for_layer(layer); - out_layer.key.assign(static_cast(max_batch_) * static_cast(max_steps) * row_elems, 0.0F); - out_layer.value.assign(static_cast(max_batch_) * static_cast(max_steps) * row_elems, 0.0F); - } - // 填充每行：已运行 slot（running）从 gpu 拷贝其 ends；新 slot 从 prefill_state 拷贝； - // 其余（dead/脏/空行）保持全零。 - for (int64_t b = 0; b < max_batch_; ++b) { - if (imported_slots_.count(b) == 0 && std::find(running_active_.begin(), running_active_.end(), b) == running_active_.end()) { - continue; // 非活跃行：全零（不拷 gpu 残留） - } - const int64_t end = (imported_slots_.count(b) > 0) - ? slots_[static_cast(b)].prefill_steps - : gpu_end_for(gpu, b); - if (end <= 0) { - continue; - } - static int splice_dbg = 0; - if (splice_dbg++ < 20) { - fprintf(stderr, "[SPLICE] row=%ld src=%s tokens=%zu end=%ld\n", b, - imported_slots_.count(b) > 0 ? "new-prefill" : "gpu-running", - slots_[static_cast(b)].request.token_ids.size(), end); - fflush(stderr); - } - const int64_t current_end = end; - for (size_t layer = 0; layer < state.layers.size(); ++layer) { - const size_t row_elems = row_elems_for_layer(layer); - const size_t copy_elems = static_cast(current_end) * row_elems; - float * dst_key = state.layers[layer].key.data() - + static_cast(b) * static_cast(max_steps) * row_elems; - float * dst_value = state.layers[layer].value.data() - + static_cast(b) * static_cast(max_steps) * row_elems; - if (imported_slots_.count(b) > 0) { - const auto & src = slots_[static_cast(b)].prefill_state.layers[layer]; - std::copy(src.key.begin(), src.key.end(), dst_key); - std::copy(src.value.begin(), src.value.end(), dst_value); - } else { - const auto & src = gpu.layers[layer]; - const size_t src_row_elems = row_elems; - const int64_t gpu_valid_steps = gpu.layers.empty() ? 0 : gpu.layers.front().valid_steps; - const size_t src_offset = static_cast(b) * static_cast(gpu_valid_steps) * src_row_elems; - std::copy( - src.key.begin() + static_cast(src_offset), - src.key.begin() + static_cast(src_offset + copy_elems), - dst_key); - std::copy( - src.value.begin() + static_cast(src_offset), - src.value.begin() + static_cast(src_offset + copy_elems), - dst_value); - } - } - } - const int64_t required_steps = compute_required_steps(max_steps); - ar_->start_decode_embeddings_batched(state, required_steps); - for (int64_t i : joining) { - imported_slots_.erase(i); - } - } else if (!joining.empty()) { - // 首次启动（无 GPU 导出，全部新 slot） - int64_t max_steps = 0; - for (int64_t i = 0; i < max_batch_; ++i) { - auto & slot = slots_[static_cast(i)]; - if (slot.prefill_done) { - max_steps = std::max(max_steps, slot.prefill_steps); - } - } - runtime::TransformerBatchedKVState state; - state.batch_size = max_batch_; - state.current_end = max_steps; - state.current_ends.assign(static_cast(max_batch_), 0); - state.layers.resize(slot_prefill_layer_count()); - for (size_t layer = 0; layer < state.layers.size(); ++layer) { - auto & out_layer = state.layers[layer]; - out_layer.valid_steps = max_steps; - const size_t row_elems = row_elems_for_layer(layer); - out_layer.key.assign(static_cast(max_batch_) * static_cast(max_steps) * row_elems, 0.0F); - out_layer.value.assign(static_cast(max_batch_) * static_cast(max_steps) * row_elems, 0.0F); - } - for (int64_t b = 0; b < max_batch_; ++b) { - auto & slot = slots_[static_cast(b)]; - if (!slot.prefill_done) { - continue; - } - state.current_ends[static_cast(b)] = slot.prefill_steps; - for (size_t layer = 0; layer < state.layers.size(); ++layer) { - const size_t row_elems = row_elems_for_layer(layer); - const size_t copy_elems = static_cast(slot.prefill_steps) * row_elems; - float * dst_key = state.layers[layer].key.data() - + static_cast(b) * static_cast(max_steps) * row_elems; - float * dst_value = state.layers[layer].value.data() - + static_cast(b) * static_cast(max_steps) * row_elems; - const auto & src = slot.prefill_state.layers[layer]; - std::copy(src.key.begin(), src.key.end(), dst_key); - std::copy(src.value.begin(), src.value.end(), dst_value); - } - } - const int64_t required_steps = compute_required_steps(max_steps); - ar_->start_decode_embeddings_batched(state, required_steps); - // 全量分支：所有 prefill slot 已 splice 进 batched KV，清空 imported 标记， - // 否则每次 tick 都会重复 rebuild（用 prefill_steps 覆盖 GPU 增量 KV）。 - imported_slots_.clear(); - } - batched_started_ = true; - dirty_rows_.clear(); - running_active_.clear(); - for (int64_t i = 0; i < max_batch_; ++i) { - auto & slot = slots_[static_cast(i)]; - if (slot.state == Slot::State::Active && slot.prefill_done) { - running_active_.push_back(i); - } - } - } - - // 计算 batched decode graph 所需 cache steps。为避免 graph 重建（昂贵）， - // 首次 build 用保守上限：max_steps + 完整 AR 预算。后续 join 若 max_steps - // 增长超预算才重建（罕见），通常复用首次 graph。 - int64_t compute_required_steps(int64_t /*max_prefill_steps*/) const { - // 固定 cache 上限：首次 build 后永不重建（避免 export 空 state / pool 顺序破坏）。 - // 预算 = 最大可能 prefill（参考+文本） + kMaxArSteps + 尾部。 - // KV 显存 = max_batch × cache_steps × layers × row_elems ≈ 可控。 - return kMaxDecodeCacheSteps; - } - - // GPU batched state 中某行的当前 end（运行中 slot 用它作为拷贝长度） - int64_t gpu_end_for(const runtime::TransformerBatchedKVState & gpu, int64_t batch) const { - if (!gpu.current_ends.empty()) { - if (static_cast(gpu.current_ends.size()) == gpu.batch_size) { - return gpu.current_ends[static_cast(batch)]; - } - } - return gpu.current_end; + // [DIAG] 轮/步日志是否开启（仅看是否有待处理请求，避免纯空闲刷屏）。 + bool diag_enabled() const { + return tick_count_ < 1000000; } size_t slot_prefill_layer_count() const { - if (max_batch_ == 0) { - return 0; - } for (int64_t i = 0; i < max_batch_; ++i) { auto & slot = slots_[static_cast(i)]; if (slot.prefill_done && !slot.prefill_state.layers.empty()) { @@ -959,11 +833,6 @@ class FireRedTTS3BatchScheduler::Impl { return 0; } - size_t row_elems_for_layer(size_t /*layer*/) const { - // FireRed TTS3 AR: kv_heads × head_dim（同 pipeline prefill state 的层元素数）。 - return static_cast(assets_->base.kv_heads * assets_->base.head_dim); - } - std::shared_ptr assets_; engine::core::ExecutionContext & execution_; bool mem_saver_ = false; @@ -978,19 +847,13 @@ class FireRedTTS3BatchScheduler::Impl { std::mutex mu_; std::condition_variable cv_; // chunk 产出 / slot 结束通知（next_chunk 等待） - std::condition_variable cv_wake_; // 调度线程唤醒（launch / stop） + std::condition_variable cv_wake_; // 调度线程唤醒（launch / release / stop） std::vector slots_; std::deque free_slots_; std::deque pending_; - bool batched_started_ = false; bool stop_ = false; std::thread scheduler_thread_; - // 已 spliced 进 batched KV 的 slot（运行中）；新 prefill 的 slot 在此集合外，需 join。 - std::vector running_active_; - // 已 prefill 但尚未 splice 进 batched KV 的 slot id - std::set imported_slots_; - // 刚结束、GPU 行残留 KV 的 slot id（下次 rebuild 时该行清零） - std::set dirty_rows_; + int64_t tick_count_ = 0; // [DIAG] }; FireRedTTS3BatchScheduler::FireRedTTS3BatchScheduler( From 5eced4ee1d447c8d385a5d77c7522bf2a9e0c409 Mon Sep 17 00:00:00 2001 From: dujing Date: Fri, 4 Sep 2026 10:47:16 +0800 Subject: [PATCH 11/12] =?UTF-8?q?fix(server):=20=E5=B9=B6=E5=8F=91?= =?UTF-8?q?=E8=AF=B7=E6=B1=82=E5=85=A8=E9=83=A8=E7=BB=91=20session=200=20?= =?UTF-8?q?=E7=9A=84=E4=B8=B2=E9=9F=B3/double-free=20=E6=A0=B9=E5=9B=A0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 根因是 SessionPoolLock 的双 index 字段 bug： - 头里 public `size_t index = 0` 与 private `size_t index_ = 0` 并存； - 构造函数只初始化 private index_，而所有调用处（sessions[pool.index]、 [REQ] 打印）都读 public index → 恒为 0； - 于是每个并发请求都借到 session 0 并绑 scheduler slot 0：新请求的 start_stream→reset() 会 abort 掉上一个还在跑的 slot → 截断在 chunk 边界、 内容互串、间歇 double-free。 修法：删掉重复的 private index_，让 public index 成为唯一字段 （构造/move/赋值/release 全用它）。 另：ensure_model_loaded_locked 是经典双重检查锁，但拿到 load_mutex_ 后 缺少二次 loaded 复查 → 并发首请求会各自完整加载一遍并重建 session 池。 补上锁内复查，确保只有首个请求真正加载。 验证：conc3 三路并发（不同文本/音色、同 seed）的输出与各自顺序 solo 逐位一致（MD5 全同）；单路不回归（audio=7.04s / WER 0.1429）。 scheduler 本身无需改动 —— 串音不在它，在 server 的 session 池借还。 Co-Authored-By: Claude --- app/server/runtime.cpp | 15 +++++++++++---- app/server/runtime.h | 5 ++++- 2 files changed, 15 insertions(+), 5 deletions(-) diff --git a/app/server/runtime.cpp b/app/server/runtime.cpp index 2c6d82adc..4997415f7 100644 --- a/app/server/runtime.cpp +++ b/app/server/runtime.cpp @@ -1766,6 +1766,13 @@ void ServerState::ensure_model_loaded_locked(LoadedModel & model) { // free_sessions 出现重复索引 —— 并发请求全部借到同一 session 下标，共享一个 // session 的 scheduler slot 与 chunk 队列，造成跨请求串音/截断。 std::unique_lock load_lock(model_load_mutex_); + // 二次检查（经典双重检查锁）：拿到锁后必须再看一眼 loaded。否则并发首请求 + // 都在锁外看到 loaded=false，排队拿锁后各自又完整加载一遍 —— 每次都 clear()+ + // 重建 session 池，多个请求各自借到新池的 index 0（全绑同一 session），造成 + // 跨请求串音 + double free。二次检查让只有第一个请求真正加载，其余直接返回。 + if (model.loaded.load(std::memory_order_acquire)) { + return; + } if (config_.max_loaded_models > 0) { evict_for_model_limit(model); // Note: if ensure_model_fits_memory() below then refuses the load, the @@ -2039,16 +2046,16 @@ struct ServerState::TimedTaskResult { }; ServerState::SessionPoolLock::SessionPoolLock(ServerState::LoadedModel & model, size_t index) - : model_(&model), index_(index) {} + : model_(&model), index(index) {} ServerState::SessionPoolLock::SessionPoolLock(SessionPoolLock && other) noexcept - : model_(other.model_), index_(other.index_) { other.model_ = nullptr; } + : model_(other.model_), index(other.index) { other.model_ = nullptr; } ServerState::SessionPoolLock & ServerState::SessionPoolLock::operator=(SessionPoolLock && other) noexcept { if (this != &other) { release(); model_ = other.model_; - index_ = other.index_; + index = other.index; other.model_ = nullptr; } return *this; @@ -2059,7 +2066,7 @@ ServerState::SessionPoolLock::~SessionPoolLock() { release(); } void ServerState::SessionPoolLock::release() { if (model_ != nullptr) { std::lock_guard lock(model_->pool_mutex); - model_->free_sessions.push_back(index_); + model_->free_sessions.push_back(index); model_ = nullptr; } } diff --git a/app/server/runtime.h b/app/server/runtime.h index 4e66cbfdc..c2d94db0b 100644 --- a/app/server/runtime.h +++ b/app/server/runtime.h @@ -61,12 +61,15 @@ class ServerState final : public IHttpHandler { 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; - size_t index_ = 0; }; struct LoadedModel { From bff5f39246880c76fb0762d838fabaffa0c09b1f Mon Sep 17 00:00:00 2001 From: dujing Date: Fri, 4 Sep 2026 13:15:13 +0800 Subject: [PATCH 12/12] =?UTF-8?q?perf(firered):=20flow=20=E8=B7=A8=20slot?= =?UTF-8?q?=20batch=20+=20round=20=E5=BC=80=E8=BD=AE=E6=94=B6=E9=9B=86?= =?UTF-8?q?=E7=AA=97=E5=8F=A3?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 动机：profile 显示单路生成 flow(DiT) 占 56% 耗时，而多路并发只 batch 了 AR(23%)， flow 每路每 patch 串行 → 聚合吞吐上不去。 改动（batch_scheduler.cpp）： - decode_round_step 改两阶段：阶段A per-slot AR 后处理(stop/backbone/dit_cond3)； 阶段B 存活 slot 同步 flow denoise —— 同一 denoise 步把各 slot 的 x_in 拼成一次 flow.run(batch=2N)，按行拆回各自 CFG 合并；阶段C per-slot patch_encode/chunk/redae。 DiT 输出每行独立(flash batch 轴)，与各 slot 单独跑逐位一致。异构 steps 回退逐 slot。 - scheduler_loop 开轮收集窗口：pending 非空且无 round participant 时 wait 8ms 聚合同刻 请求再开轮，消除 round 拆轮（首请求抢先单独 prefill → flow batch 永不触发）。 - [PROF] 计时：AR decode / flow / redae 分项累计，finish 时打印。 验证（105 实测）： - 2-slot 同轮(members=[0,2])每路输出与各自 solo MD5 逐位一致（T1 1a419f52/T2 9d7d871d）。 - 吞吐 2 文本顺序 2.54 → 2-slot 同轮 3.11（+22%）；6 路因 HTTP 到达散布仍部分拆轮， aggregate RTF 升到 ~2.6。 - flow 图按需重建(batch 变 rebuild)；试过 fixed-capacity+零填充复用但破坏 DiT 输出，弃。 config：firered_server.json 默认 num_inference_steps=5（default_request_options）。 Co-Authored-By: Claude --- app/server/firered_server.json | 5 +- src/models/fireredtts3/batch_scheduler.cpp | 312 ++++++++++++++++++--- 2 files changed, 272 insertions(+), 45 deletions(-) diff --git a/app/server/firered_server.json b/app/server/firered_server.json index aa570993c..f004f86dd 100644 --- a/app/server/firered_server.json +++ b/app/server/firered_server.json @@ -18,7 +18,10 @@ "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" } } ] -} +} \ No newline at end of file diff --git a/src/models/fireredtts3/batch_scheduler.cpp b/src/models/fireredtts3/batch_scheduler.cpp index 0a158493f..dba08c5f3 100644 --- a/src/models/fireredtts3/batch_scheduler.cpp +++ b/src/models/fireredtts3/batch_scheduler.cpp @@ -135,6 +135,8 @@ class FireRedTTS3BatchScheduler::Impl { weight_context_bytes, storage_type, false)), redae_(std::make_unique( assets_, execution_, graph_arena_bytes, weight_context_bytes, storage_type)), + // flow 图按需重建（batch 变则 rebuild）。曾试过 fixed-capacity + 零填充复用， + // 但 padding 破坏 DiT 输出（单路 MD5 就变），故保持每轮实际 batch 建图。 flow_(std::make_unique( assets_, execution_, graph_arena_bytes, weight_context_bytes, storage_type, false)), reference_voice_cache_( @@ -556,7 +558,18 @@ class FireRedTTS3BatchScheduler::Impl { ar_->start_decode_embeddings_batched(state, kMaxDecodeCacheSteps); } - // 一次 batched decode 推进本轮所有活跃行（非活跃行 inert）。 + // 一次 batched decode 推进本轮所有活跃行，并把 flow（DiT denoise）也跨 slot batch。 + // 两阶段： + // A. per-slot AR 后处理（stop 判定 / backbone_cond / dit_cond3）——host 计算，逐 slot。 + // B. 对存活 slot 同步做 flow denoise：所有 slot 的同一 denoise 步拼一次大 batch + // flow.run（行 = 各 slot 的 cfg_batch 行首尾相接），输出按 slot 拆回各自 CFG 合并。 + // DiT 每行独立 attend，故与各 slot 单独跑逐位一致。异构 steps 则逐个单跑（回退）。 + // C. per-slot：patch_encode -> next_input，累加 latents/chunk，chunk 边界 redae。 + struct FlowSurvivor { + int64_t id; + int64_t cfg_batch; // guidance_scale>0 ? 2 : 1 + std::vector dit_cond3; // [history_patches+1, dit_hidden] + }; void decode_round_step_locked(const std::vector & active) { const int64_t hidden = assets_->base.hidden_size; std::vector embeddings(static_cast(max_batch_ * hidden), 0.0F); @@ -567,7 +580,10 @@ class FireRedTTS3BatchScheduler::Impl { embeddings.begin() + static_cast(i * hidden)); active_mask[static_cast(i)] = 1; } + auto t_dec0 = std::chrono::steady_clock::now(); auto out = ar_->decode_embeddings_batched(embeddings, max_batch_, active_mask); + prof_.ar_decode_ms += std::chrono::duration( + std::chrono::steady_clock::now() - t_dec0).count(); if (diag_enabled()) { for (int64_t i : active) { const size_t base = static_cast(i) * static_cast(hidden); @@ -586,6 +602,9 @@ class FireRedTTS3BatchScheduler::Impl { } } } + // ---- 阶段 A：per-slot AR 后处理 + stop 判定，收集存活 slot ---- + std::vector survivors; + survivors.reserve(active.size()); for (int64_t i : active) { auto & slot = slots_[static_cast(i)]; if (slot.state != Slot::State::Active || !slot.prefill_done || slot.step < 1) { @@ -594,7 +613,172 @@ class FireRedTTS3BatchScheduler::Impl { std::vector row( out.hidden.begin() + static_cast(i * hidden), out.hidden.begin() + static_cast((i + 1) * hidden)); - advance_slot_locked(i, row); + if (slot.step >= kMaxArSteps) { + slot.state = Slot::State::Dead; + continue; + } + const float stop = ar_->stop(row); + if (stop >= slot.request.stop_threshold && slot.step >= 6) { + slot.state = Slot::State::Dead; + continue; + } + slot.backbone_cond.insert(slot.backbone_cond.end(), row.begin(), row.end()); + const int64_t cond_rows = static_cast(slot.backbone_cond.size()) / hidden; + auto cond3 = last_rows(slot.backbone_cond, cond_rows, hidden, assets_->base.history_patches + 1); + FlowSurvivor m; + m.id = i; + m.cfg_batch = slot.request.guidance_scale > 0.0F ? 2 : 1; + m.dit_cond3 = ar_->dit_head(cond3, assets_->base.history_patches + 1); + survivors.push_back(std::move(m)); + } + if (survivors.empty()) { + return; + } + // ---- 阶段 B/C：flow denoise（batch 或逐 slot）+ per-slot 收尾 ---- + const size_t s0_steps = slots_[static_cast(survivors[0].id)].schedule.size(); + const bool same_steps = std::all_of(survivors.begin(), survivors.end(), [&](const FlowSurvivor & m) { + return slots_[static_cast(m.id)].schedule.size() == s0_steps; + }); + if (same_steps && survivors.size() > 1) { + auto latents = batched_flow_latents_locked(survivors); + for (size_t s = 0; s < survivors.size(); ++s) { + finish_slot_patch_locked(survivors[s].id, std::move(latents[s])); + } + } else { + for (const FlowSurvivor & m : survivors) { + auto & slot = slots_[static_cast(m.id)]; + auto latent = flow_one_patch_locked(slot, m.dit_cond3, static_cast(slot.step)); + finish_slot_patch_locked(m.id, std::move(latent)); + } + } + } + + // 多 slot 同步 flow denoise，返回每 slot 的 next_latent（[patch, redae_dim]）。 + // pred 布局：flow 输出每行 = [patch, redae_dim]（graph 已 slice 掉 history token）， + // 行 stride = patch*redae_dim；行按 slot 的 cfg_batch 首尾相接。 + std::vector> batched_flow_latents_locked(const std::vector & survivors) { + const int64_t history_tokens = assets_->base.history_patches * assets_->base.patch_size; + const int64_t tokens = history_tokens + assets_->base.patch_size; + const int64_t in_channels = assets_->base.redae_dim + assets_->base.dit_hidden_size + assets_->base.speaker_dim; + const int64_t patch = assets_->base.patch_size; + const int64_t redae_dim = assets_->base.redae_dim; + const int64_t n = static_cast(survivors.size()); + std::vector> history(n); + std::vector> current(n); + std::vector base_row(n); + int64_t total_batch = 0; + for (int64_t s = 0; s < n; ++s) { + auto & slot = slots_[static_cast(survivors[s].id)]; + history[s] = last_rows( + slot.latents_gen, + static_cast(slot.latents_gen.size()) / redae_dim, + redae_dim, + history_tokens); + const uint64_t noise_offset = + sampling::torch_cuda_tensor_iterator_offset_blocks( + static_cast(patch * redae_dim), + sampling_policy_) * + static_cast(slot.step); + current[s] = sampling::generate_torch_cuda_tensor_iterator_randn( + static_cast(patch * redae_dim), + slot.request.seed, + noise_offset, + sampling_policy_, + sampling::TorchRandnPrecision::Float32); + base_row[s] = total_batch; + total_batch += survivors[s].cfg_batch; + } + const auto & schedule = slots_[static_cast(survivors[0].id)].schedule; + for (size_t i = 0; i + 1 < schedule.size(); ++i) { + std::vector x_in(static_cast(total_batch * tokens * in_channels), 0.0F); + std::vector time_all(static_cast(total_batch * 256), 0.0F); + const auto te = firered_timestep_embedding(schedule[i]); + for (int64_t s = 0; s < n; ++s) { + auto & slot = slots_[static_cast(survivors[s].id)]; + const int64_t b0 = base_row[s]; + const int64_t cfg = survivors[s].cfg_batch; + for (int64_t b = 0; b < cfg; ++b) { + const int64_t g = b0 + b; + std::copy(te.begin(), te.end(), + time_all.begin() + static_cast(g * 256)); + for (int64_t t = 0; t < tokens; ++t) { + float * row = x_in.data() + static_cast((g * tokens + t) * in_channels); + if (t < history_tokens) { + std::copy( + history[s].begin() + static_cast(t * redae_dim), + history[s].begin() + static_cast((t + 1) * redae_dim), + row); + } else { + const int64_t local = t - history_tokens; + std::copy( + current[s].begin() + static_cast(local * redae_dim), + current[s].begin() + static_cast((local + 1) * redae_dim), + row); + } + if (b == 0) { + const int64_t cond_row = t / patch; + std::copy( + survivors[s].dit_cond3.begin() + static_cast(cond_row * assets_->base.dit_hidden_size), + survivors[s].dit_cond3.begin() + static_cast((cond_row + 1) * assets_->base.dit_hidden_size), + row + redae_dim); + std::copy(slot.spk_dit.begin(), slot.spk_dit.end(), + row + redae_dim + assets_->base.dit_hidden_size); + } + } + } + } + auto t_fl0 = std::chrono::steady_clock::now(); + auto pred = flow_->run(x_in, time_all, total_batch); + prof_.flow_ms += std::chrono::duration( + std::chrono::steady_clock::now() - t_fl0).count(); + const float dt = schedule[i + 1] - schedule[i]; + const int64_t row_elems = patch * redae_dim; + for (int64_t s = 0; s < n; ++s) { + auto & slot = slots_[static_cast(survivors[s].id)]; + const int64_t b0 = base_row[s]; + const int64_t cfg = survivors[s].cfg_batch; + for (int64_t t = 0; t < patch; ++t) { + for (int64_t c = 0; c < redae_dim; ++c) { + const size_t idx = static_cast(t * redae_dim + c); + const size_t cond_off = static_cast(b0 * row_elems + idx); + float vt = pred[cond_off]; + if (cfg == 2) { + const float uncond = pred[static_cast((b0 + 1) * row_elems + idx)]; + vt = (1.0F + slot.request.guidance_scale) * vt - slot.request.guidance_scale * uncond; + } + current[s][idx] += dt * vt; + } + } + } + } + return current; + } + + // 阶段 C：一个 slot 的 AR patch 收尾（flow 产出 next_latent 后）。 + void finish_slot_patch_locked(int64_t slot_id, std::vector next_latent) { + auto & slot = slots_[static_cast(slot_id)]; + slot.latents_gen.insert(slot.latents_gen.end(), next_latent.begin(), next_latent.end()); + slot.next_input = ar_->patch_encode(next_latent); + slot.chunk_latents.insert(slot.chunk_latents.end(), next_latent.begin(), next_latent.end()); + slot.generated_patches++; + slot.step++; + prof_.n_steps++; + prof_.n_patches++; + if (slot.generated_patches >= slot.chunk_target && slot.chunk_index < slot.chunks.size()) { + auto t_red0 = std::chrono::steady_clock::now(); + auto audio = redae_->decode_incremental(slot.redae_state, slot.chunk_latents); + prof_.redae_ms += std::chrono::duration( + std::chrono::steady_clock::now() - t_red0).count(); + prof_.n_redae++; + slot.chunk_latents.clear(); + slot.chunk_index++; + if (slot.chunk_index < slot.chunks.size()) { + slot.chunk_target += slot.chunks[slot.chunk_index]; + } + if (!audio.samples.empty()) { + slot.chunk_queue.push_back(std::move(audio)); + cv_.notify_all(); + } } } @@ -627,48 +811,6 @@ class FireRedTTS3BatchScheduler::Impl { slot.step = 1; } - // 单 slot 一步 AR（step>=1：hidden 来自一次 batched decode）。 - void advance_slot_locked(int64_t slot_id, const std::vector & hidden_row) { - auto & slot = slots_[static_cast(slot_id)]; - const int64_t hidden = assets_->base.hidden_size; - - // AR 步数上限（与 pipeline.cpp 的 max_steps=400 一致）：stop 未触发时强制结束。 - if (slot.step >= kMaxArSteps) { - slot.state = Slot::State::Dead; - return; - } - const float stop = ar_->stop(hidden_row); - if (stop >= slot.request.stop_threshold && slot.step >= 6) { - slot.state = Slot::State::Dead; - return; - } - - slot.backbone_cond.insert(slot.backbone_cond.end(), hidden_row.begin(), hidden_row.end()); - const int64_t cond_rows = static_cast(slot.backbone_cond.size()) / hidden; - auto cond3 = last_rows(slot.backbone_cond, cond_rows, hidden, assets_->base.history_patches + 1); - auto dit_cond3 = ar_->dit_head(cond3, assets_->base.history_patches + 1); - const auto next_latent = flow_one_patch_locked(slot, dit_cond3, static_cast(slot.step)); - slot.latents_gen.insert(slot.latents_gen.end(), next_latent.begin(), next_latent.end()); - slot.next_input = ar_->patch_encode(next_latent); - slot.chunk_latents.insert(slot.chunk_latents.end(), next_latent.begin(), next_latent.end()); - slot.generated_patches++; - slot.step++; - - // chunk 边界（与 pipeline.cpp 相同） - if (slot.generated_patches >= slot.chunk_target && slot.chunk_index < slot.chunks.size()) { - auto audio = redae_->decode_incremental(slot.redae_state, slot.chunk_latents); - slot.chunk_latents.clear(); - slot.chunk_index++; - if (slot.chunk_index < slot.chunks.size()) { - slot.chunk_target += slot.chunks[slot.chunk_index]; - } - if (!audio.samples.empty()) { - slot.chunk_queue.push_back(std::move(audio)); - cv_.notify_all(); // 唤醒等待该 slot 的 next_chunk - } - } - } - // flow + CFG 单 patch（与 pipeline.cpp 的 flow_one_patch 一致，但使用 slot 状态）。 std::vector flow_one_patch_locked(Slot & slot, const std::vector & dit_cond3, uint64_t step_index) { const int64_t history_tokens = assets_->base.history_patches * assets_->base.patch_size; @@ -756,13 +898,21 @@ class FireRedTTS3BatchScheduler::Impl { // 有剩余 chunk_latents → decode_incremental 输出 tail（含 istft 尾部）； // 无剩余 → flush_incremental 输出尾部。二选一，避免 decode + flush 双输出。 if (!slot.chunk_latents.empty() && slot.generated_patches > 0) { + auto t_red0 = std::chrono::steady_clock::now(); auto audio = redae_->decode_incremental(slot.redae_state, slot.chunk_latents); + prof_.redae_ms += std::chrono::duration( + std::chrono::steady_clock::now() - t_red0).count(); + prof_.n_redae++; slot.chunk_latents.clear(); if (!audio.samples.empty()) { slot.chunk_queue.push_back(std::move(audio)); } } else { + auto t_red0 = std::chrono::steady_clock::now(); auto flush = redae_->flush_incremental(slot.redae_state); + prof_.redae_ms += std::chrono::duration( + std::chrono::steady_clock::now() - t_red0).count(); + prof_.n_redae++; if (!flush.samples.empty()) { slot.chunk_queue.push_back(std::move(flush)); } @@ -773,6 +923,25 @@ class FireRedTTS3BatchScheduler::Impl { slot.finished = true; // 该 slot 的 GPU 行内容由下一轮全量 import 清零；期间它不在 active（active_mask // 置 0 → 整行 -inf 且不 advance），绝不参与 decode、绝不污染其他行。 + // [PROF] 打印自上次 finish 以来的分项累计（单路时即该请求自身） + const double d_ar = prof_.ar_decode_ms - prof_last_.ar_decode_ms; + const double d_flow = prof_.flow_ms - prof_last_.flow_ms; + const double d_red = prof_.redae_ms - prof_last_.redae_ms; + fprintf(stderr, + "[PROF] slot=%ld steps=%lld patches=%lld redae=%lld | AR_decode=%.1fms flow=%.1fms " + "redae=%.1fms | ar%%=%.0f flow%%=%.0f redae%%=%.0f | per_patch_ar=%.2fms per_patch_flow=%.2fms\n", + slot_id, + (long long)(prof_.n_steps - prof_last_.n_steps), + (long long)(prof_.n_patches - prof_last_.n_patches), + (long long)(prof_.n_redae - prof_last_.n_redae), + d_ar, d_flow, d_red, + d_ar + d_flow + d_red > 0 ? 100.0 * d_ar / (d_ar + d_flow + d_red) : 0, + d_ar + d_flow + d_red > 0 ? 100.0 * d_flow / (d_ar + d_flow + d_red) : 0, + d_ar + d_flow + d_red > 0 ? 100.0 * d_red / (d_ar + d_flow + d_red) : 0, + prof_.n_patches - prof_last_.n_patches > 0 ? d_ar / (prof_.n_patches - prof_last_.n_patches) : 0, + prof_.n_patches - prof_last_.n_patches > 0 ? d_flow / (prof_.n_patches - prof_last_.n_patches) : 0); + fflush(stderr); + prof_last_ = prof_; cv_.notify_all(); } @@ -791,11 +960,40 @@ class FireRedTTS3BatchScheduler::Impl { return false; } + // 开轮收集窗口：pending 有待启动请求、且当前无任何 round 参与者在跑时，说明 + // "下一轮即将开始"。此时若立刻 tick，第一个请求会单独 prefill 成一轮（round + // 拆轮 —— 同一瞬间到达的其余请求被 gate 挡到整轮跑完才并下一轮，flow batch + // 因而失效）。解法：短暂等待，让同刻到达的请求在 pending 里聚集，再一次性开轮。 + // 已攒满（pending 能占满空闲 slot）则不等直接开。 + bool should_collect_for_round_locked() const { + if (pending_.empty()) { + return false; + } + if (any_round_participant_locked()) { + return false; // 已有轮在跑，无需收集（新 pending 等本轮 drain） + } + // pending 已占满 max_batch → 一轮已能打满，不必再等。 + return static_cast(pending_.size()) < max_batch_; + } + // 专用调度线程主循环（唯一 GPU decode 线程）。 // 持锁跑 tick（一次 batched decode 推进所有活跃行），无工作/无 pending 时等待。 void scheduler_loop() { std::unique_lock lock(mu_); + constexpr auto kRoundCollectWindow = std::chrono::microseconds(8000); while (!stop_) { + // 开轮前收集：若下一轮即将开始且还没攒满，等一小段让同刻请求加入同一轮， + // 避免 round 拆轮。collect_started_ 防连续多个 1.5ms 空等（攒满/超时即开）。 + if (should_collect_for_round_locked() && !collect_started_) { + collect_started_ = true; + fprintf(stderr, "[COLLECT] t=%lld pend=%zu free=%zu waiting %dus\n", + (long long)tick_count_, pending_.size(), free_slots_.size(), + (int)std::chrono::duration_cast(kRoundCollectWindow).count()); + fflush(stderr); + cv_wake_.wait_for(lock, kRoundCollectWindow); + collect_started_ = false; + // wait_for 醒来即继续到 tick —— pending 已尽量聚集，本轮一并 prefill + } // 记录 tick 前各 slot 的已产出 chunk 数 size_t chunks_before = 0; for (int64_t i = 0; i < max_batch_; ++i) { @@ -854,6 +1052,32 @@ class FireRedTTS3BatchScheduler::Impl { bool stop_ = false; std::thread scheduler_thread_; int64_t tick_count_ = 0; // [DIAG] + bool collect_started_ = false; // 开轮收集窗口进行中（防连续空等） + + // [PROF] 各组件累计耗时（单调度线程，串行累计；finish_slot 时打印分项） + struct ProfAccum { + double ar_decode_ms = 0; // batched AR decode（图前 + 图后 readback） + double flow_ms = 0; // flow_one_patch_locked（整函数，含 dit_head? 不含，dit_head 单独在 advance） + double redae_ms = 0; // RedAE decode_incremental / flush + double other_ms = 0; // 其余（advance 的 stop/dit_head/patch_encode/backbone 等） + int64_t n_steps = 0; // AR step 数 + int64_t n_patches = 0; // generated patches + int64_t n_redae = 0; // redae 调用次数 + } prof_; + ProfAccum prof_last_; // [PROF] 上次打印时的累计快照 + // 一段代码的计时 RAII：析构时把 elapsed_ms 累加到 acc + struct ProfScoped { + ProfAccum & acc; + double & slot; + std::chrono::steady_clock::time_point t0; + ProfScoped(ProfAccum & a, double & target) : acc(a), slot(target) { + (void)acc; + t0 = std::chrono::steady_clock::now(); + } + ~ProfScoped() { + slot += std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); + } + }; }; FireRedTTS3BatchScheduler::FireRedTTS3BatchScheduler(