From 148624808ace201465742c71eba4dab13d91f600 Mon Sep 17 00:00:00 2001 From: Suraj Kolla Date: Wed, 7 Oct 2026 00:58:57 -0700 Subject: [PATCH] Allow `use_sliced_mla_proj` with Qwix interception quantization. `use_sliced_mla_proj` slices the single `wq_b` / `wkv_b` kernel param before a plain `lax.dot_general` (see `DenseGeneral(slice_bounds=...)`). With Qwix interception quantization (`use_qwix_quantization=True`), the layers keep `quant=None`, and Qwix intercepts that dot_general by module path. So the sliced projections are still quantized in both forward and backward. The current config check rejects every quantized config, which blocks sliced MLA projections for Qwix-quantized (e.g. fp8) training. # Tests Updated `tests/unit/configs_value_test.py` (non-Qwix quantization still rejected, Qwix allowed, Qwix + batch split rejected). #### Auxiliary Test (similar to [PR#5141](https://github.com/AI-Hypercomputer/maxtext/pull/5141)) Sliced MLA (`use_sliced_mla_proj`) + Qwix: interception correct. [internal link] [full [internal link] ``` # use_sliced_mla_proj=false [QWIX] module='decoder/moe_layers/self_attention/wq_b' op=dot_general0 rule=0 --- [QWIX] module='decoder/moe_layers/self_attention/wkv_b' op=dot_general0 rule=0 --- [QWIX] module='mtp_block/mtp_layer_1/mtp_1_transformer_layer/self_attention/wq_b' op=dot_general0 rule=0 --- [QWIX] module='mtp_block/mtp_layer_1/mtp_1_transformer_layer/self_attention/wkv_b' op=dot_general0 rule=0 # use_sliced_mla_proj=true [QWIX] module='decoder/moe_layers/self_attention/wq_b' op=dot_general0 rule=0 [QWIX] module='decoder/moe_layers/self_attention/wq_b' op=dot_general1 rule=0 --- [QWIX] module='decoder/moe_layers/self_attention/wkv_b' op=dot_general0 rule=0 [QWIX] module='decoder/moe_layers/self_attention/wkv_b' op=dot_general1 rule=0 --- [QWIX] module='mtp_block/mtp_layer_1/mtp_1_transformer_layer/self_attention/wq_b' op=dot_general0 rule=0 [QWIX] module='mtp_block/mtp_layer_1/mtp_1_transformer_layer/self_attention/wq_b' op=dot_general1 rule=0 --- [QWIX] module='mtp_block/mtp_layer_1/mtp_1_transformer_layer/self_attention/wkv_b' op=dot_general0 rule=0 [QWIX] module='mtp_block/mtp_layer_1/mtp_1_transformer_layer/self_attention/wkv_b' op=dot_general1 rule=0 ``` # Checklist - [X] I have performed a self-review of my code. - [X] I have necessary comments in my code, particularly in hard-to-understand areas. - [X] I have run end-to-end tests tests and provided workload links above if applicable. - [X] I have made or will make corresponding changes to the doc if needed. PiperOrigin-RevId: 994917175 --- src/maxtext/configs/types.py | 10 ++++++++-- tests/unit/configs_value_test.py | 34 +++++++++++++++++++++++++++++--- 2 files changed, 39 insertions(+), 5 deletions(-) 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):