Explicitly pass qwix config for deepseek batch split

[shuningjin]

Description

Fix: b/489513157

Overview

PR 3182 switched deepseek batch split config to use pure JAX, which broke the existing Qwix integration that relies on the Qwix interception feature. As a workaround, PR 3319 explicitly pass the QwixRule to gmm kernels.

This PR furthers the workaround. We explicitly pass the QwixQuantization to individual layers, in a similar vein as existing AQTQuantization.

# flow1 (existing), aqt plumb
DeepSeekMoE, non Batch-Split Version (MoE & MLA): flax, aqt plumb

# flow2 (existing), qwix intercept
DeepSeekMoE, non Batch-Split Version (MoE & MLA): flax, intercept

# flow3 (goal of this PR), qwix plumb
DeepSeekMoE, Batch-Split Version (MoE & MLA): jax, qwix plumb

In summary,

• when deepseek + batch_split: we will use flow3
• any other cases: remain the same as before, either flow1 or flow2

Main Changes

• quantization.py

  • Add QwixQuantization, similar to AQTQuantization
    • implements dot_general_cls and einsum, with QwixDotGeneral and QwixEinsum
    • use qwix._src.core.dot_general_qt.dot_general_qt config, with "fp8_full" and calibration method
  • when use qwix and batch split, not use intercept (qwix.quantize_model) but return QwixQuantization

• deepseek_batchsplit.py

  • Allow dot to use quantized version quant.dot_general_cls
  • pass quant to methods

Tests

locally on v5p-8 for deepseek3-test (trim down to 2 dense layer + 2 MoE layer)

• command: https://paste.googleplex.com/5927802336575488
• log: https://paste.googleplex.com/4739506990350336
• profile: https://xprof.corp.google.com/trace_viewer/shuningjin-6983180425227296013

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• 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, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[codecov]

Codecov Report

❌ Patch coverage is 45.09804% with 28 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/models/deepseek_batchsplit.py	15.78%	16 Missing ⚠️
src/maxtext/layers/quantizations.py	62.50%	11 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!