gemma4: fix 16x audio soft-token magnitude on NPU via single-source scaling

.github/workflows/inference.yml

@@ -7,6 +7,7 @@ on:
 
 permissions:
   contents: read
+  actions: read
 
 concurrency:
   group: inference-${{ github.event.pull_request.number || github.ref }}
@@ -29,9 +30,41 @@ jobs:
             echo "approved=false" >> $GITHUB_OUTPUT
           fi
 
-  build:
+  runner-check:
     needs: gate
     if: needs.gate.outputs.approved == 'true'
+    runs-on: ubuntu-latest
+    timeout-minutes: 5
+    outputs:
+      available: ${{ steps.check.outputs.available }}
+    steps:
+      - id: check
+        env:
+          GH_TOKEN: ${{ github.token }}
+          REPO: ${{ github.repository }}
+        run: |
+          set -euo pipefail
+          status=$(curl -sS -L -w "%{http_code}" -o runners.json \
+            -H "Authorization: Bearer ${GH_TOKEN}" \
+            -H "Accept: application/vnd.github+json" \
+            -H "X-GitHub-Api-Version: 2022-11-28" \
+            "https://api.github.com/repos/${REPO}/actions/runners?per_page=100" || true)
+          if [ "$status" != "200" ]; then
+            echo "available=false" >> "$GITHUB_OUTPUT"
+            echo "::notice::Could not inspect self-hosted runners (HTTP $status); skipping inference tests instead of queueing."
+            exit 0
+          fi
+          count=$(jq '[.runners[] | select(.status == "online" and .busy == false) | select(([.labels[].name] | contains(["self-hosted", "linux", "arm64"])))] | length' runners.json)
+          if [ "$count" -gt 0 ]; then
+            echo "available=true" >> "$GITHUB_OUTPUT"
+          else
+            echo "available=false" >> "$GITHUB_OUTPUT"
+            echo "::notice::No idle self-hosted linux arm64 runner is available; skipping inference tests."
+          fi
+
+  build:
+    needs: [gate, runner-check]
+    if: needs.gate.outputs.approved == 'true' && needs.runner-check.outputs.available == 'true'
     runs-on: [self-hosted, linux, arm64]
     timeout-minutes: 60
     steps:

python/cactus/convert/cactus_adapters/tensor_io.py

@@ -3,7 +3,7 @@
 from pathlib import Path
 from typing import Optional, Dict, Any, List
 from .weight_patterns import EMBED_NAMES
-from ..model_adapters.naming import GEMMA4_WEIGHT_SCALE
+from ..model_adapters.naming import GEMMA4_WEIGHT_SCALE, gemma4_scale_factor
 
 try:
     import torch
@@ -283,19 +283,7 @@ def save_tensor_with_header(tensor, output_path, precision='INT8', transpose=Fal
             data = data / GEMMA4_WEIGHT_SCALE
 
     if model_type == 'gemma4':
-        filename = output_path.name
-        is_audio_weight = filename.startswith('audio_')
-        if any(x in filename for x in ['input_norm', 'post_attn_norm', 'pre_ffn_norm', 'post_ffn_norm',
-                                       'post_per_layer_norm', 'post_proj_norm']):
-            data = data / GEMMA4_WEIGHT_SCALE
-        elif any(x in filename for x in ['ffn_gate', 'ffn_up', 'per_layer_gate', 'moe_gate_proj', 'moe_up_proj']):
-            data = data * GEMMA4_WEIGHT_SCALE
-        elif 'router_scale' in filename:
-            data = data / GEMMA4_WEIGHT_SCALE
-        elif filename in ('token_embeddings.weights', 'output_weight.weights', 'embed_vision_proj.weights', 'embed_audio_proj.weights'):
-            data = data / GEMMA4_WEIGHT_SCALE
-        elif filename == 'output_norm.weights':
-            data = data * GEMMA4_WEIGHT_SCALE
+        data = data * gemma4_scale_factor(output_path.name)
 
     if precision in ('INT8', 'INT4'):
         filename = output_path.name

python/cactus/convert/model_adapters/naming.py

@@ -23,6 +23,8 @@ class NameMatch:
 GEMMA4_DIV_BASENAMES = {
     "token_embeddings",
     "output_weight",
+    "embed_vision_proj",
+    "embed_audio_proj",
 }
 
 

python/cactus/transpile/component_pipeline.py

@@ -1,10 +1,11 @@
 from __future__ import annotations
 
 import copy
+from contextlib import AbstractContextManager
 from dataclasses import dataclass
 from dataclasses import field
 import re
-from typing import Any
+from typing import Any, Callable
 
 import numpy as np
 import torch
@@ -30,6 +31,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_reparam: Callable[[torch.nn.Module], AbstractContextManager] | None = None
     npu_example_inputs: tuple[torch.Tensor, ...] | None = None
 
 

python/cactus/transpile/model_adapters.py

@@ -16,6 +16,7 @@
 from cactus.transpile.component_pipeline import ComponentModuleSpec
 from cactus.convert.cactus_adapters.tensor_io import CACTUS_MAGIC
 from cactus.convert.cactus_adapters.tensor_io import align_offset
+from cactus.convert.model_adapters.naming import gemma4_scale_factor
 
 
 _GEMMA4_SAFE_TEXT_MLP_PRODUCT_SCALE = 1.0 / 64.0
@@ -1001,6 +1002,30 @@ def _gemma4_compute_native_like_image_features(
     return projected
 
 
+_GEMMA4_NPU_PROJ_WEIGHTS = (
+    ("embed_audio", "embed_audio_proj.weights"),
+    ("embed_vision", "embed_vision_proj.weights"),
+)
+
+
+@contextmanager
+def _gemma4_npu_projection_reparam(module: torch.nn.Module):
+    backbone = getattr(module, "multimodal_backbone", None)
+    scaled: list[tuple[torch.nn.Linear, float]] = []
+    for attr, cactus_name in _GEMMA4_NPU_PROJ_WEIGHTS:
+        projection = getattr(getattr(backbone, attr, None), "embedding_projection", None)
+        factor = gemma4_scale_factor(cactus_name)
+        if isinstance(projection, torch.nn.Linear) and factor != 1.0:
+            scaled.append((projection, factor))
+    for projection, factor in scaled:
+        projection.weight.data.mul_(factor)
+    try:
+        yield
+    finally:
+        for projection, factor in scaled:
+            projection.weight.data.div_(factor)
+
+
 def _gemma4_compute_native_like_audio_features(
     multimodal_backbone: torch.nn.Module,
     input_features: torch.Tensor,
@@ -3745,6 +3770,7 @@ def _require(component: str) -> None:
             metadata={"family": "gemma4", "task": "multimodal_causal_lm_logits"},
             npu_module=vision_encoder_npu,
             npu_runtime_input_count=2,
+            npu_reparam=_gemma4_npu_projection_reparam,
         ))
     if "audio_encoder" in expanded_components:
         specs.append(ComponentModuleSpec(
@@ -3755,6 +3781,7 @@ def _require(component: str) -> None:
             output_keys=("audio_features",),
             graph_meta={**common_graph_meta, "component": "audio_encoder"},
             metadata={"family": "gemma4", "task": "multimodal_causal_lm_logits"},
+            npu_reparam=_gemma4_npu_projection_reparam,
         ))
     if "lm_encoder" in expanded_components:
         if image_features is None or audio_features is None:

python/cactus/transpile/npu/pipeline.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from contextlib import nullcontext
 from pathlib import Path
 
 from .audio import emit_audio_encoder_mlpackage
@@ -58,34 +59,36 @@ def _resolve(per_component: int | None, default: int | None) -> int | None:
         qbits = component_quants[component]
         qdesc = f"int{qbits}" if qbits else "fp16"
         print(f"npu.pipeline: emitting {component} quant={qdesc}")
-        if component == "source_encoder":
-            emitted = emit_fn(
-                module,
-                bundle_root,
-                example_inputs=example_inputs,
-                input_names=tuple(getattr(spec, "input_keys", ()) or ()),
-                filename=filename,
-                quantize_bits=qbits,
-            )
-        elif component == "vision_encoder":
-            names = ("x",) + tuple(spec.input_keys[1:n_runtime])
-            emitted = emit_fn(
-                module,
-                bundle_root,
-                runtime_inputs=tuple(zip(names, example_inputs[:n_runtime])),
-                baked_inputs=example_inputs[n_runtime:],
-                filename=filename,
-                quantize_bits=qbits,
-            )
-        else:
-            emitted = emit_fn(
-                module,
-                bundle_root,
-                example_input=example_inputs[0],
-                baked_inputs=example_inputs[1:],
-                filename=filename,
-                quantize_bits=qbits,
-            )
+        reparam = getattr(spec, "npu_reparam", None)
+        with (reparam(module) if reparam else nullcontext()):
+            if component == "source_encoder":
+                emitted = emit_fn(
+                    module,
+                    bundle_root,
+                    example_inputs=example_inputs,
+                    input_names=tuple(getattr(spec, "input_keys", ()) or ()),
+                    filename=filename,
+                    quantize_bits=qbits,
+                )
+            elif component == "vision_encoder":
+                names = ("x",) + tuple(spec.input_keys[1:n_runtime])
+                emitted = emit_fn(
+                    module,
+                    bundle_root,
+                    runtime_inputs=tuple(zip(names, example_inputs[:n_runtime])),
+                    baked_inputs=example_inputs[n_runtime:],
+                    filename=filename,
+                    quantize_bits=qbits,
+                )
+            else:
+                emitted = emit_fn(
+                    module,
+                    bundle_root,
+                    example_input=example_inputs[0],
+                    baked_inputs=example_inputs[1:],
+                    filename=filename,
+                    quantize_bits=qbits,
+                )
         if emitted:
             results[f"npu_{component}"] = f"components/{emitted}"
     return results