Quantize activations before EP all gather

src/maxtext/configs/base.yml

@@ -226,6 +226,8 @@ load_balance_loss_weight: 0.0 # weight for the load balance loss
 use_random_routing: false # whether to use random routing for debug/test purpose
 use_custom_sort_vjp: true # whether to use a custom VJP sort for efficient backward pass processing in sparse matmul
 use_ring_of_experts: false # whether to use ring of experts for sparse matmul expert parallelism
+quantize_before_ep_all_gather: false # whether to quantize activations before the ring-of-experts EP all-gather
+                                    # (fp8 collective + ragged sort) vs. quantizing later inside the gmm call
 num_moe_emb_chunks: 0 # number of chunks for overlapping token all-gather and GMM computation along embedding dimension
 # If true, peel the 'expert' mesh axis off the MoE dispatch/MLP batch dim so the expert GEMM
 # stays expert-parallel (AllToAll); false keeps 'expert' on the batch dim (activation_batch_moe).

src/maxtext/configs/types.py

@@ -867,6 +867,13 @@ class MoEGeneral(BaseModel):
       False,
       description="Whether to use Ring of Experts for sparse matmul expert parallelism.",
   )
+  quantize_before_ep_all_gather: bool = Field(
+      False,
+      description=(
+          "Whether to quantize activations before the Ring of Experts EP all-gather (so the "
+          "collective and ragged sort move fp8, not bf16), vs. quantizing later inside the gmm call."
+      ),
+  )
   moe_dispatch_no_expert_sharding: bool = Field(
       False,
       description=(
@@ -2841,6 +2848,16 @@ def validate_ragged_buffer_factor(self):
           "  2. Ragged sort with ring of experts (use_ring_of_experts=True AND use_ragged_sort=True)"
       )
 
+  def _validate_quantize_before_ep_all_gather(self):
+    """Validates quantize_before_ep_all_gather is used with supported settings."""
+    if self.quantize_before_ep_all_gather and not (
+        self.use_ring_of_experts and self.use_qwix_quantization and self.use_gmm_v2
+    ):
+      raise ValueError(
+          "quantize_before_ep_all_gather=True is only supported with use_ring_of_experts=True and "
+          "qwix quantization, and gmm v2 kernel"
+      )
+
   def _validate_use_te_comm_gemm_overlap(self):
     """Validates that use_te_comm_gemm_overlap is used with supported settings to enable TE Collective GEMM ops."""
     te_has_distributed_env = jax.local_device_count() == 1 and jax.distributed.is_initialized()
@@ -3948,6 +3965,8 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
       if self.use_batch_split_schedule:
         raise ValueError("GMM v2 is not supported with a batch split schedule.")
 
+    self._validate_quantize_before_ep_all_gather()
+
     for val in self.compress_ratios:
       if val != 0 and val < 4:
         raise ValueError(f"compress_ratio must be 0 (disabled) or >= 4, got {val}")

src/maxtext/kernels/megablox/ops.py

@@ -27,6 +27,8 @@
 from maxtext.layers import quantizations
 import qwix
 import qwix.pallas as qpl
+from qwix._src.core import numerics as qwix_numerics
+from qwix._src.core.qarray import call_with_generic_broadcast
 import tokamax
 
 
@@ -102,13 +104,18 @@ def gmm(
           act_calibration_method="absmax",
       )
 
+  lhs_scale = None
+  if isinstance(lhs, qpl.QArray):
+    lhs_scale = lhs.scale
+    lhs = lhs.qvalue
+
   gmm_fwd_bwd = lambda *args: _gmm_fwd(*args)[0]  # pylint: disable=C3001
   gmm_fwd_bwd = jax.custom_vjp(
       gmm_fwd_bwd,
       nondiff_argnums=(3, 4, 7, 8, 9, 10, 11, 12, 13, 14, 15),
   )
   gmm_fwd_bwd.defvjp(_gmm_fwd, functools.partial(_gmm_bwd, lhs.dtype, rhs.dtype))
-  return gmm_fwd_bwd(
+  out = gmm_fwd_bwd(
       lhs,
       rhs,
       group_sizes,
@@ -127,6 +134,9 @@ def gmm(
       use_gmm_v2,
       partial_sum,
   )
+  if lhs_scale is not None:
+    out = call_with_generic_broadcast(jnp.multiply, out, lhs_scale.astype(out.dtype))
+  return out
 
 
 # ==============================================================================
@@ -202,15 +212,7 @@ def _gmm_fwd(
     out = _fwd_run_tokamax_v1(lhs, rhs, group_sizes, preferred_element_type, transpose_rhs, use_manual_quantization)
   elif use_tokamax_backend and use_gmm_v2:
     out = _fwd_run_tokamax_v2(
-        lhs,
-        rhs,
-        group_sizes,
-        preferred_element_type,
-        tiling,
-        group_offset,
-        partial_sum,
-        transpose_rhs,
-        quantization_rule,
+        lhs, rhs, group_sizes, preferred_element_type, tiling, group_offset, partial_sum, transpose_rhs, quantization_rule
     )
   else:
     out = _fwd_run_megablox(
@@ -238,7 +240,7 @@ def _fwd_quantize_activation_and_weight(
     transpose_rhs: bool,
 ) -> tuple[jnp.ndarray | qpl.QArray, jnp.ndarray | qpl.QArray]:
   """Handles act and weight quantization for GMM forward inputs."""
-  if quantization_rule.act_qtype and not isinstance(lhs, qpl.QArray) and not use_gmm_v2:
+  if quantization_rule.act_qtype and not isinstance(lhs, qpl.QArray) and qwix_numerics.should_quantize(lhs.dtype):
     lhs = qpl.quantize(  # pyrefly: ignore[bad-assignment]
         lhs,
         quantization_rule.act_qtype,
@@ -382,24 +384,37 @@ def _fwd_run_tokamax_v2(
     rhs_operand = rhs_operand.qvalue
     rhs_scale = _fwd_prepare_rhs_scale(rhs, transpose_rhs=transpose_rhs)
 
+  lhs_operand = lhs.qvalue if isinstance(lhs, qpl.QArray) else lhs
+  maybe_quantize_lhs = not isinstance(lhs, qpl.QArray) and qwix_numerics.should_quantize(lhs_operand.dtype)
+
+  lhs_scale = _fwd_prepare_lhs_scale(quantization_rule) if maybe_quantize_lhs else None
+
   custom_fwd_tiling = gmm_v2.TileSizes(
       tile_m=tiling[0],
       tile_k=tiling[1],
       tile_n=tiling[2],
   )
 
-  return gmm_v2.gmm_v2(
-      lhs=lhs,  # pyrefly: ignore[bad-argument-type]
+  eff_pref_dtype = preferred_element_type if qwix_numerics.should_quantize(lhs_operand.dtype) else jnp.bfloat16
+
+  out = gmm_v2.gmm_v2(
+      lhs=lhs_operand,  # pyrefly: ignore[bad-argument-type]
       rhs=rhs_operand,  # pyrefly: ignore[bad-argument-type]
       group_sizes=group_sizes,
       rhs_scale=rhs_scale,
       tile_info=custom_fwd_tiling,
-      preferred_element_type=preferred_element_type,
+      preferred_element_type=eff_pref_dtype,
       partial_sum=partial_sum,
       group_offset=group_offset,
-      lhs_scale=_fwd_prepare_lhs_scale(quantization_rule),
+      maybe_quantize_lhs=maybe_quantize_lhs,
+      lhs_scale=lhs_scale,
   )
 
+  if isinstance(lhs, qpl.QArray):
+    out *= lhs.scale.astype(out.dtype)
+
+  return out
+
 
 def _fwd_run_megablox(
     lhs: jnp.ndarray,
@@ -560,7 +575,12 @@ def _bwd_prepare_inputs(
 
   # GMM2 FWD performs lhs quantization inside kernel, lhs is stored as unquantized dtype
   # in the residual tuple. In BWD, we explicitly quantize lhs.
-  if quantization_rule and quantization_rule.act_qtype and not isinstance(lhs, qpl.QArray):
+  if (
+      quantization_rule
+      and quantization_rule.act_qtype
+      and not isinstance(lhs, qpl.QArray)
+      and qwix_numerics.should_quantize(lhs.dtype)
+  ):
     lhs = qpl.quantize(  # pyrefly: ignore[bad-assignment]
         lhs,
         quantization_rule.act_qtype,

src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py

@@ -525,7 +525,12 @@ def _matmul(is_first_k_step: bool, is_last_k_step: bool):
           # 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:
+          if tiled_lhs.dtype == lhs_q_dtype:
+            # lhs block already arrives quantized, the real dequant
+            # scale is applied externally, so just pass an identity scale here.
+            block_lhs_q = block_lhs
+            block_scale = jnp.array(1.0, dtype=acc_ref.dtype)
+          elif 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)
@@ -1206,7 +1211,11 @@ def make_gmm_configs(
   )
 
   lhs_q_dtype = None
-  if maybe_quantize_lhs and rhs_cfgs.should_dequantize_after_matmul:
+  if jnp.issubdtype(lhs.dtype, jnp.integer) or jnp.issubdtype(lhs.dtype, jnp.float8_e4m3fn):
+    # lhs arrives already quantized (e.g. pre-quantized ahead-of-time by the
+    # caller): use its dtype as-is, no in-kernel quantization/scale needed.
+    lhs_q_dtype = lhs.dtype
+  elif maybe_quantize_lhs and rhs_cfgs.should_dequantize_after_matmul:
     # Choose lhs quantization dtype based on TPU hardware support.
     is_rhs_float = jnp.issubdtype(rhs_quant_dtype, jnp.floating)  # pyrefly: ignore[bad-argument-type]
     tpu_info = pltpu.get_tpu_info()
@@ -1217,10 +1226,10 @@ def make_gmm_configs(
       # floating rhs as conversion to int8 will cause numeric issues.
       is_rhs_4bits = jax.dtypes.itemsize_bits(rhs_quant_dtype) == 4  # pyrefly: ignore[bad-argument-type]
       if is_rhs_float or is_rhs_4bits:
-        lhs_q_dtype = jnp.float8_e4m3fn.dtype
+        lhs_q_dtype = jnp.float8_e4m3fn
     if tpu_info.int8_ops_per_second > 0:
       if not is_rhs_float:
-        lhs_q_dtype = jnp.int8.dtype
+        lhs_q_dtype = jnp.int8
 
   if lhs_scale is not None:
     assert lhs_q_dtype is not None, (
@@ -1241,17 +1250,21 @@ def make_gmm_configs(
       has_scale=has_lhs_scale,
   )
 
-  if out_dtype is None:
-    out_dtype = lhs.dtype
+  if out_dtype is None or jnp.issubdtype(out_dtype, jnp.float8_e4m3fn):
+    # The raw quantized-domain matmul output isn't yet rescaled -- writing it
+    # directly as fp8 would lose precision before the scale multiply happens
+    # (either inside this kernel via lhs_scale/block_scale, or externally by
+    # the caller for a pre-quantized lhs). Floor to bf16 as a safe intermediate.
+    out_dtype = jnp.bfloat16
 
   if acc_dtype is None:
     if lhs_cfgs.quant_dtype is None:
-      acc_dtype = jnp.float32.dtype
+      acc_dtype = jnp.float32
     else:
       # Input quantization requires elementwise ops which can put pressure on
       # VPUs. Using faster bf16 hardware during accumulation can help offset the
       # pressure.
-      acc_dtype = jnp.bfloat16.dtype
+      acc_dtype = jnp.bfloat16
 
   if isinstance(tile_info, TileSizes):
     tiles = tile_info

src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_tgmm_kernel.py

@@ -248,7 +248,7 @@ def make_tgmm_configs(
 
   fuse_act = None  # fuse_act has to be None in tgmm.
   if acc_dtype is None:
-    acc_dtype = jnp.float32.dtype
+    acc_dtype = jnp.float32
   if isinstance(tile_info, gmm_v2.TileSizes):
     tiles = tile_info
   else:

src/maxtext/kernels/ragged/ragged_sort.py

@@ -102,18 +102,21 @@ def _ring_ragged_sort_fwd(hidden_states_local, topk_indices_local):
     shard_output_start = group_offsets[experts_start]
     shard_output_end = group_offsets[experts_end]
 
-    if buffer_size is None or buffer_size >= num_tokens_local * topk:
-      local_buffer_size = num_tokens_local * topk
-      x = ragged_gather(
+    def _gather(indices, start, end):
+      return ragged_gather(
           hidden_states_local,
-          token_indices_sorted,
-          shard_output_start[None],
-          shard_output_end[None],
+          indices,
+          start,
+          end,
           enforce_fallback=enforce_gather_fallback,
           flops_override=gather_flops_override,
           bytes_accessed_override=gather_bytes_accessed_override,
           use_single_sparsecore=use_single_sparsecore,
       )
+
+    if buffer_size is None or buffer_size >= num_tokens_local * topk:
+      local_buffer_size = num_tokens_local * topk
+      x = _gather(token_indices_sorted, shard_output_start[None], shard_output_end[None])
     else:
       local_buffer_size = buffer_size
       # We only gather up to the available buffer size or the actual number of
@@ -128,16 +131,7 @@ def _ring_ragged_sort_fwd(hidden_states_local, topk_indices_local):
           local_buffer_size,
           axis=0,
       )
-      x = ragged_gather(
-          hidden_states_local,
-          sliced_indices,
-          jnp.int32(0)[None],
-          gather_end[None],
-          enforce_fallback=enforce_gather_fallback,
-          flops_override=gather_flops_override,
-          bytes_accessed_override=gather_bytes_accessed_override,
-          use_single_sparsecore=use_single_sparsecore,
-      )
+      x = _gather(sliced_indices, jnp.int32(0)[None], gather_end[None])
 
     out = (x, group_sizes_local, topk_argsort_revert_indices)
 
@@ -174,17 +168,11 @@ def _ring_ragged_sort_bwd(res, g_out):
     # rather than materializing a (mostly-zero) dense buffer ourselves.
     n = topk_argsort_revert_indices.shape[0]
 
-    if local_buffer_size >= n:
-      valid_rows_mask = (topk_argsort_revert_indices >= shard_output_start) & (
-          topk_argsort_revert_indices < shard_output_end
-      )
-      # The forward scatter-add over `token_indices_sorted` is equivalent to a
-      # gather-reduce: each input token has exactly `topk` contributions located
-      # at sorted positions `topk_argsort_revert_indices[t*topk:(t+1)*topk]`.
-      # `topk_weights` is set to ones because this op has no per-row weighting.
-      grad_hidden_states = ragged_gather_reduce(
+    def _gather_reduce(indices, valid_rows_mask):
+      """`topk_weights` is set to ones because this op has no per-row weighting."""
+      return ragged_gather_reduce(
           g_x,
-          topk_argsort_revert_indices,
+          indices,
           topk_weights=jnp.ones((n,), dtype=jnp.float32),
           valid_rows_mask=valid_rows_mask,
           reduce_group_size=topk,
@@ -193,6 +181,12 @@ def _ring_ragged_sort_bwd(res, g_out):
           bytes_accessed_override=gather_reduce_bytes_accessed_override,
           use_single_sparsecore=use_single_sparsecore,
       )
+
+    if local_buffer_size >= n:
+      valid_rows_mask = (topk_argsort_revert_indices >= shard_output_start) & (
+          topk_argsort_revert_indices < shard_output_end
+      )
+      grad_hidden_states = _gather_reduce(topk_argsort_revert_indices, valid_rows_mask)
     else:
       # Buffering: g_x has size `local_buffer_size` (packed).
       # The revert indices are global [0, n), but they must map to the local
@@ -207,18 +201,7 @@ def _ring_ragged_sort_bwd(res, g_out):
       # Clamp invalid indices to 0 to prevent compile-time/run-time out-of-bounds
       # in JAX. These clamped values will be ignored due to `valid_rows_mask`.
       safe_indices = jnp.where(valid_rows_mask, shifted_indices, 0)
-
-      grad_hidden_states = ragged_gather_reduce(
-          g_x,
-          safe_indices,
-          topk_weights=jnp.ones((n,), dtype=jnp.float32),
-          valid_rows_mask=valid_rows_mask,
-          reduce_group_size=topk,
-          enforce_fallback=enforce_gather_reduce_fallback,
-          flops_override=gather_reduce_flops_override,
-          bytes_accessed_override=gather_reduce_bytes_accessed_override,
-          use_single_sparsecore=use_single_sparsecore,
-      )
+      grad_hidden_states = _gather_reduce(safe_indices, valid_rows_mask)
     return grad_hidden_states, None
 
   _ring_ragged_sort.defvjp(_ring_ragged_sort_fwd, _ring_ragged_sort_bwd)