Parakeet-TDT NPU acceleration on the Apple Neural Engine

README.md

@@ -132,11 +132,11 @@ graph.hard_reset();
 
 | Device | Gemma4 Text | Gemma4 Vision | Parakeet 1.1B | RAM |
 |--------|----------|------------|---------------|-----|
-| Mac M4 Pro | - | - | - | - |
+| Mac M4 Pro | 324 / 39 | 1.2s / 48 | 0.2s (NPU) | 1385 MB |
 | iPad/Mac M3 | - | - | - | - |
 | iPhone 17 Pro | - | - | - | - |
 | iPhone 13 Mini | - | - | - | - |
-| Galaxy S26 | - | - | - | - |
+| Galaxy S26 | 248 / 21 | 17s / 16 | 4.2s | 4024 MB |
 | Pixel 10 Pro | - | - | - | - |
 | Pixel 6a | - | - | - | - |
 | Galaxy A17 5G | - | - | - | - |

cactus-engine/src/engine.h

@@ -29,7 +29,7 @@ struct NPUNamedInput {
 class NPUEncoder {
 public:
     virtual ~NPUEncoder() = default;
-    virtual bool load(const std::string& model_path) = 0;
+    virtual bool load(const std::string& model_path, const std::string& compute_units = "") = 0;
     virtual bool preallocate(
         const std::vector<int>& input_shape,
         const std::string& input_name = "x",
@@ -701,7 +701,7 @@ class Model {
     size_t last_prefill_tail_chunk_tokens() const { return last_prefill_tail_chunk_tokens_; }
     size_t last_prefill_tail_padding_tokens() const { return last_prefill_tail_padding_tokens_; }
 
-    bool load_npu_audio_encoder(const std::string& model_path);
+    bool load_npu_audio_encoder(const std::string& model_path, const std::string& compute_units = "");
     bool has_npu_audio_encoder() const { return npu_audio_encoder_ != nullptr; }
 
     bool load_npu_vision_encoder(const std::string& model_path);
@@ -860,6 +860,7 @@ class Model {
 
     std::string family_;
     std::string npu_audio_encoder_mlpackage_;
+    std::string npu_audio_compute_units_;
     std::string npu_vision_encoder_mlpackage_;
     std::string npu_source_encoder_mlpackage_;
 

cactus-engine/src/model.cpp

@@ -614,7 +614,7 @@ bool Model::init(const std::string& bundle_dir, size_t context_size,
 
     if (!npu_audio_encoder_mlpackage_.empty()) {
         std::string full_path = bundle_dir + "/" + npu_audio_encoder_mlpackage_;
-        if (!load_npu_audio_encoder(full_path)) {
+        if (!load_npu_audio_encoder(full_path, npu_audio_compute_units_)) {
             CACTUS_LOG_WARN("model", "NPU audio encoder load failed for " << full_path << "; falling back to CPU");
         }
     }
@@ -662,6 +662,9 @@ bool Model::load_manifest() {
     if (obj.count("npu_audio_encoder") && obj.at("npu_audio_encoder").is<std::string>()) {
         npu_audio_encoder_mlpackage_ = obj.at("npu_audio_encoder").get<std::string>();
     }
+    if (obj.count("npu_audio_compute_units") && obj.at("npu_audio_compute_units").is<std::string>()) {
+        npu_audio_compute_units_ = obj.at("npu_audio_compute_units").get<std::string>();
+    }
     if (obj.count("npu_vision_encoder") && obj.at("npu_vision_encoder").is<std::string>()) {
         npu_vision_encoder_mlpackage_ = obj.at("npu_vision_encoder").get<std::string>();
     }
@@ -2900,7 +2903,8 @@ std::vector<uint32_t> Model::transcribe_parakeet_tdt(const std::vector<float>& a
         const std::vector<int> out_shape = npu_audio_encoder_->get_output_shape();
         if (in_shape.size() >= 3 && out_shape.size() >= 3 &&
             in_shape[1] > 0 && in_shape[2] > 0 && out_shape[1] > 0 && out_shape[2] > 0 &&
-            static_cast<size_t>(in_shape[2]) == expected_mels) {
+            static_cast<size_t>(in_shape[2]) == expected_mels &&
+            copy_frames <= static_cast<size_t>(in_shape[1])) {
             const size_t window_frames = static_cast<size_t>(in_shape[1]);
             const size_t window_hidden = static_cast<size_t>(out_shape[1]);
             const size_t hidden_dim_npu = static_cast<size_t>(out_shape[2]);

cactus-engine/src/model_npu.cpp

@@ -10,10 +10,10 @@
 namespace cactus {
 namespace engine {
 
-bool Model::load_npu_audio_encoder(const std::string& model_path) {
+bool Model::load_npu_audio_encoder(const std::string& model_path, const std::string& compute_units) {
     auto encoder = npu::create_encoder();
     if (!encoder) return false;
-    if (!encoder->load(model_path)) return false;
+    if (!encoder->load(model_path, compute_units)) return false;
     if (!encoder->is_available()) return false;
     npu_audio_encoder_ = std::move(encoder);
     CACTUS_LOG_INFO("model", "NPU audio encoder loaded from: " << model_path);

cactus-engine/src/npu_ane.h

@@ -26,7 +26,7 @@ class ANEEncoder : public NPUEncoder {
     ANEEncoder(ANEEncoder&& other) noexcept;
     ANEEncoder& operator=(ANEEncoder&& other) noexcept;
 
-    bool load(const std::string& model_path) override;
+    bool load(const std::string& model_path, const std::string& compute_units = "") override;
 
     bool preallocate(const std::vector<int>& input_shape,
                      const std::string& input_name = "x",
@@ -68,7 +68,7 @@ class ANEEncoder : public NPUEncoder {
     ANEEncoder() = default;
     ~ANEEncoder() override = default;
 
-    bool load(const std::string&) override { return false; }
+    bool load(const std::string&, const std::string& = "") override { return false; }
 
     bool preallocate(const std::vector<int>&,
                      const std::string& = "x",

cactus-engine/src/npu_ane.mm

@@ -89,17 +89,26 @@ static void maybe_apply_compute_units_env(const char* env_name, MLComputeUnits&
     CACTUS_LOG_WARN("npu", "Ignoring invalid " << env_name << "=" << raw);
 }
 
-static MLComputeUnits resolve_compute_units_for_model(NSString* model_path, bool is_prefill) {
+static MLComputeUnits resolve_compute_units_for_model(NSString* model_path, bool is_prefill,
+                                                      const std::string& compute_units_hint) {
     MLComputeUnits units = MLComputeUnitsCPUAndNeuralEngine;
 
-    // Gemma4 multimodal encoders show better fidelity on CPU+GPU.
     if (!is_prefill && model_path_looks_like_audio_encoder(model_path)) {
         units = MLComputeUnitsCPUAndGPU;
     }
     if (!is_prefill && model_path_looks_like_vision_encoder(model_path)) {
         units = MLComputeUnitsCPUAndGPU;
     }
 
+    if (!compute_units_hint.empty()) {
+        MLComputeUnits parsed;
+        if (parse_compute_units(compute_units_hint.c_str(), parsed)) {
+            units = parsed;
+        } else {
+            CACTUS_LOG_WARN("npu", "Ignoring invalid compute units hint=" << compute_units_hint);
+        }
+    }
+
     maybe_apply_compute_units_env("CACTUS_ANE_COMPUTE_UNITS", units);
     if (is_prefill) {
         maybe_apply_compute_units_env("CACTUS_ANE_PREFILL_COMPUTE_UNITS", units);
@@ -207,7 +216,7 @@ @interface CactusANEImpl : NSObject
 @property (nonatomic, assign) BOOL hasOutputBackings;
 @property (nonatomic, assign) BOOL hasMultiOutputBackings;
 
-- (instancetype)initWithModelPath:(NSString*)path;
+- (instancetype)initWithModelPath:(NSString*)path computeUnits:(NSString*)cu;
 - (NSArray<NSNumber*>*)getInputShape;
 - (NSArray<NSNumber*>*)getOutputShape;
 - (BOOL)preallocateBuffersWithInput:(NSString*)inputName
@@ -229,14 +238,15 @@ - (size_t)predictMultiInputWithDict:(NSMutableDictionary<NSString*, MLFeatureVal
 
 @implementation CactusANEImpl
 
-- (instancetype)initWithModelPath:(NSString*)path {
+- (instancetype)initWithModelPath:(NSString*)path computeUnits:(NSString*)cu {
     self = [super init];
     if (self) {
         NSError* error = nil;
         NSURL* modelURL = resolve_or_compile_model_url(path, &error);
 
         MLModelConfiguration* config = [[MLModelConfiguration alloc] init];
-        config.computeUnits = resolve_compute_units_for_model(path, false);
+        config.computeUnits = resolve_compute_units_for_model(
+            path, false, cu ? std::string([cu UTF8String]) : std::string(""));
         CACTUS_LOG_INFO("npu",
                         "ANEEncoder compute units: "
                             << compute_units_to_string(config.computeUnits)
@@ -794,7 +804,7 @@ - (size_t)predictMultiInputWithDict:(NSMutableDictionary<NSString*, MLFeatureVal
     return *this;
 }
 
-bool ANEEncoder::load(const std::string& model_path) {
+bool ANEEncoder::load(const std::string& model_path, const std::string& compute_units) {
     @autoreleasepool {
         CACTUS_LOG_INFO("npu", "ANEEncoder loading model: " << model_path);
         NSString* path = [NSString stringWithUTF8String:model_path.c_str()];
@@ -804,7 +814,7 @@ - (size_t)predictMultiInputWithDict:(NSMutableDictionary<NSString*, MLFeatureVal
             return false;
         }
 
-        CactusANEImpl* impl = [[CactusANEImpl alloc] initWithModelPath:path];
+        CactusANEImpl* impl = [[CactusANEImpl alloc] initWithModelPath:path computeUnits:(compute_units.empty() ? nil : [NSString stringWithUTF8String:compute_units.c_str()])];
 
         if (impl && impl.model) {
             impl_ = (__bridge_retained void*)impl;

python/cactus/transpile/component_pipeline.py

@@ -30,6 +30,7 @@ class ComponentModuleSpec:
     metadata: dict[str, object] = field(default_factory=dict)
     npu_module: torch.nn.Module | None = None
     npu_runtime_input_count: int = 1
+    npu_example_inputs: tuple[torch.Tensor, ...] | None = None
 
 
 @dataclass

python/cactus/transpile/hf_model.py

@@ -506,6 +506,8 @@ def _write_component_bundle(
     for manifest_key, mlpackage_path in (npu_encoder_mlpackages or {}).items():
         if mlpackage_path:
             manifest_payload[manifest_key] = mlpackage_path
+    if family == "parakeet_tdt" and "npu_audio_encoder" in manifest_payload:
+        manifest_payload["npu_audio_compute_units"] = "CPU_AND_NE"
     _write_json(manifest_path, manifest_payload)
     return manifest_path
 

python/cactus/transpile/npu/coremltools_patches.py

@@ -14,6 +14,7 @@ def apply_all_coremltools_patches() -> None:
         _override_layer_norm_translator()
         _override_one_hot_translator()
         _register_unfold_op()
+        _register_noop_assert_ops()
         _patch_pipeline_remove_fuse_prelu()
         _patch_mb_binops_scalar_cast()
         _patch_fp16_cast_skip_layer_norm()
@@ -384,6 +385,17 @@ def one_hot(context, node):
     _TORCH_OPS_REGISTRY.set_func_by_name(one_hot, "one_hot")
 
 
+def _register_noop_assert_ops() -> None:
+    from coremltools.converters.mil.frontend.torch.torch_op_registry import (
+        _TORCH_OPS_REGISTRY,
+    )
+
+    def _noop(context, node):
+        return
+
+    _TORCH_OPS_REGISTRY.set_func_by_name(_noop, "_assert_tensor_metadata")
+
+
 def _register_unfold_op() -> None:
     from coremltools.converters.mil.frontend.torch.torch_op_registry import (
         _TORCH_OPS_REGISTRY,

python/cactus/transpile/npu/pipeline.py

@@ -47,7 +47,9 @@ def _resolve(per_component: int | None, default: int | None) -> int | None:
         if component not in _ENCODER_COMPONENTS:
             continue
         emit_fn, filename = _ENCODER_COMPONENTS[component]
-        example_inputs = tuple(getattr(spec, "example_inputs", ()) or ())
+        example_inputs = tuple(
+            getattr(spec, "npu_example_inputs", None) or getattr(spec, "example_inputs", ()) or ()
+        )
         if not example_inputs:
             print(f"npu.pipeline: {component} spec has no example inputs; skipping")
             continue