gemma4: run the audio encoder correctly and fast on the Apple Neural Engine

python/cactus/transpile/hf_model.py

@@ -506,7 +506,7 @@ 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:
+    if family in ("parakeet_tdt", "gemma4") 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/model_adapters.py

@@ -3533,6 +3533,22 @@ def _build_gemma3_causal_lm_component_specs(
     return specs
 
 
+def _gemma4_make_audio_mask_fp16_safe(model: torch.nn.Module) -> None:
+    backbone = getattr(model, "model", model)
+    audio_tower = getattr(backbone, "audio_tower", None)
+    if audio_tower is None:
+        return
+    seen: set[int] = set()
+    for mod in audio_tower.modules():
+        cfg = getattr(mod, "config", None)
+        if cfg is None or id(cfg) in seen:
+            continue
+        seen.add(id(cfg))
+        value = getattr(cfg, "attention_invalid_logits_value", None)
+        if isinstance(value, (int, float)) and value < -1e4:
+            cfg.attention_invalid_logits_value = -1e4
+
+
 def _build_gemma4_multimodal_component_specs(
     model: torch.nn.Module,
     *,
@@ -3616,6 +3632,7 @@ def _require(component: str) -> None:
         native_image_soft_token_counts=native_image_soft_token_counts,
         native_image_pool_shapes=None,
     ).eval()
+    _gemma4_make_audio_mask_fp16_safe(model)
     audio_encoder = Gemma4AudioEncoderAdapter(model, weights_dir=weights_dir).eval()
 
     image_features: torch.Tensor | None = None