From 20ecaa8b775431cee816a4512fefd566d1bd113c Mon Sep 17 00:00:00 2001 From: Noah Cylich Date: Tue, 16 Jun 2026 12:37:16 -1000 Subject: [PATCH] gemma4: fix fp16 -inf overflow in audio encoder mask, route audio to ANE The audio attention masked-fill constant (-1e9) overflows to -inf in the FLOAT16 NPU graph: fully-masked rows softmax to NaN and partial rows diverge on the Neural Engine. Clamp it to a finite fp16-safe -1e4 before export (soft-capped logits stay within +/-50, so masking is unchanged), and route the gemma4 audio encoder to CPU_AND_NE via the bundle manifest. Signed-off-by: Noah Cylich --- python/cactus/transpile/hf_model.py | 2 +- python/cactus/transpile/model_adapters.py | 17 +++++++++++++++++ 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/python/cactus/transpile/hf_model.py b/python/cactus/transpile/hf_model.py index 4151ccb4a..dc8beacd2 100644 --- a/python/cactus/transpile/hf_model.py +++ b/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 diff --git a/python/cactus/transpile/model_adapters.py b/python/cactus/transpile/model_adapters.py index b6e2cee0d..e10a3f803 100644 --- a/python/cactus/transpile/model_adapters.py +++ b/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