diff --git a/README.md b/README.md index 02203083c..dbe690029 100644 --- a/README.md +++ b/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 | - | - | - | - | diff --git a/cactus-engine/src/engine.h b/cactus-engine/src/engine.h index 9f0f6c8b3..34fd2153b 100644 --- a/cactus-engine/src/engine.h +++ b/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& 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_; diff --git a/cactus-engine/src/model.cpp b/cactus-engine/src/model.cpp index 0a8f8b490..a46b6fc61 100644 --- a/cactus-engine/src/model.cpp +++ b/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()) { npu_audio_encoder_mlpackage_ = obj.at("npu_audio_encoder").get(); } + if (obj.count("npu_audio_compute_units") && obj.at("npu_audio_compute_units").is()) { + npu_audio_compute_units_ = obj.at("npu_audio_compute_units").get(); + } if (obj.count("npu_vision_encoder") && obj.at("npu_vision_encoder").is()) { npu_vision_encoder_mlpackage_ = obj.at("npu_vision_encoder").get(); } @@ -2900,7 +2903,8 @@ std::vector Model::transcribe_parakeet_tdt(const std::vector& a const std::vector 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(in_shape[2]) == expected_mels) { + static_cast(in_shape[2]) == expected_mels && + copy_frames <= static_cast(in_shape[1])) { const size_t window_frames = static_cast(in_shape[1]); const size_t window_hidden = static_cast(out_shape[1]); const size_t hidden_dim_npu = static_cast(out_shape[2]); diff --git a/cactus-engine/src/model_npu.cpp b/cactus-engine/src/model_npu.cpp index 89cd88977..8c421dee1 100644 --- a/cactus-engine/src/model_npu.cpp +++ b/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); diff --git a/cactus-engine/src/npu_ane.h b/cactus-engine/src/npu_ane.h index c03dcc0cf..7594960bf 100644 --- a/cactus-engine/src/npu_ane.h +++ b/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& 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&, const std::string& = "x", diff --git a/cactus-engine/src/npu_ane.mm b/cactus-engine/src/npu_ane.mm index d7334067a..99638e2e0 100644 --- a/cactus-engine/src/npu_ane.mm +++ b/cactus-engine/src/npu_ane.mm @@ -89,10 +89,10 @@ 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; } @@ -100,6 +100,15 @@ static MLComputeUnits resolve_compute_units_for_model(NSString* model_path, bool 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*)getInputShape; - (NSArray*)getOutputShape; - (BOOL)preallocateBuffersWithInput:(NSString*)inputName @@ -229,14 +238,15 @@ - (size_t)predictMultiInputWithDict:(NSMutableDictionary 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, diff --git a/python/cactus/transpile/npu/pipeline.py b/python/cactus/transpile/npu/pipeline.py index e04429298..f3a277806 100644 --- a/python/cactus/transpile/npu/pipeline.py +++ b/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 diff --git a/python/cactus/transpile/tdt_runtime.py b/python/cactus/transpile/tdt_runtime.py index 7f6d2d562..01275d7ec 100644 --- a/python/cactus/transpile/tdt_runtime.py +++ b/python/cactus/transpile/tdt_runtime.py @@ -3,6 +3,7 @@ from dataclasses import dataclass from dataclasses import replace import json +import os from pathlib import Path import re from typing import Any @@ -415,6 +416,17 @@ def _relative_position_bias(query: torch.Tensor, relative_key: torch.Tensor, *, return gathered * float(scale) +def _relative_position_bias_static(query: torch.Tensor, relative_key: torch.Tensor, *, scale: float) -> torch.Tensor: + batch, heads, seq_len, _ = query.shape + pos_len = 2 * seq_len - 1 + scores = torch.matmul(query, relative_key.transpose(-1, -2)) + x = F.pad(scores, (1, 0)) + x = x.reshape(batch, heads, pos_len + 1, seq_len) + x = x[:, :, 1:, :] + x = x.reshape(batch, heads, seq_len, pos_len) + return x[:, :, :, :seq_len] * float(scale) + + class ParakeetTDTFeedForward(nn.Module): def __init__(self, config: ParakeetTDTConfig, prefix: str, state_dict: dict[str, torch.Tensor]): super().__init__() @@ -550,10 +562,12 @@ def __init__(self, config: ParakeetTDTConfig, prefix: str, state_dict: dict[str, ) self.config = config - def forward(self, x: torch.Tensor) -> torch.Tensor: + def forward(self, x: torch.Tensor, pad_mask: torch.Tensor | None = None) -> torch.Tensor: x = x.transpose(1, 2) x = self.pointwise_conv1(x) x = F.glu(x, dim=1) + if pad_mask is not None: + x = x * pad_mask x = self.depthwise_conv(x) x = self.batch_norm(x) x = _apply_activation(x, self.config.encoder_hidden_act) @@ -644,6 +658,111 @@ def forward(self, input_features: torch.Tensor) -> torch.Tensor: return x +class ParakeetTDTANESelfAttention(nn.Module): + def __init__(self, attn: "ParakeetTDTSelfAttention", seq_len: int): + super().__init__() + self.linear_q = attn.linear_q + self.linear_k = attn.linear_k + self.linear_v = attn.linear_v + self.linear_out = attn.linear_out + self.linear_pos = attn.linear_pos + self.pos_bias_u = attn.pos_bias_u + self.pos_bias_v = attn.pos_bias_v + self.config = attn.config + weight = attn.linear_pos.weight + rel_pos = _relative_position_embeddings( + seq_len=int(seq_len), + hidden_dim=self.config.hidden_dim, + device=weight.device, + dtype=weight.dtype, + ) + with torch.no_grad(): + rel_k = attn.linear_pos(rel_pos).view( + 1, 2 * int(seq_len) - 1, self.config.attention_heads, self.config.attention_head_dim + ).flip(1) + self.register_buffer("_rel_k", rel_k.detach(), persistent=False) + + def forward(self, x: torch.Tensor, key_bias: torch.Tensor | None = None) -> torch.Tensor: + batch, seq_len, _ = x.shape + num_heads = self.config.attention_heads + head_dim = self.config.attention_head_dim + + q = self.linear_q(x).view(batch, seq_len, num_heads, head_dim) + k = self.linear_k(x).view(batch, seq_len, num_heads, head_dim) + v = self.linear_v(x).view(batch, seq_len, num_heads, head_dim) + + q_u = q + self.pos_bias_u.view(1, 1, num_heads, head_dim).to(dtype=x.dtype, device=x.device) + q_v = q + self.pos_bias_v.view(1, 1, num_heads, head_dim).to(dtype=x.dtype, device=x.device) + + rel_k_heads = self._rel_k.to(dtype=x.dtype, device=x.device).permute(0, 2, 1, 3) + rel_bias = _relative_position_bias_static( + q_v.permute(0, 2, 1, 3), + rel_k_heads, + scale=self.config.attention_scale, + ) + if key_bias is not None: + rel_bias = rel_bias + key_bias + attn = F.scaled_dot_product_attention( + q_u.permute(0, 2, 1, 3), + k.permute(0, 2, 1, 3), + v.permute(0, 2, 1, 3), + attn_mask=rel_bias, + dropout_p=0.0, + is_causal=False, + ) + attn = attn.permute(0, 2, 1, 3).reshape(batch, seq_len, self.config.hidden_dim) + return self.linear_out(attn) + + +class ParakeetTDTANEEncoderLayer(nn.Module): + def __init__(self, layer: "ParakeetTDTEncoderLayer", seq_len: int): + super().__init__() + self.feed_forward1 = layer.feed_forward1 + self.feed_forward2 = layer.feed_forward2 + self.self_attn = ParakeetTDTANESelfAttention(layer.self_attn, seq_len) + self.conv = layer.conv + self.norm_feed_forward1 = layer.norm_feed_forward1 + self.norm_self_att = layer.norm_self_att + self.norm_conv = layer.norm_conv + self.norm_feed_forward2 = layer.norm_feed_forward2 + self.norm_out = layer.norm_out + self.config = layer.config + + def forward(self, x: torch.Tensor, conv_mask: torch.Tensor, key_bias: torch.Tensor) -> torch.Tensor: + x = x + 0.5 * self.feed_forward1(self.norm_feed_forward1(x), activation=self.config.encoder_hidden_act) + x = x + self.self_attn(self.norm_self_att(x), key_bias) + x = x + self.conv(self.norm_conv(x), conv_mask) + x = x + 0.5 * self.feed_forward2(self.norm_feed_forward2(x), activation=self.config.encoder_hidden_act) + return self.norm_out(x) + + +class ParakeetTDTANEEncoder(nn.Module): + def __init__(self, encoder: "ParakeetTDTEncoder", seq_len: int): + super().__init__() + self.pre_encode = encoder.pre_encode + self.layers = nn.ModuleList( + [ParakeetTDTANEEncoderLayer(layer, seq_len) for layer in encoder.layers] + ) + factor = max(1, int(encoder.layers[0].config.subsampling_factor)) + self._subsample_stages = max(0, factor.bit_length() - 1) + + def _frame_masks(self, input_features: torch.Tensor, seq_len: int) -> tuple[torch.Tensor, torch.Tensor]: + valid = (input_features.abs().sum(-1) > 0).to(input_features.dtype).sum(-1, keepdim=True) + for _ in range(self._subsample_stages): + valid = torch.floor((valid - 1.0) / 2.0) + 1.0 + positions = torch.arange(seq_len, device=input_features.device, dtype=input_features.dtype) + enc_valid = (positions.view(1, 1, seq_len) < valid.view(-1, 1, 1)).to(input_features.dtype) + key_bias = (1.0 - enc_valid).unsqueeze(1) * -1e4 + return enc_valid, key_bias + + def forward(self, input_features: torch.Tensor) -> torch.Tensor: + x = self.pre_encode(input_features) + conv_mask, key_bias = self._frame_masks(input_features, x.shape[1]) + for layer in self.layers: + x = layer(x, conv_mask, key_bias) + return x + + class ParakeetTDTDecoderCell(nn.Module): def __init__(self, hidden_dim: int, prefix: str, state_dict: dict[str, torch.Tensor]): super().__init__() @@ -868,6 +987,22 @@ def build_parakeet_tdt_component_specs( "task": "tdt_transcription", "adapter_family": "parakeet_tdt", } + + factor = max(1, int(model.config.subsampling_factor)) + npu_frames = int(os.environ.get("CACTUS_NPU_PARAKEET_FRAMES", "3000")) + if npu_frames <= 0 or npu_frames % factor != 0: + raise ValueError(f"CACTUS_NPU_PARAKEET_FRAMES must be a positive multiple of {factor}, got {npu_frames}") + npu_input_features = torch.zeros( + (int(input_features.shape[0]), npu_frames, int(input_features.shape[-1])), + device=input_features.device, + dtype=input_features.dtype, + ) + valid_frames = min(int(input_features.shape[1]), npu_frames) + npu_input_features[:, :valid_frames, :] = input_features[:, :valid_frames, :] + with torch.no_grad(): + ane_seq_len = int(model.encoder.pre_encode(npu_input_features).shape[1]) + ane_encoder = ParakeetTDTANEEncoder(model.encoder, ane_seq_len) + return [ ComponentModuleSpec( component="audio_encoder", @@ -877,6 +1012,9 @@ def build_parakeet_tdt_component_specs( output_keys=("encoder_hidden_states",), graph_meta={**common_graph_meta, "component": "audio_encoder"}, metadata={"family": "parakeet_tdt", "task": "tdt_transcription"}, + npu_module=ane_encoder, + npu_example_inputs=(npu_input_features,), + npu_runtime_input_count=1, ), ComponentModuleSpec( component="decoder",