Allow use_sliced_mla_proj with Qwix interception quantization.

src/maxtext/configs/types.py

@@ -5758,8 +5758,14 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
       raise ValueError("`share_kv_projections` is not compatible with `fused_qkv`.")
     if self.share_kv_projections and self.attention_type == "mla":
       raise ValueError("`share_kv_projections` is not compatible with `attention_type='mla'`.")
-    if self.use_sliced_mla_proj and (self.quantization or self.use_qwix_quantization):
-      raise ValueError("`use_sliced_mla_proj` is not supported with quantization.")
+    if self.use_sliced_mla_proj and self.quantization:
+      # Qwix interception still quantizes the sliced dot_general; module-level quant
+      # (AQT / Fp8 / TE / batch split) would silently fall back to the unsliced path.
+      if not self.use_qwix_quantization or self.use_batch_split_schedule:
+        raise ValueError(
+            "`use_sliced_mla_proj` with quantization is only supported via Qwix interception "
+            "(use_qwix_quantization=True, use_batch_split_schedule=False)."
+        )
 
     if self.use_manual_quantization and not self.use_batch_split_schedule:
       raise ValueError("manual quantization is only used when `use_batch_split_schedule=True`.")

tests/unit/configs_value_test.py

@@ -1161,16 +1161,44 @@ def test_indexer_cutoff_threshold_remat_policy(self):
     with self.assertRaises(ValueError):
       pyconfig.initialize(argv_invalid)
 
-  def test_sliced_mla_proj_disallows_quantization(self):
-    """Tests that use_sliced_mla_proj=True is incompatible with quantization."""
+  def test_sliced_mla_proj_disallows_non_qwix_quantization(self):
+    """Tests that use_sliced_mla_proj=True is incompatible with module-level (non-Qwix) quantization."""
     argv = [
         "",
         _BASE_CONFIG_PATH,
         "run_name=test",
         "use_sliced_mla_proj=true",
         "quantization=int8",
     ]
-    with self.assertRaises(pydantic.ValidationError):
+    with self.assertRaisesRegex(pydantic.ValidationError, "only supported via Qwix interception"):
+      pyconfig.initialize(argv)
+
+  def test_sliced_mla_proj_allows_qwix_quantization(self):
+    """Tests that use_sliced_mla_proj=True is allowed with Qwix interception quantization."""
+    argv = [
+        "",
+        _BASE_CONFIG_PATH,
+        "run_name=test",
+        "use_sliced_mla_proj=true",
+        "quantization=int8",
+        "use_qwix_quantization=true",
+    ]
+    config = pyconfig.initialize(argv)
+    self.assertTrue(config.use_sliced_mla_proj)
+    self.assertTrue(config.use_qwix_quantization)
+
+  def test_sliced_mla_proj_disallows_qwix_quantization_with_batch_split(self):
+    """Tests that use_sliced_mla_proj=True is incompatible with batch-split Qwix quantization."""
+    argv = [
+        "",
+        _BASE_CONFIG_PATH,
+        "run_name=test",
+        "use_sliced_mla_proj=true",
+        "quantization=fp8_full",
+        "use_qwix_quantization=true",
+        "use_batch_split_schedule=true",
+    ]
+    with self.assertRaisesRegex(pydantic.ValidationError, "only supported via Qwix interception"):
       pyconfig.initialize(argv)
 
   def test_serve_fp8_weight_accepts_valid_config(self):