From 39bf166a70ddee72e18b8354c409891e0c774002 Mon Sep 17 00:00:00 2001 From: Shuwen-Fang Date: Thu, 23 Jul 2026 21:05:35 +0000 Subject: [PATCH] add --- .../kernels/ragged/ragged_gather_reduce_v2.py | 84 ++++++++++--- src/maxtext/kernels/ragged/ragged_sort.py | 119 ++++++++++-------- src/maxtext/layers/moe.py | 34 +++-- 3 files changed, 153 insertions(+), 84 deletions(-) diff --git a/src/maxtext/kernels/ragged/ragged_gather_reduce_v2.py b/src/maxtext/kernels/ragged/ragged_gather_reduce_v2.py index b6ccc1e281..427888dac8 100644 --- a/src/maxtext/kernels/ragged/ragged_gather_reduce_v2.py +++ b/src/maxtext/kernels/ragged/ragged_gather_reduce_v2.py @@ -263,12 +263,35 @@ def _calculate_col_chunk_size(col_size: int, num_simd_lanes: int) -> int: return 128 +def _compute_sorted_by_validity( + valid_rows_mask_2d: jax.Array, + num_row_partitions: int, + row_chunk_size: int, +) -> jax.Array: + """Computes row index permutation sorted by validity for SparseCore tiling.""" + row_partition_size = valid_rows_mask_2d.shape[1] + # Stable sort of a boolean key is a stable partition: valid rows keep their + # relative order and move ahead of the invalid ones. + sorted_by_validity = jnp.argsort(~valid_rows_mask_2d, descending=False, stable=True, axis=-1) + sorted_by_validity += jnp.arange(num_row_partitions)[:, None] * row_partition_size + + pad_to = _align_to(row_partition_size, row_chunk_size) + if pad_to > row_partition_size: + sorted_by_validity = jnp.pad( + sorted_by_validity, + ((0, 0), (0, pad_to - row_partition_size)), + constant_values=0, + ) + return sorted_by_validity.reshape(-1).astype(jnp.int32) + + def _preprocess( valid_rows_mask: jax.Array, reduce_group_size: int, num_row_partitions: int, num_simd_lanes: int, row_chunk_size: int, + precomputed_sorted_by_validity: jax.Array | None = None, ) -> tuple[jax.Array, jax.Array, jax.Array]: """Sorts valid source rows to the front of each row partition. @@ -279,22 +302,12 @@ def _preprocess( ``num_simd_lanes`` so the kernel can load it as a single vector. mask: per output group, whether the group has any valid source row. """ - row_partition_size = valid_rows_mask.shape[0] // num_row_partitions valid_rows_mask_2d = valid_rows_mask.reshape(num_row_partitions, -1) - # Stable sort of a boolean key is a stable partition: valid rows keep their - # relative order and move ahead of the invalid ones. - sorted_by_validity = jnp.argsort(~valid_rows_mask_2d, descending=False, stable=True, axis=-1) - sorted_by_validity += jnp.arange(num_row_partitions)[:, None] * row_partition_size - - pad_to = _align_to(row_partition_size, row_chunk_size) - if pad_to > row_partition_size: - sorted_by_validity = jnp.pad( - sorted_by_validity, - ((0, 0), (0, pad_to - row_partition_size)), - constant_values=0, - ) - sorted_by_validity = sorted_by_validity.reshape(-1) + if precomputed_sorted_by_validity is not None: + sorted_by_validity = precomputed_sorted_by_validity + else: + sorted_by_validity = _compute_sorted_by_validity(valid_rows_mask_2d, num_row_partitions, row_chunk_size) num_src_rows_per_row_partition = jnp.pad( jnp.sum(valid_rows_mask_2d, axis=-1).astype(jnp.int32), @@ -302,7 +315,7 @@ def _preprocess( ) mask = jnp.any(valid_rows_mask.reshape(-1, reduce_group_size), axis=-1) return ( - sorted_by_validity.astype(jnp.int32), + sorted_by_validity, num_src_rows_per_row_partition, mask, ) @@ -624,9 +637,40 @@ def col_loop(col_compute_offset): ) -@functools.partial( - jax.jit, static_argnames=("reduce_group_size", "enforce_fallback", "flops_override", "bytes_accessed_override") -) +def get_sorted_by_validity( + hidden_size: int, + indices_size: int, + valid_rows_mask: jax.Array, + reduce_group_size: int, +) -> jax.Array | None: + """Computes sorted_by_validity row index permutation from valid_rows_mask.""" + if jax.devices()[0].platform != "tpu": + return None + sc_info = pltpu.get_tpu_info().sparse_core + if sc_info is None: + return None + + num_simd_lanes = sc_info.num_lanes + num_lanes = pltpu.get_tpu_info().num_lanes + num_cores = sc_info.num_cores * sc_info.num_subcores + + num_column_partitions = _calculate_num_column_partitions( + hidden_size, indices_size, num_cores, num_lanes, num_simd_lanes + ) + num_row_partitions = num_cores // num_column_partitions + _, row_chunk_size = _calculate_row_tiling(indices_size, num_simd_lanes, num_row_partitions) + + padded_input_size = _align_to(indices_size, num_row_partitions * reduce_group_size) + padded_valid_rows_mask = jnp.pad( + valid_rows_mask, + (0, padded_input_size - indices_size), + constant_values=False, + ) + + valid_rows_mask_2d = padded_valid_rows_mask.reshape(num_row_partitions, -1) + return _compute_sorted_by_validity(valid_rows_mask_2d, num_row_partitions, row_chunk_size) + + def ragged_gather_reduce( x: jax.Array, indices: jax.Array, @@ -636,6 +680,7 @@ def ragged_gather_reduce( enforce_fallback: bool = False, flops_override: int = -1, bytes_accessed_override: int = -1, + precomputed_sorted_by_validity: jax.Array | None = None, ) -> jax.Array: """Gathers ``x`` by ``indices``, weights and masks, then reduces by group. @@ -645,6 +690,8 @@ def ragged_gather_reduce( topk_weights: 1-D per-row weights, ``(input_size,)``. valid_rows_mask: 1-D bool mask of valid gathered rows, ``(input_size,)``. reduce_group_size: number of consecutive rows summed into one output row. + precomputed_sorted_by_validity: optional pre-computed sorted_by_validity array from + ``get_sorted_by_validity``. Returns: Reduced output, ``(input_size // reduce_group_size, hidden_size)``. @@ -707,6 +754,7 @@ def ragged_gather_reduce( num_row_partitions, num_simd_lanes, row_chunk_size, + precomputed_sorted_by_validity=precomputed_sorted_by_validity, ) # Step 4: Launch the SparseCore kernel. diff --git a/src/maxtext/kernels/ragged/ragged_sort.py b/src/maxtext/kernels/ragged/ragged_sort.py index 639d3172b7..4330c18298 100644 --- a/src/maxtext/kernels/ragged/ragged_sort.py +++ b/src/maxtext/kernels/ragged/ragged_sort.py @@ -17,7 +17,7 @@ import jax import jax.numpy as jnp from maxtext.kernels.ragged.ragged_gather import ragged_gather -from maxtext.kernels.ragged.ragged_gather_reduce_v2 import ragged_gather_reduce +from maxtext.kernels.ragged.ragged_gather_reduce_v2 import get_sorted_by_validity, ragged_gather_reduce def ring_ragged_sort( @@ -34,6 +34,7 @@ def ring_ragged_sort( gather_reduce_flops_override=-1, gather_bytes_accessed_override=-1, gather_reduce_bytes_accessed_override=-1, + precomputed_sorted_by_validity=None, ): """Ragged-gather variant for AG-RS Expert Parallelism token routing. @@ -62,6 +63,7 @@ def ring_ragged_sort( ep_name: ``str`` identifying the expert parallel axis name. ep_size: scalar ``int`` representing the expert parallel mesh size. buffer_size: optional scalar ``int`` representing the size of the local buffer. + precomputed_sorted_by_validity: optional pre-computed sorted_by_validity array. Returns: A tuple containing: @@ -72,13 +74,41 @@ def ring_ragged_sort( - 1D tensor ``topk_argsort_revert_indices`` for inverse routing. """ + def _compute_valid_rows_mask(topk_argsort_revert_indices, shard_output_start, shard_output_end, local_buffer_size): + """Computes valid rows mask and indices for ragged gather / gather reduce operations.""" + # Restrict to the [start, end) source range via a validity bitmask. The + # ragged kernel packs valid rows to the front of each row-partition and + # only iterates over the populated prefix, so we hand it the mask directly + # 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 + ) + indices = topk_argsort_revert_indices + else: + # Buffering: g_x has size `local_buffer_size` (packed). + # The revert indices are global [0, n), but they must map to the local + # packed g_x buffer. + shifted_indices = topk_argsort_revert_indices - shard_output_start + local_num_tokens = shard_output_end - shard_output_start + # We only reduce gradients from the valid portion of the local buffer. + limit = jnp.minimum(local_num_tokens, local_buffer_size) + # Mask out tokens that were not gathered (either because they belong to + # other shards, or they exceeded the local buffer size). + valid_rows_mask = (shifted_indices >= 0) & (shifted_indices < limit) + # 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`. + indices = jnp.where(valid_rows_mask, shifted_indices, 0) + return valid_rows_mask, indices + @jax.custom_vjp - def _ring_ragged_sort(hidden_states_local, topk_indices_local): + def _ring_ragged_sort(hidden_states_local, topk_indices_local, precomputed_sorted_by_validity): """Sort and gather activations to different EP shards.""" - return _ring_ragged_sort_fwd(hidden_states_local, topk_indices_local)[0] + return _ring_ragged_sort_fwd(hidden_states_local, topk_indices_local, precomputed_sorted_by_validity)[0] @jax.named_scope("ragged-sort-fwd") - def _ring_ragged_sort_fwd(hidden_states_local, topk_indices_local): + def _ring_ragged_sort_fwd(hidden_states_local, topk_indices_local, precomputed_sorted_by_validity): """Sort and gather activations forward pass.""" num_tokens_local = hidden_states_local.shape[0] @@ -136,7 +166,17 @@ def _ring_ragged_sort_fwd(hidden_states_local, topk_indices_local): bytes_accessed_override=gather_bytes_accessed_override, ) - out = (x, group_sizes_local, topk_argsort_revert_indices) + n = topk_argsort_revert_indices.shape[0] + sorted_by_validity_cache = None + if precomputed_sorted_by_validity is None and not enforce_gather_reduce_fallback: + valid_rows_mask_fwd, _ = _compute_valid_rows_mask( + topk_argsort_revert_indices, shard_output_start, shard_output_end, local_buffer_size + ) + sorted_by_validity_cache = get_sorted_by_validity(hidden_states_local.shape[1], n, valid_rows_mask_fwd, topk) + else: + sorted_by_validity_cache = precomputed_sorted_by_validity + + out = (x, group_sizes_local, topk_argsort_revert_indices, sorted_by_validity_cache) res = ( topk_argsort_revert_indices, @@ -144,6 +184,7 @@ def _ring_ragged_sort_fwd(hidden_states_local, topk_indices_local): shard_output_end, local_buffer_size, hidden_states_local.shape, + sorted_by_validity_cache, ) return out, res @@ -163,62 +204,30 @@ def _ring_ragged_sort_bwd(res, g_out): shard_output_end, local_buffer_size, _, + sorted_by_validity_cache, ) = res - g_x, _, _ = g_out - # Restrict to the [start, end) source range via a validity bitmask. The - # ragged kernel packs valid rows to the front of each row-partition and - # only iterates over the populated prefix, so we hand it the mask directly - # rather than materializing a (mostly-zero) dense buffer ourselves. - n = topk_argsort_revert_indices.shape[0] + g_x, _, _, _ = g_out - 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( - g_x, - topk_argsort_revert_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, - ) - else: - # Buffering: g_x has size `local_buffer_size` (packed). - # The revert indices are global [0, n), but they must map to the local - # packed g_x buffer. - shifted_indices = topk_argsort_revert_indices - shard_output_start - local_num_tokens = shard_output_end - shard_output_start - # We only reduce gradients from the valid portion of the local buffer. - limit = jnp.minimum(local_num_tokens, local_buffer_size) - # Mask out tokens that were not gathered (either because they belong to - # other shards, or they exceeded the local buffer size). - valid_rows_mask = (shifted_indices >= 0) & (shifted_indices < limit) - # 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) + valid_rows_mask, indices = _compute_valid_rows_mask( + topk_argsort_revert_indices, shard_output_start, shard_output_end, local_buffer_size + ) - 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, - ) - return grad_hidden_states, None + grad_hidden_states = ragged_gather_reduce( + g_x, + indices, + topk_weights=jnp.ones((topk_argsort_revert_indices.shape[0],), 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, + precomputed_sorted_by_validity=sorted_by_validity_cache, + ) + return grad_hidden_states, None, None _ring_ragged_sort.defvjp(_ring_ragged_sort_fwd, _ring_ragged_sort_bwd) - return _ring_ragged_sort(hidden_states_local, topk_indices_local) + return _ring_ragged_sort(hidden_states_local, topk_indices_local, precomputed_sorted_by_validity) def ring_ragged_unsort( diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index 7737734b09..0da0fe6486 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -70,6 +70,7 @@ class RouteMetadata: # Shape [num_ep, num_ep]. all_gather of reshaped_group_sizes across EP shards. # [i, j] = number of tokens from batch shard i sent to expert shard j. all_shards_group_sizes: Optional[jax.Array] + precomputed_sorted_by_validity: Optional[jax.Array] = None @struct.dataclass @@ -846,11 +847,12 @@ def permute( self, inputs, gate_logits, - pre_bias_logits, + pre_bias_logits=None, use_custom_sort_vjp=True, rngs=None, roll_to_expert_id=None, input_ids=None, + precomputed_sorted_by_validity=None, ): """Permute tokens to group by expert to fit gmm call.""" # reshape inputs (batch, sequence, emb) to (batch * sequence, emb) @@ -888,6 +890,7 @@ def permute( # local prefix of valid rows. use_ragged_in_permute = self.config.use_ragged_sort and self.config.use_ring_of_experts buffer_size = None + sorted_by_validity = None if use_ragged_in_permute: topk_indices_2d = jnp.reshape(selected_experts, (bsz_times_seq_len, selected_experts.shape[2])) # roll_to_expert_id is not directly used in the kernel, ep axis id is directly called @@ -902,8 +905,7 @@ def permute( ) else: buffer_size = None - - sorted_inputs, group_size, sorted_selected_experts = ring_ragged_sort( + sorted_inputs, group_size, sorted_selected_experts, sorted_by_validity = ring_ragged_sort( inputs_2d, topk_indices_2d, self.config.num_experts, @@ -917,6 +919,7 @@ def permute( gather_reduce_flops_override=self.config.ragged_gather_reduce_cost_estimate_flops, gather_bytes_accessed_override=self.config.ragged_gather_cost_estimate_bytes_accessed, gather_reduce_bytes_accessed_override=self.config.ragged_gather_reduce_cost_estimate_bytes_accessed, + precomputed_sorted_by_validity=precomputed_sorted_by_validity, ) else: flatten_selected_experts = jnp.ravel(selected_experts) @@ -971,6 +974,7 @@ def permute( lb_loss, bias_updates, local_group_size, + sorted_by_validity, ) def unpermute( @@ -1648,7 +1652,9 @@ def get_routed_moe_shardings(is_batch_sharded_by_expert, has_input_ids): ) = get_routed_moe_shardings(is_batch_sharded_by_expert, input_ids is not None) w0_pspec, w1_pspec, wo_pspec = maybe_aqt_partition(w0_kernel, w0_pspec, w1_kernel, w1_pspec, wo_kernel, wo_pspec) - def roe_ag_and_route(x, logits, pre_bias_logits, num_ep, expert_shard_id, rngs, input_ids=None): + def roe_ag_and_route( + x, logits, pre_bias_logits, num_ep, expert_shard_id, rngs, input_ids=None, precomputed_sorted_by_validity=None + ): # The ring-of-experts strategy first duplicates the inputs to all # expert shards, and then routes within each shard. @@ -1668,6 +1674,7 @@ def roe_ag_and_route(x, logits, pre_bias_logits, num_ep, expert_shard_id, rngs, lb_loss, bias_updates, local_group_sizes, + sorted_by_validity, ) = self.permute( x, logits, @@ -1676,6 +1683,7 @@ def roe_ag_and_route(x, logits, pre_bias_logits, num_ep, expert_shard_id, rngs, roll_to_expert_id=num_experts_per_shard * expert_shard_id, rngs=rngs, input_ids=input_ids, + precomputed_sorted_by_validity=precomputed_sorted_by_validity, ) return ( x, @@ -1694,6 +1702,7 @@ def roe_ag_and_route(x, logits, pre_bias_logits, num_ep, expert_shard_id, rngs, all_shards_group_sizes=None, reshaped_group_sizes=None, ), + sorted_by_validity, ) def ra2a_and_route(x, logits, pre_bias_logits, num_ep, expert_shard_id, rngs, input_ids=None): @@ -1709,6 +1718,7 @@ def ra2a_and_route(x, logits, pre_bias_logits, num_ep, expert_shard_id, rngs, in lb_loss, bias_updates, local_group_sizes, + _, ) = self.permute( x, logits, @@ -1801,7 +1811,7 @@ def route(x, logits, pre_bias_logits, rngs, input_ids=None): expert_shard_id = jax.lax.axis_index(self._expert_parallelism_name) if num_ep > 1 else 0 if self.config.use_ring_of_experts: - return roe_ag_and_route( + routed_x, route_out, route_meta, _ = roe_ag_and_route( x, logits, pre_bias_logits, @@ -1810,6 +1820,7 @@ def route(x, logits, pre_bias_logits, rngs, input_ids=None): rngs, input_ids=input_ids, ) + return routed_x, route_out, route_meta else: return ra2a_and_route( x, @@ -2029,7 +2040,7 @@ def moe_emb_chunking( ) first_x_unrouted = jax.lax.dynamic_slice_in_dim(x, 0, chunk_dim, axis=2) - cur_x_chunk, routing, route_metadata = roe_ag_and_route( + cur_x_chunk, routing, route_metadata, sorted_by_validity = roe_ag_and_route( first_x_unrouted, logits, pre_bias_logits, @@ -2046,7 +2057,7 @@ def moe_emb_chunking( partial_sum1 = jnp.zeros((cur_x_chunk.shape[0], w1.shape[-1]), dtype=cur_x_chunk.dtype) def scan_fn(carry, _): - cur_x, ps0, ps1, chunk_idx = carry + cur_x, ps0, ps1, chunk_idx, precomputed_sorted_by_validity = carry cur_w0 = jax.lax.dynamic_slice_in_dim(w0, chunk_idx * chunk_dim, chunk_dim, axis=1) cur_w1 = jax.lax.dynamic_slice_in_dim(w1, chunk_idx * chunk_dim, chunk_dim, axis=1) @@ -2063,7 +2074,7 @@ def scan_fn(carry, _): partial_accum1=ps1, ) next_x_unrouted = jax.lax.dynamic_slice_in_dim(x, (chunk_idx + 1) * chunk_dim, chunk_dim, axis=2) - next_x, _, _ = roe_ag_and_route( + next_x, _, _, _ = roe_ag_and_route( next_x_unrouted, logits, pre_bias_logits, @@ -2071,12 +2082,13 @@ def scan_fn(carry, _): expert_shard_id, chunk_rngs_key, input_ids=sharded_input_ids, + precomputed_sorted_by_validity=precomputed_sorted_by_validity, ) - return (next_x, next_ps0, next_ps1, chunk_idx + 1), None + return (next_x, next_ps0, next_ps1, chunk_idx + 1, precomputed_sorted_by_validity), None - (last_x_chunk, ps0, ps1, _), _ = jax.lax.scan( + (last_x_chunk, ps0, ps1, _, _), _ = jax.lax.scan( scan_fn, - (cur_x_chunk, partial_sum0, partial_sum1, 0), + (cur_x_chunk, partial_sum0, partial_sum1, 0, sorted_by_validity), None, length=self.config.num_moe_emb_chunks - 1, )