From 0bf1938a28022685af31c8480f1b698f875e76e9 Mon Sep 17 00:00:00 2001 From: Shuning Jin Date: Thu, 29 Oct 2026 18:48:56 +0000 Subject: [PATCH 1/4] MoE: native FP8 QAT for routed-expert GMMs (FP8 forward + FP8-weight backward dLHS), and fix the dropless full-buffer gather index Adds a native-QAT path on top of fake QAT (fp8_moe_fake_quant + fp8_fake_quant_hoist=gather8): - fp8_moe_native_fwd: the wi/wo forward GMMs run FP8 x FP8 on the MXU from the gathered FP8 weights and per-channel scales; gmm_v2 quantizes the activation in the kernel (fp8_moe_native_fwd_lhs_quant, fp8_moe_native_fwd_acc_dtype). - fp8_moe_native_bwd_dlhs: the backward dLHS GMMs read the FP8 weights and dequantize each tile in VMEM, fused into the first matmul of each block (rhs_contract_dequant); bit-identical to the XLA dequant backward. Unsupported tiles fall back to the XLA dequant and log a warning. - Fp8NativeWeight(anchor, codes, scale) carries the weight through shard_map; the weight gradient is returned as the anchor's cotangent (straight-through into the master weight). - A RoutedMoE call with fp8_moe_native_fwd but no gather8 absmax (outside the layer scan) raises instead of silently running fake QAT. - ragged_sort: clamp the masked (-1) indices of the full-buffer ring gather-reduces; the unclamped rows halted the SparseCore on dropless replay. The clamp applies on every full-buffer ring gather-reduce, not only with moe_routing_index_math=count; it leaves outputs unchanged. - gmm_v2: new lhs quantization kwargs (lhs_quant_block_size, lhs_f32_scale, lhs_hoist_quant, rhs_scale_first) and GmmConfigs fields; rejects a negative lhs_quant_block_size. - Tests: fp8_moe_native_test (flags, config validator, outside-scan guard, dequantize_fp8_codes bit-identical to dequantize_weight_codes_ste, lhs_quant config values match ops._LHS_QUANT_MODES); ragged_sort_full_buffer_indices_test (cpu_only: it passes unclamped indices on purpose). All new flags default to off; with them off, the lowered program is unchanged (lowered TPU StableHLO + Mosaic kernels compared: 12 GMM-op programs and 3 model fwd+grad programs identical). --- src/maxtext/configs/base.yml | 13 + src/maxtext/configs/types.py | 83 +++ src/maxtext/kernels/megablox/__init__.py | 4 + src/maxtext/kernels/megablox/ops.py | 608 +++++++++++++++++- .../pallas_mosaic_tpu_v2_gmm_kernel.py | 345 ++++++++-- src/maxtext/kernels/ragged/ragged_sort.py | 14 +- src/maxtext/layers/moe.py | 175 ++++- tests/unit/fp8_moe_native_test.py | 204 ++++++ .../ragged_sort_full_buffer_indices_test.py | 185 ++++++ 9 files changed, 1575 insertions(+), 56 deletions(-) create mode 100644 tests/unit/fp8_moe_native_test.py create mode 100644 tests/unit/ragged_sort_full_buffer_indices_test.py diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 2e05c57bc1..c8bfc68002 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -347,6 +347,19 @@ fp8_moe_fake_quant: false # once for the scanned layer stack, before the scan. gather8: also all-gather the fp8 codes and dequantize after the # gather. Env PARAMBOLE_FQ_HOIST (0, absmax, 1 = gather8) selects a mode while this is off. fp8_fake_quant_hoist: "off" +# With fp8_moe_fake_quant and fp8_fake_quant_hoist=gather8, run the routed-expert wi/wo forward GMMs in FP8 (gathered +# codes + per-channel scales into gmm_v2, activation quantized in the kernel). Backward: bf16 DLHS / TGMM at the +# dequantized weight, unless fp8_moe_native_bwd_dlhs dequantizes it inside the DLHS kernel. +# Needs sparse_matmul, moe_wi_gmm_fused_swiglu, moe_wo_gmm_fused_dswiglu, use_tokamax_gmm and use_gmm_v2, and runs only +# inside the decoder's layer scan (where gather8's absmax is attached): a routed-MoE call outside it raises. +fp8_moe_native_fwd: false +# gmm_v2 accumulator across K blocks for fp8_moe_native_fwd: "" (kernel default, bfloat16), "bfloat16" or "float32". +fp8_moe_native_fwd_acc_dtype: "" +# In-kernel lhs quantization scheme for fp8_moe_native_fwd ("" = default block512_bf16; see types.py). +fp8_moe_native_fwd_lhs_quant: "" +# fp8_moe_native_fwd: the backward DLHS GMMs read the fp8 weight codes and dequantize them in VMEM inside gmm_v2 +# (bit-identical; no XLA dequantization of the expert weights). Needs moe_gmm_v2_dlhs_transpose_rhs and fixed tiles. +fp8_moe_native_bwd_dlhs: false # Trainer side: the rollout serves fp8_moe experts, so weight sync quantizes them (and emits their scales). rollout_fp8_moe: false load_balance_loss_weight: 0.0 # weight for the load balance loss diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index bf988eebb0..2a7556c527 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -1937,6 +1937,89 @@ def _off_for_false_fp8_fake_quant_hoist(v: Any) -> Any: _off_for_false_fp8_fake_quant_hoist ) + fp8_moe_native_fwd: bool = Field( + False, + description=( + "Native FP8 forward for the routed experts. With fp8_moe_fake_quant and fp8_fake_quant_hoist=gather8," + " the wi and wo forward GMMs run on the MXU in float8_e4m3fn: the gathered weight codes and their" + " per-channel scales feed gmm_v2 directly (no dequantization), and gmm_v2 quantizes the bf16 activation" + " in the kernel with a dynamic per-row absmax scale (see fp8_moe_native_fwd_lhs_quant). The backward runs" + " bf16 DLHS / TGMM at the dequantized weight (straight-through weight gradient), unless" + " fp8_moe_native_bwd_dlhs moves the DLHS dequantization into the kernel. Runs only inside the decoder's" + " layer scan, where gather8's absmax is attached; a RoutedMoE call outside it (unscanned layers, KV-cache" + " decode) raises instead of silently running fake QAT." + ), + ) + fp8_moe_native_fwd_acc_dtype: Literal["", "bfloat16", "float32"] = Field( + "", + description=( + "fp8_moe_native_fwd: gmm_v2's accumulator across K blocks. '' keeps the kernel default (bfloat16 when the" + " lhs is quantized in the kernel, as in the rollout's kernel); 'float32' accumulates in float32." + ), + ) + fp8_moe_native_fwd_lhs_quant: Literal[ + "", + "block512_bf16", + "block512_bf16_rhs1", + "block512_f32", + "block512_f32_hoist", + "fullk_bf16", + "fullk_bf16_hoist", + "fullk_f32", + "fullk_f32_hoist", + ] = Field( + "", + description=( + "fp8_moe_native_fwd: in-kernel lhs (activation) quantization scheme. 'block512_*': one scale per row and" + " 512-column block ('' / 'block512_bf16' is gmm_v2's default); 'fullk_*': one scale per row over the full" + " K dim. '*_bf16': bf16 absmax and reciprocal; '*_f32': float32 max(absmax, 1e-12) / 448 and a true" + " division (with fullk, the rollout's FP8 MoE kernel's scheme). 'block512_bf16_rhs1': the weight scale is" + " applied to the block accumulator before the activation scale. '*_hoist': each lhs block is quantized once" + " per tile instead of per mxu-wide n slice; same values and, since the compiler deduplicates the repeated" + " quantization, the same speed." + ), + ) + fp8_moe_native_bwd_dlhs: bool = Field( + False, + description=( + "fp8_moe_native_fwd: the backward DLHS GMMs (wi, and wo's fused-dSwiGLU DLHS, including the ring-of-experts" + " one) read the float8 weight codes and dequantize each weight block in VMEM inside gmm_v2, with the same" + " expression (codes.astype(float32) * scale).astype(bf16), instead of an XLA dequantization of every" + " expert weight before the kernel. Bit-identical results; the weight-gradient TGMMs are unchanged. Needs" + " moe_gmm_v2_dlhs_transpose_rhs and fixed tiles (use_gmm_v2_heuristic_tiling=False)." + ), + ) + + @model_validator(mode="after") + def validate_fp8_moe_native_fwd(self) -> "MoEGeneral": + """Rejects fp8_moe_native_fwd / fp8_moe_native_bwd_dlhs without the config they require. + + The FP8 GMMs replace the gather8 dequantization on the fused gmm_v2 SwiGLU / dSwiGLU path only, and the + in-kernel DLHS dequantization needs the transposed DLHS rhs and fixed tiles. + """ + if self.fp8_moe_native_fwd: + _reject_unless( + "fp8_moe_native_fwd=True", + [ + (self.fp8_moe_fake_quant, "fp8_moe_fake_quant=True"), + (self.fp8_fake_quant_hoist == "gather8", "fp8_fake_quant_hoist=gather8"), + (self.sparse_matmul, "sparse_matmul=True"), + (self.use_tokamax_gmm and self.use_gmm_v2, "use_tokamax_gmm=True and use_gmm_v2=True"), + (self.moe_wi_gmm_fused_swiglu, "moe_wi_gmm_fused_swiglu=True"), + (self.moe_wo_gmm_fused_dswiglu, "moe_wo_gmm_fused_dswiglu=True"), + ], + ) + if self.fp8_moe_native_bwd_dlhs: + _reject_unless( + "fp8_moe_native_bwd_dlhs=True", + [ + (self.fp8_moe_native_fwd, "fp8_moe_native_fwd=True"), + (self.moe_gmm_v2_dlhs_transpose_rhs, "moe_gmm_v2_dlhs_transpose_rhs=True"), + (not self.use_gmm_v2_heuristic_tiling, "use_gmm_v2_heuristic_tiling=False"), + ], + ) + return self + rollout_fp8_moe: bool = Field( False, description=( diff --git a/src/maxtext/kernels/megablox/__init__.py b/src/maxtext/kernels/megablox/__init__.py index 15c55a2384..1e9c0d5fe9 100644 --- a/src/maxtext/kernels/megablox/__init__.py +++ b/src/maxtext/kernels/megablox/__init__.py @@ -15,9 +15,13 @@ from maxtext.kernels.megablox.ops import ( TcUnsortCfg, + dequantize_fp8_codes, gmm, gmm_dswiglu, gmm_dswiglu_bwd, + gmm_dswiglu_fp8, + gmm_fp8_forward, gmm_silu, + gmm_silu_fp8, gmm_silu_pre_act_tile_n, ) diff --git a/src/maxtext/kernels/megablox/ops.py b/src/maxtext/kernels/megablox/ops.py index 4981e1ce70..f8a6be3b3f 100644 --- a/src/maxtext/kernels/megablox/ops.py +++ b/src/maxtext/kernels/megablox/ops.py @@ -19,6 +19,7 @@ import dataclasses import functools from typing import List, Literal, NamedTuple, Tuple +from absl import logging import jax import jax.numpy as jnp from maxtext.kernels.megablox import backend @@ -1716,6 +1717,8 @@ def gmm_dswiglu_bwd( use_gmm_v2_heuristic_tiling: bool = False, use_dlhs_transpose_rhs: bool = False, dact_pre_act_tile_n: int | None = None, + rhs_fp8_scale: jnp.ndarray | None = None, + rhs_dtype: jnp.dtype | None = None, ): """The backward of `lhs_act @ rhs` (per group), lhs_act = silu(gate) * up, with the SwiGLU backward fused. @@ -1728,11 +1731,25 @@ def gmm_dswiglu_bwd( gate_up of gmm_silu(fuse_bwd_dswiglu=True), gmm_silu_pre_act_tile_n(its tiling, 2n): the epilogue then reads it in that kernel's per-tile [gate_j | up_j] layout. d[gate | up] is contiguous either way. - Only the bf16 tokamax gmm_v2 path, as gmm_silu: no quantization, partial sums or explicit weight gathers. + Only the tokamax gmm_v2 path, as gmm_silu: no quantized activations, partial sums or explicit weight gathers. The + weight is bf16, or FP8 codes with rhs_fp8_scale: + + rhs_fp8_scale (fp8_moe_native_bwd_dlhs): rhs is float8_e4m3fn codes with this per-channel [g, 1, n] scale, i.e. + the weight dequantize_fp8_codes(rhs, rhs_fp8_scale, rhs_dtype). With use_dlhs_transpose_rhs and fixed tiles the + DLHS kernel dequantizes it in VMEM (_dlhs_fp8_native_kwargs); otherwise it is dequantized here first. The results + are bit-identical to passing the dequantized weight; drhs is in rhs_dtype. Returns: (d_gate_up, dot, drhs); dot is None without return_dact_dot. """ + native_dlhs = {} + if rhs_fp8_scale is not None: + rhs_dtype = jnp.dtype(preferred_element_type if rhs_dtype is None else rhs_dtype) + if _use_fp8_native_dlhs(use_dlhs_transpose_rhs, use_gmm_v2_heuristic_tiling, tiling, rhs.shape): + native_dlhs = _dlhs_fp8_native_kwargs(rhs_fp8_scale, rhs_dtype) + else: + rhs = dequantize_fp8_codes(rhs, rhs_fp8_scale, rhs_dtype) + drhs_dtype = rhs_dtype if rhs_fp8_scale is not None else rhs.dtype if use_gmm_v2_heuristic_tiling: dlhs_tiling = gmm_v2.calculate_tiling else: @@ -1763,6 +1780,7 @@ def gmm_dswiglu_bwd( dact_row_scale=dact_row_scale, return_dact_dot=return_dact_dot, dact_pre_act_tile_n=dact_pre_act_tile_n, + **native_dlhs, ) drhs = _drhs_run_tokamax_v2( dout, @@ -1770,7 +1788,7 @@ def gmm_dswiglu_bwd( group_sizes, group_offset, rhs.shape[0], - rhs.dtype, + drhs_dtype, tiling, use_gmm_v2_heuristic_tiling, partial_sum=None, @@ -1874,3 +1892,589 @@ def _gmm_dswiglu_bwd(preferred_element_type, tiling, heuristic, dlhs_transpose, _gmm_dswiglu.defvjp(_gmm_dswiglu_fwd, _gmm_dswiglu_bwd) + + +# ============================================================================== +# fp8_moe_native_fwd: native FP8 forward for the routed experts +# ============================================================================== +# +# With fp8_moe_fake_quant + fp8_fake_quant_hoist=gather8 the trainer gathers each routed-expert weight as +# float8_e4m3fn codes with a per-(expert, output channel) float32 scale, dequantizes them to bf16 and runs the +# GMMs in bf16. The ops below take the codes and the scale instead and run the forward GMMs on the MXU in FP8, as +# the FP8 rollout does: the weight stays in FP8 (its scale is applied to the block accumulator, whose dtype is +# acc_dtype) and the bf16 activation is quantized to float8_e4m3fn inside the kernel with a dynamic per-row absmax +# scale (per 512-column block by default; lhs_quant_mode selects other schemes). +# +# The backward runs the same bf16 DLHS / TGMM as gmm_silu(fuse_bwd_dswiglu=True) and gmm_dswiglu, at the weight +# dequantized with dequantize_weight_codes_ste's expression, (codes.astype(float32) * scale).astype(bf16), which is +# bit-identical to the bf16 weight the bf16 path multiplies by. With native_bwd_dlhs the DLHS dequantizes the codes +# inside the kernel instead (see below). Only the codes and the scale are saved for the backward (1 byte per weight +# element instead of 2). +# +# `rhs` is the gradient anchor: the bf16 weight the bf16 path would multiply by (in the trainer, the output of +# dequantize_weight_codes_ste). Its value is not read, so XLA removes the forward dequantization; the weight +# gradient is returned as its cotangent, and dequantize_weight_codes_ste carries it to the master weight +# (straight-through, as today). `rhs_codes` and `rhs_scale` get no gradient. + + +def _check_per_channel_scale(scale: jnp.ndarray, name: str): + """Raises unless `scale` is a per-(expert, output channel) [g, 1, n] weight scale.""" + if scale.ndim != 3 or scale.shape[1] != 1: + raise ValueError(f"{name}: expected a per-channel [g, 1, n] scale, got {scale.shape}.") + + +def _fp8_rhs_scale(scale: jnp.ndarray) -> jnp.ndarray: + """quantize_weight_codes_per_channel's [g, 1, n] scale as gmm_v2's rhs_scale [g, 1, 1, n] (one K block).""" + _check_per_channel_scale(scale, "fp8_moe_native_fwd") + return scale.astype(jnp.float32)[:, :, None, :] + + +def dequantize_fp8_codes(codes: jnp.ndarray, scale: jnp.ndarray, dtype: jnp.dtype) -> jnp.ndarray: + """The weight of `codes` in `dtype`: dequantize_weight_codes_ste's forward expression (also used in the kernel).""" + return gmm_v2.dequantize_fp8_codes(codes, scale, dtype) + + +# fp8_moe_native_bwd_dlhs: the backward DLHS (dout @ W^T) reads the codes and dequantizes them inside gmm_v2 +# (rhs_contract_scale): per weight block, in VMEM, with dequantize_fp8_codes' expression, cached across the m tiles +# of an expert. The matmul is the same bf16 one, so the DLHS is bit-identical to the one at the dequantized weight, +# without the XLA dequantization of every expert weight (and its bf16 copy in HBM). The TGMM (weight gradient) does +# not read the weight and is unchanged. Only with the DLHS rhs transpose (the codes' [g, k, n] layout is the +# kernel's transposed [g, size_n, size_k] rhs, the scale varying along its contraction dim) and fixed tiles (the +# heuristic tiling would size its tiles for an fp8 rhs). + + +def _use_fp8_native_dlhs(dlhs_transpose: bool, heuristic: bool, tiling, rhs_shape) -> bool: + """Whether the DLHS can dequantize the [g, size_n, size_k] codes in the kernel; otherwise they are dequantized first. + + Besides the rhs transpose and fixed tiles, gmm_v2's rhs_contract_scale needs whole [tile_n, tile_k] weight blocks + with tile_n % 32 == 0 and tile_k % 128 == 0 (tiling[4], tiling[5] are the DLHS tile_k, tile_n). Called only when + fp8_moe_native_bwd_dlhs is requested, so a fallback is logged. + """ + if not dlhs_transpose or heuristic: + supported = False + else: + tile_k, tile_n = tiling[4], tiling[5] + size_n, size_k = rhs_shape[1], rhs_shape[2] + supported = tile_n % 32 == 0 and tile_k % 128 == 0 and size_n % tile_n == 0 and size_k % tile_k == 0 + if not supported: + logging.warning( + "fp8_moe_native_bwd_dlhs: the DLHS falls back to an XLA weight dequantization (dlhs_transpose=%s," + " heuristic_tiling=%s, tiling=%s, rhs_shape=%s); it needs the DLHS rhs transpose, fixed tiles and whole" + " [tile_n, tile_k] blocks with tile_n %% 32 == 0 and tile_k %% 128 == 0.", + dlhs_transpose, + heuristic, + tiling, + rhs_shape, + ) + return supported + + +def _dlhs_fp8_native_kwargs(rhs_scale: jnp.ndarray, rhs_dtype) -> dict: + """gmm_v2 kwargs of a transpose_rhs DLHS whose rhs is the codes of the per-channel [g, 1, n] rhs_scale.""" + _check_per_channel_scale(rhs_scale, "fp8_moe_native_bwd_dlhs") + return {"rhs_contract_scale": rhs_scale, "rhs_dequant_dtype": jnp.dtype(rhs_dtype)} + + +def _check_fp8_rhs(rhs, rhs_codes, rhs_scale, name): + """Raises unless rhs_codes are float8_e4m3fn of the anchor's shape with a matching [g, 1, n] scale.""" + if rhs.shape != rhs_codes.shape: + raise ValueError(f"{name}: anchor {rhs.shape} and codes {rhs_codes.shape} must have the same shape.") + if rhs_codes.dtype != jnp.float8_e4m3fn: + raise ValueError(f"{name}: rhs_codes must be float8_e4m3fn, got {rhs_codes.dtype}.") + if rhs_scale.shape != (rhs_codes.shape[0], 1, rhs_codes.shape[2]): + raise ValueError(f"{name}: rhs_scale must be [g, 1, n] = {(rhs_codes.shape[0], 1, rhs_codes.shape[2])}.") + + +def _acc_kwargs(acc_dtype): + """gmm_v2's acc_dtype kwarg, passed only when set: None keeps its default (bf16 when lhs is quantized).""" + return {} if acc_dtype is None else {"acc_dtype": jnp.dtype(acc_dtype)} + + +_LHS_QUANT_MODES: dict[str, dict] = { + "": {}, + "block512_bf16": {}, + "block512_bf16_rhs1": {"rhs_scale_first": True}, + "block512_f32": {"lhs_f32_scale": True}, + "block512_f32_hoist": {"lhs_f32_scale": True, "lhs_hoist_quant": True}, + "fullk_bf16": {"lhs_quant_block_size": 0}, + "fullk_bf16_hoist": {"lhs_quant_block_size": 0, "lhs_hoist_quant": True}, + "fullk_f32": {"lhs_quant_block_size": 0, "lhs_f32_scale": True}, + "fullk_f32_hoist": {"lhs_quant_block_size": 0, "lhs_f32_scale": True, "lhs_hoist_quant": True}, +} + + +def _lhs_quant_kwargs(lhs_quant_mode: str) -> dict: + """gmm_v2's in-kernel lhs quantization kwargs for a fp8_moe_native_fwd_lhs_quant mode.""" + if lhs_quant_mode not in _LHS_QUANT_MODES: + raise ValueError(f"Unknown lhs_quant_mode {lhs_quant_mode!r}; expected one of {list(_LHS_QUANT_MODES)}.") + return _LHS_QUANT_MODES[lhs_quant_mode] + + +def gmm_silu_fp8( + lhs: jnp.ndarray, + rhs: jnp.ndarray, + rhs_codes: jnp.ndarray, + rhs_scale: jnp.ndarray, + group_sizes: jnp.ndarray, + preferred_element_type: jnp.dtype, + tiling: tuple[int, int, int, int, int, int, int, int, int], + group_offset: jnp.ndarray | int | None = None, + use_gmm_v2_heuristic_tiling: bool = False, + use_dlhs_transpose_rhs: bool = False, + fuse_dlhs_scale: bool = False, + fuse_3d_dlhs_ps: bool = False, + fwd_rhs_buffer_count: int | None = None, + acc_dtype: jnp.dtype | None = None, + lhs_quant_mode: str = "", + native_bwd_dlhs: bool = False, +) -> tuple[jnp.ndarray, jnp.ndarray]: + """gmm_silu(fuse_bwd_dswiglu=True) with an FP8 forward: returns (act, gate_up). + + rhs: the bf16 [g, k, 2n] gradient anchor (not read); rhs_codes: its float8_e4m3fn codes; rhs_scale: [g, 1, 2n]. + Same contract as gmm_silu(fuse_bwd_dswiglu=True): the caller stops act's gradient and routes it through + gate_up. + + acc_dtype: gmm_v2's accumulator across K blocks (None: the kernel default). lhs_quant_mode: the in-kernel lhs + quantization scheme, a key of _LHS_QUANT_MODES (fp8_moe_native_fwd_lhs_quant). native_bwd_dlhs + (fp8_moe_native_bwd_dlhs): the backward DLHS dequantizes the codes inside the kernel (same values). The other + arguments are gmm_silu's. + """ + if isinstance(lhs, qpl.QArray) or lhs.ndim != 2 or rhs.ndim != 3: + raise NotImplementedError("gmm_silu_fp8: plain 2D lhs and [g, k, 2n] rhs only.") + _check_fp8_rhs(rhs, rhs_codes, rhs_scale, "gmm_silu_fp8") + return _gmm_silu_fp8( + lhs, + rhs, + rhs_codes, + rhs_scale, + group_sizes, + group_offset, + preferred_element_type, + tuple(tiling), + use_gmm_v2_heuristic_tiling, + use_dlhs_transpose_rhs, + fuse_dlhs_scale, + fuse_3d_dlhs_ps, + fwd_rhs_buffer_count, + None if acc_dtype is None else jnp.dtype(acc_dtype), + jnp.dtype(rhs.dtype), + lhs_quant_mode, + bool(native_bwd_dlhs), + ) + + +def _gmm_silu_fp8_kernel( + lhs, + rhs_codes, + rhs_scale, + group_sizes, + group_offset, + preferred_element_type, + tiling, + heuristic, + return_pre_act, + rhs_buffer_count, + acc_dtype, + lhs_quant_mode="", +): + """_gmm_silu_kernel with FP8 weights (codes + per-channel scale) and in-kernel FP8 lhs quantization.""" + if heuristic: + tile_info = gmm_v2.calculate_tiling + else: + tile_n = gmm_silu_pre_act_tile_n(tiling, rhs_codes.shape[-1]) + tile_info = gmm_v2.TileSizes(tile_m=tiling[0], tile_k=tiling[1], tile_n=tile_n) + buffering = {} if rhs_buffer_count is None else {"rhs_buffer_count": rhs_buffer_count} + return gmm_v2.gmm_v2( + lhs=lhs, + rhs=rhs_codes, + group_sizes=group_sizes, + rhs_scale=_fp8_rhs_scale(rhs_scale), + tile_info=tile_info, + preferred_element_type=preferred_element_type, + group_offset=group_offset, + maybe_quantize_lhs=True, + fuse_act="silu", + return_pre_act=return_pre_act, + **buffering, + **_acc_kwargs(acc_dtype), + **_lhs_quant_kwargs(lhs_quant_mode), + ) + + +@functools.partial(jax.custom_vjp, nondiff_argnums=(6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16)) +def _gmm_silu_fp8( + lhs, + rhs, + rhs_codes, + rhs_scale, + group_sizes, + group_offset, + preferred_element_type, + tiling, + heuristic, + dlhs_transpose, + fuse_scale, + fuse_3d, + fwd_rhs_buffer_count, + acc_dtype, + rhs_dtype, + lhs_quant_mode, + native_bwd_dlhs, +): + """gmm_silu_fp8 when not differentiated: the activation, and a zero gate_up placeholder XLA removes.""" + del rhs, dlhs_transpose, fuse_scale, fuse_3d, rhs_dtype, native_bwd_dlhs + act = _gmm_silu_fp8_kernel( + lhs, + rhs_codes, + rhs_scale, + group_sizes, + group_offset, + preferred_element_type, + tiling, + heuristic, + return_pre_act=False, + rhs_buffer_count=fwd_rhs_buffer_count, + acc_dtype=acc_dtype, + lhs_quant_mode=lhs_quant_mode, + ) + return act, jnp.zeros((lhs.shape[0], rhs_codes.shape[-1]), dtype=preferred_element_type) + + +def _gmm_silu_fp8_fwd( + lhs, + rhs, + rhs_codes, + rhs_scale, + group_sizes, + group_offset, + preferred_element_type, + tiling, + heuristic, + dlhs_transpose, + fuse_scale, + fuse_3d, + fwd_rhs_buffer_count, + acc_dtype, + rhs_dtype, + lhs_quant_mode, + native_bwd_dlhs, +): + """Differentiated: (act, gate_up) from the FP8 kernel; saves the codes and scale, not the bf16 weight.""" + del rhs, dlhs_transpose, fuse_scale, fuse_3d, rhs_dtype, native_bwd_dlhs + act, gate_up = _gmm_silu_fp8_kernel( + lhs, + rhs_codes, + rhs_scale, + group_sizes, + group_offset, + preferred_element_type, + tiling, + heuristic, + return_pre_act="concat", + rhs_buffer_count=fwd_rhs_buffer_count, + acc_dtype=acc_dtype, + lhs_quant_mode=lhs_quant_mode, + ) + return (act, gate_up), (lhs, rhs_codes, rhs_scale, group_sizes, group_offset) + + +def _gmm_silu_fp8_bwd( + preferred_element_type, + tiling, + heuristic, + dlhs_transpose, + fuse_scale, + fuse_3d, + fwd_rhs_buffer_count, + acc_dtype, + rhs_dtype, + lhs_quant_mode, + native_bwd_dlhs, + residual, + grad_out, +): + """_gmm_silu_dswiglu_bwd at the dequantized bf16 weight; the weight gradient goes to the anchor.""" + del acc_dtype, lhs_quant_mode + lhs, rhs_codes, rhs_scale, group_sizes, group_offset = residual + if native_bwd_dlhs and _use_fp8_native_dlhs(dlhs_transpose, heuristic, tiling, rhs_codes.shape): + # fp8_moe_native_bwd_dlhs: _gmm_silu_dswiglu_bwd's _gmm_bwd (no quantization, partial sums or 3D layouts: + # _dlhs_run_tokamax_v2 and _compute_drhs reduce to these two calls) with the codes dequantized in the DLHS. + _, d_gate_up = grad_out + grad = d_gate_up.astype(preferred_element_type) + dlhs = gmm_v2.gmm_v2( + lhs=grad, + rhs=rhs_codes, + group_sizes=group_sizes, + rhs_scale=None, + partial_sum=None, + tile_info=gmm_v2.TileSizes(tile_m=tiling[3], tile_k=tiling[4], tile_n=tiling[5]), + preferred_element_type=lhs.dtype, + group_offset=group_offset, + lhs_scale=None, + transpose_rhs=True, + maybe_quantize_lhs=True, + zero_initialize=True, + out_is_3d=False, + **_dlhs_fp8_native_kwargs(rhs_scale, rhs_dtype), + ) + drhs = _drhs_run_tokamax_v2( + grad, + lhs, + group_sizes, + group_offset, + rhs_codes.shape[0], + rhs_dtype, + tiling, + heuristic, + partial_sum=None, + ) + return dlhs, drhs, None, None, None, None + rhs = dequantize_fp8_codes(rhs_codes, rhs_scale, rhs_dtype) + dlhs, drhs, _, _ = _gmm_silu_dswiglu_bwd( + preferred_element_type, + tiling, + heuristic, + dlhs_transpose, + fuse_scale, + fuse_3d, + fwd_rhs_buffer_count, + (lhs, rhs, group_sizes, group_offset), + grad_out, + ) + return dlhs, drhs, None, None, None, None + + +_gmm_silu_fp8.defvjp(_gmm_silu_fp8_fwd, _gmm_silu_fp8_bwd) + + +def gmm_dswiglu_fp8( + lhs_act: jnp.ndarray, + gate_up: jnp.ndarray, + rhs: jnp.ndarray, + rhs_codes: jnp.ndarray, + rhs_scale: jnp.ndarray, + group_sizes: jnp.ndarray, + preferred_element_type: jnp.dtype, + tiling: tuple[int, int, int, int, int, int, int, int, int], + group_offset: jnp.ndarray | int | None = None, + use_gmm_v2_heuristic_tiling: bool = False, + use_dlhs_transpose_rhs: bool = False, + dact_pre_act_tile_n: int | None = None, + acc_dtype: jnp.dtype | None = None, + lhs_quant_mode: str = "", + native_bwd_dlhs: bool = False, +) -> jnp.ndarray: + """gmm_dswiglu with an FP8 forward. rhs: the bf16 [g, k, n] anchor (not read); rhs_codes / rhs_scale [g, 1, n]. + + Same contract as gmm_dswiglu: the backward returns d[gate | up] and a zero cotangent for lhs_act. + + acc_dtype, lhs_quant_mode and native_bwd_dlhs: as in gmm_silu_fp8. The other arguments are gmm_dswiglu's. + """ + _check_fp8_rhs(rhs, rhs_codes, rhs_scale, "gmm_dswiglu_fp8") + return _gmm_dswiglu_fp8( + lhs_act, + gate_up, + rhs, + rhs_codes, + rhs_scale, + group_sizes, + group_offset, + preferred_element_type, + tuple(tiling), + use_gmm_v2_heuristic_tiling, + use_dlhs_transpose_rhs, + dact_pre_act_tile_n, + None if acc_dtype is None else jnp.dtype(acc_dtype), + jnp.dtype(rhs.dtype), + lhs_quant_mode, + bool(native_bwd_dlhs), + ) + + +def _gmm_fp8_fwd_kernel( + lhs, + rhs_codes, + rhs_scale, + group_sizes, + group_offset, + preferred_element_type, + tiling, + heuristic, + acc_dtype, + lhs_quant_mode="", + rhs_buffer_count=None, +): + """gmm's tokamax gmm_v2 forward with FP8 weights and in-kernel FP8 lhs quantization.""" + if heuristic: + tile_info = gmm_v2.calculate_tiling + else: + tile_info = gmm_v2.TileSizes(tile_m=tiling[0], tile_k=tiling[1], tile_n=tiling[2]) + buffering = {} if rhs_buffer_count is None else {"rhs_buffer_count": rhs_buffer_count} + return gmm_v2.gmm_v2( + lhs=lhs, + rhs=rhs_codes, + group_sizes=group_sizes, + rhs_scale=_fp8_rhs_scale(rhs_scale), + tile_info=tile_info, + preferred_element_type=preferred_element_type, + group_offset=group_offset, + maybe_quantize_lhs=True, + **buffering, + **_acc_kwargs(acc_dtype), + **_lhs_quant_kwargs(lhs_quant_mode), + ) + + +@functools.partial(jax.custom_vjp, nondiff_argnums=(7, 8, 9, 10, 11, 12, 13, 14, 15)) +def _gmm_dswiglu_fp8( + lhs_act, + gate_up, + rhs, + rhs_codes, + rhs_scale, + group_sizes, + group_offset, + preferred_element_type, + tiling, + heuristic, + dlhs_transpose, + pre_act_tile_n, + acc_dtype, + rhs_dtype, + lhs_quant_mode, + native_bwd_dlhs, +): + """gmm_dswiglu_fp8 when not differentiated: the FP8 forward.""" + del gate_up, rhs, dlhs_transpose, pre_act_tile_n, rhs_dtype, native_bwd_dlhs + return _gmm_fp8_fwd_kernel( + lhs_act, + rhs_codes, + rhs_scale, + group_sizes, + group_offset, + preferred_element_type, + tiling, + heuristic, + acc_dtype, + lhs_quant_mode, + ) + + +def _gmm_dswiglu_fp8_fwd( + lhs_act, + gate_up, + rhs, + rhs_codes, + rhs_scale, + group_sizes, + group_offset, + preferred_element_type, + tiling, + heuristic, + dlhs_transpose, + pre_act_tile_n, + acc_dtype, + rhs_dtype, + lhs_quant_mode, + native_bwd_dlhs, +): + """Differentiated: the FP8 forward; saves the FP8 weight and scale (not a bf16 copy) for the backward.""" + del rhs, dlhs_transpose, pre_act_tile_n, rhs_dtype, native_bwd_dlhs + out = _gmm_fp8_fwd_kernel( + lhs_act, + rhs_codes, + rhs_scale, + group_sizes, + group_offset, + preferred_element_type, + tiling, + heuristic, + acc_dtype, + lhs_quant_mode, + ) + return out, (lhs_act, gate_up, rhs_codes, rhs_scale, group_sizes, group_offset) + + +def _gmm_dswiglu_fp8_bwd( + preferred_element_type, + tiling, + heuristic, + dlhs_transpose, + pre_act_tile_n, + acc_dtype, + rhs_dtype, + lhs_quant_mode, + native_bwd_dlhs, + residual, + grad, +): + """_gmm_dswiglu_bwd at the dequantized bf16 weight; the weight gradient goes to the anchor.""" + del acc_dtype, lhs_quant_mode + lhs_act, gate_up, rhs_codes, rhs_scale, group_sizes, group_offset = residual + if native_bwd_dlhs: + # fp8_moe_native_bwd_dlhs: _gmm_dswiglu_bwd with the codes (gmm_dswiglu_bwd dequantizes them in the DLHS). + d_gate_up, _, drhs = gmm_dswiglu_bwd( + grad, + lhs_act, + rhs_codes, + group_sizes, + group_offset, + preferred_element_type, + tiling, + gate_up, + use_gmm_v2_heuristic_tiling=heuristic, + use_dlhs_transpose_rhs=dlhs_transpose, + dact_pre_act_tile_n=pre_act_tile_n, + rhs_fp8_scale=rhs_scale, + rhs_dtype=rhs_dtype, + ) + return jnp.zeros_like(lhs_act), d_gate_up, drhs, None, None, None, None + rhs = dequantize_fp8_codes(rhs_codes, rhs_scale, rhs_dtype) + d_lhs_act, d_gate_up, drhs, _, _ = _gmm_dswiglu_bwd( + preferred_element_type, + tiling, + heuristic, + dlhs_transpose, + pre_act_tile_n, + (lhs_act, gate_up, rhs, group_sizes, group_offset), + grad, + ) + return d_lhs_act, d_gate_up, drhs, None, None, None, None + + +_gmm_dswiglu_fp8.defvjp(_gmm_dswiglu_fp8_fwd, _gmm_dswiglu_fp8_bwd) + + +def gmm_fp8_forward( + lhs: jnp.ndarray, + rhs_codes: jnp.ndarray, + rhs_scale: jnp.ndarray, + group_sizes: jnp.ndarray, + preferred_element_type: jnp.dtype, + tiling: tuple[int, int, int, int, int, int, int, int, int], + group_offset: jnp.ndarray | int | None = None, + use_gmm_v2_heuristic_tiling: bool = False, + acc_dtype: jnp.dtype | None = None, + lhs_quant_mode: str = "", + fwd_rhs_buffer_count: int | None = None, +) -> jnp.ndarray: + """The FP8 forward of gmm (no VJP): for callers that own the backward (ring_ragged_unsort_down_projection). + + rhs_codes: float8_e4m3fn [g, k, n]; rhs_scale: [g, 1, n]. acc_dtype and lhs_quant_mode: as in gmm_silu_fp8; + fwd_rhs_buffer_count: gmm_v2's rhs_buffer_count. The other arguments are gmm's. + """ + if rhs_codes.dtype != jnp.float8_e4m3fn: + raise ValueError(f"gmm_fp8_forward: rhs_codes must be float8_e4m3fn, got {rhs_codes.dtype}.") + return _gmm_fp8_fwd_kernel( + lhs, + rhs_codes, + rhs_scale, + group_sizes, + group_offset, + preferred_element_type, + tuple(tiling), + use_gmm_v2_heuristic_tiling, + None if acc_dtype is None else jnp.dtype(acc_dtype), + lhs_quant_mode, + fwd_rhs_buffer_count, + ) diff --git a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py index a41b24ffc6..912301e966 100644 --- a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py +++ b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py @@ -386,6 +386,20 @@ class GmmConfigs: # starts one grid step before its first use instead of two, and one weight tile # of VMEM is freed (with fuse_act, a gate and an up tile). Values are unchanged. rhs_buffer_count: int | None = None + # In-kernel lhs quantization: compute the per-row block scale in float32 as max(absmax, eps) / dtype_max and divide + # by it, instead of the bf16 absmax and reciprocal. + lhs_f32_scale: bool = False + # In-kernel lhs quantization: quantize each lhs k block once per tile, before the n loop, instead of inside it. + lhs_hoist_quant: bool = False + # Quantized matmul: apply the rhs (sub-channel) scale to the block accumulator before the lhs scale. + rhs_scale_first: bool = False + # rhs_contract_scale (fp8_moe_native_bwd_dlhs): rhs holds float8 codes [size_group, size_n, size_k] + # (transpose_rhs) whose scale varies along the contraction dim K ([size_group, 1, size_k] float32, read in + # the rhs scale slot). Each rhs block is dequantized in VMEM to rhs_dequant_dtype with + # (codes.astype(float32) * scale).astype(rhs_dequant_dtype) into a scratch tile, cached across the grid + # steps that read the same block (the m tiles of a group), and the unquantized matmul reads that tile. + rhs_contract_dequant: bool = False + rhs_dequant_dtype: jnp.dtype | None = None @property def lhs_is_3d(self) -> bool: @@ -516,6 +530,12 @@ def rhs_scale_index_map(self, n_id: jax.Array, gm_id: jax.Array, k_id: jax.Array b_tile_id = b_row // self.cfgs.num_quant_blocks_per_tile_k return (group_id, b_tile_id, 0, n_id) + def rhs_contract_scale_index_map(self, n_id: jax.Array, gm_id: jax.Array, k_id: jax.Array): + """rhs_contract_dequant: the [1, tile_k] block of the group's [size_group, 1, size_k] contraction scale.""" + del n_id + group_id = self.metadata_ref.gm_id_to_group_id[gm_id] + return (group_id, 0, k_id) + def out_index_map(self, n_id: jax.Array, gm_id: jax.Array, _: jax.Array): """Calculates index map for the output tensor.""" is_last_gm = gm_id == (pl.num_programs(self.cfgs.gm_axis) - 1) @@ -609,6 +629,12 @@ def _ordered(fn): (None, cfgs.num_quant_blocks_per_tile_k, 1, cfgs.tiles.tile_n), _ordered(index_map.rhs_scale_index_map), ) + elif cfgs.rhs_contract_dequant: + # The contraction scale rides in the rhs scale slot (has_scale is False: nothing else reads it). + rhs_scale_block_spec = pl.BlockSpec( + (None, 1, cfgs.tiles.tile_k), + _ordered(index_map.rhs_contract_scale_index_map), + ) if cfgs.has_partial_sum and cfgs.out_is_3d: ps_d0 = cfgs.out_size_n // num_lanes if cfgs.out_window else cfgs.out_tile_d0 @@ -673,6 +699,7 @@ def inner_kernel( partial_out_ref: jax.Array, # [size_lhs_sublane, tile_n] acc_ref: jax.Array, # [tile_m, tile_n] metadata_ref: MetadataRef, + rhs_cache_refs: tuple[jax.Array, jax.Array] | None = None, *, cfgs: GmmConfigs, ): @@ -701,6 +728,9 @@ def inner_kernel( the activation is applied. With cfgs.fuse_dact, tiled_ps_ref is a DactInputRef and the accumulator becomes [d_gate | d_up] before the regular store path; with cfgs.return_dact_dot the out refs are (out, dot) pairs. + With cfgs.rhs_contract_dequant, rhs_cache_refs is (the [tile_n, tile_k] + dequantized weight tile, the SMEM key of the block it holds); see + _rhs_cache_miss. """ pre_act_refs = () @@ -708,7 +738,7 @@ def inner_kernel( (tiled_out_ref, *tiled_pre_act_refs), (partial_out_ref, *partial_pre_act_refs) = tiled_out_ref, partial_out_ref pre_act_refs = tuple(zip(tiled_pre_act_refs, partial_pre_act_refs)) - def _matmul(is_first_k_step: bool, is_last_k_step: bool): + def _matmul(is_first_k_step: bool, is_last_k_step: bool, rhs_fill: bool = False): tpu_info = pltpu.get_tpu_info() mxu_size = tpu_info.mxu_column_size @@ -728,13 +758,28 @@ def _matmul(is_first_k_step: bool, is_last_k_step: bool): tiled_lhs = relayout.load_3d_as_2d(tiled_lhs_ref.value, cfgs.tiles.tile_m, cfgs.lhs_tile_d0) else: tiled_lhs = tiled_lhs_ref.get_value().reshape(-1, cfgs.tiles.tile_k)[...] - tiled_rhs = tiled_rhs_ref.get_weight(transpose_rhs=cfgs.transpose_rhs) + rhs_fill_refs = None + if cfgs.rhs_contract_dequant and rhs_fill: + # The first grid step of this weight block: rhs_block dequantizes each [mxu, tile_k] piece + # of the codes right before its matmul (so the VPU work can overlap the MXU work of the previous piece) + # and stores it in the cache tile, which the later steps of the block read. + rhs_fill_refs = (tiled_rhs_ref.weight, tiled_rhs_ref.scale, rhs_cache_refs[0]) # pyrefly: ignore + tiled_rhs = None + elif cfgs.rhs_contract_dequant: + # This block's weight, dequantized in VMEM by the block's first step: the same bf16 tile the bf16 path DMAs, + # so the matmul below is that path's. + tiled_rhs = rhs_cache_refs[0][...] # pyrefly: ignore[unsupported-operation] + else: + tiled_rhs = tiled_rhs_ref.get_weight(transpose_rhs=cfgs.transpose_rhs) # When rhs is packed (quantized dtype packed into uint32), unpack it # back to the original dtype using pltpu.bitcast which operates on K # axis. This expands the K dimension back to tile_k. if cfgs.rhs_cfgs.should_bitcast: tiled_rhs = pltpu.bitcast(tiled_rhs, cfgs.rhs_cfgs.dtype) - rhs_tile_n = tiled_rhs.shape[0 if cfgs.transpose_rhs else 1] + if tiled_rhs is None: + rhs_tile_n = cfgs.tiles.tile_n + else: + rhs_tile_n = tiled_rhs.shape[0 if cfgs.transpose_rhs else 1] # This should only be taken in the case where we don't requantize # the scales and thus we need to dequantize inside VMEM to avoid small @@ -765,6 +810,15 @@ def _matmul(is_first_k_step: bool, is_last_k_step: bool): tiled_lhs = jnp.where(mask_lhs, tiled_lhs, 0) def rhs_block(start_k: int, end_k: int, start_n: int, end_n: int): + if rhs_fill_refs is not None: + # rhs_contract_dequant (transpose_rhs) fill step: dequantize this piece, keep it for the block's later steps. + codes_ref, scale_ref, cache_ref = rhs_fill_refs + end_k = min(end_k, cfgs.tiles.tile_k) + block = dequantize_fp8_codes( + codes_ref[start_n:end_n, start_k:end_k], scale_ref[:, start_k:end_k], cache_ref.dtype + ) + cache_ref[start_n:end_n, start_k:end_k] = block + return block if cfgs.transpose_rhs: return tiled_rhs[start_n:end_n, start_k:end_k] return tiled_rhs[start_k:end_k, start_n:end_n] @@ -852,45 +906,57 @@ def rhs_block(start_k: int, end_k: int, start_n: int, end_n: int): lhs_scale = tiled_lhs_ref.get_scale().astype(acc_ref.dtype) lhs_scale_inv = 1.0 / lhs_scale - # Without n outer loop, result of quantized matmul becomes available only - # at the last iteration of the loop. This means [tile_m, tile_n] value - # needs to be stored until the last iteration. By adding n outer loop, - # result of [tile_m, mxu_size] becomes available at the end of every k - # inner loop which can be used to pipeline subsequent VPU or VST ops with - # MXU ops for the next [tile_m, mxu_size]. + def _quantize_lhs_block(block_lhs): + """Quantizes one [tile_m, k block] of lhs: (quantized block, its [tile_m, 1] or [1, 1] scale).""" + if should_use_external_scale: + assert lhs_scale is not None + assert lhs_scale_inv is not None + return jnp.clip(block_lhs * lhs_scale_inv, -dtype_max, dtype_max).astype(lhs_q_dtype), lhs_scale + if cfgs.lhs_f32_scale: + x_f32 = block_lhs.astype(jnp.float32) + block_abs_max = jnp.max(jnp.abs(x_f32), axis=1, keepdims=True) + block_scale = jnp.maximum(block_abs_max, _LHS_F32_SCALE_EPS) / dtype_max + return (x_f32 / block_scale).astype(lhs_q_dtype), block_scale + block_abs_max = jnp.max(jnp.abs(block_lhs), axis=1, keepdims=True) + block_scale = block_abs_max / dtype_max + # If block_scale=0, 1 / block_scale is NaN or Inf, which breaks the downcast to the quantized dtype; use 0. + block_scale_inv = jnp.where(block_scale == 0, 0, 1 / block_scale) + return (block_lhs * block_scale_inv).astype(lhs_q_dtype), block_scale + + k_ranges = [ + (start_k, min(cfgs.tiles.tile_k, start_k + q_block_size)) # pyrefly: ignore[unsupported-operation] + for start_k in range(0, cfgs.tiles.tile_k, q_block_size) # pyrefly: ignore[bad-argument-type] + ] + hoisted_lhs_blocks = None + if cfgs.lhs_hoist_quant: + hoisted_lhs_blocks = [_quantize_lhs_block(tiled_lhs[:, start_k:end_k]) for start_k, end_k in k_ranges] + + def _apply_rhs_scale(acc, start_k, start_n, end_n): + """Applies the rhs sub-channel scale of the quant block at start_k, if it is applied after the matmul.""" + if not cfgs.rhs_cfgs.should_dequantize_after_matmul: + return acc + b_id = start_k // cfgs.rhs_cfgs.quant_block_size # pyrefly: ignore[unsupported-operation] + rhs_scale_slice = tiled_rhs_ref.get_scale() + return acc * rhs_scale_slice[b_id, :, start_n:end_n].astype(acc_ref.dtype) + + # Without the n outer loop, the quantized matmul's result is available only at the last k iteration, so the + # whole [tile_m, tile_n] value is kept until then. With it, each [tile_m, mxu_size] result is ready at the end + # of its k loop, and the following VPU / VST work pipelines with the MXU work of the next one. for start_n in range(0, rhs_tile_n, mxu_size): end_n = min(rhs_tile_n, start_n + mxu_size) col_size = end_n - start_n acc_n = jnp.zeros((cfgs.tiles.tile_m, col_size), dtype=acc_ref.dtype) - for start_k in range(0, cfgs.tiles.tile_k, q_block_size): # pyrefly: ignore[bad-argument-type] - end_k = min(cfgs.tiles.tile_k, start_k + q_block_size) # pyrefly: ignore[unsupported-operation] - - block_lhs = tiled_lhs[:, start_k:end_k] + for k_index, (start_k, end_k) in enumerate(k_ranges): block_rhs = rhs_block(start_k, end_k, start_n, end_n) - - # Perform lhs quantization. Note that for every block_lhs, - # same computation will be performed tiles_n//mxu_size times. - # But we can let compiler perform CSE and avoid recomputation. - if should_use_external_scale: - assert lhs_scale is not None - assert lhs_scale_inv is not None - block_lhs_q = jnp.clip(block_lhs * lhs_scale_inv, -dtype_max, dtype_max).astype(lhs_q_dtype) - block_scale = lhs_scale # [1, 1] + if hoisted_lhs_blocks is not None: + block_lhs_q, block_scale = hoisted_lhs_blocks[k_index] else: - block_abs_max = jnp.max(jnp.abs(block_lhs), axis=1, keepdims=True) - block_scale = block_abs_max / dtype_max - - # If block_scale=0, it will cause division by zero and return either - # NaN or Inf. Since this can cause numeric issue when downcasting to - # quantized value, we convert them into 0. - block_scale_inv = jnp.where(block_scale == 0, 0, 1 / block_scale) - # Convert lhs into quantized dtype. - block_lhs_q = (block_lhs * block_scale_inv).astype(lhs_q_dtype) - - # Unlike unquantized path, compiler may not perform implicit type - # conversion due to numeric concerns. As this can cause unsupported - # matmul error, explicit type conversion is performed. + # The same block is quantized for every n iteration; the compiler's CSE removes the recomputation. + block_lhs_q, block_scale = _quantize_lhs_block(tiled_lhs[:, start_k:end_k]) + + # Unlike the unquantized path, the compiler may not convert the operand types implicitly (numerics), which + # gives an unsupported-matmul error, so the conversion is explicit. if not tpu_info.is_matmul_supported(lhs_q_dtype, block_rhs.dtype): block_rhs = block_rhs.astype(lhs_q_dtype) @@ -908,13 +974,10 @@ def rhs_block(start_k: int, end_k: int, start_n: int, end_n: int): preferred_element_type=preferred_element_type, ).astype(acc_ref.dtype) - block_acc *= block_scale.astype(acc_ref.dtype) - - # Apply rhs subchannel scale per quant block. - if cfgs.rhs_cfgs.should_dequantize_after_matmul: - b_id = start_k // cfgs.rhs_cfgs.quant_block_size # pyrefly: ignore[unsupported-operation] - rhs_scale_slice = tiled_rhs_ref.get_scale() - block_acc *= rhs_scale_slice[b_id, :, start_n:end_n].astype(acc_ref.dtype) + if cfgs.rhs_scale_first: + block_acc = _apply_rhs_scale(block_acc, start_k, start_n, end_n) * block_scale.astype(acc_ref.dtype) + else: + block_acc = _apply_rhs_scale(block_acc * block_scale.astype(acc_ref.dtype), start_k, start_n, end_n) acc_n += block_acc acc_list.append(acc_n) @@ -1031,6 +1094,32 @@ def matmul_last(): num_k = pl.num_programs(2) k_id = pl.program_id(2) + if cfgs.rhs_contract_dequant: + # The first step of each weight block dequantizes it piece by piece inside its matmul, so the VPU work overlaps + # the MXU work (see _matmul's rhs_fill); the other steps read the cache tile. + cache_miss = _rhs_cache_miss(rhs_cache_refs, metadata_ref, cfgs) + is_first = k_id == 0 + is_last = k_id == (num_k - 1) + + def _variants(fill: bool): + suffix = "_rhs_fill" if fill else "" + return [ + jax.named_scope(name + suffix)(functools.partial(_matmul, first, last, rhs_fill=fill)) + for name, first, last in ( + ("matmul_first_last", True, True), + ("matmul_first", True, False), + ("matmul_last", False, True), + ("matmul", False, False), + ) + ] + + def _dispatch(fill: bool): + first_last, first, last, mid = _variants(fill) + lax.cond(is_first, lambda: lax.cond(is_last, first_last, first), lambda: lax.cond(is_last, last, mid)) + + lax.cond(cache_miss, lambda: _dispatch(True), lambda: _dispatch(False)) + return + is_first_k_step = k_id == 0 is_last_k_step = k_id == (num_k - 1) @@ -1049,6 +1138,46 @@ def matmul_last(): ) +# The default k block of the in-kernel lhs quantization (one scale per row and block). +_DEFAULT_LHS_QUANT_BLOCK_SIZE = 512 +# lhs_f32_scale: the floor of a row block's absmax, so an all-zero block gets a finite scale. +_LHS_F32_SCALE_EPS = 1e-12 + + +def dequantize_fp8_codes(codes: jax.Array, scale: jax.Array, dtype: jnp.dtype) -> jax.Array: + """The weight of float8 `codes` in `dtype`: (codes.astype(float32) * scale).astype(dtype). + + dequantize_weight_codes_ste's forward expression. ops.dequantize_fp8_codes calls it on whole weights and the + rhs_contract_dequant path on pieces of an rhs block (codes [r, c], scale [1, c]), so both are bit-identical. + """ + return (codes.astype(jnp.float32) * scale).astype(dtype) + + +def _rhs_block_key(metadata_ref: MetadataRef, cfgs: GmmConfigs) -> jax.Array: + """rhs_contract_dequant: an int32 id of this grid step's rhs block (group, n tile, k tile).""" + num_n = pl.num_programs(cfgs.n_axis) + num_k = pl.num_programs(2) + group_id = metadata_ref.gm_id_to_group_id[pl.program_id(cfgs.gm_axis)] + return (group_id * num_n + pl.program_id(cfgs.n_axis)) * num_k + pl.program_id(2) + + +def _rhs_cache_miss( + rhs_cache_refs: tuple[jax.Array, jax.Array] | None, metadata_ref: MetadataRef, cfgs: GmmConfigs +) -> jax.Array: + """rhs_contract_dequant: whether the cache tile holds another rhs block than this step's. + + The rhs block index is (group, n tile, k tile), so with a single n and k tile (the MoE DLHS tiles) the cache is + filled once per group and every m tile of the group reuses it. Side effect: records this step's block key in the + SMEM key, since the step that misses fills the cache tile. + """ + assert rhs_cache_refs is not None + key_ref = rhs_cache_refs[1] + key = _rhs_block_key(metadata_ref, cfgs) + miss = key != key_ref[0] + key_ref[0] = key + return miss + + def _mask_group_rows(acc: jax.Array, metadata_ref: MetadataRef, cfgs: GmmConfigs): """Zeroes the rows of `acc` that do not belong to the current gm tile's group. @@ -1323,6 +1452,8 @@ def kernel_main( pre_act_partial_refs: tuple[jax.Array, ...] = (), # return_pre_act: their partial-row scratch # Per extra output, the VMEM zero buffer of its zero-init DMAs when its dtype is not out_ref's (None: zero_ref). extra_zero_refs: tuple[jax.Array | None, ...] = (), + # rhs_contract_dequant: (the dequantized weight tile [tile_n, tile_k], its SMEM block key int32[1]). + rhs_cache_refs: tuple[jax.Array, jax.Array] | None = None, ): """Entry point for GMM kernel. @@ -1459,6 +1590,10 @@ def kernel_main( if pre_act_out_refs: out_in = (out_in, *(r.reshape(-1, cfgs.dims.size_lhs_sublane, *r.shape[1:]) for r in pre_act_out_refs)) scratches[0] = (partial_out_ref, *pre_act_partial_refs) + if cfgs.rhs_contract_dequant: + assert rhs_cache_refs is not None + rhs_cache_refs[1][0] = jnp.int32(-1) # no block dequantized yet + scratches.append(rhs_cache_refs) pipeline_fn(lhs_in, rhs_ref, ps_in, out_in, scratches=scratches) if cfgs.zero_init: @@ -1486,6 +1621,7 @@ def kernel_main_pre_act( partial_up_ref: jax.Array, # [size_lhs_sublane, tile_n] *, cfgs: GmmConfigs, + rhs_cache_refs: tuple[jax.Array, jax.Array] | None = None, ): """kernel_main for return_pre_act: pallas_call passes the gate and up outputs right after act.""" kernel_main( @@ -1503,6 +1639,7 @@ def kernel_main_pre_act( cfgs=cfgs, pre_act_out_refs=(gate_out_ref, up_out_ref), pre_act_partial_refs=(partial_gate_ref, partial_up_ref), + rhs_cache_refs=rhs_cache_refs, ) @@ -1523,6 +1660,7 @@ def kernel_main_single_extra_out( extra_zero_ref: jax.Array | None, *, cfgs: GmmConfigs, + rhs_cache_refs: tuple[jax.Array, jax.Array] | None = None, ): """kernel_main for one extra output: [gate | up] (return_pre_act="concat") or the dot (return_dact_dot). @@ -1545,6 +1683,7 @@ def kernel_main_single_extra_out( pre_act_out_refs=(extra_out_ref,), pre_act_partial_refs=(partial_extra_ref,), extra_zero_refs=(extra_zero_ref,), + rhs_cache_refs=rhs_cache_refs, ) @@ -1797,6 +1936,16 @@ def get_scope_name(cfgs: GmmConfigs) -> str: scope_name += f"-preact_tn_{cfgs.dact_pre_act_tile_n}" if cfgs.rhs_buffer_count is not None: scope_name += f"-rhsbuf_{cfgs.rhs_buffer_count}" + if cfgs.lhs_cfgs.quant_dtype is not None and cfgs.lhs_cfgs.quant_block_size != _DEFAULT_LHS_QUANT_BLOCK_SIZE: + scope_name += f"-lqbs_{cfgs.lhs_cfgs.quant_block_size}" + if cfgs.lhs_f32_scale: + scope_name += "-lf32s" + if cfgs.lhs_hoist_quant: + scope_name += "-lhoist" + if cfgs.rhs_scale_first: + scope_name += "-rhs1" + if cfgs.rhs_contract_dequant: + scope_name += "-rhs_cdq" return scope_name @@ -1826,8 +1975,17 @@ def make_gmm_configs( return_dact_dot: bool = False, dact_pre_act_tile_n: int | None = None, rhs_buffer_count: int | None = None, + lhs_quant_block_size: int | None = None, + lhs_f32_scale: bool = False, + lhs_hoist_quant: bool = False, + rhs_scale_first: bool = False, + rhs_contract_scale: jax.Array | None = None, + rhs_dequant_dtype: jnp.dtype | None = None, ): """Fills the GMM config for the GMM kernel.""" + if lhs_quant_block_size is not None and (not isinstance(lhs_quant_block_size, int) or lhs_quant_block_size < 0): + # A negative block size would make the quantized k loop empty (an all-zero output). + raise ValueError(f"lhs_quant_block_size must be None or an int >= 0; got {lhs_quant_block_size!r}.") dims = validate_inputs( lhs, @@ -1892,7 +2050,9 @@ def make_gmm_configs( # scale value. Since this operation is very memory intensive, we use a # block size that is small enough to minimize memory overhead but large # enough to minimize compute overhead of quantization. - quant_block_size=512, + quant_block_size=( + dims.size_k if lhs_quant_block_size == 0 else (lhs_quant_block_size or _DEFAULT_LHS_QUANT_BLOCK_SIZE) + ), dtype=lhs.dtype, has_scale=has_lhs_scale, is_3d=lhs.ndim == 3, @@ -2002,6 +2162,29 @@ def make_gmm_configs( if rhs_buffer_count is not None and (not isinstance(rhs_buffer_count, int) or rhs_buffer_count < 2): # One buffer would serialize every weight DMA with the matmul that reads it. raise ValueError(f"rhs_buffer_count must be None or an int >= 2; got {rhs_buffer_count!r}.") + rhs_contract_dequant = rhs_contract_scale is not None + if rhs_contract_dequant: + # The rhs is float8 codes dequantized in VMEM before an unquantized matmul (see GmmConfigs). + if not transpose_rhs or rhs_scale is not None or rhs_bias is not None or fuse_act is not None: + raise NotImplementedError("rhs_contract_scale needs transpose_rhs and no rhs_scale, rhs_bias or fuse_act.") + if lhs_cfgs.is_3d or lhs_scale is not None or lhs_cfgs.quant_dtype is not None: + raise NotImplementedError("rhs_contract_scale needs a 2D, unscaled, unquantized lhs.") + if jax.dtypes.itemsize_bits(rhs.dtype) != 8 or not jnp.issubdtype(rhs.dtype, jnp.floating): + raise ValueError(f"rhs_contract_scale needs float8 rhs codes; got {rhs.dtype}.") + if rhs_contract_scale.shape != (dims.size_group, 1, dims.size_k): + raise ValueError( + f"rhs_contract_scale must have shape {(dims.size_group, 1, dims.size_k)}; got {rhs_contract_scale.shape}." + ) + if tiles.tile_n % 32 or tiles.tile_k % num_lanes or dims.size_n % tiles.tile_n or dims.size_k % tiles.tile_k: + # 32: a float8 sublane tile (the dequantization runs on 32-row pieces); whole tiles only, so every scale and + # code block is in bounds. + raise NotImplementedError( + f"rhs_contract_scale needs tile_n % 32 == 0, tile_k % {num_lanes} == 0 and tiles that divide size_n and" + f" size_k; got {tiles=}, {dims=}." + ) + rhs_dequant_dtype = jnp.dtype(lhs.dtype if rhs_dequant_dtype is None else rhs_dequant_dtype) + elif rhs_dequant_dtype is not None: + raise ValueError("rhs_dequant_dtype requires rhs_contract_scale.") return GmmConfigs( dims=dims, @@ -2021,6 +2204,11 @@ def make_gmm_configs( return_dact_dot=return_dact_dot, dact_pre_act_tile_n=dact_pre_act_tile_n, rhs_buffer_count=rhs_buffer_count, + lhs_f32_scale=lhs_f32_scale, + lhs_hoist_quant=lhs_hoist_quant, + rhs_scale_first=rhs_scale_first, + rhs_contract_dequant=rhs_contract_dequant, + rhs_dequant_dtype=rhs_dequant_dtype, ) @@ -2037,6 +2225,15 @@ def get_metadata(cfgs: GmmConfigs) -> dict[str, str | int | float]: del cfgs_dict["dact_pre_act_tile_n"] # the contiguous [gate | up] reader keeps its metadata too if cfgs.rhs_buffer_count is None: del cfgs_dict["rhs_buffer_count"] # the default rhs buffering keeps its metadata too + if not cfgs.lhs_f32_scale: + del cfgs_dict["lhs_f32_scale"] + if not cfgs.lhs_hoist_quant: + del cfgs_dict["lhs_hoist_quant"] + if not cfgs.rhs_scale_first: + del cfgs_dict["rhs_scale_first"] + if not cfgs.rhs_contract_dequant: + del cfgs_dict["rhs_contract_dequant"] + del cfgs_dict["rhs_dequant_dtype"] ret = {} for path, val in jax.tree_util.tree_leaves_with_path(cfgs_dict): key = jax.tree_util.keystr(path, simple=True, separator=".") @@ -2063,6 +2260,11 @@ def get_metadata(cfgs: GmmConfigs) -> dict[str, str | int | float]: "return_dact_dot", "dact_pre_act_tile_n", "rhs_buffer_count", + "lhs_quant_block_size", + "lhs_f32_scale", + "lhs_hoist_quant", + "rhs_scale_first", + "rhs_dequant_dtype", ] ) def gmm_v2( @@ -2092,6 +2294,12 @@ def gmm_v2( return_dact_dot: bool = False, dact_pre_act_tile_n: int | None = None, rhs_buffer_count: int | None = None, + lhs_quant_block_size: int | None = None, + lhs_f32_scale: bool = False, + lhs_hoist_quant: bool = False, + rhs_scale_first: bool = False, + rhs_contract_scale: jax.Array | None = None, # [size_group, 1, size_k] + rhs_dequant_dtype: jnp.dtype | None = None, ) -> jax.Array | tuple[jax.Array, ...]: """GMM kernel implemented with emit_pipeline. @@ -2161,6 +2369,28 @@ def gmm_v2( weight block's DMA one grid step ahead instead of two. With a single K tile (tile_k == size_k) the weight block changes only at group and n-tile boundaries. The values are unchanged. + lhs_quant_block_size: With in-kernel lhs quantization, the k block that + gets one scale per row: None for the default (512), 0 for the whole K + dim, or a positive int. + lhs_f32_scale: With in-kernel lhs quantization, compute each block scale + in float32 as max(absmax, 1e-12) / dtype_max and divide by it, instead + of the bf16 absmax and reciprocal. + lhs_hoist_quant: With in-kernel lhs quantization, quantize each lhs block + once per tile instead of per mxu-wide n slice. The values are unchanged. + rhs_scale_first: With a quantized matmul, apply the rhs sub-channel scale + to the block accumulator before the lhs scale. + rhs_contract_scale: Optional. With transpose_rhs, rhs holds float8 codes + whose scale varies along the contraction dim: [size_group, 1, size_k] + (a per-output-channel weight scale, read by the DLHS of a weight-codes + forward). Each rhs block is dequantized in VMEM, + (codes.astype(float32) * scale).astype(rhs_dequant_dtype), into a + scratch tile that is reused while the block index is unchanged (with + one n and k tile, every m tile of a group), and the matmul is the + unquantized one on that tile: the output is bit-identical to passing + the dequantized weight. No rhs_scale, rhs_bias, fuse_act, lhs_scale or + 3D lhs; tile_n % 32 == 0, tile_k % 128 == 0, whole tiles. + rhs_dequant_dtype: With rhs_contract_scale, the dtype of the dequantized + weight (None: lhs.dtype). Returns: Output of shape [size_m, size_n] (or [size_m, size_n // 128, 128]); with @@ -2208,6 +2438,12 @@ def gmm_v2( return_dact_dot=return_dact_dot, dact_pre_act_tile_n=dact_pre_act_tile_n, rhs_buffer_count=rhs_buffer_count, + lhs_quant_block_size=lhs_quant_block_size, + lhs_f32_scale=lhs_f32_scale, + lhs_hoist_quant=lhs_hoist_quant, + rhs_scale_first=rhs_scale_first, + rhs_contract_scale=rhs_contract_scale, + rhs_dequant_dtype=rhs_dequant_dtype, ) dims = cfgs.dims tiles = cfgs.tiles @@ -2228,6 +2464,11 @@ def gmm_v2( if rhs_bias is not None: rhs_bias = rhs_bias.astype(jnp.float32) rhs_bias_spec = pl.BlockSpec(memory_space=pltpu.HBM) + if cfgs.rhs_contract_dequant: + # The contraction scale rides in the rhs scale slot (see generate_block_specs). + assert rhs_contract_scale is not None and rhs_scale is None + rhs_scale = rhs_contract_scale.astype(jnp.float32) + rhs_scale_spec = pl.BlockSpec(memory_space=pltpu.HBM) # Initialize scratch shapes. max_num_gm = dims.size_group + pl.cdiv(dims.size_m, tiles.tile_m) - 1 @@ -2303,6 +2544,12 @@ def gmm_v2( ) dot_zero_scratch = pltpu.VMEM((dot_zero_tile_m, num_lanes), cfgs.acc_dtype) if cfgs.zero_init else None scratch_shapes += [pltpu.VMEM((dims.size_lhs_sublane, num_lanes), cfgs.acc_dtype), dot_zero_scratch] + if cfgs.rhs_contract_dequant: + # Last: the dequantized weight tile and the SMEM key of the rhs block it holds (kernel_body peels them off). + scratch_shapes += [ + pltpu.VMEM((tiles.tile_n, tiles.tile_k), cfgs.rhs_dequant_dtype), + pltpu.SMEM((1,), jnp.int32), + ] aligned_n = align_to(cfgs.out_size_n, num_lanes) if cfgs.out_is_3d: @@ -2358,8 +2605,18 @@ def gmm_v2( pallas_out_shape = out_init pallas_out_specs = out_spec + kernel_body = functools.partial(kernel_fn, cfgs=cfgs) + if cfgs.rhs_contract_dequant: + base_kernel_body = kernel_body + + def kernel_body_with_rhs_cache(*refs): + *refs, rhs_cache_ref, rhs_cache_key_ref = refs + base_kernel_body(*refs, rhs_cache_refs=(rhs_cache_ref, rhs_cache_key_ref)) + + kernel_body = kernel_body_with_rhs_cache + out = pl.pallas_call( - functools.partial(kernel_fn, cfgs=cfgs), + kernel_body, out_shape=pallas_out_shape, grid_spec=pltpu.PrefetchScalarGridSpec( num_scalar_prefetch=2, diff --git a/src/maxtext/kernels/ragged/ragged_sort.py b/src/maxtext/kernels/ragged/ragged_sort.py index ccc3fb1c7e..28f483703f 100644 --- a/src/maxtext/kernels/ragged/ragged_sort.py +++ b/src/maxtext/kernels/ragged/ragged_sort.py @@ -85,6 +85,16 @@ def _local_slot_positions(topk_indices_local, num_experts, shard_idx, ep_size, b return pos.reshape(-1) +def _clamp_masked_indices(indices, valid_rows_mask): + """Sets the masked rows' gather indices to 0, as the packed branches do. + + With index_math="count", the inverse permutation is -1 on other shards' slots. The SparseCore gather-reduce kernel + reads the input at the index of every lane of its last partial row block, masked lanes included, without bounds + checks, so -1 is an out-of-bounds read. The masked lanes are discarded, so the output is unchanged. + """ + return jnp.where(valid_rows_mask, indices, 0) + + def ring_ragged_sort( hidden_states_local, topk_indices_local, @@ -295,7 +305,7 @@ def _ring_ragged_sort_bwd(res, g_out): # `topk_weights` is set to ones because this op has no per-row weighting. grad_hidden_states = ragged_gather_reduce( g_x, - topk_argsort_revert_indices, + _clamp_masked_indices(topk_argsort_revert_indices, valid_rows_mask), topk_weights=jnp.ones((n,), dtype=jnp.float32), valid_rows_mask=valid_rows_mask, reduce_group_size=topk, @@ -411,7 +421,7 @@ def _ring_unsort_combine( ) return ragged_gather_reduce( sorted_tokens_local, - topk_argsort_revert_indices, + _clamp_masked_indices(topk_argsort_revert_indices, valid_rows_mask), topk_weights=topk_weights_flat, valid_rows_mask=valid_rows_mask, reduce_group_size=topk, diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index 6a9e1543ba..574c4a03bd 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -888,6 +888,59 @@ class Fp8HoistedAbsmax(nnx.Variable): """One layer's per-channel absmax of its routed-expert weights, computed for the whole scanned stack at once.""" +@dataclasses.dataclass(frozen=True) +class Fp8NativeWeight: + """One gathered routed-expert weight as the native FP8 GMMs take it (fp8_moe_native_fwd). + + Attributes: + anchor: the bf16 weight the gather8 path would multiply by (dequantize_weight_codes_ste's output). The FP8 GMMs + never read its value, so XLA removes its computation; they return the weight gradient as its cotangent, and + dequantize_weight_codes_ste carries it to the master weight (straight-through). + codes: the float8_e4m3fn weight codes [g, k, n]. + scale: the float32 per-(expert, output channel) scale [g, 1, n]. + """ + + anchor: jax.Array + codes: jax.Array + scale: jax.Array + + @classmethod + def pspec(cls, weight_pspec): + """The shard_map in_spec of an Fp8NativeWeight whose anchor and codes are laid out as `weight_pspec`.""" + if weight_pspec is None: + return None + return cls(weight_pspec, weight_pspec, P(weight_pspec[0], None, weight_pspec[2])) + + @classmethod + def concat_last(cls, first, second): + """[first | second] along the output channel (the prefused wi), for split w0 / w1.""" + return cls( + jnp.concatenate([first.anchor, second.anchor], axis=-1), + jnp.concatenate([first.codes, second.codes], axis=-1), + jnp.concatenate([first.scale, second.scale], axis=-1), + ) + + +jax.tree_util.register_dataclass(Fp8NativeWeight, data_fields=["anchor", "codes", "scale"], meta_fields=[]) + + +def fp8_native_fwd_acc_dtype(config) -> jnp.dtype | None: + """fp8_moe_native_fwd_acc_dtype as a dtype, or None for gmm_v2's default.""" + name = getattr(config, "fp8_moe_native_fwd_acc_dtype", "") + return jnp.dtype(name) if name else None + + +def fp8_native_fwd_kwargs(config, with_bwd_dlhs: bool = True) -> dict: + """The FP8 GMM ops' kwargs from the fp8_moe_native_fwd_* (and, with_bwd_dlhs, fp8_moe_native_bwd_dlhs) config.""" + kwargs = { + "acc_dtype": fp8_native_fwd_acc_dtype(config), + "lhs_quant_mode": getattr(config, "fp8_moe_native_fwd_lhs_quant", ""), + } + if with_bwd_dlhs: + kwargs["native_bwd_dlhs"] = getattr(config, "fp8_moe_native_bwd_dlhs", False) + return kwargs + + class GateLogit(nnx.Module): """A layer used to compute gate logits, allowing to return the pre bias values for DeepSeek routing.""" @@ -4052,12 +4105,38 @@ def gmm_up_or_silu( inputs, padding_amount = max_utils.maybe_pad(x, fwd_tile("wi_tile_fwd_batch_seq")) # Weights are stored as (G,K,2N): w0/w1 are adjacent slices (the same concatenate as gmm_up), or w0 is already # the whole [w0 | w1] (moe_unsplit_prefused_wi). - w_fused = w0 if w1 is None else jnp.concatenate([w0, w1], axis=-1) + if w1 is None: + w_fused = w0 + elif isinstance(w0, Fp8NativeWeight): + w_fused = Fp8NativeWeight.concat_last(w0, w1) + else: + w_fused = jnp.concatenate([w0, w1], axis=-1) # wi_tile_fwd_rhs_buffer_count: the fused forward kernel's weight pipeline depth, applied only when that GMM # runs one K tile (its weight block then changes only at expert and n-tile boundaries); otherwise, e.g. with # the eval tiles, the kernel's default buffering. rhs_buffers = self.config.wi_tile_fwd_rhs_buffer_count one_k_tile = not self.config.use_gmm_v2_heuristic_tiling and wi_tile_size[1] >= inputs.shape[-1] + if isinstance(w_fused, Fp8NativeWeight): + # fp8_moe_native_fwd: the FP8 forward of the same fused GMM (the validator requires fuse_dswiglu_in_wo_gmm). + if not fuse_dswiglu_in_wo_gmm: + raise NotImplementedError("fp8_moe_native_fwd requires moe_wo_gmm_fused_dswiglu=True.") + act, gate_up = mblx.gmm_silu_fp8( + inputs.astype(self.dtype), + w_fused.anchor.astype(self.dtype), + w_fused.codes, + w_fused.scale, + gmm_operands["group_sizes"], + preferred_element_type=self.dtype, + tiling=wi_tile_size, + group_offset=gmm_operands["group_offset"], + use_gmm_v2_heuristic_tiling=self.config.use_gmm_v2_heuristic_tiling, + use_dlhs_transpose_rhs=self.config.moe_gmm_v2_dlhs_transpose_rhs, + fuse_dlhs_scale=self.config.moe_accumulate_wi_dlhs, + fuse_3d_dlhs_ps=self.config.moe_3d_dlhs_partial_sum, + fwd_rhs_buffer_count=rhs_buffers if rhs_buffers and one_k_tile else None, + **fp8_native_fwd_kwargs(self.config), + ) + return (act[: x.shape[0]], gate_up[: x.shape[0]]) if padding_amount else (act, gate_up) out = mblx.gmm_silu( inputs.astype(self.dtype), w_fused.astype(self.dtype), @@ -4133,6 +4212,24 @@ def _run_wo_gmm_dswiglu( raise NotImplementedError(f"moe_wo_gmm_fused_dswiglu does not support {', '.join(unsupported)}.") padded_h, padding_amount = max_utils.maybe_pad(intermediate_layer, fwd_tile("wo_tile_fwd_batch_seq")) padded_gate_up = jnp.pad(gate_up, ((0, padding_amount), (0, 0))) if padding_amount > 0 else gate_up + if isinstance(wo, Fp8NativeWeight): + # fp8_moe_native_fwd: the FP8 forward of the same GMM; the backward is gmm_dswiglu's. + out = mblx.gmm_dswiglu_fp8( + padded_h.astype(self.dtype), + padded_gate_up.astype(self.dtype), + wo.anchor.astype(self.dtype), + wo.codes, + wo.scale, + gmm_fn.keywords["group_sizes"], + preferred_element_type=self.dtype, + tiling=wo_tile_size, + group_offset=gmm_fn.keywords["group_offset"], + use_gmm_v2_heuristic_tiling=self.config.use_gmm_v2_heuristic_tiling, + use_dlhs_transpose_rhs=self.config.moe_gmm_v2_dlhs_transpose_rhs, + dact_pre_act_tile_n=_dswiglu_pre_act_tile_n(gate_up), + **fp8_native_fwd_kwargs(self.config), + ) + return out[: intermediate_layer.shape[0]] if padding_amount > 0 else out out = mblx.gmm_dswiglu( padded_h.astype(self.dtype), padded_gate_up.astype(self.dtype), @@ -4491,6 +4588,24 @@ def _ring_down_project_and_unpermute( wo_gather_axes, wo_tile_size = get_wo_gmm_params() def down_projection(h, kernel, group_sizes, expert_assignments, group_offset): + if isinstance(kernel, Fp8NativeWeight): + # fp8_moe_native_fwd: the FP8 forward. No VJP needed: ring_ragged_unsort_down_projection owns the + # backward (down_projection_bwd_dswiglu below). Rows padded to the forward tile, as gmm() pads them. + del expert_assignments + num_rows = h.shape[0] + h, padding_amount = max_utils.maybe_pad(h, fwd_tile("wo_tile_fwd_batch_seq")) + out = mblx.gmm_fp8_forward( + h.astype(self.dtype), + kernel.codes, + kernel.scale, + group_sizes, + preferred_element_type=self.dtype, + tiling=wo_tile_size, + group_offset=group_offset, + use_gmm_v2_heuristic_tiling=self.config.use_gmm_v2_heuristic_tiling, + **fp8_native_fwd_kwargs(self.config, with_bwd_dlhs=False), + ) + return out[:num_rows] if padding_amount > 0 else out return gmm( h, kernel, @@ -4508,6 +4623,16 @@ def down_projection_bwd_dswiglu( with the SwiGLU backward in the DLHS epilogue (megablox gmm_dswiglu_bwd): (d[gate | up], , dwo). """ del expert_assignments + fp8_kernel = kernel if isinstance(kernel, Fp8NativeWeight) else None + use_fp8_dlhs = fp8_kernel is not None and getattr(self.config, "fp8_moe_native_bwd_dlhs", False) + fp8_dlhs = {} + if use_fp8_dlhs: + # fp8_moe_native_bwd_dlhs: the codes go to the backward, whose DLHS dequantizes them in VMEM (same values). + kernel = fp8_kernel.codes + fp8_dlhs = {"rhs_fp8_scale": fp8_kernel.scale, "rhs_dtype": self.dtype} + elif fp8_kernel is not None: + # fp8_moe_native_fwd: the bf16 backward at the dequantized weight (bit-identical to the gather8 one). + kernel = mblx.dequantize_fp8_codes(fp8_kernel.codes, fp8_kernel.scale, self.dtype) num_rows = gathered.shape[0] # Rows padded to the forward tile, as gmm() pads them; the padding is in no group. gathered, padding_amount = max_utils.maybe_pad(gathered, fwd_tile("wo_tile_fwd_batch_seq")) @@ -4518,7 +4643,7 @@ def down_projection_bwd_dswiglu( d_gate_up, dot, grad_wo = mblx.gmm_dswiglu_bwd( gathered.astype(self.dtype), weighted_h.astype(self.dtype), - kernel.astype(self.dtype), + kernel if use_fp8_dlhs else kernel.astype(self.dtype), group_sizes, group_offset, self.dtype, @@ -4529,10 +4654,18 @@ def down_projection_bwd_dswiglu( use_gmm_v2_heuristic_tiling=self.config.use_gmm_v2_heuristic_tiling, use_dlhs_transpose_rhs=self.config.moe_gmm_v2_dlhs_transpose_rhs, dact_pre_act_tile_n=_dswiglu_pre_act_tile_n(dact_gate_up), + **fp8_dlhs, ) if padding_amount > 0: # dot is an array here (return_dact_dot=True). d_gate_up, dot = d_gate_up[:num_rows], dot[:num_rows] # pylint: disable=unsubscriptable-object + if fp8_kernel is not None: + # The cotangent of the Fp8NativeWeight: the weight gradient to the anchor; the codes and the scale get none. + grad_wo = Fp8NativeWeight( + grad_wo.astype(fp8_kernel.anchor.dtype), + jnp.zeros_like(fp8_kernel.codes), + jnp.zeros_like(fp8_kernel.scale), + ) return d_gate_up, dot, grad_wo operands = get_gmm_operands_for_local_experts(intermediate_layer, routing, route_metadata) @@ -4907,6 +5040,16 @@ def _cn(t): # Not every mesh axis: shard_map rejects axes an enclosing vmap mapped away ("stage", "diloco"). input_mesh_axes = {axis for spec in in_specs if spec is not None for dim in spec for axis in mesh_axes_for_dim(dim)} overflow_mesh_axes = tuple(axis for axis in self.mesh.axis_names if axis in input_mesh_axes) + # fp8_moe_native_fwd: the expert weights enter as Fp8NativeWeight (anchor, codes, scale); see below. + use_fp8_native_fwd = getattr(self.config, "fp8_moe_native_fwd", False) and fp8_fake_quant_at_gather is not None + if use_fp8_native_fwd: + weight_specs = slice(3, 6) # w0_pspec, w1_pspec, wo_pspec in in_specs above + assert in_specs[weight_specs] == (w0_pspec, w1_pspec, wo_pspec), "in_specs order changed" + native_weight_specs = tuple( + spec if kernel is None else Fp8NativeWeight.pspec(spec) + for kernel, spec in ((w0_kernel, w0_pspec), (w1_kernel, w1_pspec), (wo_kernel, wo_pspec)) + ) + in_specs = in_specs[: weight_specs.start] + native_weight_specs + in_specs[weight_specs.stop :] @functools.partial( jax.shard_map, @@ -5537,13 +5680,21 @@ def _fsdp_all_gather(w, pspec): if fp8_dequantize_after_gather is not None: # After the barrier, so the barrier orders the (cheaper) code gathers and the dequantization can fuse # into the GEMM operand prologue. + def _after_gather(codes, deq): + if deq is None: + return None + master_weight, scale, grad_pspec = deq + w_deq = quantizations.dequantize_weight_codes_ste( + master_weight, codes, scale, grad_transform=functools.partial(self._maybe_shard_with_pspec, pspec=grad_pspec) + ) + if not use_fp8_native_fwd: + return w_deq + # fp8_moe_native_fwd: the GMMs read the codes and the scale; w_deq is only the gradient anchor (its value is + # never read, so XLA removes the dequantization from the forward). + return Fp8NativeWeight(w_deq, codes, scale) + w0_kernel, w1_kernel, wo_kernel = ( - None - if deq is None - else quantizations.dequantize_weight_codes_ste( - deq[0], codes, deq[1], grad_transform=functools.partial(self._maybe_shard_with_pspec, pspec=deq[2]) - ) - for codes, deq in zip((w0_kernel, w1_kernel, wo_kernel), fp8_dequantize_after_gather) + _after_gather(codes, deq) for codes, deq in zip((w0_kernel, w1_kernel, wo_kernel), fp8_dequantize_after_gather) ) route_and_compute_outputs = sparse_matmul_route_and_compute( @@ -6664,6 +6815,14 @@ def _maybe_native_fused_weight(kernel, scale): if wo_absmax is not None and cfg.sparse_matmul and fp8_fake_quant_hoist_mode(cfg) == "gather8": fp8_fake_quant_at_gather = (w0_absmax, w1_absmax, wo_absmax) else: + if getattr(cfg, "fp8_moe_native_fwd", False): + # The validator guarantees gather8 and sparse_matmul, so the absmax is missing: this call is outside the + # decoder's layer scan. Without the codes, sparse_matmul would quietly run fake QAT. + raise ValueError( + "fp8_moe_native_fwd=True needs the gather8 absmax that the decoder's layer scan attaches" + " (stacked_fp8_fake_quant_absmax); this RoutedMoE call has none (unscanned layers, KV-cache decode," + " or a Qwix-intercepted call), so it would silently run fake QAT instead of the FP8 forward." + ) w0_kernel = quantizations.fake_quantize_weight_per_channel(w0_kernel, absmax=w0_absmax) if w1_kernel is not None: # None: moe_unsplit_prefused_wi, and the line above rounded both halves at once w1_kernel = quantizations.fake_quantize_weight_per_channel(w1_kernel, absmax=w1_absmax) diff --git a/tests/unit/fp8_moe_native_test.py b/tests/unit/fp8_moe_native_test.py new file mode 100644 index 0000000000..3cdfe8ddf7 --- /dev/null +++ b/tests/unit/fp8_moe_native_test.py @@ -0,0 +1,204 @@ +# 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. +"""fp8_moe_native_fwd / fp8_moe_native_bwd_dlhs on CPU: config validation, the outside-layer-scan guard, and the +dequantization and lhs-quant-mode contracts the native path relies on. +""" + +import sys +import typing + +from flax import nnx +from flax.linen import partitioning as nn_partitioning +import jax +import jax.numpy as jnp +from jax.sharding import Mesh +import numpy as np +import pytest + +from maxtext.configs import pyconfig +from maxtext.configs import types as config_types +from maxtext.kernels.megablox import ops as megablox_ops +from maxtext.layers import quantizations +from maxtext.utils import maxtext_utils +from maxtext.utils import maxtext_utils_nnx +from maxtext.utils import model_creation_utils +from tests.utils.test_helpers import get_test_config_path + +pytestmark = pytest.mark.cpu_only + +_MOE_MLP = 128 +_SEQ = 64 +_VOCAB = 512 +_EXPERTS = 8 +_TOP_K = 2 + +# A shrunk Qwen3.5 MoE model on one device, with the fake-QAT gather8 path the native flags build on. +_MODEL = { + "model_name": "qwen3.5-397b-a17b", + "override_model_config": True, + "base_num_decoder_layers": 4, + "base_emb_dim": 256, + "base_num_query_heads": 4, + "base_num_kv_heads": 2, + "head_dim": 64, + "vocab_size": _VOCAB, + "base_mlp_dim": _MOE_MLP, + "base_moe_mlp_dim": _MOE_MLP, + "num_experts": _EXPERTS, + "num_experts_per_tok": _TOP_K, + "gdn_num_key_heads": 2, + "gdn_num_value_heads": 4, + "gdn_key_head_dim": 32, + "gdn_value_head_dim": 32, + "gdn_chunk_size": 16, + "use_gdn_kernel": False, + "use_mrope": False, + "max_target_length": _SEQ, + "max_prefill_predict_length": _SEQ, + "per_device_batch_size": 2, + "dtype": "bfloat16", + "weight_dtype": "bfloat16", + "attention": "dot_product", + "use_tokamax_splash": False, + "prefuse_moe_weights": True, + "moe_unsplit_prefused_wi": True, + "megablox": True, + "sparse_matmul": True, + "fp8_moe_fake_quant": True, + "fp8_fake_quant_hoist": "gather8", + "ragged_buffer_factor": -1.0, + "scan_layers": True, + "enable_dropout": False, +} + +# What fp8_moe_native_fwd needs on top of the fake-QAT gather8 model. +_NATIVE = { + "fp8_moe_native_fwd": True, + "use_tokamax_gmm": True, + "use_gmm_v2": True, + "moe_wi_gmm_fused_swiglu": True, + "moe_wo_gmm_fused_dswiglu": True, + "wo_tile_dlhs_mlp_dim": _MOE_MLP, +} + +_NATIVE_BWD_DLHS = {"fp8_moe_native_bwd_dlhs": True, "moe_gmm_v2_dlhs_transpose_rhs": True} + + +def _cfg(**overrides): + kwargs = { + "run_name": "fp8_moe_native", + "enable_checkpointing": False, + "log_config": False, + "skip_jax_distributed_system": True, + **_MODEL, + **overrides, + } + return pyconfig.initialize([sys.argv[0], get_test_config_path()], **kwargs) + + +def test_flags_are_defined_and_off_by_default(): + for key in ( + "fp8_moe_native_fwd", + "fp8_moe_native_fwd_acc_dtype", + "fp8_moe_native_fwd_lhs_quant", + "fp8_moe_native_bwd_dlhs", + ): + assert key in set(config_types.MoEGeneral.model_fields), key + cfg = _cfg() + assert not cfg.fp8_moe_native_fwd and not cfg.fp8_moe_native_bwd_dlhs + assert cfg.fp8_moe_native_fwd_acc_dtype == "" and cfg.fp8_moe_native_fwd_lhs_quant == "" + + +def test_native_config_is_accepted(): + cfg = _cfg( + **_NATIVE, **_NATIVE_BWD_DLHS, fp8_moe_native_fwd_lhs_quant="fullk_f32", fp8_moe_native_fwd_acc_dtype="float32" + ) + assert cfg.fp8_moe_native_fwd and cfg.fp8_moe_native_bwd_dlhs + + +@pytest.mark.parametrize( + "override,requirement", + [ + ({"fp8_moe_fake_quant": False}, "fp8_moe_fake_quant=True"), + ({"fp8_fake_quant_hoist": "absmax"}, "fp8_fake_quant_hoist=gather8"), + ({"use_gmm_v2": False}, "use_tokamax_gmm=True and use_gmm_v2=True"), + ], +) +def test_native_fwd_requirement_is_enforced(override, requirement): + with pytest.raises(Exception, match=requirement): + _cfg(**{**_NATIVE, **override}) + + +@pytest.mark.parametrize( + "override,requirement", + [ + ({"fp8_moe_native_fwd": False}, "fp8_moe_native_fwd=True"), + ({"moe_gmm_v2_dlhs_transpose_rhs": False}, "moe_gmm_v2_dlhs_transpose_rhs=True"), + ({"use_gmm_v2_heuristic_tiling": True}, "use_gmm_v2_heuristic_tiling=False"), + ], +) +def test_native_bwd_dlhs_requirement_is_enforced(override, requirement): + with pytest.raises(Exception, match=requirement): + _cfg(**{**_NATIVE, **_NATIVE_BWD_DLHS, **override}) + + +def _lower_forward(cfg): + """Builds the model and lowers its forward on CPU (no compile).""" + mesh = Mesh(maxtext_utils.create_device_mesh(cfg), cfg.mesh_axes) + create_fn, abstract_model = model_creation_utils.create_nnx_abstract_model(cfg, mesh=mesh) + model = maxtext_utils_nnx.create_nnx_sharded_model(abstract_model, create_fn, mesh=mesh) + batch = int(cfg.per_device_batch_size * jax.device_count()) + rng = np.random.default_rng(0) + tokens = jnp.asarray(rng.integers(1, _VOCAB, size=(batch, _SEQ), dtype=np.int32)) + positions = jnp.broadcast_to(jnp.arange(_SEQ, dtype=jnp.int32), (batch, _SEQ)) + segments = jnp.ones((batch, _SEQ), jnp.int32) + with jax.set_mesh(mesh), nn_partitioning.axis_rules(cfg.logical_axis_rules): + graphdef, params, rest = nnx.split(model, nnx.Param, ...) + + def forward(p): + return nnx.merge(graphdef, p, rest, copy=True)( + tokens, positions, decoder_segment_ids=segments, enable_dropout=False + ) + + jax.jit(forward).lower(params) + + +def test_native_fwd_outside_the_layer_scan_raises(): + # Unscanned layers get no gather8 absmax; without the guard the layer would quietly run fake QAT. + with pytest.raises(ValueError, match="needs the gather8 absmax"): + _lower_forward(_cfg(**_NATIVE, scan_layers=False)) + + +def test_fake_quant_gather8_outside_the_layer_scan_still_lowers(): + # Control: the same unscanned model without fp8_moe_native_fwd lowers (fake QAT rounds the weights in the layer). + _lower_forward(_cfg(scan_layers=False)) + + +def test_dequantize_fp8_codes_matches_the_gather8_dequantization(): + # The native backward multiplies by dequantize_fp8_codes; fake QAT's gather8 path by dequantize_weight_codes_ste. + # The native path claims bit-identical gradients, so the two expressions must agree bit for bit. + rng = np.random.default_rng(0) + weight = jnp.asarray(rng.standard_normal((4, 64, 32)) * 0.02, jnp.bfloat16) + codes, scale = quantizations.quantize_weight_codes_per_channel(weight) + native = megablox_ops.dequantize_fp8_codes(codes, scale, jnp.bfloat16) + gather8 = quantizations.dequantize_weight_codes_ste(weight, codes, scale) + assert native.dtype == gather8.dtype == jnp.bfloat16 + np.testing.assert_array_equal(np.asarray(native).view(np.uint16), np.asarray(gather8).view(np.uint16)) + + +def test_lhs_quant_modes_match_the_config_values(): + # A value the config accepts but ops does not know would fail only at trace time. + field = dict(config_types.MoEGeneral.model_fields)["fp8_moe_native_fwd_lhs_quant"] + config_values = set(typing.get_args(field.annotation)) + assert config_values == set(megablox_ops._LHS_QUANT_MODES) # pylint: disable=protected-access diff --git a/tests/unit/ragged_sort_full_buffer_indices_test.py b/tests/unit/ragged_sort_full_buffer_indices_test.py new file mode 100644 index 0000000000..714b85cd8d --- /dev/null +++ b/tests/unit/ragged_sort_full_buffer_indices_test.py @@ -0,0 +1,185 @@ +# 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. +"""Full-buffer (dropless) ring-of-experts gather-reduces receive in-bounds indices. + +With moe_routing_index_math="count", ring_ragged_sort's inverse permutation is -1 on the slots of other shards. The +full-buffer branches of `_ring_unsort_combine` and of ring_ragged_sort's backward must clamp those masked lanes before +`ragged_gather_reduce`: its SparseCore kernel reads every lane of the last row block without bounds checks. + +On CPU, `ragged_gather_reduce` uses its jnp fallback, so a spy on it checks the arguments of every call instead of the +TPU fault. The tests are CPU-only: the last one deliberately passes the unclamped indices. +""" + +from unittest import mock + +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +from maxtext.kernels.ragged import ragged_gather_reduce_v2 +from maxtext.kernels.ragged import ragged_sort + +pytestmark = pytest.mark.cpu_only + +_HIDDEN = 64 +_GATHER_REDUCE_KWARGS = { + "enforce_fallback": True, + "flops_override": -1, + "bytes_accessed_override": -1, + "use_single_sparsecore": False, +} + + +def _bits_equal(a, b): + a, b = np.asarray(a), np.asarray(b) + if a.shape != b.shape or a.dtype != b.dtype: + return False + return bool(np.array_equal(a.reshape(-1).view(np.uint8), b.reshape(-1).view(np.uint8))) + + +def _routing(seed, tokens, experts, topk): + k = jax.random.split(jax.random.PRNGKey(seed), 3) + ids = jax.lax.top_k(jax.random.normal(k[0], (tokens, experts)), topk)[1].astype(jnp.int32) + w = jax.random.uniform(k[1], (tokens * topk,), jnp.float32, 0.25, 0.3) + return ids, w, k[2] + + +def _maps(ids, experts, ep, shard, index_math): + """This shard's window and the inverse permutation ring_ragged_sort builds for `index_math`.""" + flat = ids.reshape(-1) + local = experts // ep + offs = jnp.concatenate([jnp.zeros((1,), jnp.int32), jnp.cumsum(jax.nn.one_hot(flat, experts, dtype=jnp.int32).sum(0))]) + start, end = offs[shard * local], offs[(shard + 1) * local] + if index_math == "count": + pos = ragged_sort._local_slot_positions(ids, experts, shard, ep) # pylint: disable=protected-access + revert = jnp.where(pos >= 0, start + pos, -1).astype(jnp.int32) + else: + revert = jnp.argsort(jnp.argsort(flat, stable=True)).astype(jnp.int32) + return start, end, revert + + +class _Spy: + """Wraps ragged_sort.ragged_gather_reduce; records (min index, max index, rows) of every call at run time.""" + + def __init__(self, declamp=False): + self.calls = [] + self.declamp = declamp + self.real = ragged_sort.ragged_gather_reduce + + def __call__(self, x, indices, *args, **kwargs): + if self.declamp: + # The arguments before the fix: the masked lanes of a full buffer carry the raw inverse (-1 under "count"). + indices = jnp.where(kwargs["valid_rows_mask"], indices, -1) + rows = x.shape[0] + jax.debug.callback(lambda lo, hi: self.calls.append((int(lo), int(hi), rows)), jnp.min(indices), jnp.max(indices)) + return self.real(x, indices, *args, **kwargs) + + def out_of_bounds(self): + return [c for c in self.calls if c[0] < 0 or c[1] >= c[2]] + + +_COMBINE_CASES = [(ep, m) for ep in (4, 8) for m in ("count", "sort")] + + +@pytest.mark.parametrize("ep,index_math", _COMBINE_CASES, ids=[f"ep{ep}-{m}" for ep, m in _COMBINE_CASES]) +def test_full_buffer_combine_in_bounds_and_unchanged(ep, index_math): + tokens, experts, topk = 256, 2 * ep, 4 + n = tokens * topk + for seed in range(3): + ids, w, key = _routing(seed, tokens, experts, topk) + x = jax.random.normal(key, (n, _HIDDEN), jnp.float32).astype(jnp.bfloat16) # full buffer: rows at global positions + for shard in range(ep): + start, end, revert = _maps(ids, experts, ep, shard, index_math) + if index_math == "count": + assert int(jnp.sum(revert < 0)) > 0, "count must leave other shards' slots at -1" + spy = _Spy() + with mock.patch.object(ragged_sort, "ragged_gather_reduce", spy): + out = jax.jit( + lambda x, r, w, s, e: ragged_sort._ring_unsort_combine( # pylint: disable=protected-access + x, r, w, s, e, topk, _GATHER_REDUCE_KWARGS + ) + )(x, revert, w, start, end) + out.block_until_ready() + jax.effects_barrier() + assert len(spy.calls) == 1, spy.calls + assert not spy.out_of_bounds(), f"seed{seed} shard{shard}: indices out of [0, {n}): {spy.calls}" + # The call before the fix: the raw inverse. The masked lanes are discarded, so the output is the same bits. + valid = (revert >= start) & (revert < end) + before = ragged_gather_reduce_v2.ragged_gather_reduce( + x=x, + indices=revert, + topk_weights=w, + valid_rows_mask=valid, + reduce_group_size=topk, + enforce_fallback=True, + ) + assert _bits_equal(out, before), f"seed{seed} shard{shard}: output changed" + # float64 reference: each token's routing-weighted sum over its slots in this shard's window. + x_np, revert_np, w_np, valid_np = (np.asarray(a) for a in (x.astype(jnp.float32), revert, w, valid)) + gathered = x_np[np.where(valid_np, revert_np, 0)] * w_np[:, None] + ref = np.where(valid_np[:, None], gathered, 0.0).reshape((tokens, topk, _HIDDEN)).sum(1) + np.testing.assert_allclose(np.asarray(out, np.float32), ref, rtol=2e-2, atol=2e-2) + + +def _op(ep, index_math, tokens=128, topk=4, seed=0): + """ring_ragged_sort (full buffer) -> ring_ragged_unsort per shard, vmapped over "ep": (out, dx, dw).""" + experts = 2 * ep + local = experts // ep + ids, w, key = _routing(seed, tokens, experts, topk) + k = jax.random.split(key, 2) + x = jax.random.normal(k[0], (tokens, _HIDDEN), jnp.float32).astype(jnp.bfloat16) + cot = jax.random.normal(k[1], (ep, tokens, _HIDDEN), jnp.float32).astype(jnp.bfloat16) + + def per_shard(x, w, cot): + def f(x, w): + sorted_x, gs, revert, perm = ragged_sort.ring_ragged_sort( + x, ids, experts, topk, "ep", ep, buffer_size=None, index_math=index_math + ) + assert sorted_x.shape[0] == tokens * topk # the full buffer + return ragged_sort.ring_ragged_unsort( + sorted_x * 2, gs, revert, topk, local, "ep", w, topk_argsort_indices=perm, index_math=index_math + ) + + out, vjp = jax.vjp(f, x, w) + dx, dw = vjp(cot) + return out, dx, dw + + return jax.jit(jax.vmap(per_shard, in_axes=(None, None, 0), axis_name="ep"))(x, w, cot) + + +@pytest.mark.parametrize("ep", [4, 8]) +def test_full_buffer_op_in_bounds_and_count_matches_sort(ep): + results = {} + for index_math in ("count", "sort"): + spy = _Spy() + with mock.patch.object(ragged_sort, "ragged_gather_reduce", spy): + results[index_math] = jax.block_until_ready(_op(ep, index_math)) + jax.effects_barrier() + # Per shard: one combine (forward) and one sort backward. + assert len(spy.calls) == 2 * ep, spy.calls + assert not spy.out_of_bounds(), f"{index_math}: indices out of bounds: {spy.out_of_bounds()}" + for name, a, b in zip(("out", "dx", "dw"), results["count"], results["sort"]): + assert _bits_equal(a, b), f"ep{ep}: {name} differs between count and sort" + + +def test_unclamped_indices_are_flagged(): + """The arguments before the fix (-1 on the masked lanes) fail the bounds check, in both call sites.""" + ep = 4 + spy = _Spy(declamp=True) + with mock.patch.object(ragged_sort, "ragged_gather_reduce", spy): + jax.block_until_ready(_op(ep, "count")) + jax.effects_barrier() + assert len(spy.calls) == 2 * ep, spy.calls + assert len(spy.out_of_bounds()) == 2 * ep, spy.calls From 10fdc4fa92733f40fef75542cbec3a3a3c602bad Mon Sep 17 00:00:00 2001 From: Shuning Jin Date: Sat, 31 Oct 2026 00:18:20 +0000 Subject: [PATCH 2/4] Address review: fixed native-forward numerics, fixed-tile requirement, multi-K DLHS doc; split out the ragged_sort fix - Remove fp8_moe_native_fwd_lhs_quant and fp8_moe_native_fwd_acc_dtype: fp8_moe_native_fwd always quantizes the activation with one float32 absmax scale per row over the full K dim and accumulates in float32 (the setting that measured closest to the rollout's FP8 MoE kernel and fastest); the other modes stay ops-level arguments for benchmarks and tests. Lowered programs are unchanged vs the old explicit fullk_f32 + float32 config. - fp8_moe_native_fwd requires use_gmm_v2_heuristic_tiling=False: the heuristic tiling budgets a bf16 accumulator when the lhs is quantized in the kernel, so with the float32 accumulator its tiles could exceed VMEM. - _rhs_cache_miss: document that the DLHS dequant cache only hits with a single K tile; with more, each step dequantizes inside its matmul, unchanged outputs, still faster than the XLA dequant fallback. - Move the ragged_sort full-buffer index clamp and its test to a separate PR. --- src/maxtext/configs/base.yml | 14 +- src/maxtext/configs/types.py | 39 +--- src/maxtext/kernels/megablox/ops.py | 4 +- .../pallas_mosaic_tpu_v2_gmm_kernel.py | 8 +- src/maxtext/kernels/ragged/ragged_sort.py | 14 +- src/maxtext/layers/moe.py | 16 +- tests/unit/fp8_moe_native_test.py | 33 ++-- .../ragged_sort_full_buffer_indices_test.py | 185 ------------------ 8 files changed, 45 insertions(+), 268 deletions(-) delete mode 100644 tests/unit/ragged_sort_full_buffer_indices_test.py diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index c8bfc68002..0a40440180 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -348,15 +348,13 @@ fp8_moe_fake_quant: false # gather. Env PARAMBOLE_FQ_HOIST (0, absmax, 1 = gather8) selects a mode while this is off. fp8_fake_quant_hoist: "off" # With fp8_moe_fake_quant and fp8_fake_quant_hoist=gather8, run the routed-expert wi/wo forward GMMs in FP8 (gathered -# codes + per-channel scales into gmm_v2, activation quantized in the kernel). Backward: bf16 DLHS / TGMM at the -# dequantized weight, unless fp8_moe_native_bwd_dlhs dequantizes it inside the DLHS kernel. -# Needs sparse_matmul, moe_wi_gmm_fused_swiglu, moe_wo_gmm_fused_dswiglu, use_tokamax_gmm and use_gmm_v2, and runs only -# inside the decoder's layer scan (where gather8's absmax is attached): a routed-MoE call outside it raises. +# codes + per-channel scales into gmm_v2; the activation is quantized in the kernel with one float32 absmax scale per +# row over the full K dim, and accumulated in float32). Backward: bf16 DLHS / TGMM at the dequantized weight, unless +# fp8_moe_native_bwd_dlhs dequantizes it inside the DLHS kernel. +# Needs sparse_matmul, moe_wi_gmm_fused_swiglu, moe_wo_gmm_fused_dswiglu, use_tokamax_gmm, use_gmm_v2 and fixed tiles +# (use_gmm_v2_heuristic_tiling=false), and runs only inside the decoder's layer scan (where gather8's absmax is +# attached): a routed-MoE call outside it raises. fp8_moe_native_fwd: false -# gmm_v2 accumulator across K blocks for fp8_moe_native_fwd: "" (kernel default, bfloat16), "bfloat16" or "float32". -fp8_moe_native_fwd_acc_dtype: "" -# In-kernel lhs quantization scheme for fp8_moe_native_fwd ("" = default block512_bf16; see types.py). -fp8_moe_native_fwd_lhs_quant: "" # fp8_moe_native_fwd: the backward DLHS GMMs read the fp8 weight codes and dequantize them in VMEM inside gmm_v2 # (bit-identical; no XLA dequantization of the expert weights). Needs moe_gmm_v2_dlhs_transpose_rhs and fixed tiles. fp8_moe_native_bwd_dlhs: false diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 2a7556c527..eb3c797129 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -1943,42 +1943,14 @@ def _off_for_false_fp8_fake_quant_hoist(v: Any) -> Any: "Native FP8 forward for the routed experts. With fp8_moe_fake_quant and fp8_fake_quant_hoist=gather8," " the wi and wo forward GMMs run on the MXU in float8_e4m3fn: the gathered weight codes and their" " per-channel scales feed gmm_v2 directly (no dequantization), and gmm_v2 quantizes the bf16 activation" - " in the kernel with a dynamic per-row absmax scale (see fp8_moe_native_fwd_lhs_quant). The backward runs" + " in the kernel with one dynamic float32 absmax scale per row over the full K dim (the rollout FP8 MoE kernel's" + " decode scheme) and accumulates in float32. The backward runs" " bf16 DLHS / TGMM at the dequantized weight (straight-through weight gradient), unless" " fp8_moe_native_bwd_dlhs moves the DLHS dequantization into the kernel. Runs only inside the decoder's" " layer scan, where gather8's absmax is attached; a RoutedMoE call outside it (unscanned layers, KV-cache" " decode) raises instead of silently running fake QAT." ), ) - fp8_moe_native_fwd_acc_dtype: Literal["", "bfloat16", "float32"] = Field( - "", - description=( - "fp8_moe_native_fwd: gmm_v2's accumulator across K blocks. '' keeps the kernel default (bfloat16 when the" - " lhs is quantized in the kernel, as in the rollout's kernel); 'float32' accumulates in float32." - ), - ) - fp8_moe_native_fwd_lhs_quant: Literal[ - "", - "block512_bf16", - "block512_bf16_rhs1", - "block512_f32", - "block512_f32_hoist", - "fullk_bf16", - "fullk_bf16_hoist", - "fullk_f32", - "fullk_f32_hoist", - ] = Field( - "", - description=( - "fp8_moe_native_fwd: in-kernel lhs (activation) quantization scheme. 'block512_*': one scale per row and" - " 512-column block ('' / 'block512_bf16' is gmm_v2's default); 'fullk_*': one scale per row over the full" - " K dim. '*_bf16': bf16 absmax and reciprocal; '*_f32': float32 max(absmax, 1e-12) / 448 and a true" - " division (with fullk, the rollout's FP8 MoE kernel's scheme). 'block512_bf16_rhs1': the weight scale is" - " applied to the block accumulator before the activation scale. '*_hoist': each lhs block is quantized once" - " per tile instead of per mxu-wide n slice; same values and, since the compiler deduplicates the repeated" - " quantization, the same speed." - ), - ) fp8_moe_native_bwd_dlhs: bool = Field( False, description=( @@ -1994,8 +1966,8 @@ def _off_for_false_fp8_fake_quant_hoist(v: Any) -> Any: def validate_fp8_moe_native_fwd(self) -> "MoEGeneral": """Rejects fp8_moe_native_fwd / fp8_moe_native_bwd_dlhs without the config they require. - The FP8 GMMs replace the gather8 dequantization on the fused gmm_v2 SwiGLU / dSwiGLU path only, and the - in-kernel DLHS dequantization needs the transposed DLHS rhs and fixed tiles. + The FP8 GMMs replace the gather8 dequantization on the fused gmm_v2 SwiGLU / dSwiGLU path only and need fixed + tiles (float32 accumulator); the in-kernel DLHS dequantization also needs the transposed DLHS rhs. """ if self.fp8_moe_native_fwd: _reject_unless( @@ -2007,6 +1979,9 @@ def validate_fp8_moe_native_fwd(self) -> "MoEGeneral": (self.use_tokamax_gmm and self.use_gmm_v2, "use_tokamax_gmm=True and use_gmm_v2=True"), (self.moe_wi_gmm_fused_swiglu, "moe_wi_gmm_fused_swiglu=True"), (self.moe_wo_gmm_fused_dswiglu, "moe_wo_gmm_fused_dswiglu=True"), + # gmm_v2's heuristic tiling budgets a bfloat16 accumulator when the lhs is quantized in the kernel, + # but the native forward accumulates in float32, so heuristic tiles could exceed VMEM. + (not self.use_gmm_v2_heuristic_tiling, "use_gmm_v2_heuristic_tiling=False"), ], ) if self.fp8_moe_native_bwd_dlhs: diff --git a/src/maxtext/kernels/megablox/ops.py b/src/maxtext/kernels/megablox/ops.py index f8a6be3b3f..e47ff0b38b 100644 --- a/src/maxtext/kernels/megablox/ops.py +++ b/src/maxtext/kernels/megablox/ops.py @@ -2004,7 +2004,7 @@ def _acc_kwargs(acc_dtype): def _lhs_quant_kwargs(lhs_quant_mode: str) -> dict: - """gmm_v2's in-kernel lhs quantization kwargs for a fp8_moe_native_fwd_lhs_quant mode.""" + """gmm_v2's in-kernel lhs quantization kwargs for an lhs_quant_mode (fp8_moe_native_fwd uses "fullk_f32").""" if lhs_quant_mode not in _LHS_QUANT_MODES: raise ValueError(f"Unknown lhs_quant_mode {lhs_quant_mode!r}; expected one of {list(_LHS_QUANT_MODES)}.") return _LHS_QUANT_MODES[lhs_quant_mode] @@ -2035,7 +2035,7 @@ def gmm_silu_fp8( gate_up. acc_dtype: gmm_v2's accumulator across K blocks (None: the kernel default). lhs_quant_mode: the in-kernel lhs - quantization scheme, a key of _LHS_QUANT_MODES (fp8_moe_native_fwd_lhs_quant). native_bwd_dlhs + quantization scheme, a key of _LHS_QUANT_MODES (fp8_moe_native_fwd passes "fullk_f32"). native_bwd_dlhs (fp8_moe_native_bwd_dlhs): the backward DLHS dequantizes the codes inside the kernel (same values). The other arguments are gmm_silu's. """ diff --git a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py index 912301e966..8a2ce1ad56 100644 --- a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py +++ b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py @@ -1166,8 +1166,12 @@ def _rhs_cache_miss( ) -> jax.Array: """rhs_contract_dequant: whether the cache tile holds another rhs block than this step's. - The rhs block index is (group, n tile, k tile), so with a single n and k tile (the MoE DLHS tiles) the cache is - filled once per group and every m tile of the group reuses it. Side effect: records this step's block key in the + The rhs block index is (group, n tile, k tile), and k is the innermost grid axis. With a single k tile (the MoE + DLHS tiles) the cache is filled once per (group, n tile) and every m tile of the group reuses it, for any number + of n tiles (the DLHS grid is (n, gm, k)). With two or more k tiles consecutive steps need different blocks, so the + cache never hits and every step dequantizes its block inside the matmul; outputs are unchanged, and it still beats + dequantizing the whole weight in XLA first (TPU7x, production DLHS shapes, 2-4 k tiles: 0.84-0.91x its time), + because the kernel reads 1-byte codes instead of a bf16 copy. Side effect: records this step's block key in the SMEM key, since the step that misses fills the cache tile. """ assert rhs_cache_refs is not None diff --git a/src/maxtext/kernels/ragged/ragged_sort.py b/src/maxtext/kernels/ragged/ragged_sort.py index 28f483703f..ccc3fb1c7e 100644 --- a/src/maxtext/kernels/ragged/ragged_sort.py +++ b/src/maxtext/kernels/ragged/ragged_sort.py @@ -85,16 +85,6 @@ def _local_slot_positions(topk_indices_local, num_experts, shard_idx, ep_size, b return pos.reshape(-1) -def _clamp_masked_indices(indices, valid_rows_mask): - """Sets the masked rows' gather indices to 0, as the packed branches do. - - With index_math="count", the inverse permutation is -1 on other shards' slots. The SparseCore gather-reduce kernel - reads the input at the index of every lane of its last partial row block, masked lanes included, without bounds - checks, so -1 is an out-of-bounds read. The masked lanes are discarded, so the output is unchanged. - """ - return jnp.where(valid_rows_mask, indices, 0) - - def ring_ragged_sort( hidden_states_local, topk_indices_local, @@ -305,7 +295,7 @@ def _ring_ragged_sort_bwd(res, g_out): # `topk_weights` is set to ones because this op has no per-row weighting. grad_hidden_states = ragged_gather_reduce( g_x, - _clamp_masked_indices(topk_argsort_revert_indices, valid_rows_mask), + topk_argsort_revert_indices, topk_weights=jnp.ones((n,), dtype=jnp.float32), valid_rows_mask=valid_rows_mask, reduce_group_size=topk, @@ -421,7 +411,7 @@ def _ring_unsort_combine( ) return ragged_gather_reduce( sorted_tokens_local, - _clamp_masked_indices(topk_argsort_revert_indices, valid_rows_mask), + topk_argsort_revert_indices, topk_weights=topk_weights_flat, valid_rows_mask=valid_rows_mask, reduce_group_size=topk, diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index 574c4a03bd..1513e00382 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -924,18 +924,16 @@ def concat_last(cls, first, second): jax.tree_util.register_dataclass(Fp8NativeWeight, data_fields=["anchor", "codes", "scale"], meta_fields=[]) -def fp8_native_fwd_acc_dtype(config) -> jnp.dtype | None: - """fp8_moe_native_fwd_acc_dtype as a dtype, or None for gmm_v2's default.""" - name = getattr(config, "fp8_moe_native_fwd_acc_dtype", "") - return jnp.dtype(name) if name else None +# fp8_moe_native_fwd's in-kernel activation quantization and accumulator: one float32 absmax scale per row over the +# full K dim, float32 accumulation. Closest to the rollout's FP8 MoE kernel and fastest of the measured gmm_v2 +# variants; the others stay selectable as ops-level arguments for benchmarks and tests. +_FP8_NATIVE_FWD_LHS_QUANT_MODE = "fullk_f32" +_FP8_NATIVE_FWD_ACC_DTYPE = jnp.dtype(jnp.float32) def fp8_native_fwd_kwargs(config, with_bwd_dlhs: bool = True) -> dict: - """The FP8 GMM ops' kwargs from the fp8_moe_native_fwd_* (and, with_bwd_dlhs, fp8_moe_native_bwd_dlhs) config.""" - kwargs = { - "acc_dtype": fp8_native_fwd_acc_dtype(config), - "lhs_quant_mode": getattr(config, "fp8_moe_native_fwd_lhs_quant", ""), - } + """The FP8 GMM ops' kwargs for fp8_moe_native_fwd (and, with_bwd_dlhs, fp8_moe_native_bwd_dlhs).""" + kwargs = {"acc_dtype": _FP8_NATIVE_FWD_ACC_DTYPE, "lhs_quant_mode": _FP8_NATIVE_FWD_LHS_QUANT_MODE} if with_bwd_dlhs: kwargs["native_bwd_dlhs"] = getattr(config, "fp8_moe_native_bwd_dlhs", False) return kwargs diff --git a/tests/unit/fp8_moe_native_test.py b/tests/unit/fp8_moe_native_test.py index 3cdfe8ddf7..35bde4f268 100644 --- a/tests/unit/fp8_moe_native_test.py +++ b/tests/unit/fp8_moe_native_test.py @@ -11,12 +11,11 @@ # 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. -"""fp8_moe_native_fwd / fp8_moe_native_bwd_dlhs on CPU: config validation, the outside-layer-scan guard, and the -dequantization and lhs-quant-mode contracts the native path relies on. +"""fp8_moe_native_fwd / fp8_moe_native_bwd_dlhs on CPU: config validation, the outside-layer-scan guard, the +dequantization contract the native path relies on, and the fixed activation-quantization / accumulator settings. """ import sys -import typing from flax import nnx from flax.linen import partitioning as nn_partitioning @@ -29,6 +28,7 @@ from maxtext.configs import pyconfig from maxtext.configs import types as config_types from maxtext.kernels.megablox import ops as megablox_ops +from maxtext.layers import moe from maxtext.layers import quantizations from maxtext.utils import maxtext_utils from maxtext.utils import maxtext_utils_nnx @@ -108,22 +108,14 @@ def _cfg(**overrides): def test_flags_are_defined_and_off_by_default(): - for key in ( - "fp8_moe_native_fwd", - "fp8_moe_native_fwd_acc_dtype", - "fp8_moe_native_fwd_lhs_quant", - "fp8_moe_native_bwd_dlhs", - ): + for key in ("fp8_moe_native_fwd", "fp8_moe_native_bwd_dlhs"): assert key in set(config_types.MoEGeneral.model_fields), key cfg = _cfg() assert not cfg.fp8_moe_native_fwd and not cfg.fp8_moe_native_bwd_dlhs - assert cfg.fp8_moe_native_fwd_acc_dtype == "" and cfg.fp8_moe_native_fwd_lhs_quant == "" def test_native_config_is_accepted(): - cfg = _cfg( - **_NATIVE, **_NATIVE_BWD_DLHS, fp8_moe_native_fwd_lhs_quant="fullk_f32", fp8_moe_native_fwd_acc_dtype="float32" - ) + cfg = _cfg(**_NATIVE, **_NATIVE_BWD_DLHS) assert cfg.fp8_moe_native_fwd and cfg.fp8_moe_native_bwd_dlhs @@ -133,6 +125,8 @@ def test_native_config_is_accepted(): ({"fp8_moe_fake_quant": False}, "fp8_moe_fake_quant=True"), ({"fp8_fake_quant_hoist": "absmax"}, "fp8_fake_quant_hoist=gather8"), ({"use_gmm_v2": False}, "use_tokamax_gmm=True and use_gmm_v2=True"), + # The heuristic tiling budgets a bf16 accumulator; the native forward accumulates in f32. + ({"use_gmm_v2_heuristic_tiling": True}, "use_gmm_v2_heuristic_tiling=False"), ], ) def test_native_fwd_requirement_is_enforced(override, requirement): @@ -197,8 +191,11 @@ def test_dequantize_fp8_codes_matches_the_gather8_dequantization(): np.testing.assert_array_equal(np.asarray(native).view(np.uint16), np.asarray(gather8).view(np.uint16)) -def test_lhs_quant_modes_match_the_config_values(): - # A value the config accepts but ops does not know would fail only at trace time. - field = dict(config_types.MoEGeneral.model_fields)["fp8_moe_native_fwd_lhs_quant"] - config_values = set(typing.get_args(field.annotation)) - assert config_values == set(megablox_ops._LHS_QUANT_MODES) # pylint: disable=protected-access +def test_native_fwd_uses_full_k_f32_activation_quant_and_f32_accumulation(): + # One float32 absmax scale per row over the full K dim and a float32 accumulator, with no config knob: the plain + # fp8_moe_native_fwd=True must get the numerics closest to the rollout's FP8 MoE kernel. + kwargs = moe.fp8_native_fwd_kwargs(_cfg(**_NATIVE)) + assert kwargs["lhs_quant_mode"] == "fullk_f32" + assert kwargs["lhs_quant_mode"] in megablox_ops._LHS_QUANT_MODES # pylint: disable=protected-access + assert kwargs["acc_dtype"] == jnp.float32 + assert not kwargs["native_bwd_dlhs"] diff --git a/tests/unit/ragged_sort_full_buffer_indices_test.py b/tests/unit/ragged_sort_full_buffer_indices_test.py deleted file mode 100644 index 714b85cd8d..0000000000 --- a/tests/unit/ragged_sort_full_buffer_indices_test.py +++ /dev/null @@ -1,185 +0,0 @@ -# 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. -"""Full-buffer (dropless) ring-of-experts gather-reduces receive in-bounds indices. - -With moe_routing_index_math="count", ring_ragged_sort's inverse permutation is -1 on the slots of other shards. The -full-buffer branches of `_ring_unsort_combine` and of ring_ragged_sort's backward must clamp those masked lanes before -`ragged_gather_reduce`: its SparseCore kernel reads every lane of the last row block without bounds checks. - -On CPU, `ragged_gather_reduce` uses its jnp fallback, so a spy on it checks the arguments of every call instead of the -TPU fault. The tests are CPU-only: the last one deliberately passes the unclamped indices. -""" - -from unittest import mock - -import jax -import jax.numpy as jnp -import numpy as np -import pytest - -from maxtext.kernels.ragged import ragged_gather_reduce_v2 -from maxtext.kernels.ragged import ragged_sort - -pytestmark = pytest.mark.cpu_only - -_HIDDEN = 64 -_GATHER_REDUCE_KWARGS = { - "enforce_fallback": True, - "flops_override": -1, - "bytes_accessed_override": -1, - "use_single_sparsecore": False, -} - - -def _bits_equal(a, b): - a, b = np.asarray(a), np.asarray(b) - if a.shape != b.shape or a.dtype != b.dtype: - return False - return bool(np.array_equal(a.reshape(-1).view(np.uint8), b.reshape(-1).view(np.uint8))) - - -def _routing(seed, tokens, experts, topk): - k = jax.random.split(jax.random.PRNGKey(seed), 3) - ids = jax.lax.top_k(jax.random.normal(k[0], (tokens, experts)), topk)[1].astype(jnp.int32) - w = jax.random.uniform(k[1], (tokens * topk,), jnp.float32, 0.25, 0.3) - return ids, w, k[2] - - -def _maps(ids, experts, ep, shard, index_math): - """This shard's window and the inverse permutation ring_ragged_sort builds for `index_math`.""" - flat = ids.reshape(-1) - local = experts // ep - offs = jnp.concatenate([jnp.zeros((1,), jnp.int32), jnp.cumsum(jax.nn.one_hot(flat, experts, dtype=jnp.int32).sum(0))]) - start, end = offs[shard * local], offs[(shard + 1) * local] - if index_math == "count": - pos = ragged_sort._local_slot_positions(ids, experts, shard, ep) # pylint: disable=protected-access - revert = jnp.where(pos >= 0, start + pos, -1).astype(jnp.int32) - else: - revert = jnp.argsort(jnp.argsort(flat, stable=True)).astype(jnp.int32) - return start, end, revert - - -class _Spy: - """Wraps ragged_sort.ragged_gather_reduce; records (min index, max index, rows) of every call at run time.""" - - def __init__(self, declamp=False): - self.calls = [] - self.declamp = declamp - self.real = ragged_sort.ragged_gather_reduce - - def __call__(self, x, indices, *args, **kwargs): - if self.declamp: - # The arguments before the fix: the masked lanes of a full buffer carry the raw inverse (-1 under "count"). - indices = jnp.where(kwargs["valid_rows_mask"], indices, -1) - rows = x.shape[0] - jax.debug.callback(lambda lo, hi: self.calls.append((int(lo), int(hi), rows)), jnp.min(indices), jnp.max(indices)) - return self.real(x, indices, *args, **kwargs) - - def out_of_bounds(self): - return [c for c in self.calls if c[0] < 0 or c[1] >= c[2]] - - -_COMBINE_CASES = [(ep, m) for ep in (4, 8) for m in ("count", "sort")] - - -@pytest.mark.parametrize("ep,index_math", _COMBINE_CASES, ids=[f"ep{ep}-{m}" for ep, m in _COMBINE_CASES]) -def test_full_buffer_combine_in_bounds_and_unchanged(ep, index_math): - tokens, experts, topk = 256, 2 * ep, 4 - n = tokens * topk - for seed in range(3): - ids, w, key = _routing(seed, tokens, experts, topk) - x = jax.random.normal(key, (n, _HIDDEN), jnp.float32).astype(jnp.bfloat16) # full buffer: rows at global positions - for shard in range(ep): - start, end, revert = _maps(ids, experts, ep, shard, index_math) - if index_math == "count": - assert int(jnp.sum(revert < 0)) > 0, "count must leave other shards' slots at -1" - spy = _Spy() - with mock.patch.object(ragged_sort, "ragged_gather_reduce", spy): - out = jax.jit( - lambda x, r, w, s, e: ragged_sort._ring_unsort_combine( # pylint: disable=protected-access - x, r, w, s, e, topk, _GATHER_REDUCE_KWARGS - ) - )(x, revert, w, start, end) - out.block_until_ready() - jax.effects_barrier() - assert len(spy.calls) == 1, spy.calls - assert not spy.out_of_bounds(), f"seed{seed} shard{shard}: indices out of [0, {n}): {spy.calls}" - # The call before the fix: the raw inverse. The masked lanes are discarded, so the output is the same bits. - valid = (revert >= start) & (revert < end) - before = ragged_gather_reduce_v2.ragged_gather_reduce( - x=x, - indices=revert, - topk_weights=w, - valid_rows_mask=valid, - reduce_group_size=topk, - enforce_fallback=True, - ) - assert _bits_equal(out, before), f"seed{seed} shard{shard}: output changed" - # float64 reference: each token's routing-weighted sum over its slots in this shard's window. - x_np, revert_np, w_np, valid_np = (np.asarray(a) for a in (x.astype(jnp.float32), revert, w, valid)) - gathered = x_np[np.where(valid_np, revert_np, 0)] * w_np[:, None] - ref = np.where(valid_np[:, None], gathered, 0.0).reshape((tokens, topk, _HIDDEN)).sum(1) - np.testing.assert_allclose(np.asarray(out, np.float32), ref, rtol=2e-2, atol=2e-2) - - -def _op(ep, index_math, tokens=128, topk=4, seed=0): - """ring_ragged_sort (full buffer) -> ring_ragged_unsort per shard, vmapped over "ep": (out, dx, dw).""" - experts = 2 * ep - local = experts // ep - ids, w, key = _routing(seed, tokens, experts, topk) - k = jax.random.split(key, 2) - x = jax.random.normal(k[0], (tokens, _HIDDEN), jnp.float32).astype(jnp.bfloat16) - cot = jax.random.normal(k[1], (ep, tokens, _HIDDEN), jnp.float32).astype(jnp.bfloat16) - - def per_shard(x, w, cot): - def f(x, w): - sorted_x, gs, revert, perm = ragged_sort.ring_ragged_sort( - x, ids, experts, topk, "ep", ep, buffer_size=None, index_math=index_math - ) - assert sorted_x.shape[0] == tokens * topk # the full buffer - return ragged_sort.ring_ragged_unsort( - sorted_x * 2, gs, revert, topk, local, "ep", w, topk_argsort_indices=perm, index_math=index_math - ) - - out, vjp = jax.vjp(f, x, w) - dx, dw = vjp(cot) - return out, dx, dw - - return jax.jit(jax.vmap(per_shard, in_axes=(None, None, 0), axis_name="ep"))(x, w, cot) - - -@pytest.mark.parametrize("ep", [4, 8]) -def test_full_buffer_op_in_bounds_and_count_matches_sort(ep): - results = {} - for index_math in ("count", "sort"): - spy = _Spy() - with mock.patch.object(ragged_sort, "ragged_gather_reduce", spy): - results[index_math] = jax.block_until_ready(_op(ep, index_math)) - jax.effects_barrier() - # Per shard: one combine (forward) and one sort backward. - assert len(spy.calls) == 2 * ep, spy.calls - assert not spy.out_of_bounds(), f"{index_math}: indices out of bounds: {spy.out_of_bounds()}" - for name, a, b in zip(("out", "dx", "dw"), results["count"], results["sort"]): - assert _bits_equal(a, b), f"ep{ep}: {name} differs between count and sort" - - -def test_unclamped_indices_are_flagged(): - """The arguments before the fix (-1 on the masked lanes) fail the bounds check, in both call sites.""" - ep = 4 - spy = _Spy(declamp=True) - with mock.patch.object(ragged_sort, "ragged_gather_reduce", spy): - jax.block_until_ready(_op(ep, "count")) - jax.effects_barrier() - assert len(spy.calls) == 2 * ep, spy.calls - assert len(spy.out_of_bounds()) == 2 * ep, spy.calls From 22ebd030453b83d54077f8ce43f929cfcefb7438 Mon Sep 17 00:00:00 2001 From: Shuning Jin Date: Sat, 31 Oct 2026 08:49:19 +0000 Subject: [PATCH 3/4] Enforce single K tile for full-K dynamic LHS quantization in fp8_moe_native_fwd --- src/maxtext/configs/types.py | 78 ++++++++++++++++++- .../pallas_mosaic_tpu_v2_gmm_kernel.py | 11 ++- tests/unit/fp8_moe_native_test.py | 69 +++++++++++++++- 3 files changed, 155 insertions(+), 3 deletions(-) diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index eb3c797129..547dc7a6d4 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -1982,7 +1982,8 @@ def validate_fp8_moe_native_fwd(self) -> "MoEGeneral": # gmm_v2's heuristic tiling budgets a bfloat16 accumulator when the lhs is quantized in the kernel, # but the native forward accumulates in float32, so heuristic tiles could exceed VMEM. (not self.use_gmm_v2_heuristic_tiling, "use_gmm_v2_heuristic_tiling=False"), - ], + ] + + self._fp8_moe_native_fwd_tile_requirements(), ) if self.fp8_moe_native_bwd_dlhs: _reject_unless( @@ -2156,6 +2157,81 @@ def _moe_mlp_dim_if_unsharded(self) -> int | None: return None return 2 ** get_individual_scales(scale)[2] * base + def _moe_emb_dim_if_unsharded(self) -> int | None: + """The routed experts' input dim (wi's K dim) if no mesh axis shards it in the sparse_matmul GMMs, else None.""" + if self.use_2d_fsdp_sharding or getattr(self, "num_moe_emb_chunks", 0) > 0: + return None + axes = {"tensor"} + axes.update(axes_for_logical(self.logical_axis_rules or [], "activation_embed_moe")) + for axis in axes: + if getattr(self, f"ici_{axis}_parallelism", None) != 1 or getattr(self, f"dcn_{axis}_parallelism", None) != 1: + return None + if self.moe_expert_input_dim > 0: + return self.moe_expert_input_dim + base, scale = self.base_emb_dim, self.global_parameter_scale + if base <= 0 or scale <= 0 or scale & (scale - 1): + return None + return 2 ** get_individual_scales(scale)[0] * base + + def _fp8_moe_native_fwd_tile_requirements(self) -> list[tuple[bool, str]]: + """Tile rules for fp8_moe_native_fwd's full-K (lhs_quant_block_size=0) forward GMMs, checked at config time. + + gmm_v2 quantizes one [tile_m, tile_k] lhs slice per K step, so full-K dynamic scaling needs tile_k == size_k in + both forward GMMs (wi contracts emb_dim, wo contracts the MoE MLP dim); make_gmm_configs enforces the same rule + when traced. When separate eval logical axis rules are set, RoutedMoE's fwd_tile reads eval_wi_tile_fwd_embed_dim + and eval_wo_tile_fwd_mlp_dim in eval_step, so those are checked too. + """ + if self.use_gmm_v2_heuristic_tiling: + return [] + eval_rule = getattr(self, "custom_mesh_and_rule_for_eval", CustomRule.DEFAULT) + eval_axis_rules = getattr(self, "logical_axis_rules_for_eval", None) + separate_eval_rules = eval_rule not in ( + CustomRule.DEFAULT, + getattr(self, "custom_mesh_and_rule", CustomRule.DEFAULT), + ) or bool(eval_axis_rules and eval_axis_rules != self.logical_axis_rules) + reqs: list[tuple[bool, str]] = [] + emb = self._moe_emb_dim_if_unsharded() + if emb is not None: + reqs.append( + ( + self.wi_tile_fwd_embed_dim == emb, + f"wi_tile_fwd_embed_dim={emb}, the embed dim (one K tile for full-K activation quantization;" + f" got {self.wi_tile_fwd_embed_dim})", + ) + ) + if separate_eval_rules and all( + getattr(self, f"ici_{ax}_parallelism", None) == 1 and getattr(self, f"dcn_{ax}_parallelism", None) == 1 + for ax in axes_for_logical(eval_axis_rules or [], "activation_embed_moe") + ): + reqs.append( + ( + self.eval_wi_tile_fwd_embed_dim == emb, + f"eval_wi_tile_fwd_embed_dim={emb} when evaluation uses separate logical axis rules" + f" (got {self.eval_wi_tile_fwd_embed_dim})", + ) + ) + n = self._moe_mlp_dim_if_unsharded() + if n is not None: + reqs.append( + ( + self.wo_tile_fwd_mlp_dim == n, + f"wo_tile_fwd_mlp_dim={n}, the MoE MLP dim (one K tile for full-K activation quantization;" + f" got {self.wo_tile_fwd_mlp_dim})", + ) + ) + if separate_eval_rules and all( + getattr(self, f"ici_{ax}_parallelism", None) == 1 and getattr(self, f"dcn_{ax}_parallelism", None) == 1 + for ax in axes_for_logical(eval_axis_rules or [], "mlp_no_fsdp") + ): + reqs.append( + ( + self.eval_wo_tile_fwd_mlp_dim == n, + f"eval_wo_tile_fwd_mlp_dim={n} when evaluation uses separate logical axis rules" + f" (got {self.eval_wo_tile_fwd_mlp_dim})", + ) + ) + return reqs + @model_validator(mode="after") def validate_moe_roe_packed_routing_gather(self) -> "MoEGeneral": """The packed gather replaces the per-chunk routing gathers of the pipelined local-routing chunk loop.""" diff --git a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py index 8a2ce1ad56..412473b8f8 100644 --- a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py +++ b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py @@ -2079,6 +2079,15 @@ def make_gmm_configs( else: tiles = tile_info(dims, lhs_cfgs, rhs_cfgs, vmem_limit_bytes, fuse_act, partial_sum is not None) + if not has_lhs_scale and lhs_cfgs.quant_dtype is not None and lhs_quant_block_size == 0 and tiles.tile_k != dims.size_k: + # inner_kernel quantizes one [tile_m, tile_k] lhs tile per K step in blocks of min(tile_k, size_k): tile_k < size_k + # would compute one scale per K tile instead of one over size_k, and tile_k > size_k would step past size_k and read + # out-of-bounds rhs_scale blocks. + raise NotImplementedError( + "lhs_quant_block_size=0 (full-K dynamic lhs quantization) requires a single K tile (tile_k == size_k);" + f" got {tiles=}, {dims=}." + ) + num_lanes = pltpu.get_tpu_info().num_lanes if lhs_cfgs.is_3d: if tiles.tile_k % num_lanes or dims.size_k % tiles.tile_k: @@ -2375,7 +2384,7 @@ def gmm_v2( boundaries. The values are unchanged. lhs_quant_block_size: With in-kernel lhs quantization, the k block that gets one scale per row: None for the default (512), 0 for the whole K - dim, or a positive int. + dim (requires tile_k == size_k), or a positive int. lhs_f32_scale: With in-kernel lhs quantization, compute each block scale in float32 as max(absmax, 1e-12) / dtype_max and divide by it, instead of the bf16 absmax and reciprocal. diff --git a/tests/unit/fp8_moe_native_test.py b/tests/unit/fp8_moe_native_test.py index 35bde4f268..a3db896789 100644 --- a/tests/unit/fp8_moe_native_test.py +++ b/tests/unit/fp8_moe_native_test.py @@ -16,6 +16,8 @@ """ import sys +import types +from unittest import mock from flax import nnx from flax.linen import partitioning as nn_partitioning @@ -28,6 +30,7 @@ from maxtext.configs import pyconfig from maxtext.configs import types as config_types from maxtext.kernels.megablox import ops as megablox_ops +from maxtext.kernels.megablox import pallas_mosaic_tpu_v2_gmm_kernel as gmm_v2 from maxtext.layers import moe from maxtext.layers import quantizations from maxtext.utils import maxtext_utils @@ -37,6 +40,7 @@ pytestmark = pytest.mark.cpu_only +_EMB = 256 _MOE_MLP = 128 _SEQ = 64 _VOCAB = 512 @@ -48,7 +52,7 @@ "model_name": "qwen3.5-397b-a17b", "override_model_config": True, "base_num_decoder_layers": 4, - "base_emb_dim": 256, + "base_emb_dim": _EMB, "base_num_query_heads": 4, "base_num_kv_heads": 2, "head_dim": 64, @@ -89,6 +93,8 @@ "use_gmm_v2": True, "moe_wi_gmm_fused_swiglu": True, "moe_wo_gmm_fused_dswiglu": True, + "wi_tile_fwd_embed_dim": _EMB, + "wo_tile_fwd_mlp_dim": _MOE_MLP, "wo_tile_dlhs_mlp_dim": _MOE_MLP, } @@ -117,6 +123,13 @@ def test_flags_are_defined_and_off_by_default(): def test_native_config_is_accepted(): cfg = _cfg(**_NATIVE, **_NATIVE_BWD_DLHS) assert cfg.fp8_moe_native_fwd and cfg.fp8_moe_native_bwd_dlhs + eval_cfg = _cfg( + **_NATIVE, + custom_mesh_and_rule_for_eval="pure-fsdp", + eval_wi_tile_fwd_embed_dim=_EMB, + eval_wo_tile_fwd_mlp_dim=_MOE_MLP, + ) + assert eval_cfg.fp8_moe_native_fwd @pytest.mark.parametrize( @@ -127,6 +140,22 @@ def test_native_config_is_accepted(): ({"use_gmm_v2": False}, "use_tokamax_gmm=True and use_gmm_v2=True"), # The heuristic tiling budgets a bf16 accumulator; the native forward accumulates in f32. ({"use_gmm_v2_heuristic_tiling": True}, "use_gmm_v2_heuristic_tiling=False"), + # Full-K dynamic lhs quantization requires one K tile (tile_k == size_k) in wi, wo and separate-rule eval. + ({"base_emb_dim": 4096, "wi_tile_fwd_embed_dim": 1024}, "wi_tile_fwd_embed_dim=4096"), + ({"wi_tile_fwd_embed_dim": 2 * _EMB}, f"wi_tile_fwd_embed_dim={_EMB}"), + ({"wo_tile_fwd_mlp_dim": _MOE_MLP // 2}, f"wo_tile_fwd_mlp_dim={_MOE_MLP}"), + ( + {"ici_expert_parallelism": 2, "custom_mesh_and_rule_for_eval": "ep-as-cp"}, + f"eval_wi_tile_fwd_embed_dim={_EMB}", + ), + ( + { + "ici_expert_parallelism": 2, + "custom_mesh_and_rule_for_eval": "ep-as-cp", + "eval_wi_tile_fwd_embed_dim": _EMB, + }, + f"eval_wo_tile_fwd_mlp_dim={_MOE_MLP}", + ), ], ) def test_native_fwd_requirement_is_enforced(override, requirement): @@ -199,3 +228,41 @@ def test_native_fwd_uses_full_k_f32_activation_quant_and_f32_accumulation(): assert kwargs["lhs_quant_mode"] in megablox_ops._LHS_QUANT_MODES # pylint: disable=protected-access assert kwargs["acc_dtype"] == jnp.float32 assert not kwargs["native_bwd_dlhs"] + + +def test_gmm_v2_full_k_lhs_quant_requires_single_k_tile(): + # lhs_quant_block_size=0 means one dynamic scale per row over all size_k columns; inner_kernel slices one tile_k at a + # time, so make_gmm_configs must reject tile_k != size_k instead of computing per-tile scales. + lhs = jnp.zeros((512, 4096), jnp.bfloat16) + rhs = jnp.zeros((4, 4096, 2048), jnp.float8_e4m3fn) + rhs_scale = jnp.ones((4, 1, 1, 2048), jnp.float32) + group_sizes = jnp.array([128, 128, 128, 128], jnp.int32) + group_offset = jnp.array([0], jnp.int32) + fake_tpu = types.SimpleNamespace( + fp8_ops_per_second=1, int8_ops_per_second=1, num_lanes=128, mxu_column_size=256, get_sublane_tiling=lambda _: 8 + ) + + def make(tile_k): + return gmm_v2.make_gmm_configs( + lhs, + rhs, + rhs_scale, + None, + None, + group_sizes, + group_offset, + tile_info=gmm_v2.TileSizes(tile_m=512, tile_k=tile_k, tile_n=1024), + vmem_limit_bytes=None, + out_dtype=jnp.bfloat16, + acc_dtype=jnp.float32, + maybe_quantize_lhs=True, + zero_initialize=True, + lhs_quant_block_size=0, + lhs_f32_scale=True, + ) + + with mock.patch.object(gmm_v2.pltpu, "get_tpu_info", return_value=fake_tpu): + assert make(4096).lhs_cfgs.quant_block_size == 4096 + for bad_tile_k in (1024, 8192): + with pytest.raises(NotImplementedError, match="tile_k == size_k"): + make(bad_tile_k) From 77cc0acfb716c6a8f48a98c94eb3abd6715c660d Mon Sep 17 00:00:00 2001 From: Shuning Jin Date: Mon, 2 Nov 2026 21:45:19 +0000 Subject: [PATCH 4/4] Reject split K for fp8_moe_native_fwd; cache rhs_contract_dequant blocks only with one K tile - fp8_moe_native_fwd scales each activation row over the whole K dim, so reject a split K instead of skipping the tile check: tensor / tensor_sequence / autoregressive (or any mlp_no_fsdp axis) for wo, tensor / activation_embed_moe for wi, and num_moe_emb_chunks > 0. RoutedMoE also checks K in the traced w0/wo partition specs (custom_mesh_and_rule, parallelism -1). - gmm_v2 rhs_contract_dequant: with one K tile dispatch only the first_last fill/cached pair; with two or more K tiles drop the never-hit VMEM cache tile and SMEM key and dequantize each piece straight into the matmul (first/mid/last only). Bit-identical to the XLA-dequant DLHS at 1/2/4 K tiles on tpu7x-8. --- src/maxtext/configs/types.py | 35 ++++++ .../pallas_mosaic_tpu_v2_gmm_kernel.py | 113 ++++++++++-------- src/maxtext/layers/moe.py | 25 ++++ tests/unit/fp8_moe_native_test.py | 31 +++++ 4 files changed, 152 insertions(+), 52 deletions(-) diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 547dc7a6d4..47d18bdc38 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -1983,6 +1983,7 @@ def validate_fp8_moe_native_fwd(self) -> "MoEGeneral": # but the native forward accumulates in float32, so heuristic tiles could exceed VMEM. (not self.use_gmm_v2_heuristic_tiling, "use_gmm_v2_heuristic_tiling=False"), ] + + self._fp8_moe_native_fwd_unsplit_k_requirements() + self._fp8_moe_native_fwd_tile_requirements(), ) if self.fp8_moe_native_bwd_dlhs: @@ -2173,6 +2174,40 @@ def _moe_emb_dim_if_unsharded(self) -> int | None: return None return 2 ** get_individual_scales(scale)[0] * base + def _fp8_moe_native_fwd_unsplit_k_requirements(self) -> list[tuple[bool, str]]: + """fp8_moe_native_fwd scales each activation row over the whole K dim, so neither forward GMM may see a K slice. + + wi contracts the embed dim and wo the MoE MLP dim. A mesh axis that splits one of them in the expert shard_map, or + num_moe_emb_chunks, would make each shard or chunk quantize its slice with its own scale, and make_gmm_configs + cannot tell: it only sees the local size_k. The axes are the ones _moe_emb_dim_if_unsharded and + _moe_mlp_dim_if_unsharded read, under the training rules and separate eval rules. Parallelism -1 is filled in from + the device count later and custom_mesh_and_rule replaces logical_axis_rules after this validator, so RoutedMoE + checks the traced weight partition specs too. + """ + rules = [self.logical_axis_rules or []] + eval_axis_rules = getattr(self, "logical_axis_rules_for_eval", None) + if eval_axis_rules: + rules.append(eval_axis_rules) + + def split_by(axes: set[str], logical: str) -> list[str]: + for rule in rules: + axes = axes | set(axes_for_logical(rule, logical)) + return [ + f"{kind}_{axis}_parallelism={degree}" + for axis in sorted(axes) + for kind in ("ici", "dcn") + if (degree := getattr(self, f"{kind}_{axis}_parallelism", None)) is not None and degree > 1 + ] + + emb_split = split_by({"tensor"}, "activation_embed_moe") + mlp_split = split_by({"tensor", "tensor_sequence", "autoregressive"}, "mlp_no_fsdp") + emb_chunks = getattr(self, "num_moe_emb_chunks", 0) + return [ + (not emb_split, f"an unsplit embed dim, wi's K (got {', '.join(emb_split)})"), + (not mlp_split, f"an unsplit MoE MLP dim, wo's K (got {', '.join(mlp_split)})"), + (emb_chunks <= 0, f"num_moe_emb_chunks=0, an unchunked wi K (got {emb_chunks})"), + ] + def _fp8_moe_native_fwd_tile_requirements(self) -> list[tuple[bool, str]]: """Tile rules for fp8_moe_native_fwd's full-K (lhs_quant_block_size=0) forward GMMs, checked at config time. diff --git a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py index 412473b8f8..4efb70ca73 100644 --- a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py +++ b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py @@ -396,11 +396,25 @@ class GmmConfigs: # rhs_contract_scale (fp8_moe_native_bwd_dlhs): rhs holds float8 codes [size_group, size_n, size_k] # (transpose_rhs) whose scale varies along the contraction dim K ([size_group, 1, size_k] float32, read in # the rhs scale slot). Each rhs block is dequantized in VMEM to rhs_dequant_dtype with - # (codes.astype(float32) * scale).astype(rhs_dequant_dtype) into a scratch tile, cached across the grid - # steps that read the same block (the m tiles of a group), and the unquantized matmul reads that tile. + # (codes.astype(float32) * scale).astype(rhs_dequant_dtype) inside the unquantized matmul that reads it. With a + # single K tile it is also kept in a scratch tile for the later grid steps that read the same block (the m tiles of + # a group); see rhs_dequant_cache. rhs_contract_dequant: bool = False rhs_dequant_dtype: jnp.dtype | None = None + @property + def rhs_dequant_cache(self) -> bool: + """rhs_contract_dequant with one K tile: the dequantized rhs block is cached in VMEM across its m tiles. + + The rhs block index is (group, n tile, k tile), and k is the innermost grid axis. With one K tile, consecutive + m tiles of a group read the same block, so it is dequantized once per (group, n tile). With two or more K tiles, + consecutive steps read different blocks, a cache would never hit, and every step dequantizes its block inside the + matmul with no cache tile (make_gmm_configs requires tile_k to divide size_k). That still beats an XLA + dequantization of the whole weight first, because the kernel reads 1-byte codes instead of a bf16 copy (TPU7x, + production DLHS shapes, 2-4 k tiles: 0.84-0.91x its time, measured when these steps also wrote an unused cache tile). + """ + return self.rhs_contract_dequant and self.tiles.tile_k == self.dims.size_k + @property def lhs_is_3d(self) -> bool: return self.lhs_cfgs.is_3d @@ -728,7 +742,7 @@ def inner_kernel( the activation is applied. With cfgs.fuse_dact, tiled_ps_ref is a DactInputRef and the accumulator becomes [d_gate | d_up] before the regular store path; with cfgs.return_dact_dot the out refs are (out, dot) pairs. - With cfgs.rhs_contract_dequant, rhs_cache_refs is (the [tile_n, tile_k] + With cfgs.rhs_dequant_cache, rhs_cache_refs is (the [tile_n, tile_k] dequantized weight tile, the SMEM key of the block it holds); see _rhs_cache_miss. """ @@ -760,14 +774,15 @@ def _matmul(is_first_k_step: bool, is_last_k_step: bool, rhs_fill: bool = False) tiled_lhs = tiled_lhs_ref.get_value().reshape(-1, cfgs.tiles.tile_k)[...] rhs_fill_refs = None if cfgs.rhs_contract_dequant and rhs_fill: - # The first grid step of this weight block: rhs_block dequantizes each [mxu, tile_k] piece - # of the codes right before its matmul (so the VPU work can overlap the MXU work of the previous piece) - # and stores it in the cache tile, which the later steps of the block read. - rhs_fill_refs = (tiled_rhs_ref.weight, tiled_rhs_ref.scale, rhs_cache_refs[0]) # pyrefly: ignore + # This step dequantizes its weight block: rhs_block dequantizes each [mxu, tile_k] piece of the codes right + # before its matmul (so the VPU work can overlap the MXU work of the previous piece). With rhs_dequant_cache it + # also stores the piece in the cache tile, which the later steps of the block read. + cache_ref = rhs_cache_refs[0] if cfgs.rhs_dequant_cache else None # pyrefly: ignore[unsupported-operation] + rhs_fill_refs = (tiled_rhs_ref.weight, tiled_rhs_ref.scale, cache_ref) # pyrefly: ignore tiled_rhs = None elif cfgs.rhs_contract_dequant: - # This block's weight, dequantized in VMEM by the block's first step: the same bf16 tile the bf16 path DMAs, - # so the matmul below is that path's. + # rhs_dequant_cache hit: this block's weight, dequantized in VMEM by the block's first step: the same bf16 tile + # the bf16 path DMAs, so the matmul below is that path's. tiled_rhs = rhs_cache_refs[0][...] # pyrefly: ignore[unsupported-operation] else: tiled_rhs = tiled_rhs_ref.get_weight(transpose_rhs=cfgs.transpose_rhs) @@ -811,13 +826,15 @@ def _matmul(is_first_k_step: bool, is_last_k_step: bool, rhs_fill: bool = False) def rhs_block(start_k: int, end_k: int, start_n: int, end_n: int): if rhs_fill_refs is not None: - # rhs_contract_dequant (transpose_rhs) fill step: dequantize this piece, keep it for the block's later steps. + # rhs_contract_dequant (transpose_rhs) fill step: dequantize this piece; with the cache, keep it for the + # block's later steps. codes_ref, scale_ref, cache_ref = rhs_fill_refs end_k = min(end_k, cfgs.tiles.tile_k) block = dequantize_fp8_codes( - codes_ref[start_n:end_n, start_k:end_k], scale_ref[:, start_k:end_k], cache_ref.dtype + codes_ref[start_n:end_n, start_k:end_k], scale_ref[:, start_k:end_k], cfgs.rhs_dequant_dtype ) - cache_ref[start_n:end_n, start_k:end_k] = block + if cache_ref is not None: + cache_ref[start_n:end_n, start_k:end_k] = block return block if cfgs.transpose_rhs: return tiled_rhs[start_n:end_n, start_k:end_k] @@ -1095,29 +1112,26 @@ def matmul_last(): k_id = pl.program_id(2) if cfgs.rhs_contract_dequant: - # The first step of each weight block dequantizes it piece by piece inside its matmul, so the VPU work overlaps - # the MXU work (see _matmul's rhs_fill); the other steps read the cache tile. - cache_miss = _rhs_cache_miss(rhs_cache_refs, metadata_ref, cfgs) - is_first = k_id == 0 - is_last = k_id == (num_k - 1) - - def _variants(fill: bool): - suffix = "_rhs_fill" if fill else "" - return [ - jax.named_scope(name + suffix)(functools.partial(_matmul, first, last, rhs_fill=fill)) - for name, first, last in ( - ("matmul_first_last", True, True), - ("matmul_first", True, False), - ("matmul_last", False, True), - ("matmul", False, False), - ) - ] - - def _dispatch(fill: bool): - first_last, first, last, mid = _variants(fill) - lax.cond(is_first, lambda: lax.cond(is_last, first_last, first), lambda: lax.cond(is_last, last, mid)) - - lax.cond(cache_miss, lambda: _dispatch(True), lambda: _dispatch(False)) + # Every step that dequantizes a weight block does it piece by piece inside its matmul, so the VPU work overlaps + # the MXU work (see _matmul's rhs_fill). + if cfgs.rhs_dequant_cache: + # One K tile: every step is the first and the last K step. The step that starts a block fills the cache tile; + # the block's later m tiles read it. + fill = jax.named_scope("matmul_first_last_rhs_fill")(functools.partial(_matmul, True, True, rhs_fill=True)) + cached = jax.named_scope("matmul_first_last")(functools.partial(_matmul, True, True)) + lax.cond(_rhs_cache_miss(rhs_cache_refs, metadata_ref, cfgs), fill, cached) + return + # Two or more K tiles: consecutive steps read different blocks, so every step dequantizes its own (no cache), and + # no step is both the first and the last K step. + first, last, mid = ( + jax.named_scope(name + "_rhs_fill")(functools.partial(_matmul, first_k, last_k, rhs_fill=True)) + for name, first_k, last_k in ( + ("matmul_first", True, False), + ("matmul_last", False, True), + ("matmul", False, False), + ) + ) + lax.cond(k_id == 0, first, lambda: lax.cond(k_id == (num_k - 1), last, mid)) return is_first_k_step = k_id == 0 @@ -1154,27 +1168,22 @@ def dequantize_fp8_codes(codes: jax.Array, scale: jax.Array, dtype: jnp.dtype) - def _rhs_block_key(metadata_ref: MetadataRef, cfgs: GmmConfigs) -> jax.Array: - """rhs_contract_dequant: an int32 id of this grid step's rhs block (group, n tile, k tile).""" + """rhs_dequant_cache: an int32 id of this grid step's rhs block (group, n tile); there is one k tile.""" num_n = pl.num_programs(cfgs.n_axis) - num_k = pl.num_programs(2) group_id = metadata_ref.gm_id_to_group_id[pl.program_id(cfgs.gm_axis)] - return (group_id * num_n + pl.program_id(cfgs.n_axis)) * num_k + pl.program_id(2) + return group_id * num_n + pl.program_id(cfgs.n_axis) def _rhs_cache_miss( rhs_cache_refs: tuple[jax.Array, jax.Array] | None, metadata_ref: MetadataRef, cfgs: GmmConfigs ) -> jax.Array: - """rhs_contract_dequant: whether the cache tile holds another rhs block than this step's. - - The rhs block index is (group, n tile, k tile), and k is the innermost grid axis. With a single k tile (the MoE - DLHS tiles) the cache is filled once per (group, n tile) and every m tile of the group reuses it, for any number - of n tiles (the DLHS grid is (n, gm, k)). With two or more k tiles consecutive steps need different blocks, so the - cache never hits and every step dequantizes its block inside the matmul; outputs are unchanged, and it still beats - dequantizing the whole weight in XLA first (TPU7x, production DLHS shapes, 2-4 k tiles: 0.84-0.91x its time), - because the kernel reads 1-byte codes instead of a bf16 copy. Side effect: records this step's block key in the - SMEM key, since the step that misses fills the cache tile. + """rhs_dequant_cache: whether the cache tile holds another rhs block than this step's. + + The cache is filled once per (group, n tile) and every m tile of the group reuses it, for any number of n tiles (the + DLHS grid is (n, gm, k)). Side effect: records this step's block key in the SMEM key, since the step that misses + fills the cache tile. """ - assert rhs_cache_refs is not None + assert cfgs.rhs_dequant_cache and rhs_cache_refs is not None key_ref = rhs_cache_refs[1] key = _rhs_block_key(metadata_ref, cfgs) miss = key != key_ref[0] @@ -1456,7 +1465,7 @@ def kernel_main( pre_act_partial_refs: tuple[jax.Array, ...] = (), # return_pre_act: their partial-row scratch # Per extra output, the VMEM zero buffer of its zero-init DMAs when its dtype is not out_ref's (None: zero_ref). extra_zero_refs: tuple[jax.Array | None, ...] = (), - # rhs_contract_dequant: (the dequantized weight tile [tile_n, tile_k], its SMEM block key int32[1]). + # rhs_dequant_cache: (the dequantized weight tile [tile_n, tile_k], its SMEM block key int32[1]). rhs_cache_refs: tuple[jax.Array, jax.Array] | None = None, ): """Entry point for GMM kernel. @@ -1594,7 +1603,7 @@ def kernel_main( if pre_act_out_refs: out_in = (out_in, *(r.reshape(-1, cfgs.dims.size_lhs_sublane, *r.shape[1:]) for r in pre_act_out_refs)) scratches[0] = (partial_out_ref, *pre_act_partial_refs) - if cfgs.rhs_contract_dequant: + if cfgs.rhs_dequant_cache: assert rhs_cache_refs is not None rhs_cache_refs[1][0] = jnp.int32(-1) # no block dequantized yet scratches.append(rhs_cache_refs) @@ -2557,7 +2566,7 @@ def gmm_v2( ) dot_zero_scratch = pltpu.VMEM((dot_zero_tile_m, num_lanes), cfgs.acc_dtype) if cfgs.zero_init else None scratch_shapes += [pltpu.VMEM((dims.size_lhs_sublane, num_lanes), cfgs.acc_dtype), dot_zero_scratch] - if cfgs.rhs_contract_dequant: + if cfgs.rhs_dequant_cache: # Last: the dequantized weight tile and the SMEM key of the rhs block it holds (kernel_body peels them off). scratch_shapes += [ pltpu.VMEM((tiles.tile_n, tiles.tile_k), cfgs.rhs_dequant_dtype), @@ -2619,7 +2628,7 @@ def gmm_v2( pallas_out_specs = out_spec kernel_body = functools.partial(kernel_fn, cfgs=cfgs) - if cfgs.rhs_contract_dequant: + if cfgs.rhs_dequant_cache: base_kernel_body = kernel_body def kernel_body_with_rhs_cache(*refs): diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index 1513e00382..efaffbd920 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -939,6 +939,25 @@ def fp8_native_fwd_kwargs(config, with_bwd_dlhs: bool = True) -> dict: return kwargs +def fp8_native_fwd_k_splits(mesh_shape, wi_pspec, wo_pspec, num_moe_emb_chunks: int) -> list[str]: + """What splits the K dim of the native FP8 forward GMMs, as seen by the expert shard_map (empty: nothing). + + fp8_moe_native_fwd quantizes each activation row with one scale over K: dim 1 of the [experts, K, N] weights (wi: + the embed dim, wo: the MoE MLP dim; the tokens enter the shard_map with the embed dim unsharded). A split K would + give each shard or chunk its own scale, and gmm_v2 only sees the local size_k. The config validator rejects the + splits it can see; this checks the traced partition specs (custom_mesh_and_rule, parallelism -1). + """ + splits = [] + for name, dim, pspec in (("wi", "embed", wi_pspec), ("wo", "MoE MLP", wo_pspec)): + axes = [axis for axis in mesh_axes_for_dim(pspec[1]) if mesh_shape.get(axis, 1) > 1] if pspec is not None else [] + if axes: + degree = math.prod(mesh_shape[axis] for axis in axes) + splits.append(f"{name}'s K ({dim} dim) sharded {degree}-way over {tuple(axes)}") + if num_moe_emb_chunks > 0: + splits.append(f"wi's K in num_moe_emb_chunks={num_moe_emb_chunks} chunks") + return splits + + class GateLogit(nnx.Module): """A layer used to compute gate logits, allowing to return the pre bias values for DeepSeek routing.""" @@ -5041,6 +5060,12 @@ def _cn(t): # fp8_moe_native_fwd: the expert weights enter as Fp8NativeWeight (anchor, codes, scale); see below. use_fp8_native_fwd = getattr(self.config, "fp8_moe_native_fwd", False) and fp8_fake_quant_at_gather is not None if use_fp8_native_fwd: + k_splits = fp8_native_fwd_k_splits(self.mesh.shape, w0_pspec, wo_pspec, self.config.num_moe_emb_chunks) + if k_splits: + raise ValueError( + "fp8_moe_native_fwd=True quantizes each activation row with one scale over the whole K dim of the forward" + f" GMMs, but here: {'; '.join(k_splits)}." + ) weight_specs = slice(3, 6) # w0_pspec, w1_pspec, wo_pspec in in_specs above assert in_specs[weight_specs] == (w0_pspec, w1_pspec, wo_pspec), "in_specs order changed" native_weight_specs = tuple( diff --git a/tests/unit/fp8_moe_native_test.py b/tests/unit/fp8_moe_native_test.py index a3db896789..e178d58cc6 100644 --- a/tests/unit/fp8_moe_native_test.py +++ b/tests/unit/fp8_moe_native_test.py @@ -156,6 +156,14 @@ def test_native_config_is_accepted(): }, f"eval_wo_tile_fwd_mlp_dim={_MOE_MLP}", ), + # One activation scale per row over the whole K: no mesh axis or emb chunking may split wi's or wo's K. + # (moe_unsplit_prefused_wi rejects a sharded MLP dim and emb chunks on its own; off, so only this check fires.) + ( + {"ici_tensor_sequence_parallelism": 2, "moe_unsplit_prefused_wi": False}, + "an unsplit MoE MLP dim, wo's K", + ), + ({"ici_tensor_parallelism": 2, "moe_unsplit_prefused_wi": False}, "an unsplit embed dim, wi's K"), + ({"num_moe_emb_chunks": 2, "moe_unsplit_prefused_wi": False}, "num_moe_emb_chunks=0"), ], ) def test_native_fwd_requirement_is_enforced(override, requirement): @@ -208,6 +216,29 @@ def test_fake_quant_gather8_outside_the_layer_scan_still_lowers(): _lower_forward(_cfg(scan_layers=False)) +def test_native_fwd_k_splits_reads_the_weight_k_dim(): + # Dim 1 of the [experts, K, N] weights is K (wi: embed, wo: MoE MLP); size-1 mesh axes split nothing. + spec = jax.sharding.PartitionSpec + mesh_shape = {"expert": 4, "fsdp": 2, "tensor": 2, "tensor_sequence": 1} + wi, wo = spec("expert", None, "tensor"), spec("expert", "tensor", None) + assert moe.fp8_native_fwd_k_splits(mesh_shape, wi, wo, 0) == ["wo's K (MoE MLP dim) sharded 2-way over ('tensor',)"] + assert moe.fp8_native_fwd_k_splits( + mesh_shape, spec("expert", ("tensor_sequence", "tensor", "fsdp"), None), spec("expert", "tensor_sequence", None), 2 + ) == [ + "wi's K (embed dim) sharded 4-way over ('tensor', 'fsdp')", + "wi's K in num_moe_emb_chunks=2 chunks", + ] + assert not moe.fp8_native_fwd_k_splits({**mesh_shape, "tensor": 1}, wi, wo, 0) + + +def test_native_fwd_k_split_guard_runs_before_the_gmms(): + # custom_mesh_and_rule and parallelism -1 are only resolved in the layer, so RoutedMoE checks the traced specs. + with mock.patch.object(moe, "fp8_native_fwd_k_splits", return_value=["wo's K split for the test"]) as k_splits: + with pytest.raises(ValueError, match="one scale over the whole K dim.*wo's K split for the test"): + _lower_forward(_cfg(**_NATIVE)) + assert k_splits.called + + def test_dequantize_fp8_codes_matches_the_gather8_dequantization(): # The native backward multiplies by dequantize_fp8_codes; fake QAT's gather8 path by dequantize_weight_codes_ste. # The native path claims bit-identical gradients, so the two expressions must agree bit for bit.