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 new file mode 100644 index 000000000..f004f86dd --- /dev/null +++ b/app/server/firered_server.json @@ -0,0 +1,27 @@ +{ + "host": "0.0.0.0", + "port": 8007, + "backend": "cuda", + "device": 0, + "threads": 4, + "lazy_load": true, + "busy_timeout_ms": 0, + "models": [ + { + "id": "firered-base", + "family": "fireredtts3", + "path": "/data/megastore/Projects/DuJing/models/FireRedTTS3-Base-GGUF/fireredtts3-base-q8_0.gguf", + "model_spec_override": "/data/megastore/Projects/DuJing/code/audio.cpp/model_specs/fireredtts3.json", + "task": "clon", + "mode": "streaming", + "instance_count": 3, + "session_options": { + "fireredtts3.reference_cache_slots": "8", + "fireredtts3.chunk_sizes": "3,12,12,12,12,12,12,12,12,12,12,12,12,12,12,12" + }, + "default_request_options": { + "num_inference_steps": "5" + } + } + ] +} \ No newline at end of file diff --git a/app/server/runtime.cpp b/app/server/runtime.cpp index 74d1e3807..4997415f7 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,18 +1756,22 @@ 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.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_); + // 二次检查(经典双重检查锁):拿到锁后必须再看一眼 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); @@ -1785,6 +1797,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 +1825,32 @@ 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)); } + 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.session = std::move(session); model.offline = offline; model.streaming = streaming; model.loaded.store(true); @@ -1881,6 +1915,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(); } @@ -2010,6 +2045,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 +2124,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 +2154,23 @@ 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)); + 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) { if (!timed_result.ttft_ms.has_value() && stream_event_has_output(event)) { @@ -2076,8 +2181,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 +2989,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 +3025,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 +3125,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 +3156,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..c2d94db0b 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,30 @@ class ServerState final : public IHttpHandler { LiveIngestLimits live_ingest_limits(const HttpRequest & request) const override; private: + struct LoadedModel; + + // session 池借用锁(RAII):借一个空闲 session 实例,析构归还。 + class SessionPoolLock { + public: + SessionPoolLock() = default; + SessionPoolLock(LoadedModel & model, size_t index); + SessionPoolLock(SessionPoolLock && other) noexcept; + SessionPoolLock & operator=(SessionPoolLock && other) noexcept; + SessionPoolLock(const SessionPoolLock &) = delete; + SessionPoolLock & operator=(const SessionPoolLock &) = delete; + ~SessionPoolLock(); + + // 借到的 session 下标。public 且是唯一存储:构造函数/move/赋值都写这里, + // release() 也读这里归还。曾有个 private index_ 与 public index 并存, + // 构造函数只写 index_ 而调用处全读 public index → 恒为 0 → 所有并发请求 + // 都绑 session 0(跨请求串音/截断/double free 根因)。 + size_t index = 0; + + private: + void release(); + LoadedModel * model_ = nullptr; + }; + struct LoadedModel { struct RuntimeVoicePreset { std::optional voice_id; @@ -54,10 +82,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 +117,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/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/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..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 @@ -76,6 +79,24 @@ class RedAeCodecRuntime { std::vector encode(const std::vector & audio_24k); 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(DecodeState & state); + // 解码一批 latent(chunk),返回该块对应的音频(float32 24k mono)。 + // 内部用 decoder Qwen 的 KV 缓存跨块保持上下文,并用增量 iSTFT 逐块输出。 + runtime::AudioBuffer decode_incremental(DecodeState & state, const std::vector & latents); + // flush 增量 iSTFT 尾部样本(生成结束时调用)。 + runtime::AudioBuffer flush_incremental(DecodeState & state); + void release_runtime_graphs(); private: 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..d18cf757b 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, @@ -90,9 +98,16 @@ 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); // Snapshot of the batched decode KV cache (host vectors), suitable for // replication and re-import via start_decode_*_batched with a different @@ -104,6 +119,9 @@ class QwenCausalDecodeRuntime { int64_t decode_valid_steps() const noexcept; void release_runtime_graphs(); + // [DIAG] 每 batch 行的当前解码结束位置(member_ends_ 拷贝;未启动则空)。 + std::vector 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 4231fcaeb..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, - int64_t visible_prefix_steps, - int64_t current_slot); + const std::vector & member_ends, + const std::vector & cache_slots, + const std::vector * active_mask = nullptr); } // 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/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/include/engine/models/fireredtts3/ar.h b/include/engine/models/fireredtts3/ar.h index 4da49135a..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 @@ -35,9 +36,24 @@ 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); + // --- 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, + 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 new file mode 100644 index 000000000..7793b05b8 --- /dev/null +++ b/include/engine/models/fireredtts3/batch_scheduler.h @@ -0,0 +1,107 @@ +#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 { + +// 真·并发调度器(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, 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; + // 代次:launch 时 +1,release_slot 时再 +1。句柄失效/幂等判定用。 + uint64_t epoch = 0; + }; + + // 句柄:{slot id, 代次}。持有句柄才能读块/归还/终止。 + struct SlotHandle { + int64_t id = -1; + uint64_t epoch = 0; + bool valid() const noexcept { return id >= 0; } + }; + + 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),返回句柄。槽池耗尽返回无效句柄(id == -1)。 + SlotHandle launch(const FireRedTTS3BaseRequest & request, const std::vector & chunk_patches); + // 取下一音频块(驱动调度轮次直到该 slot 产出块或结束)。 + 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(); + + // 槽池容量(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/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..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" @@ -31,6 +32,13 @@ class FireRedRedAeRuntime { std::vector encode(const std::vector & audio_24k); engine::runtime::AudioBuffer decode(const std::vector & latents); + + // 增量解码(流式):重置 + 逐块解码 + 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(); private: diff --git a/include/engine/models/fireredtts3/session.h b/include/engine/models/fireredtts3/session.h index 95cb9d9c8..59836ac66 100644 --- a/include/engine/models/fireredtts3/session.h +++ b/include/engine/models/fireredtts3/session.h @@ -4,20 +4,24 @@ #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 { 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 +36,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 +56,23 @@ 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 句柄(无效 = 无)。 + FireRedTTS3BatchScheduler::SlotHandle scheduler_slot_; + bool scheduler_enabled_ = 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_; + // 流式方法互斥(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 f34f27686..d24ec36e5 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,21 @@ "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" + }, + { + "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/audio/istft_graph.cpp b/src/framework/audio/istft_graph.cpp index 700854291..4bcc22e63 100644 --- a/src/framework/audio/istft_graph.cpp +++ b/src/framework/audio/istft_graph.cpp @@ -211,6 +211,159 @@ 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 {}; + } + // 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(); + 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 +438,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..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; @@ -643,6 +644,76 @@ class RedAeRuntime { return audio; } + // --- 增量解码(流式)--- 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(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"); + } + 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 (!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); + st.dec_state = std::move(decoder.state); + } else { + // 后续块:用 decode_embedding 逐 qwen frame 解码,复用前序 KV + 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( + 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()); + } + } + st.dec_qwen_frames += qwen_frames; + + auto spec = encoder_out_.istft_head(decoder_hidden, qwen_frames); + + // 增量 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)); + st.inc_istft = std::make_unique(cfg); + st.inc_istft_frames = qwen_frames; + } + auto chunk_audio = st.inc_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(RedAeCodecRuntime::DecodeState & st) { + runtime::AudioBuffer audio; + audio.sample_rate = static_cast(config_.sample_rate); + audio.channels = 1; + if (st.inc_istft) { + audio.samples = st.inc_istft->finish_incremental(); + } + return audio; + } + void release_graphs() { encoder_in_.release_graph(); encoder_out_.release_graph(); @@ -766,6 +837,18 @@ runtime::AudioBuffer RedAeCodecRuntime::decode(const std::vector & latent return impl_->runtime.decode(latents); } +void RedAeCodecRuntime::decode_reset(RedAeCodecRuntime::DecodeState & state) { + impl_->runtime.decode_reset(state); +} + +runtime::AudioBuffer RedAeCodecRuntime::decode_incremental(RedAeCodecRuntime::DecodeState & state, const std::vector & latents) { + return impl_->runtime.decode_incremental(state, latents); +} + +runtime::AudioBuffer RedAeCodecRuntime::flush_incremental(RedAeCodecRuntime::DecodeState & state) { + return impl_->runtime.flush_incremental(state); +} + void RedAeCodecRuntime::release_runtime_graphs() { impl_->runtime.release_graphs(); } 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 8d0e65ad9..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 @@ -203,11 +204,11 @@ 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, + const std::vector * active_mask) { 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, active_mask); return; } if (tensor == nullptr) { @@ -219,11 +220,8 @@ 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"); + 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); @@ -232,14 +230,24 @@ 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 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); 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; @@ -482,6 +490,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, @@ -572,6 +628,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) { @@ -591,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"); @@ -607,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 { @@ -623,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(); @@ -648,7 +719,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; @@ -800,13 +874,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)); @@ -814,6 +889,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; } @@ -1178,11 +1260,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( @@ -1299,31 +1382,39 @@ 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 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()); + 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)); + 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(), 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_, 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_, + 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_); @@ -1348,7 +1439,14 @@ class QwenCausalDecodeRuntime::Impl { out.hidden.size() * sizeof(float)); round_readback(out.hidden, config_); } - batched_decode_cache_.advance_after_direct_append(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; } @@ -1476,6 +1574,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_; @@ -1523,6 +1623,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; @@ -1547,6 +1648,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, @@ -1593,14 +1701,19 @@ 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); } 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 { @@ -1619,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 94ca0050c..d31ea4278 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" @@ -7,6 +8,7 @@ #include #include +#include #include #include #include @@ -502,8 +504,9 @@ 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, + const std::vector * active_mask) { if (tensor == nullptr) { throw std::runtime_error("write_qwen_batched_cached_step_mask requires a tensor"); } @@ -511,11 +514,8 @@ 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"); + 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); @@ -526,6 +526,18 @@ void write_qwen_batched_cached_step_mask( } 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)]; std::fill( scratch.begin() + static_cast(offset), scratch.begin() + static_cast(offset + row_size), @@ -533,7 +545,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 a50b7191c..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" @@ -861,8 +862,41 @@ 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]) { + // 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]; + 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, 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, 0, 1, 3, 2), + core::TensorShape::from_dims({batch, 1, x_heads, 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..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" @@ -260,6 +261,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 +276,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 +320,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 +362,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 +385,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, 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/ar.cpp b/src/models/fireredtts3/ar.cpp index f340c59ea..f9263a4d9 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); } @@ -1011,6 +1016,28 @@ 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, + 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(); + } + void release_graphs() { if (token_embedding_) { token_embedding_->release_graph(); @@ -1106,6 +1133,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) { @@ -1116,6 +1150,30 @@ 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, + 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(); +} + void FireRedArRuntime::release_graphs() { impl_->release_graphs(); } diff --git a/src/models/fireredtts3/batch_scheduler.cpp b/src/models/fireredtts3/batch_scheduler.cpp new file mode 100644 index 000000000..dba08c5f3 --- /dev/null +++ b/src/models/fireredtts3/batch_scheduler.cpp @@ -0,0 +1,1129 @@ +#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 + +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 复用):所有请求的 prefill 都 padding 到该长度, +// 避免因 steps 变化重建 prefill graph 破坏 CUDA pool 逆序 free 约束。 +constexpr int64_t kMaxPrefillSteps = 1024; +// batched decode KV cache 固定上限(首次 build 后永不重建): +// 最大可能 prefill(参考 + 长文本 ~300) + kMaxArSteps(400) + 尾部余量。 +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 图按需重建(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_( + 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(); + } + } + + SlotHandle launch(const FireRedTTS3BaseRequest & request, const std::vector & chunk_patches) { + std::lock_guard lock(mu_); + if (free_slots_.empty()) { + 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 epoch=%llu tokens=%zu\n", + slot_id, (unsigned long long)epoch, request.token_ids.size()); + fflush(stderr); + cv_wake_.notify_all(); + return SlotHandle{slot_id, epoch}; + } + + engine::runtime::AudioBuffer next_chunk(const SlotHandle & handle) { + std::unique_lock lock(mu_); + // 首行所有权校验(在读 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); + } + 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。 + 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 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; + 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 (...) { + } + release_slot(handle); + throw; + } + release_slot(handle); + 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(); + } + + // owner 在排空结束(next_chunk 见空 / reset / generate drain 完)时显式归还 slot。 + // 只回收 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_) { + return; + } + auto & slot = slots_[static_cast(handle.id)]; + if (slot.epoch != handle.epoch) { + return; // 已 release 或复用 → 幂等 + } + 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()) { + free_slots_.push_back(handle.id); + } + 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,调度线程收尾。 + 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) { + // 尚未被领进 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(); + } + +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 + 单路径 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); + 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); + } + + // 是否存在某 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) { + 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; + } + + // ---- 调度主循环(锁内)---- + // 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() { + 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); + } + // 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::Dead || slot.state == Slot::State::Failed) { + finish_slot_locked(i); + } + } + } + + // 开一轮: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) { + const auto & slot = slots_[static_cast(i)]; + if (slot.state == Slot::State::Active && slot.prefill_done) { + members.push_back(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 += ","; + } + } + 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); + } + // 全量干净 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); + } + } + + // 组装全量 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); + } + } + // start_decode_embeddings_batched:graph 未建则建(固定 cache),已建则仅 import + // (重写整张 cache tensor —— 本轮成员行写入,非成员行清零)。 + ar_->start_decode_embeddings_batched(state, kMaxDecodeCacheSteps); + } + + // 一次 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); + std::vector active_mask(static_cast(max_batch_), 0); + for (int64_t i : active) { + auto & slot = slots_[static_cast(i)]; + std::copy(slot.next_input.begin(), slot.next_input.end(), + 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); + 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); + } + } + } + // ---- 阶段 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) { + continue; + } + std::vector row( + out.hidden.begin() + static_cast(i * hidden), + out.hidden.begin() + static_cast((i + 1) * hidden)); + 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(); + } + } + } + + // 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; + 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; + } + + // 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(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) { + // 有剩余 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)); + } + } + } + // Failed slot 的 error 由 owner next_chunk rethrow;此处仅置 finished。 + slot.state = Slot::State::Idle; + 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(); + } + + // 是否有待启动请求或待收尾 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; + } + + // 开轮收集窗口: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) { + chunks_before += slots_[static_cast(i)].chunk_queue.size(); + } + tick_locked(); + // 若本 tick 产出了新 chunk,短暂让出锁:让 next_chunk(cv_.wait)能抢到锁 + // 取走 chunk,实现真正的增量流式(否则调度线程持锁持续 tick,一次性 + // 生成完整个 round,流式退化成离线)。 + 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)); + } + if (!has_work_locked()) { + cv_wake_.wait(lock); + } + } + } + + // [DIAG] 轮/步日志是否开启(仅看是否有待处理请求,避免纯空闲刷屏)。 + bool diag_enabled() const { + return tick_count_ < 1000000; + } + + size_t slot_prefill_layer_count() const { + 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; + } + + 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 / release / stop) + std::vector slots_; + std::deque free_slots_; + std::deque pending_; + 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( + 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; + +FireRedTTS3BatchScheduler::SlotHandle FireRedTTS3BatchScheduler::launch( + const FireRedTTS3BaseRequest & request, + const std::vector & chunk_patches) { + return impl_->launch(request, chunk_patches); +} + +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) { + 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/pipeline.cpp b/src/models/fireredtts3/pipeline.cpp index d914c4e73..9b7161a09 100644 --- a/src/models/fireredtts3/pipeline.cpp +++ b/src/models/fireredtts3/pipeline.cpp @@ -256,6 +256,147 @@ 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; + FireRedRedAeRuntime::DecodeState redae_state; + + 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(redae_state, 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(redae_state, chunk_latents); + chunk_latents.clear(); + finished = true; + return std::move(tail); + } + auto flush = owner.redae_->flush_incremental(redae_state); + 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(redae_state); + started = true; + } + }; + auto session = std::make_unique(*this, request, chunk_patches); + return session; + } + void release_graphs() { if (redae_) { redae_->release_graphs(); @@ -855,6 +996,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..63e15afcd 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(FireRedRedAeRuntime::DecodeState & st) { + runtime_.decode_reset(st); + } + + runtime::AudioBuffer decode_incremental(FireRedRedAeRuntime::DecodeState & st, const std::vector & latents) { + return runtime_.decode_incremental(st, latents); + } + + runtime::AudioBuffer flush_incremental(FireRedRedAeRuntime::DecodeState & st) { + return runtime_.flush_incremental(st); + } + void release_graphs() { runtime_.release_runtime_graphs(); } @@ -157,6 +169,18 @@ engine::runtime::AudioBuffer FireRedRedAeRuntime::decode(const std::vectordecode(latents); } +void FireRedRedAeRuntime::decode_reset(DecodeState & st) { + impl_->decode_reset(st); +} + +engine::runtime::AudioBuffer FireRedRedAeRuntime::decode_incremental(DecodeState & st, const std::vector & latents) { + return impl_->decode_incremental(st, latents); +} + +engine::runtime::AudioBuffer FireRedRedAeRuntime::flush_incremental(DecodeState & st) { + return impl_->flush_incremental(st); +} + void FireRedRedAeRuntime::release_graphs() { impl_->release_graphs(); } diff --git a/src/models/fireredtts3/session.cpp b/src/models/fireredtts3/session.cpp index 79313eb3c..328111b35 100644 --- a/src/models/fireredtts3/session.cpp +++ b/src/models/fireredtts3/session.cpp @@ -8,8 +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 { @@ -21,6 +26,81 @@ constexpr const char * kInstructName = "FireRedTTS3 Instruct"; constexpr int64_t kDefaultTextChunkSize = 600; constexpr size_t kDefaultReferenceCacheSlots = 4; +// ---- 共享 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> 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, + std::shared_ptr context, + 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 = scheduler_by_assets.find(assets.get()); + if (it != scheduler_by_assets.end()) { + if (auto existing = it->second.lock()) { + return existing; + } + scheduler_by_assets.erase(it); + } + const FireRedTTS3Assets * key = assets.get(); + auto scheduler = std::make_shared( + std::move(assets), *context, graph_arena_bytes, helper_graph_arena_bytes, + weight_context_bytes, storage_type, reference_cache_slots, mem_saver, max_batch); + scheduler_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; } @@ -44,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; @@ -216,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); @@ -233,9 +335,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, @@ -273,6 +373,26 @@ 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 用共享 context 构造(与基类 borrow 的是同一个)。 + scheduler_ = scheduler_registry().get_or_create( + assets_, + scheduler_registry().get_or_create_context(assets_, options.backend), + 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_, @@ -337,15 +457,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; @@ -353,8 +480,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); @@ -365,6 +494,151 @@ 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) { + 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"); + } + reset(); + initialize_stream_request(request); + auto base = make_base_request(*tokenizer_, request); + if (scheduler_enabled_ && scheduler_) { + // 共享 scheduler:launch 一个 slot(并发请求共享 GPU batch decode) + scheduler_slot_ = scheduler_->launch(base, stream_chunk_patches_); + if (!scheduler_slot_.valid()) { + 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"); + } + runtime::AudioBuffer audio; + if (scheduler_enabled_ && scheduler_) { + if (!scheduler_slot_.valid()) { + return std::nullopt; + } + audio = scheduler_->next_chunk(scheduler_slot_); + if (audio.samples.empty()) { + // 流结束:显式归还 slot(此后该 id 才可能被复用),并失效本句柄, + // 杜绝"本 session 尚未完全 drain / 仍在持句柄时 slot 被其他请求复用"。 + scheduler_->release_slot(scheduler_slot_); + scheduler_slot_ = FireRedTTS3BatchScheduler::SlotHandle{}; + 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; + 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(); + 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_->release_slot(scheduler_slot_); + scheduler_slot_ = FireRedTTS3BatchScheduler::SlotHandle{}; + } + 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;