feat(cosyvoice3): run the flow DiT estimator on the Rockchip NPU

CMakeLists.txt

@@ -2473,6 +2473,7 @@ audiocpp_add_model(cosyvoice3
         src/models/cosyvoice3/flow.cpp
         src/models/cosyvoice3/frontend.cpp
         src/models/cosyvoice3/hift.cpp
+        src/models/cosyvoice3/rknn_flow.cpp
         src/models/cosyvoice3/session.cpp
         src/models/cosyvoice3/tokenizer_text.cpp
     INCLUDES

docs/build/rockchip-npu.md

@@ -366,6 +366,7 @@ the configurable components:
 | Qwen3 `speech_decoder` | `qwen3_tts.exec_plan` | `ggml`, `rknn` | RKNN: required `main`, optional `first` |
 | sanoTTS `decoder` | `sanotts.exec_plan` | `default`, `ggml`, `neon` | None |
 | Inflect v2 `decoder` | `inflect_v2.exec_plan` | `ggml`, `rknn` | RKNN: required `main` |
+| CosyVoice3 `flow` | `cosyvoice3.exec_plan` | `ggml`, `rknn` | RKNN: required `main` (comma-separated ladder of per-length engines) |
 
 These options also work with older GGUFs whose embedded specs do not declare
 `exec_plan`; the model validates the plan without requiring a spec override.

docs/models/cosyvoice3.md

@@ -89,6 +89,72 @@ audiocpp_cli \
 |---|---|---:|---|
 | `cosyvoice3.reference_cache_slots` | integer >= 0 | `4` | Prepared reference-audio cache slots. |
 | `cosyvoice3.mem_saver` | bool | `false` | Release cached runtime graphs after request phases. |
+| `cosyvoice3.exec_plan` | JSON object | `{}` | Component execution plans. The only component is `flow`; see the RKNN section. |
+
+## Rockchip NPU (RKNN) estimator
+
+The flow decoder is the expensive half of this family: the DiT estimator runs once per Euler step
+(`num_inference_steps`, default 10) over the whole utterance, and on an RK3588 it is ~79% of
+synthesis. `cosyvoice3.exec_plan` voices that one stage through `librknnrt` and leaves the AR LLM
+and the HiFT vocoder on the host, which is the split the reference RKNN port of this model
+measured best: a quantized NPU LLM degrades the speech tokens, and a quantized HiFT adds
+high-frequency artifacts.
+
+RKNN has no runtime graph builder, so each engine is compiled offline for one fixed sequence
+length. Build them on an x86_64 host from the same checkpoint the GGUF was converted from:
+
+```bash
+python3 tools/rknn/prepare_cosyvoice3_rknn_flow.py \
+    --estimator-onnx flow.decoder.estimator.fp32.onnx \
+    --seq-len 200 --seq-len 500 --seq-len 900 --seq-len 1100 \
+    --output-dir models/CosyVoice3-RKNN
+```
+
+Then point the session at the ladder (order does not matter; the smallest engine that fits is
+used, and an utterance longer than the largest engine runs the CPU graph). The ladder travels as
+the `flow` component's `main` artifact, comma-separated — see
+[component execution plans](../build/rockchip-npu.md#component-execution-plans):
+
+```bash
+audiocpp_cli --task clon --family cosyvoice3 --model models/CosyVoice3-GGUF/cosyvoice3-q8_0.gguf \
+    --backend cpu --session-option 'cosyvoice3.exec_plan={"flow":{"impl":"rknn","artifacts":{"main":"models/CosyVoice3-RKNN/flow_estimator_seq200.rknn,models/CosyVoice3-RKNN/flow_estimator_seq900.rknn"}}}' \
+    --text "..." --voice-ref assets/resources/b.wav --reference-text "..." --out out.wav
+```
+
+Build with `-DENGINE_ENABLE_RKNN=ON` (plus `ENGINE_RKNN_LIBRARY` / `ENGINE_RKNN_INCLUDE_DIR`
+pointing at the board's `librknnrt`); a plan selecting `impl: rknn` throws at load in a build
+without it rather than falling back silently. Each engine is ~700 MB and an utterance is padded up
+to the engine's
+length, so the ladder is a time trade: the estimator costs about `S^2`, and the padding is exact
+(measured bit-identical to an unpadded run), not a quality one.
+
+### What it measures on an RK3588
+
+Measured on an Orange Pi 5 Plus (RK3588, 32 GB, `librknnrt` 2.3.2, `--backend cpu --threads 4`,
+10 Euler steps, same text/seed/reference on both paths), `dit_ms` for the whole ODE loop:
+
+| mel frames | engine (pad) | ggml CPU | RKNN |
+|---|---|---|---|
+| 102 | seq200 (98) | 8.48 s | 11.49 s |
+| 212 | seq500 (288) | 19.14 s | 57.10 s |
+| 102 | seq110 (8) | 8.48 s | 6.14 s |
+| 212 | seq250 (38) | 19.29 s | 18.51 s |
+| 402 | seq400 (2) | 38.87 s | 38.84 s |
+
+The estimator on the NPU is at parity with four A76 cores running the ggml graph, so the padding is
+what decides: a coarse ladder loses by 3x, a dense one breaks even above 200 frames and wins 1.38x
+only at the shortest lengths. Peak RSS is 3.8-6.8 GB against 2.3-2.6 GB.
+
+Two ceilings bound the ladder. The flow's frame count includes the prompt mel, so a 4-second
+reference puts every utterance above 700 frames. And the driver's IOVA arena is 32-bit on this
+kernel: one engine asks 0.7 GB at seq110 and 4.71 GB at seq900, and above about 4 GB `rknn_init`
+fails with `-4` (`RKNPU: failed to allocate IOVA: -12`), which is why the ladder is loaded lazily,
+one engine resident. Engines above ~700 frames are therefore unusable on this board, and long
+utterances — where the estimator is 79% of synthesis — run on the ggml graph.
+
+The fp16 cost of the graph, measured on the PC simulator against the fp32 ONNX (seq200, 106 real
+frames): snr 44.0 dB, tail snr 60.5 dB, `max_abs 0.5544`, mean `-3.83689` against `-3.83806`. There
+is no systematic mel bias, so no offset knob is needed.
 
 Legacy `cosyvoice3.conv_weight_type` remains accepted but has no consumer and
 does not change convolution storage. The implementation uses the shared

include/engine/models/cosyvoice3/flow.h

@@ -31,7 +31,8 @@ class CosyVoice3DiTFlowRuntime {
         engine::core::ExecutionContext & execution,
         size_t graph_arena_bytes,
         size_t weight_context_bytes,
-        engine::assets::TensorStorageType storage_type);
+        engine::assets::TensorStorageType storage_type,
+        const std::vector<std::string> & rknn_engine_paths = {});
     ~CosyVoice3DiTFlowRuntime();
 
     CosyVoice3DiTFlowRuntime(const CosyVoice3DiTFlowRuntime &) = delete;

include/engine/models/cosyvoice3/rknn_flow.h

@@ -0,0 +1,63 @@
+#pragma once
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <vector>
+
+namespace engine::models::cosyvoice3 {
+
+// Rockchip NPU (RKNN) backend for the flow-matching DiT estimator: the ODE step that the ggml
+// DiT graph in flow.cpp computes. The estimator is the flow decoder's velocity field, and it is
+// the stage the reference RKNN port keeps on the NPU (the AR LLM and the HiFT vocoder stay on
+// CPU there, for the same reasons this repository keeps its sidecars on the decoder half).
+//
+// RKNN has no dynamic shapes, so each engine is compiled offline for one fixed sequence length
+// (tools/rknn/prepare_cosyvoice3_rknn_flow.py) and a ladder of them is registered here. An utterance of
+// `frames` mel frames runs on the smallest engine that fits it: `x`/`mu`/`cond` zero-padded on
+// the right, `mask` 1 on the real frames and 0 in the pad, output cropped back. Measured on the
+// host, a padded run is bit-identical to an exact run at the real length (355 dB, max_abs 0.0),
+// so the ladder is a time trade, not a quality one: the estimator costs ~S^2, and every frame of
+// pad is paid for ten times, once per ODE step.
+//
+// The engine takes the raw timestep `t` (its time-embedding MLP is inside the graph) and a frame
+// `mask`; the ggml graph takes the timestep embedding and has no mask. Both are supplied here
+// from the same ODE loop, so the two backends see the same step.
+class CosyVoice3RknnFlow {
+public:
+    // Registers the ladder without loading any of it: one engine asks the driver for 0.7-3.3 GB
+    // of NPU memory (measured, seq200 to seq1100), so a ladder cannot be resident at once. The
+    // length an engine covers is read from its file name (`flow_estimator_seq900.rknn`, the
+    // producer's convention) and validated against the engine's own attributes when it loads.
+    // `mel_channels` is the checkpoint's, checked the same way.
+    CosyVoice3RknnFlow(const std::vector<std::string> & engine_paths, int64_t mel_channels);
+    ~CosyVoice3RknnFlow();
+
+    CosyVoice3RknnFlow(const CosyVoice3RknnFlow &) = delete;
+    CosyVoice3RknnFlow & operator=(const CosyVoice3RknnFlow &) = delete;
+
+    // Longest utterance the ladder covers, in mel frames. A longer one has no engine and has to
+    // run on the ggml graph, which is why the caller keeps it.
+    int64_t max_frames() const;
+    int64_t engine_count() const;
+    // Mel channels the ladder was built for.
+    int64_t mel_channels() const;
+
+    // One ODE step. `x`, `mu` and `cond` are channel-major [2, mel_channels, frames] and `spks`
+    // is [2, mel_channels], with the classifier-free-guidance uncond branch zeroed in `mu`,
+    // `cond` and `spks` exactly as the ggml path builds it. `t` is the raw timestep. Returns
+    // dphi_dt as [2, mel_channels, frames]: the cond branch first, then uncond.
+    std::vector<float> run(
+        const float * x,
+        const float * mu,
+        const float * cond,
+        const float * spks,
+        float t,
+        int64_t frames);
+
+private:
+    struct Impl;
+    std::unique_ptr<Impl> impl_;
+};
+
+}  // namespace engine::models::cosyvoice3

model_specs/cosyvoice3.json

@@ -174,6 +174,13 @@
         "description": "Release cached runtime graphs after each request to reduce peak VRAM; default false.",
         "required": false,
         "default": false
+      },
+      {
+        "name": "exec_plan",
+        "type": "string",
+        "description": "JSON object of component execution plans. Component: flow. Each plan accepts impl, backend, artifacts, and params. impl is default/ggml/rknn. GGML keeps the session backend. RKNN requires artifacts.main, a comma-separated ladder of estimator engines compiled one per sequence length (tools/rknn/prepare_cosyvoice3_rknn_flow.py), and ENGINE_ENABLE_RKNN; omit backend. An utterance longer than the largest engine runs the CPU graph. Neither implementation takes params.",
+        "required": false,
+        "default": "{}"
       }
     ],
     "load": []

src/models/cosyvoice3/flow.cpp

@@ -14,6 +14,7 @@
 #include "engine/framework/modules/primitive_modules.h"
 #include "engine/framework/modules/structural_modules.h"
 #include "engine/framework/modules/weight_binding.h"
+#include "engine/models/cosyvoice3/rknn_flow.h"
 
 #include <ggml-alloc.h>
 #include <ggml.h>
@@ -269,6 +270,10 @@ std::shared_ptr<CosyDiTFlowWeights> load_flow_weights(
         *weights->store, source, "pre_lookahead_layer.conv1", storage_type, c.flow_hidden_size, c.flow_mel_channels, c.pre_lookahead_len + 1, true);
     weights->pre_lookahead_conv2 = binding::conv1d_from_source(
         *weights->store, source, "pre_lookahead_layer.conv2", storage_type, c.flow_mel_channels, c.flow_hidden_size, 3, true);
+    // The ggml DiT weights are loaded in both paths, deliberately: the RKNN ladder covers a fixed
+    // set of lengths, and an utterance longer than the largest engine runs the ggml graph, so the
+    // copy is not spare memory. Skipping it would trade a fallback that works for one that throws
+    // `weight is invalid` the first time the ladder runs out.
     weights->input_projection = binding::linear_from_source(
         *weights->store, source, "decoder.estimator.input_embed.proj", storage_type, c.flow_hidden_size, 320, true);
     weights->conv_pos_1.weight = weights->store->load_tensor(
@@ -514,7 +519,7 @@ class DiTGraph {
         const std::vector<float> & mu,
         const std::vector<float> & cond,
         const std::vector<float> & speaker,
-        const std::vector<float> & time_embedding,
+        float t,
         int64_t frames) {
         if (frames <= 0) {
             throw std::runtime_error("CosyVoice3 DiT frames must be positive");
@@ -524,10 +529,15 @@ class DiTGraph {
         if (static_cast<int64_t>(x.size()) != batch * channels * frames ||
             static_cast<int64_t>(mu.size()) != batch * channels * frames ||
             static_cast<int64_t>(cond.size()) != batch * channels * frames ||
-            static_cast<int64_t>(speaker.size()) != batch * channels ||
-            static_cast<int64_t>(time_embedding.size()) != batch * kTimeEmbeddingSize) {
+            static_cast<int64_t>(speaker.size()) != batch * channels) {
             throw std::runtime_error("CosyVoice3 DiT input size mismatch");
         }
+        // The graph takes the timestep embedding, not the timestep: the sinusoidal expansion is
+        // host-side here and inside the graph in the RKNN engine. Both backends get `t`.
+        const auto temb = timestep_embedding(t);
+        std::vector<float> time_embedding(static_cast<size_t>(batch) * kTimeEmbeddingSize);
+        std::copy(temb.begin(), temb.end(), time_embedding.begin());
+        std::copy(temb.begin(), temb.end(), time_embedding.begin() + kTimeEmbeddingSize);
         ensure(frames);
         ggml_backend_tensor_set(x_, x.data(), 0, x.size() * sizeof(float));
         ggml_backend_tensor_set(mu_, mu.data(), 0, mu.size() * sizeof(float));
@@ -728,16 +738,29 @@ class CosyVoice3DiTFlowRuntime::Impl {
         core::ExecutionContext & execution,
         size_t graph_arena_bytes,
         size_t weight_context_bytes,
-        engine::assets::TensorStorageType storage_type)
+        engine::assets::TensorStorageType storage_type,
+        const std::vector<std::string> & rknn_engine_paths)
         : assets_(std::move(assets)),
           execution_(execution),
-          weights_(load_flow_weights(*assets_, execution_, weight_context_bytes, storage_type)),
+          weights_(load_flow_weights(
+              *assets_,
+              execution_,
+              weight_context_bytes,
+              storage_type)),
           condition_(execution_, weights_, assets_->config, graph_arena_bytes),
           dit_(execution_, weights_, assets_->config, graph_arena_bytes),
           graph_arena_bytes_(graph_arena_bytes) {
         if (assets_ == nullptr) {
             throw std::runtime_error("CosyVoice3 flow runtime requires assets");
         }
+        if (!rknn_engine_paths.empty()) {
+            // The engines' own channel count is checked against the checkpoint when an engine
+            // loads, so nothing to compare here: registering the ladder loads none of it.
+            rknn_ = std::make_unique<CosyVoice3RknnFlow>(
+                rknn_engine_paths, assets_->config.flow_mel_channels);
+            debug::trace_log_scalar("cosyvoice3.flow.dit.rknn.engines", rknn_->engine_count());
+            debug::trace_log_scalar("cosyvoice3.flow.dit.rknn.max_frames", rknn_->max_frames());
+        }
     }
 
     CosyVoice3FlowOutput generate(const CosyVoice3FlowRequest & request) {
@@ -788,15 +811,20 @@ class CosyVoice3DiTFlowRuntime::Impl {
         std::copy(speaker_cond.begin(), speaker_cond.end(), spks.begin());
 
         const auto schedule = cosine_time_schedule(request.num_inference_steps);
+        // The NPU ladder covers a fixed range of lengths; past its largest engine the ggml graph
+        // is the only backend that has a shape for the utterance, so the fallback is per call.
+        const bool on_npu = rknn_ != nullptr && total_frames <= rknn_->max_frames();
+        if (rknn_ != nullptr && !on_npu) {
+            debug::timing_log_scalar(
+                "cosyvoice3.flow.dit.rknn.fallback_frames", static_cast<int64_t>(total_frames));
+        }
         const auto dit_start = std::chrono::steady_clock::now();
         for (int64_t step = 1; step < static_cast<int64_t>(schedule.size()); ++step) {
             const float t = schedule[static_cast<size_t>(step - 1)];
             const float dt = schedule[static_cast<size_t>(step)] - t;
-            auto temb = timestep_embedding(t);
-            std::vector<float> time_batched(static_cast<size_t>(2 * kTimeEmbeddingSize));
-            std::copy(temb.begin(), temb.end(), time_batched.begin());
-            std::copy(temb.begin(), temb.end(), time_batched.begin() + kTimeEmbeddingSize);
-            auto pred = dit_.run(x_batched, mu, cond, spks, time_batched, total_frames);
+            auto pred = on_npu
+                ? rknn_->run(x_batched.data(), mu.data(), cond.data(), spks.data(), t, total_frames)
+                : dit_.run(x_batched, mu, cond, spks, t, total_frames);
             const size_t branch = static_cast<size_t>(c.flow_mel_channels * total_frames);
             for (size_t i = 0; i < branch; ++i) {
                 const float guided = (1.0F + kInferenceCfgRate) * pred[i] - kInferenceCfgRate * pred[branch + i];
@@ -832,6 +860,7 @@ class CosyVoice3DiTFlowRuntime::Impl {
     std::shared_ptr<CosyDiTFlowWeights> weights_;
     ConditionGraph condition_;
     DiTGraph dit_;
+    std::unique_ptr<CosyVoice3RknnFlow> rknn_;
     size_t graph_arena_bytes_ = 0;
 };
 
@@ -840,13 +869,15 @@ CosyVoice3DiTFlowRuntime::CosyVoice3DiTFlowRuntime(
     engine::core::ExecutionContext & execution,
     size_t graph_arena_bytes,
     size_t weight_context_bytes,
-    engine::assets::TensorStorageType storage_type)
+    engine::assets::TensorStorageType storage_type,
+    const std::vector<std::string> & rknn_engine_paths)
     : impl_(std::make_unique<Impl>(
           std::move(assets),
           execution,
           graph_arena_bytes,
           weight_context_bytes,
-          storage_type)) {}
+          storage_type,
+          rknn_engine_paths)) {}
 
 CosyVoice3DiTFlowRuntime::~CosyVoice3DiTFlowRuntime() = default;