Merge branch 'main' into draft/int3-qmoe-export

Author: titaiwangms

src/mobius/integrations/ort_genai/auto_export.py

@@ -1356,6 +1356,7 @@ def _write_audio_processor_config(
         # 128-dim log-mel frame. Reproduced natively by the ort-extensions
         # ``Gemma4Audio`` op with ``type="raw_frames"`` (pad to a whole number of
         # frames, reshape to (num_tokens, 640)).
+        # Disable the decoder's Whisper-compatible 30-second truncation.
         samples_per_token = getattr(audio, "hidden_size", None) or 640
         processor = {
             "feature_extraction": {
@@ -1364,6 +1365,7 @@ def _write_audio_processor_config(
                         "operation": {
                             "name": "audio_decoder",
                             "type": "AudioDecoder",
+                            "attrs": {"max_samples": 0},
                         }
                     },
                     {
@@ -1384,7 +1386,7 @@ def _write_audio_processor_config(
         proc_filename = "audio_feature_extraction.json"
     elif model_type in _GEMMA4_MODEL_TYPES:
         # Gemma4 USM-style 128-dim log-mel spectrogram via the ort-extensions
-        # ``Gemma4Audio`` op with ``type="log_mel"``.
+        # ``Gemma4LogMel`` op.
         # OrtxCreateSpeechFeatureExtractor requires the feature_extraction.sequence format.
         processor = {
             "feature_extraction": {
@@ -1397,10 +1399,9 @@ def _write_audio_processor_config(
                     },
                     {
                         "operation": {
-                            "name": "gemma4_audio",
-                            "type": "Gemma4Audio",
+                            "name": "gemma4_log_mel",
+                            "type": "Gemma4LogMel",
                             "attrs": {
-                                "type": "log_mel",
                                 "feature_size": 128,
                                 "sampling_rate": 16000,
                                 "frame_length_ms": 20.0,

src/mobius/integrations/ort_genai/auto_export_test.py

@@ -1003,6 +1003,7 @@ def test_audio_gemma4_unified_writes_raw_frames(self, tmp_path):
 
         seq = data["feature_extraction"]["sequence"]
         assert seq[0]["operation"]["type"] == "AudioDecoder"
+        assert seq[0]["operation"]["attrs"] == {"max_samples": 0}
         op = seq[1]["operation"]
         assert op["name"] == "gemma4_audio"
         assert op["type"] == "Gemma4Audio"
@@ -1029,9 +1030,8 @@ def test_audio_gemma4_writes_feature_extraction_json(self, tmp_path):
         seq = data["feature_extraction"]["sequence"]
         assert len(seq) == 2
         assert seq[0]["operation"]["type"] == "AudioDecoder"
-        assert seq[1]["operation"]["type"] == "Gemma4Audio"
+        assert seq[1]["operation"]["type"] == "Gemma4LogMel"
         attrs = seq[1]["operation"]["attrs"]
-        assert attrs["type"] == "log_mel"
         assert attrs["feature_size"] == 128
         assert attrs["sampling_rate"] == 16000
         assert attrs["frame_length_ms"] == 20.0  # noqa: RUF069
@@ -1967,11 +1967,11 @@ class FakeConfig:
         # First op: AudioDecoder
         op0 = seq[0]["operation"]
         assert op0["type"] == "AudioDecoder"
+        assert "max_samples" not in op0.get("attrs", {})
 
-        # Second op: Gemma4Audio (type=log_mel) with expected attrs
+        # Second op: Gemma4LogMel with expected attrs
         op1 = seq[1]["operation"]
-        assert op1["type"] == "Gemma4Audio"
-        assert op1["attrs"]["type"] == "log_mel"
+        assert op1["type"] == "Gemma4LogMel"
         assert op1["attrs"]["feature_size"] == 128
         assert op1["attrs"]["sampling_rate"] == 16000
         assert op1["attrs"]["mel_floor"] == 0.001  # noqa: RUF069
@@ -4198,6 +4198,11 @@ def test_gemma4_unified_native_processor_contract(self, tmp_path):
 
         with open(result["audio_processor"], encoding="utf-8") as f:
             audio_processor = json.load(f)
+        assert audio_processor["feature_extraction"]["sequence"][0]["operation"] == {
+            "name": "audio_decoder",
+            "type": "AudioDecoder",
+            "attrs": {"max_samples": 0},
+        }
         audio_op = audio_processor["feature_extraction"]["sequence"][1]["operation"]
         assert audio_op == {
             "name": "gemma4_audio",

src/mobius/models/gemma4.py

@@ -13,7 +13,8 @@
 - Dual head_dim: local sliding-window layers use ``config.head_dim``; global
   full-attention layers use ``config.global_head_dim``.
 - Dual RoPE: different ``rope_theta`` and ``partial_rotary_factor`` per layer type.
-- Per-layer input gating (disabled when ``hidden_size_per_layer_input == 0``).
+- Per-layer input gating with decoder-matched quantized projections
+  (disabled when ``hidden_size_per_layer_input == 0``).
 - Vision encoder: pre-patchified input ``[B, N, 3*P^2]`` with 2D position lookup,
   bidirectional attention, 4-norm structure, and scale-then-project pooling.
 - Vision projector: scale-free RMSNorm -> Linear (matches ``embed_vision`` weights).
@@ -1577,10 +1578,11 @@ def __init__(self, config: Gemma4Config, layer_idx: int):
 
         self._per_layer_dim = config.hidden_size_per_layer_input
         if self._per_layer_dim > 0:
-            self.per_layer_input_gate = Linear(
+            linear_class = _text_linear_class(config) or Linear
+            self.per_layer_input_gate = linear_class(
                 config.hidden_size, self._per_layer_dim, bias=False
             )
-            self.per_layer_projection = Linear(
+            self.per_layer_projection = linear_class(
                 self._per_layer_dim, config.hidden_size, bias=False
             )
             self.post_per_layer_input_norm = RMSNorm(