diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index 5ea638f426..e649f86212 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -3363,12 +3363,23 @@ def fused_moe_matmul( self.config.decoder_block not in (ctypes.DecoderBlockType.LLAMA4, ctypes.DecoderBlockType.GEMMA4) ) - output_2d = fused_moe_func( + # The fused kernel quantizes for itself: expert weights go in pre-quantized with + # per-block scales (when the qwix rule covers the grouped matmul) and the kernel + # quantizes the activations in-kernel, accumulating in f32. The call runs outside + # qwix's interception: qwix only guards pallas_call, while the kernel is launched + # through pl.kernel, so under a QtProvider its tiled matmuls would otherwise be + # fake-quantized into fp8 x fp8 with a bf16 accumulator, which Mosaic rejects. + rule = quantizations.get_fused_moe_rule() + quantized_w1, w1_scale = quantizations.quantize_weight_for_fused_moe(fused_kernel, rule) + quantized_w2, w2_scale = quantizations.quantize_weight_for_fused_moe(wo_kernel, rule) + fused_moe = quantizations.without_qwix_interception(fused_moe_func) + + output_2d = fused_moe( hidden_states=hidden_states, - w1=fused_kernel, - w2=wo_kernel, - w1_scale=None, - w2_scale=None, + w1=quantized_w1, + w2=quantized_w2, + w1_scale=w1_scale, + w2_scale=w2_scale, w1_bias=None, w2_bias=None, gating_output=gating_output, diff --git a/src/maxtext/layers/quantizations.py b/src/maxtext/layers/quantizations.py index 0ef3aeca55..d604b98cb0 100644 --- a/src/maxtext/layers/quantizations.py +++ b/src/maxtext/layers/quantizations.py @@ -31,6 +31,7 @@ from qwix._src.core import numerics from qwix._src.core import dot_general_qt from qwix._src.core import sparsity +from qwix._src import interception as qwix_interception import jax import jax.numpy as jnp @@ -66,6 +67,7 @@ def _safe_find_param(x, ptq_array_type=None): from maxtext.configs.types import TeCommGemmOverlapPolicy from maxtext.common.common_types import DType, Config from maxtext.inference.kvcache import KVQuant +from maxtext.utils import max_logging # Params used to define mixed precision quantization configs DEFAULT = "__default__" # default config @@ -970,6 +972,74 @@ def maybe_quantize_model(model, config): return model +# Quantized weight dtypes that tpu-inference's fused MoE kernel (gmm_v2) takes with +# per-block scales and dequantizes in-kernel. Any other qtype keeps the expert weights +# unquantized on the fused path. +FUSED_MOE_KERNEL_WEIGHT_QTYPES = (jnp.dtype(jnp.float8_e4m3fn), jnp.dtype(jnp.int8)) + + +def get_fused_moe_rule() -> qwix.QuantizationRule | None: + """Returns the qwix rule that governs the fused MoE grouped matmul, if any. + + The fused kernel is the same grouped matmul that MaxText's own megablox / ragged_dot + paths implement, so it takes its rule from the "gmm" op exactly like they do. Rules + that only list "dot_general" (the plain fp8 / int8 recipes) leave the experts + unquantized on every MoE path, including this one. + """ + return qpl.get_current_rule("gmm") + + +def quantize_weight_for_fused_moe( + kernel: jax.Array, rule: qwix.QuantizationRule | None +) -> tuple[jax.Array, jax.Array | None]: + """Quantizes an [E, K, N] expert weight for tpu-inference's fused MoE kernel. + + The kernel contracts over K, dequantizes after the matmul with scales laid out as + [E, num_blocks, 1, N] (blocks along K), and quantizes the activations itself with a + 32-bit accumulator. Scale granularity follows the qwix rule: channelwise over (E, N), + plus blockwise along K when the rule sets tile_size. + + Returns the (possibly unchanged) weight and its scale, or None when the rule does not + ask for a weight dtype the kernel supports. + """ + + weight_qtype = getattr(rule, "weight_qtype", None) if rule is not None else None + if weight_qtype is None: + return kernel, None + qtype = jnp.dtype(weight_qtype) + if qtype not in FUSED_MOE_KERNEL_WEIGHT_QTYPES: + max_logging.log(f"fused MoE kernel does not take {qtype} weights; keeping expert weights in {kernel.dtype}.") + return kernel, None + if kernel.ndim != 3: + raise ValueError(f"fused MoE expert weights must be [E, K, N], got {kernel.shape}") + + tile_size = getattr(rule, "tile_size", None) + tiled_axes = {1: tile_size} if tile_size else {} + cal_method = getattr(rule, "weight_calibration_method", None) + quant_kwargs = {"calibration_method": cal_method} if cal_method is not None else {} + + quantized = qpl.quantize( + kernel, + qtype, + channelwise_axes=(0, 2), + tiled_axes=tiled_axes, + scale_dtype=jnp.float32, + **quant_kwargs, + ) + scale = jnp.expand_dims(quantized.scale, axis=2) + return quantized.qvalue, scale + + +def without_qwix_interception(fn: Callable) -> Callable: + """Returns `fn` wrapped so that qwix's op interception is off while it runs. + + For ops that do their own quantization (tpu-inference's fused MoE), the qwix rule is + applied once up front to their inputs and the op itself is opaque to qwix: neither + the routing math around the kernel nor anything traced inside it gets rewritten. + """ + return qwix_interception.disable_interceptions(fn) + + def _cast_reduced_from(arr, reduced_arr): aval = jax.typeof(reduced_arr) # In shard map diff --git a/tests/unit/fused_moe_qwix_test.py b/tests/unit/fused_moe_qwix_test.py new file mode 100644 index 0000000000..776d256f89 --- /dev/null +++ b/tests/unit/fused_moe_qwix_test.py @@ -0,0 +1,264 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# pylint: disable=missing-class-docstring,missing-function-docstring,unbalanced-tuple-unpacking + +"""Tests for the qwix boundary around tpu-inference's fused MoE kernel.""" + +import sys +import types +import unittest +from unittest import mock + +from flax import nnx +import jax +import jax.numpy as jnp +import numpy as np +import qwix + +from maxtext.common import common_types as ctypes +from maxtext.layers import moe +from maxtext.layers import quantizations + + +def _assert_error_within(actual, ref, bound): + """Every |actual - ref| stays under `bound`, which broadcasts against the arrays.""" + err = np.abs(np.asarray(actual - ref)) + np.testing.assert_array_less(err, np.broadcast_to(np.asarray(bound) + 1e-6, err.shape)) + + +def _fp8_rule(**kwargs): + return qwix.QtRule( + module_path=".*", + weight_qtype=jnp.float8_e4m3fn, + act_qtype=jnp.float8_e4m3fn, + op_names=("dot_general", "gmm", "ragged_dot"), + **kwargs, + ) + + +class QuantizeWeightForFusedMoeTest(unittest.TestCase): + """quantize_weight_for_fused_moe produces what the fused kernel takes.""" + + def setUp(self): + super().setUp() + self.num_experts, self.in_dim, self.out_dim = 4, 256, 128 + self.kernel = jax.random.normal(jax.random.key(0), (self.num_experts, self.in_dim, self.out_dim), jnp.bfloat16) + + def test_no_rule_keeps_weight(self): + weight, scale = quantizations.quantize_weight_for_fused_moe(self.kernel, None) + self.assertIs(weight, self.kernel) + self.assertIsNone(scale) + + def test_rule_without_weight_qtype_keeps_weight(self): + rule = qwix.QtRule(module_path=".*", act_qtype=jnp.float8_e4m3fn) + weight, scale = quantizations.quantize_weight_for_fused_moe(self.kernel, rule) + self.assertIs(weight, self.kernel) + self.assertIsNone(scale) + + def test_unsupported_qtype_keeps_weight(self): + rule = qwix.QtRule(module_path=".*", weight_qtype=jnp.float8_e5m2) + weight, scale = quantizations.quantize_weight_for_fused_moe(self.kernel, rule) + self.assertIs(weight, self.kernel) + self.assertIsNone(scale) + + def test_fp8_channelwise(self): + weight, scale = quantizations.quantize_weight_for_fused_moe(self.kernel, _fp8_rule()) + self.assertEqual(weight.dtype, jnp.float8_e4m3fn) + self.assertEqual(weight.shape, self.kernel.shape) + # one scale per (expert, output column), laid out as the kernel expects + self.assertEqual(scale.shape, (self.num_experts, 1, 1, self.out_dim)) + self.assertEqual(scale.dtype, jnp.float32) + dequantized = weight.astype(jnp.float32) * scale[:, 0] + ref = self.kernel.astype(jnp.float32) + # fp8 e4m3 keeps 3 mantissa bits: the error is bounded by half an ulp of the column max + col_max = jnp.max(jnp.abs(ref), axis=1, keepdims=True) + _assert_error_within(dequantized, ref, col_max / 16) + + def test_fp8_blockwise_from_tile_size(self): + weight, scale = quantizations.quantize_weight_for_fused_moe(self.kernel, _fp8_rule(tile_size=64)) + self.assertEqual(weight.dtype, jnp.float8_e4m3fn) + self.assertEqual(scale.shape, (self.num_experts, self.in_dim // 64, 1, self.out_dim)) + blocks = weight.astype(jnp.float32).reshape(self.num_experts, self.in_dim // 64, 64, self.out_dim) + dequantized = (blocks * scale).reshape(self.kernel.shape) + ref = self.kernel.astype(jnp.float32) + block_max = jnp.max(jnp.abs(ref.reshape(blocks.shape)), axis=2, keepdims=True) + _assert_error_within(dequantized.reshape(blocks.shape), ref.reshape(blocks.shape), block_max / 16) + + def test_int8_channelwise(self): + rule = qwix.QtRule(module_path=".*", weight_qtype=jnp.int8, op_names=("gmm",)) + weight, scale = quantizations.quantize_weight_for_fused_moe(self.kernel, rule) + self.assertEqual(weight.dtype, jnp.int8) + self.assertEqual(scale.shape, (self.num_experts, 1, 1, self.out_dim)) + dequantized = weight.astype(jnp.float32) * scale[:, 0] + ref = self.kernel.astype(jnp.float32) + col_max = jnp.max(jnp.abs(ref), axis=1, keepdims=True) + _assert_error_within(dequantized, ref, col_max / 127) + + def test_rejects_non_3d_weight(self): + with self.assertRaises(ValueError): + quantizations.quantize_weight_for_fused_moe(self.kernel[0], _fp8_rule()) + + def test_rule_with_missing_attributes(self): + """getattr fallback handles rules missing tile_size or calibration_method.""" + + class MinimalRule: + weight_qtype = jnp.float8_e4m3fn + + weight, scale = quantizations.quantize_weight_for_fused_moe(self.kernel, MinimalRule()) + self.assertEqual(weight.dtype, jnp.float8_e4m3fn) + self.assertEqual(scale.shape, (self.num_experts, 1, 1, self.out_dim)) + + +class WithoutQwixInterceptionTest(unittest.TestCase): + """The boundary hides the qwix rule (and its op rewriting) from whatever runs inside.""" + + def test_rule_visible_outside_and_hidden_inside(self): + seen = {} + + class Layer(nnx.Module): + + def __init__(self): + self.w = nnx.Param(jnp.ones((32, 64), jnp.bfloat16)) + + def __call__(self, x): + seen["outside"] = quantizations.get_fused_moe_rule() + + def opaque(y): + seen["inside"] = quantizations.get_fused_moe_rule() + return jnp.dot(y, self.w[...], preferred_element_type=jnp.float32) + + out = quantizations.without_qwix_interception(opaque)(x) + seen["after"] = quantizations.get_fused_moe_rule() + return out + + x = jax.random.normal(jax.random.key(0), (8, 32), jnp.bfloat16) + layer = qwix.quantize_model(Layer(), qwix.QtProvider([_fp8_rule()]), x) + jaxpr = str(jax.make_jaxpr(layer)(x)) + + self.assertIsNotNone(seen["outside"]) + self.assertIsNone(seen["inside"]) + self.assertIsNotNone(seen["after"]) + # the dot inside the boundary is left alone: no fp8 casts anywhere in the trace + self.assertNotIn("float8", jaxpr) + + +class FusedMoeMatmulTest(unittest.TestCase): + """RoutedMoE.fused_moe_matmul hands the kernel pre-quantized weights, outside qwix.""" + + num_experts, top_k, emb_dim, mlp_dim, tokens = 4, 2, 64, 128, 16 + + def _fake_tpu_inference(self, calls): + """Stand-in for the tpu-inference modules fused_moe_matmul imports.""" + + def fused_moe_func(**kwargs): + calls.append(dict(kwargs, rule_inside=quantizations.get_fused_moe_rule())) + return jnp.zeros((kwargs["hidden_states"].shape[0], self.emb_dim), jnp.bfloat16) + + envs = types.SimpleNamespace( + ENABLE_RS_KERNEL=False, USE_GMM_FUSED_RS_KERNEL=False, ONEHOT_MOE_PERMUTE_THRESHOLD=0, VLLM_MOE_CHUNK_SIZE=0 + ) + pkg = types.ModuleType("tpu_inference") + pkg.envs = envs + layers = types.ModuleType("tpu_inference.layers") + common = types.ModuleType("tpu_inference.layers.common") + gmm = types.ModuleType("tpu_inference.layers.common.fused_moe_gmm") + gmm.fused_moe_func = fused_moe_func + return { + "tpu_inference": pkg, + "tpu_inference.envs": envs, + "tpu_inference.layers": layers, + "tpu_inference.layers.common": common, + "tpu_inference.layers.common.fused_moe_gmm": gmm, + } + + def _make_layer(self, calls): + test = self + config = types.SimpleNamespace( + mlp_activations=("silu",), + routed_score_func="softmax", + norm_topk_prob=True, + decoder_block=ctypes.DecoderBlockType.MIXTRAL, + ) + + class Layer(nnx.Module): + """Drives RoutedMoE.fused_moe_matmul with just the attributes it reads.""" + + def __init__(self): + keys = jax.random.split(jax.random.key(0), 3) + self.w0 = nnx.Param(jax.random.normal(keys[0], (test.num_experts, test.emb_dim, test.mlp_dim), jnp.bfloat16)) + self.w1 = nnx.Param(jax.random.normal(keys[1], (test.num_experts, test.emb_dim, test.mlp_dim), jnp.bfloat16)) + self.wo = nnx.Param(jax.random.normal(keys[2], (test.num_experts, test.mlp_dim, test.emb_dim), jnp.bfloat16)) + self.config = config + self.num_experts = test.num_experts + self.num_experts_per_tok = test.top_k + self.mesh = None + + def get_expert_parallelism_size(self): + return 1 + + def __call__(self, inputs, gate_logits): + with mock.patch.dict(sys.modules, test._fake_tpu_inference(calls)): + out, _, _ = moe.RoutedMoE.fused_moe_matmul(self, inputs, gate_logits, self.wo[...], self.w0[...], self.w1[...]) + return out + + return Layer() + + def _inputs(self): + inputs = jax.random.normal(jax.random.key(1), (1, self.tokens, self.emb_dim), jnp.bfloat16) + gate_logits = jax.random.normal(jax.random.key(2), (1, self.tokens, self.num_experts), jnp.float32) + return inputs, gate_logits + + def test_unquantized_without_qwix(self): + calls = [] + inputs, gate_logits = self._inputs() + self._make_layer(calls)(inputs, gate_logits) + (call,) = calls + self.assertEqual(call["w1"].dtype, jnp.bfloat16) + self.assertEqual(call["w1"].shape, (self.num_experts, self.emb_dim, 2 * self.mlp_dim)) + self.assertIsNone(call["w1_scale"]) + self.assertIsNone(call["w2_scale"]) + + def test_fp8_rule_prequantizes_weights_outside_qwix(self): + calls = [] + inputs, gate_logits = self._inputs() + layer = qwix.quantize_model(self._make_layer(calls), qwix.QtProvider([_fp8_rule()]), inputs, gate_logits) + calls.clear() + layer(inputs, gate_logits) + (call,) = calls + # weights arrive quantized, with scales in the kernel's [E, blocks, 1, N] layout + self.assertEqual(call["w1"].dtype, jnp.float8_e4m3fn) + self.assertEqual(call["w2"].dtype, jnp.float8_e4m3fn) + self.assertEqual(call["w1_scale"].shape, (self.num_experts, 1, 1, 2 * self.mlp_dim)) + self.assertEqual(call["w2_scale"].shape, (self.num_experts, 1, 1, self.emb_dim)) + # and the kernel itself runs outside qwix's interception + self.assertIsNone(call["rule_inside"]) + + def test_dot_general_only_rule_keeps_experts_unquantized(self): + calls = [] + inputs, gate_logits = self._inputs() + rule = qwix.QtRule( + module_path=".*", weight_qtype=jnp.float8_e4m3fn, act_qtype=jnp.float8_e4m3fn, op_names=("dot_general",) + ) + layer = qwix.quantize_model(self._make_layer(calls), qwix.QtProvider([rule]), inputs, gate_logits) + calls.clear() + layer(inputs, gate_logits) + (call,) = calls + self.assertEqual(call["w1"].dtype, jnp.bfloat16) + self.assertIsNone(call["w1_scale"]) + self.assertIsNone(call["rule_inside"]) + + +if __name__ == "__main__": + unittest.main()