[copybara-service]

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)

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

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!