diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 73c4028f09..6d164091ff 100644 --- a/src/maxtext/configs/types.py +++ b/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`.") diff --git a/tests/unit/configs_value_test.py b/tests/unit/configs_value_test.py index 2d618df899..7f8589489c 100644 --- a/tests/unit/configs_value_test.py +++ b/tests/unit/configs_value_test.py @@ -1161,8 +1161,8 @@ 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, @@ -1170,7 +1170,35 @@ def test_sliced_mla_proj_disallows_quantization(self): "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):