diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 2e05c57bc1..0a40440180 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -347,6 +347,17 @@ 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; 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 +# 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..47d18bdc38 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -1937,6 +1937,66 @@ 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 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_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 need fixed + tiles (float32 accumulator); the in-kernel DLHS dequantization also needs the transposed DLHS rhs. + """ + 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"), + # 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_unsplit_k_requirements() + + self._fp8_moe_native_fwd_tile_requirements(), + ) + 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=( @@ -2098,6 +2158,115 @@ 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_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. + + 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/__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..e47ff0b38b 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 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] + + +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 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. + """ + 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..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 @@ -386,6 +386,34 @@ 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) 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: @@ -516,6 +544,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 +643,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 +713,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 +742,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_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. """ pre_act_refs = () @@ -708,7 +752,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 +772,29 @@ 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: + # 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: + # 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) # 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 +825,17 @@ 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; 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], cfgs.rhs_dequant_dtype + ) + 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] return tiled_rhs[start_k:end_k, start_n:end_n] @@ -852,45 +923,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 +991,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 +1111,29 @@ def matmul_last(): num_k = pl.num_programs(2) k_id = pl.program_id(2) + if cfgs.rhs_contract_dequant: + # 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 is_last_k_step = k_id == (num_k - 1) @@ -1049,6 +1152,45 @@ 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_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) + 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) + + +def _rhs_cache_miss( + rhs_cache_refs: tuple[jax.Array, jax.Array] | None, metadata_ref: MetadataRef, cfgs: GmmConfigs +) -> jax.Array: + """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 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] + 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 +1465,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_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. @@ -1459,6 +1603,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_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) pipeline_fn(lhs_in, rhs_ref, ps_in, out_in, scratches=scratches) if cfgs.zero_init: @@ -1486,6 +1634,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 +1652,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 +1673,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 +1696,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 +1949,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 +1988,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 +2063,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, @@ -1915,6 +2088,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: @@ -2002,6 +2184,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 +2226,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 +2247,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 +2282,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 +2316,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 +2391,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 (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. + 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 +2460,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 +2486,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 +2566,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_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), + pltpu.SMEM((1,), jnp.int32), + ] aligned_n = align_to(cfgs.out_size_n, num_lanes) if cfgs.out_is_3d: @@ -2358,8 +2627,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_dequant_cache: + 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/layers/moe.py b/src/maxtext/layers/moe.py index 6a9e1543ba..efaffbd920 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -888,6 +888,76 @@ 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=[]) + + +# 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 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 + + +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.""" @@ -4052,12 +4122,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 +4229,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 +4605,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 +4640,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 +4660,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 +4671,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 +5057,22 @@ 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: + 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( + 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 +5703,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 +6838,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..e178d58cc6 --- /dev/null +++ b/tests/unit/fp8_moe_native_test.py @@ -0,0 +1,299 @@ +# 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, the +dequantization contract the native path relies on, and the fixed activation-quantization / accumulator settings. +""" + +import sys +import types +from unittest import mock + +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.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 +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 + +_EMB = 256 +_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": _EMB, + "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, + "wi_tile_fwd_embed_dim": _EMB, + "wo_tile_fwd_mlp_dim": _MOE_MLP, + "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_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 + + +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( + "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"), + # 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}", + ), + # 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): + 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_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. + 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_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"] + + +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)