feat(quantization): add native FP8 inference (serve_fp8_weight) for DenseGeneral and GMM v2

[Shuwen-Fang]

Description

Implements native FP8 inference (quantization="serve_fp8_weight") for Qwen3.5 (and other FP8-quantized architectures) to execute forward passes directly using quantized FP8 weights without on-the-fly dequantization (b/557372531).

Previously, weights stored/loaded in FP8 were dequantized on-the-fly to BF16 prior to compute, causing substantial memory bandwidth and runtime overhead during inference and RL rollouts. This change allows weights to stay in FP8 and compute directly via native FP8 GEMMs across DenseGeneral (supporting per-tensor, per-channel, and block-wise scaling) and RoutedMoE (dispatching directly to tokamax v2 FP8 GMM kernel).

---

Benchmark & Performance Results

Evaluated on TPU v6e-8 (8 chips, GhostLite) using the per-tensor quantized qwen3.5-35b-a3b-fp8 checkpoint (/dev/shm/maxtext_qwen3.5_35b_fp8_pertensor_v3), running a standalone autoregressive decode forward loop simulating RL rollout forward inference (batch size = 1, sequence length = 64 tokens, scan_layers=True, attention="dot_product").

Performance Summary

Metric	Baseline (quantization="")	Diff (quantization="serve_fp8_weight")	Delta / Speedup
Average Forward Step Time (15 steps)	471.59 ms (±3.07 ms)	259.80 ms (±3.25 ms)	1.81x faster (-45.0% latency)
Step 0 (JIT Compilation)	5.12 s	1.32 s	3.88x faster compile
Output Logits Finite	Yes (all finite)	Yes (all finite)	Parity verified

XProf Trace Viewer Links

• Baseline Session (quantization=""):
  http://xprof.corp.google.com/?session_id=shuwenf-5004636990561406008
  Session ID: shuwenf-5004636990561406008
• Diff Session (quantization="serve_fp8_weight"):
  http://xprof.corp.google.com/?session_id=shuwenf-16130608832799170934
  Session ID: shuwenf-16130608832799170934

Persisted GCS Traces:

• Baseline: gs://test-maxtext-output/shuwenf/xprof/20260911_run/xprof/baseline/
• Diff: gs://test-maxtext-output/shuwenf/xprof/20260911_run/xprof/diff/

---

Tests

Unit Tests

• python -m unittest tests/unit/native_fp8_dot_general_test.py (23/23 passed in 42.8s):
  • Per-tensor, per-channel, and block-wise numerical verification against independent references across all Qwen3.5 weight families (self_attn, linear_attn, MoE expert gate/up/down).
  • Multi-axis output kernel rank-matching tests (e.g. fused QKV (embed, heads, head_dim)).
  • DequantizeWeightPerChannelTest: square matrix 2D per-channel scaling & MoE 3D per-expert scaling.
  • MoENativeFp8PerChannelTest: GMM v2 native FP8 execution with per-channel qpl.QArray.
• pytest tests/unit/moe_test.py -k "SparseMoeNativeGmmPerChannelAxisTest" (1/1 passed in 43.9s):
  • Validates K-axis (input-channel) scale correctly falls back to dequantize in RoutedMoE.
• python tests/unit/configs_value_test.py (4/4 passed in 7.2s):
  • Validates serve_fp8_weight configuration, FP8 weight requirement, Qwix mutual exclusivity, and fused MoE warning.
• python -m unittest tests/unit/nnx_decoders_test.py (53/53 passed in 91.9s):
  • Full decoder forward pass and configuration suite.
• Pre-commit hooks clean: codespell, pyink, pylint (10.00/10), and yamllint all passed.

Numerical Parity & Alignment (Native FP8 OFF vs. ON)

Evaluated full NNXDecoder forward pass across 128 tokens (batch_size=4, seq_len=32, vocab_size=1024) on TPU v6e comparing Post-PR Native OFF (quantization="" / dequantize) vs. Post-PR Native ON (quantization="serve_fp8_weight"):

Metric	Post-PR Native OFF vs. ON
Mean $D_{\text{KL}}$	$4.35 \times 10^{-3}\text{ nats}$ (Max: $1.96 \times 10^{-2}$)
Top-1 Token Match Rate	89.06% (114 / 128 tokens)
Top-5 Candidate Overlap	78.64%
Logit Cosine Similarity	0.9956 (99.56%)
SQNR	20.60 dB

BUGS: b/557372531

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 and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed.

[gemini-code-assist]

Code Review

This pull request introduces a native FP8 inference feature (quantization="serve_fp8_weight") to load and execute FP8 checkpoints directly on TPU without on-the-fly dequantization, yielding significant performance improvements. It adds native_fp8_dot_general utilizing qwix and updates DenseGeneral and MoE layers to bypass standard dequantization when this mode is active. Feedback on the changes highlights two critical runtime issues: a dimension mismatch in moe.py where per_expert_scale is incorrectly reshaped to 4D instead of 3D for a 3D kernel, and a potential ValueError in quantizations.py when unconditionally reshaping kernel_scale when its size is greater than 1.

[gemini-code-assist]

Unconditionally reshaping kernel_scale to (1,) * quantized_kernel.ndim when kernel_scale.ndim != quantized_kernel.ndim will raise a ValueError if kernel_scale has a size greater than 1 (e.g., a 1D per-channel or per-expert scale of size > 1). We should only reshape to all-1s if kernel_scale.size == 1. If kernel_scale.size > 1 and its leading dimensions match the kernel, we should expand it by appending trailing dimensions of size 1 to match the kernel's rank.

Suggested change

	if kernel_scale.ndim != quantized_kernel.ndim:
	kernel_scale = kernel_scale.reshape((1,) * quantized_kernel.ndim)
	if kernel_scale.ndim != quantized_kernel.ndim:
	if kernel_scale.size == 1:
	kernel_scale = kernel_scale.reshape((1,) * quantized_kernel.ndim)
	elif kernel_scale.ndim < quantized_kernel.ndim and quantized_kernel.shape[: kernel_scale.ndim] == kernel_scale.shape:
	kernel_scale = kernel_scale.reshape(kernel_scale.shape + (1,) * (quantized_kernel.ndim - kernel_scale.ndim))