Exclude RMSNormGated from 4bit weight quantization (#277)

python/src/coreai_models/export/presets.py

@@ -33,6 +33,7 @@
     "coreai_models.primitives.macos.rope.RoPE": None,
     "coreai_models.primitives.macos.rms_norm.RMSNorm": None,
     "coreai_models.primitives.macos.rms_norm.RMSNormPlusOne": None,
+    "coreai_models.primitives.macos.rms_norm.RMSNormGated": None,
 }
 
 # Embedding modules excluded from iOS palettization.

python/tests/test_model_units/test_primitives/test_macos/test_rms_norm.py

@@ -132,6 +132,27 @@ def test_rms_norm_plus_one(self):
         torch.testing.assert_close(out, ref_out, rtol=1e-3, atol=1e-3)
 
 
+@pytest.mark.skipif(not HAS_COREAI, reason="coreai-torch not available")
+class TestRMSNormGatedPresetExclusion:
+    """RMSNormGated must be excluded from 4bit weight quantization (#277)."""
+
+    def test_4bit_preset_excludes_gated_rms_norm(self):
+        from coreai_models.export.presets import get_preset
+
+        quant_config = get_preset("4bit")["torch_quantization_config"]
+        gated_path = f"{CoreaiTorchRMSNormGated.__module__}.{CoreaiTorchRMSNormGated.__qualname__}"
+
+        assert quant_config["module_type_configs"].get(gated_path, "missing") is None
+
+        # The exclusion is load-bearing: the global weight spec blocks along axis 1,
+        # which is out of bounds for RMSNormGated's rank-1 weight.
+        weight_spec = quant_config["global_config"]["op_state_spec"]["weight"]
+        global_axis = weight_spec["granularity"]["axis"]
+        weight = CoreaiTorchRMSNormGated(dim=8).weight
+        assert weight.ndim == 1
+        assert global_axis >= weight.ndim
+
+
 # =============================================================================
 # Functional-parity tests
 # =============================================================================