[FIX] Fix MoE GMM FP8 accumulator verification error and bump qwix requirement (Maxtext)

[sierraisland]

Description

This PR fixes a fatal Pallas/Mosaic MLIR verification error ('tpu.matmul' op Expected matmul acc to be 32-bit) encountered when running the MaxText vLLM benchmark with FP8 MoE models (e.g. Qwen/Qwen3.5-35B-A3B) and bumps the Qwix requirement to prevent dependency downgrade conflicts.

Root Cause

1. Pallas MLIR Verifier Accumulator Check: When running the MaxText benchmark via the vLLM MaxText adapter (MaxTextForCausalLM / MODEL_IMPL_TYPE=flax_nnx) on the vllm_rpa attention path, maxtext.layers.moe routes MoE execution to tpu_inference.layers.common.fused_moe_gmm.fused_moe_func. TPU MXU hardware requires FP8 matrix multiplications (f8E4M3FN × f8E4M3FN) to accumulate into 32-bit (f32). In recent JAX (0.11+) and libtpu releases, the Mosaic MLIR compiler enforces this check and fails compilation if the accumulator operand is 16-bit.
2. Qwix Interception Bleed into Pallas Kernels: Because 'use_qwix_quantization': True is set in MaxText's quantization configuration, Qwix's thread-local dot_general interception inadvertently intercepted jnp.matmul inside the hand-written Pallas gmm_v2 kernel, dropping preferred_element_type and defaulting the accumulator type to bfloat16.
3. Qwix Version Pinning in requirements.txt: tpu-inference/requirements.txt pinned qwix==0.1.2, which downgraded Qwix upon installation and broke MaxText compatibility (MaxText HEAD imports qwix._src.core.sparsity which requires qwix>=0.1.8, causing register_maxtext_vllm_adapter to fail to load on benchmark server startup).

Changes

1. Disable Qwix Interception inside GMM: Wrapped gmm_v2 inside gmm_wrapper with disable_interceptions to prevent Qwix from rewriting internal Pallas matmuls during MaxText MoE execution.
2. Explicit 32-bit Accumulator: Explicitly passed acc_dtype=jnp.float32 in gmm_wrapper to ensure the TPU MXU uses a 32-bit accumulator for quantized GMM kernels.
3. Bump Qwix Dependency: Updated qwix==0.1.2 to qwix>=0.1.8 in requirements.txt to align with MaxText's sparsity requirements and avoid adapter plugin import failures.

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:

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have made or will make corresponding changes to any relevant documentation.

[github-actions]

Description

Start with a short description of what the PR does and how this is a change from
the past.

The rest of the description includes relevant details and context, examples:

• why is this change being made,
• the problem being solved and any relevant context,
• why this is a good solution,
• some information about the specific implementation,
• shortcomings of the solution and possible future improvements.

If the change fixes a Github issue, please include a link, e.g.,:
FIXES: #123456

Tests

Please describe how you tested this change, and include any instructions and/or
commands to reproduce.

Checklist

Before submitting this PR, please make sure:

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have made or will make corresponding changes to any relevant documentation.

[wenxindongwork]

With this change now all GMM accumulates in f32, i think this is will impact performance for our native model runs.

The real issue is that maxtext models do not properly support fp8 weights, so the fix really should exist on the maxtext side.