Keep qwix out of tpu-inference's fused MoE kernel and pre-quantize weights

[sierraisland]

Description

This PR addresses the compile-time Mosaic accumulator crash when serving MaxText MoE models via vLLM (attention=vllm_rpa) under use_qwix_quantization=True with an FP8 recipe:

'tpu.matmul' op Expected matmul acc to be 32-bit ... (vector<512x512xf8E4M3FN>, vector<512x256xf8E4M3FN>, vector<512x256xbf16>)

Context & Prior PRs

1. vllm-project/tpu-inference#3507:
   Initially, the attempt is in tpu-inference by wrapping gmm_v2 with disable_interceptions and forcing acc_dtype=jnp.float32 globally. However, as noted during review, hardcoding f32 accumulator inside tpu-inference impacts performance for native models. The root cause lies on the MaxText side where qwix leaks into the Pallas/tokamax kernel body and weights were unquantized.
2. AI-Hypercomputer/maxtext#5114:
   @wenxindongwork proposed fixing this at the caller site in MaxText by disabling qwix interception around the fused kernel invocation and pre-quantizing weights with blockwise/channelwise scales via quantize_weight_for_fused_moe.

Changes in this PR

This PR incorporates the solution from #5114 and further hardens it based on review feedback:

• Disable qwix interception around fused_moe_func:
  Wraps the call in qwix.disable_interceptions so that tokamax's tiled matmuls launched via pl.kernel are not rewritten into fake-quantized bf16 accumulators.
• Pre-quantize weights outside the kernel:
  When the active rule covers gmm, weights are pre-quantized to float8_e4m3fn (or int8) with [E, blocks, 1, N] scales using qpl.quantize.
• Robust attribute access (Addressing review comment in Keep qwix out of tpu-inference's fused MoE kernel; pre-quantize its weights #5114):
  Uses getattr(rule, ...) with fallback defaults for weight_qtype, tile_size, and weight_calibration_method to prevent potential AttributeErrors when custom rules or test mocks are provided.
• Unit Tests:
  Includes full unit tests in tests/unit/fused_moe_qwix_test.py, including verification of fallback behavior for minimal/mock rules.

Tests

• Setup: Ran on TPU v7x-8 using the MaxText vLLM benchmark setup:
  • Model: Qwen/Qwen3.5-35B-A3B
  • Model Architecture: MaxTextForCausalLM (MODEL_IMPL_TYPE=flax_nnx)
  • MaxText Config: quantization='fp8_e4m3', use_qwix_quantization=True, attention='vllm_rpa'
• Verification:
  • MaxText adapter plugin loaded and registered successfully on startup.
  • Successfully ran MaxText benchmark workload including chunked prefill (2048 tokens) without any MLIRError or EngineCore crashes.

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.

[gemini-code-assist]

Code Review

This pull request implements weight quantization for the fused MoE kernel under the qwix framework, adding helper functions to retrieve quantization rules, quantize expert weights, and bypass qwix interception during execution. It also introduces unit tests to verify these changes. The review feedback suggests renaming the quantized weight variables in moe.py for clarity and using jnp.expand_dims instead of jnp.reshape in quantizations.py to make the dimension expansion more explicit.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!