diff --git a/tpu_inference/kernels/experimental/fused_moe/README.md b/tpu_inference/kernels/experimental/fused_moe/README.md new file mode 100644 index 0000000000..7f45b8eb03 --- /dev/null +++ b/tpu_inference/kernels/experimental/fused_moe/README.md @@ -0,0 +1,54 @@ +# Fused MoE kernels (experimental) + +Two TPU Pallas kernels for expert-parallel (EP) Mixture-of-Experts. They demonstrate +fusing the EP collective communication directly into the GMM compute. + +## `fused_moe/` — full fused MoE kernel (ready for usage) + +`fused_moe_func_rs` runs the entire EP MoE in a single Pallas call: + +``` +gather -> GMM1 -> activation -> GMM2 -> ICI A2A +``` + +The A2A is fused into the kernel, replacing the usual post-kernel +all-reduce/psum in EP MoE. It supports bf16 and quantized (fp8) weights, an +optional post-expert RMSNorm, and sequence-parallel in/out (toggle with +`TPU_MOE_ENABLE_SP`). This is the production-oriented kernel — use it as the MoE +layer entry point. + +```python +from tpu_inference.kernels.experimental.fused_moe import fused_moe_func_rs + +out = fused_moe_func_rs( + hidden_states, w1, w2, w1_scale, w2_scale, w1_bias, w2_bias, + gating_output, topk, renormalize, mesh, activation, scoring_fn, +) +``` + +## `gmm_fused/` — AG+GMM1 kernel (example) + +`gmm_v2_ag_gmm1` is a smaller, instructional kernel: a push-based all-gather +fused with `GMM1 + activation`, driven by a precomputed per-round send schedule +(`per_round_schedule.py`). It illustrates how to overlap and fuse the EP +all-gather into the grouped matmul — the same idea the full fused kernel above +builds on. It is provided as a reference/example rather than a complete MoE +layer; the paired GMM2 + ICI reduce-scatter kernels live in the same module. + +## Future work + +- **Fuse the upstream all-gather** into the kernel so it overlaps with compute, + using a persistent token cache (each token crosses to a receiver at most once) + and neighbor-relay routing (origin → nearest → next-nearest, link-local hops). + Biggest win at decode with large hidden size. The `gmm_fused/` example is a + first step. +- **Offload random-access work to SparseCore** — the routed-token gather, per-row + all-to-all, and cache lookups — keeping the dense matmul on the TensorCore. And we need to use both Tensorcore and sparsecore in one kernel. This could make the current MoE kernel even >20% faster in prefill/large_batch + +## Notes + +- Device-specific tuned block-size tables were removed for this release; + `tuned_block_sizes.py` returns caller-supplied defaults. Retune per device + for best performance. +- These kernels target TPU and require a multi-device mesh to exercise the EP + collectives. diff --git a/tpu_inference/kernels/experimental/fused_moe/__init__.py b/tpu_inference/kernels/experimental/fused_moe/__init__.py new file mode 100644 index 0000000000..9e98bc465d --- /dev/null +++ b/tpu_inference/kernels/experimental/fused_moe/__init__.py @@ -0,0 +1,27 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# 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 +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Full fused EP MoE kernel. + +Exposes ``fused_moe_func_rs``: EP MoE with ICI reduce-scatter fused into a +single ``gmm_fused_rs`` Pallas kernel (gather -> GMM1 -> act -> GMM2 -> RS). +""" + +from .fused_moe_rs import (expert_parallel_gmm_rs, fused_moe_func_rs, + moe_gmm_local_rs_nodedup) + +__all__ = [ + "fused_moe_func_rs", + "expert_parallel_gmm_rs", + "moe_gmm_local_rs_nodedup", +] diff --git a/tpu_inference/kernels/experimental/fused_moe/fused_moe_rs.py b/tpu_inference/kernels/experimental/fused_moe/fused_moe_rs.py new file mode 100644 index 0000000000..824ffd3a3b --- /dev/null +++ b/tpu_inference/kernels/experimental/fused_moe/fused_moe_rs.py @@ -0,0 +1,528 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# 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 +# +# http://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. +"""EP MoE with ICI reduce-scatter fused into the ``gmm_fused_rs`` kernel. + +A single Pallas call performs: gather -> GMM1 -> activation -> GMM2 -> ICI +reduce-scatter. Only the nodedup path is provided. +""" + +import functools +import os + +import jax +from jax import numpy as jnp +from jax.sharding import Mesh, NamedSharding +from jax.sharding import PartitionSpec as P + +from tpu_inference.layers.common.sharding import ShardingAxisName +from tpu_inference.utils import get_mesh_shape_product + +from .gmm_fused_rs_nodedup import _select_fused_rs_block_sizes +from .gmm_fused_rs_nodedup import gmm_v2_fused_rs as gmm_v2_fused_rs_nodedup +from .gmm_v2_gather_scatter import _recover_quant_block_size + +EXPERT = ShardingAxisName.EXPERT +MLP_DATA = ShardingAxisName.MLP_DATA + + +def _flatten_partition_axes(*axis_specs): + axes = () + for spec in axis_specs: + if spec is None: + continue + if isinstance(spec, tuple): + axes += _flatten_partition_axes(*spec) + else: + axes += (spec, ) + return axes[0] if len(axes) == 1 else axes + + +def combine_partition_axes(*axis_specs): + axes = _flatten_partition_axes(*axis_specs) + if axes is None or isinstance(axes, str): + return axes + combined = () + for a in axes: + if a not in combined: + combined += (a, ) + return combined[0] if len(combined) == 1 else combined + + +def get_moe_expert_axis(mesh, default_axis=EXPERT): + """Runtime FFN expert axis spanning the physical EP mesh axes.""" + axes = tuple(a for a in ("attn_dp", "attn_dp_expert", "expert", "model") + if a in mesh.shape) + return _flatten_partition_axes(*axes) if axes else default_axis + + +def apply_scoring_fn(scoring_fn: str, x: jax.Array) -> jax.Array: + if scoring_fn == "softmax": + return jax.nn.softmax(x, axis=-1) + if scoring_fn == "sigmoid": + return jax.nn.sigmoid(x) + raise NotImplementedError(f"unsupported scoring function: {scoring_fn}") + + +def enabled_tpu_sp() -> bool: + """Sequence-parallel MoE. Defaults ON; set ``TPU_MOE_ENABLE_SP=0`` for OFF.""" + return os.environ.get("TPU_MOE_ENABLE_SP", "1") != "0" + + +def _routing_and_topk( + gating_output, + scoring_fn, + topk, + renormalize, + dtype, + mesh, + *, + expert_bias: jax.Array | None = None, + route_scale: float = 1.0, + router_output_multiplier: float | None = None, + router_score_division_eps: float | None = None, +): + if router_output_multiplier is not None: + gating_output = gating_output * router_output_multiplier + + scores = apply_scoring_fn(scoring_fn, gating_output) + scores = jax.lax.with_sharding_constraint( + scores, NamedSharding(mesh, P(MLP_DATA, None))) + + if expert_bias is not None: + # Bias shifts selection but not the final weights. + scores_for_topk = scores + expert_bias + _, topk_indices = jax.lax.top_k(scores_for_topk, k=topk) + topk_weights = jnp.take_along_axis(scores, topk_indices, axis=1) + else: + topk_weights, topk_indices = jax.lax.top_k(scores, k=topk) + + if renormalize: + denom = topk_weights.sum(axis=-1, keepdims=True) + if router_score_division_eps is not None: + denom = denom + router_score_division_eps + topk_weights = topk_weights / denom + + topk_weights = (topk_weights * route_scale).astype(dtype) + return topk_weights, topk_indices + + +# Per-row scalar-prefetch arrays scale as 4 * size_m bytes each and must fit the +# kernel's SMEM budget; guard at the empirical overflow threshold. +_FUSED_RS_MAX_SAFE_SIZE_M = 262144 + + +def _assert_fused_rs_smem_safe(size_m: int) -> None: + assert size_m < _FUSED_RS_MAX_SAFE_SIZE_M, ( + f"gmm_v2_fused_rs SMEM OOM at size_m={size_m} " + f"(>= {_FUSED_RS_MAX_SAFE_SIZE_M}). Reduce tokens * top_k or shard more." + ) + + +_FP8_OUTPUT_COMM_ENV_VALUES = frozenset( + ("fp8", "float8_e4m3fn", "jnp.float8_e4m3fn")) + + +def _enable_fp8_output_comm_from_env() -> bool: + value = os.environ.get("TPU_MOE_FP8_OUTPUT_COMM", "") + return value.lower() in _FP8_OUTPUT_COMM_ENV_VALUES + + +def _all_gather_token_hidden( + token_hidden: jax.Array, + *, + axis_name, + fp8_enabled: bool, +) -> jax.Array: + """All-gather token shards, optionally quantizing the payload to FP8.""" + if not fp8_enabled: + return jax.lax.all_gather(token_hidden, + axis_name=axis_name, + axis=0, + tiled=True) + + with jax.named_scope("moe_fp8_post_gather"): + out_dtype = token_hidden.dtype + token_hidden_f32 = token_hidden.astype(jnp.float32) + fp8_max = jnp.array(jnp.finfo(jnp.float8_e4m3fn).max, + dtype=jnp.float32) + absmax = jnp.max(jnp.abs(token_hidden_f32), axis=-1, keepdims=True) + scale = jnp.maximum(absmax, jnp.array(1e-6, + dtype=jnp.float32)) / fp8_max + token_hidden_fp8 = jnp.clip(token_hidden_f32 / scale, -fp8_max, + fp8_max).astype(jnp.float8_e4m3fn) + gathered_fp8 = jax.lax.all_gather(token_hidden_fp8, + axis_name=axis_name, + axis=0, + tiled=True) + gathered_scale = jax.lax.all_gather(scale, + axis_name=axis_name, + axis=0, + tiled=True) + return (gathered_fp8.astype(jnp.float32) * + gathered_scale).astype(out_dtype) + + +def _compute_rs_routing(topk_indices, *, num_experts, topk): + """Inline routing: lhs_indices, group_sizes, output_indices, topk_slot_indices. + + Uses integer arithmetic and one-hot+sum (not gathers / bincount) to keep the + computation cheap and fusable. dtype stays int32 for scalar-prefetch refs. + """ + topk_indices_flat = topk_indices.flatten() + topk_argsort_indices = jnp.argsort(topk_indices_flat) + expert_ids = jnp.arange(num_experts, dtype=jnp.int32) + group_sizes = jnp.sum( + (topk_indices_flat[:, None] == expert_ids[None, :]).astype(jnp.int32), + axis=0, + ) + lhs_indices = topk_argsort_indices // topk + topk_slot_indices = topk_argsort_indices % topk + output_indices = lhs_indices + return lhs_indices, group_sizes, output_indices, topk_slot_indices + + +def moe_gmm_local_rs_nodedup( + hidden_states_local: jax.Array, + w1: jax.Array, + w1_scale: jax.Array | None, + w1_bias: jax.Array | None, + w2: jax.Array, + w2_scale: jax.Array | None, + w2_bias: jax.Array | None, + w1_global_scale: jax.Array | None, + w2_global_scale: jax.Array | None, + group_offset: jax.Array, + topk_weights: jax.Array, + topk_indices: jax.Array, + post_expert_norm_weight_input: jax.Array | None = None, + *, + activation: str, + topk: int, + ep_size: int, + ep_axis_name=EXPERT, + has_post_norm: bool = False, + sp_enabled: bool = True, + fp8_post_gather: bool = False, +) -> jax.Array: + """Per-chip MoE body: ICI direct-write per row, then weighted top_k reduce.""" + num_tokens = hidden_states_local.shape[0] + hidden_size = w2.shape[-1] + chunk_size = num_tokens // ep_size + num_experts = w1.shape[0] * ep_size # global num_experts + num_local_experts = w1.shape[0] + my_id = jax.lax.axis_index(ep_axis_name) + + # Routing inlined here so it fuses with the kernel pipeline. + lhs_indices, group_sizes, output_indices, topk_slot_indices = _compute_rs_routing( + topk_indices, num_experts=num_experts, topk=topk) + size_m = lhs_indices.shape[0] + _assert_fused_rs_smem_safe(size_m) + + # Local row range [local_start, local_end) for this chip's experts. + go_val = group_offset[0] + expert_idx = jnp.arange(num_experts, dtype=jnp.int32) + local_start = jnp.sum(jnp.where(expert_idx < go_val, group_sizes, 0)) + local_end = local_start + jnp.sum( + jnp.where( + jnp.logical_and(expert_idx >= go_val, expert_idx + < go_val + num_local_experts), + group_sizes, + 0, + )) + + # Rows from other chips destined for me (dest == my_id and not local). + send_dest_chips = output_indices // chunk_size + rows = jnp.arange(size_m, dtype=jnp.int32) + row_is_mine = jnp.logical_and(rows >= local_start, rows < local_end) + to_me_remote = jnp.logical_and(send_dest_chips == my_id, + jnp.logical_not(row_is_mine)) + my_recv_count = jnp.sum(jnp.where(to_me_remote, 1, 0)) + total_recv_count = jnp.array([my_recv_count], dtype=jnp.int32) + + # tile_m here MUST match the value the kernel selects internally, otherwise + # max_num_gm under-counts and the kernel's final gather DMA is left unawaited. + block_sizes = _select_fused_rs_block_sizes( + size_m=size_m, + size_k1=w1.shape[1], + size_n1=w1.shape[2], + size_k2=w2.shape[1], + size_n2=w2.shape[2], + size_group=num_local_experts, + size_lhs_group=group_sizes.shape[0], + ep_size=ep_size, + out_dtype=hidden_states_local.dtype, + w1_dtype=w1.dtype, + w2_dtype=w2.dtype, + is_quantized=w1_scale is not None, + quant_block_size=(_recover_quant_block_size( + w1.shape[1], w1_scale.shape[1]) if w1_scale is not None else None), + act_fn=activation, + fp8_direct_write=fp8_post_gather, + ) + tile_m = block_sizes.tile_m + max_num_gm = jnp.array(num_experts + (size_m + tile_m - 1) // tile_m - 1, + dtype=jnp.int32) + + out_buf = gmm_v2_fused_rs_nodedup( + hidden_states_local, + w1, + w2, + group_sizes, + lhs_indices, + output_indices, + w1_scale=w1_scale, + w2_scale=w2_scale, + w1_global_scale=w1_global_scale, + w2_global_scale=w2_global_scale, + w1_bias=w1_bias, + w2_bias=w2_bias, + act_fn=activation, + output_size=num_tokens, + group_offset=group_offset, + topk_indices=topk_slot_indices, + ep_size=ep_size, + ep_axis_name=ep_axis_name, + max_num_gm=max_num_gm, + total_recv_count=total_recv_count, + top_k=topk, + fp8_direct_write=fp8_post_gather, + ) + + # SP off: topk_weights is replicated; slice this chip's shard to match out_3d. + local_topk_weights = ( + topk_weights if sp_enabled else jax.lax.dynamic_slice_in_dim( + topk_weights, my_id * chunk_size, chunk_size, axis=0)) + out_3d = out_buf.reshape(chunk_size, topk, hidden_size) + post_expert_norm_weight = post_expert_norm_weight_input if has_post_norm else None + if post_expert_norm_weight is not None: + norm_size = post_expert_norm_weight.shape[0] # unpadded hidden_size + pnw_raw = post_expert_norm_weight.astype(jnp.float32) + 1.0 + if hidden_size > norm_size: + # Zero padded columns so they don't affect variance. + col_idx = jnp.arange(hidden_size, dtype=jnp.int32) + valid_mask = (col_idx < norm_size)[None, None, :] + out_f32 = out_3d.astype(jnp.float32) * valid_mask + pnw = jnp.concatenate([ + pnw_raw, + jnp.zeros(hidden_size - norm_size, dtype=jnp.float32) + ]) + else: + out_f32 = out_3d.astype(jnp.float32) + pnw = pnw_raw + out_f32 = out_f32 * pnw[None, None, :] + var = jnp.sum(out_f32**2, axis=-1, keepdims=True) / norm_size + out_3d = (out_f32 * jax.lax.rsqrt(var + 1e-8)).astype(out_3d.dtype) + + token_hidden = jnp.sum(out_3d * local_topk_weights[:, :, None], axis=1) + + if sp_enabled: + # Kernel reduce-scatter is the SP exit; output stays token-sharded. + return token_hidden + # SP off: gather the per-chip token shard back to the replicated batch. + return _all_gather_token_hidden( + token_hidden, + axis_name=ep_axis_name, + fp8_enabled=fp8_post_gather, + ) + + +def expert_parallel_gmm_rs( + hidden_states: jax.Array, + w1: jax.Array, + w1_scale: jax.Array | None, + w1_bias: jax.Array | None, + w2: jax.Array, + w2_scale: jax.Array | None, + w2_bias: jax.Array | None, + topk_weights: jax.Array, + topk_indices: jax.Array, + *, + activation: str, + topk: int, + mesh: Mesh, + post_expert_norm_weight: jax.Array | None = None, + w1_global_scale: jax.Array | None = None, + w2_global_scale: jax.Array | None = None, + fp8_post_gather: bool = False, +) -> jax.Array: + """shard_map driver: routing runs per-chip so it fuses with the kernel.""" + expert_axis = get_moe_expert_axis(mesh, EXPERT) + ep_size = get_mesh_shape_product(mesh, expert_axis) + ep_p_spec = P(expert_axis) + data_p_spec = P(MLP_DATA) + # SP off: hidden replicated in/out with an explicit all-gather in the body. + sp_enabled = enabled_tpu_sp() + hidden_in_spec = data_p_spec if sp_enabled else P() + moe_out_spec = (P(combine_partition_axes(MLP_DATA, expert_axis)) + if sp_enabled else P()) + fp8_post_gather = ((not sp_enabled) and w1_scale is not None + and w2_scale is not None and + (fp8_post_gather or _enable_fp8_output_comm_from_env())) + # SP off: replicate routing tensors too; the body slices its local chunk. + topk_w_spec = (P(combine_partition_axes(MLP_DATA, expert_axis), None) + if sp_enabled else P()) + topk_i_spec = P(MLP_DATA, None) if sp_enabled else P() + num_experts = w1.shape[0] + num_experts_per_shard = num_experts // ep_size + group_offset = jnp.arange(0, num_experts, num_experts_per_shard) + + w1_scale_spec = None if w1_scale is None else ep_p_spec + w1_bias_spec = None if w1_bias is None else ep_p_spec + w2_scale_spec = None if w2_scale is None else ep_p_spec + w2_bias_spec = None if w2_bias is None else ep_p_spec + w1_gs_spec = None if w1_global_scale is None else ep_p_spec + w2_gs_spec = None if w2_global_scale is None else ep_p_spec + + _has_pn_rs = post_expert_norm_weight is not None + _pnw_rs = (post_expert_norm_weight if post_expert_norm_weight is not None + else jnp.zeros((1, ), jnp.bfloat16)) + result = jax.shard_map( + functools.partial( + moe_gmm_local_rs_nodedup, + activation=activation, + topk=topk, + ep_size=ep_size, + ep_axis_name=expert_axis, + has_post_norm=_has_pn_rs, + sp_enabled=sp_enabled, + fp8_post_gather=fp8_post_gather, + ), + mesh=mesh, + in_specs=( + hidden_in_spec, + ep_p_spec, # w1 + w1_scale_spec, + w1_bias_spec, + ep_p_spec, # w2 + w2_scale_spec, + w2_bias_spec, + w1_gs_spec, + w2_gs_spec, + ep_p_spec, # group_offset + topk_w_spec, + topk_i_spec, + P(), # post_expert_norm_weight + ), + out_specs=moe_out_spec, + check_vma=False, + )( + hidden_states, + w1, + w1_scale, + w1_bias, + w2, + w2_scale, + w2_bias, + w1_global_scale, + w2_global_scale, + group_offset, + topk_weights, + topk_indices, + _pnw_rs, + ) + + return result + + +@functools.partial( + jax.jit, + static_argnames=( + "topk", + "renormalize", + "mesh", + "activation", + "scoring_fn", + "fp8_post_gather", + ), +) +def fused_moe_func_rs( + hidden_states: jax.Array, + w1: jax.Array, + w2: jax.Array, + w1_scale: jax.Array | None, + w2_scale: jax.Array | None, + w1_bias: jax.Array | None, + w2_bias: jax.Array | None, + gating_output: jax.Array | None, + topk: int, + renormalize: bool, + mesh: Mesh, + activation: str, + scoring_fn: str, + post_expert_norm_weight: jax.Array | None = None, + topk_weights: jax.Array | None = None, + topk_indices: jax.Array | None = None, + fp8_post_gather: bool = False, +) -> jax.Array: + """EP MoE with ICI reduce-scatter fused in kernel (gmm_fused_rs). + + Uses caller-supplied ``topk_weights``/``topk_indices`` when both are given; + otherwise computes top-k from ``gating_output``. Then runs the fused kernel + (gather -> GMM1 -> act -> GMM2 -> ICI reduce-scatter) and reduces over top_k. + """ + num_tokens, hidden_size = hidden_states.shape + global_num_experts, padded_hidden_size, _ = w1.shape + dtype = hidden_states.dtype + + assert (num_tokens * topk) % 16 == 0 + if topk_weights is not None and topk_indices is not None: + # Honor pre-computed routing; do not recompute from gating_output. + topk_weights = jax.lax.with_sharding_constraint( + topk_weights, NamedSharding(mesh, P(MLP_DATA, None))) + topk_indices = jax.lax.with_sharding_constraint( + topk_indices, NamedSharding(mesh, P(MLP_DATA, None))) + else: + assert gating_output is not None, ( + "fused_moe_func_rs: either pre-computed topk_weights+topk_indices " + "or gating_output must be provided.") + assert gating_output.shape == (num_tokens, global_num_experts) + topk_weights, topk_indices = _routing_and_topk(gating_output, + scoring_fn, topk, + renormalize, dtype, + mesh) + + # Pad hidden_states to w1's K dimension if needed. + if padded_hidden_size != hidden_size: + hidden_states = jnp.pad(hidden_states, + ((0, 0), + (0, padded_hidden_size - hidden_size))) + + result = expert_parallel_gmm_rs( + hidden_states, + w1, + w1_scale, + w1_bias, + w2, + w2_scale, + w2_bias, + topk_weights, + topk_indices, + activation=activation, + topk=topk, + mesh=mesh, + post_expert_norm_weight=post_expert_norm_weight, + fp8_post_gather=fp8_post_gather, + ) + + return result[:num_tokens, :hidden_size] + + +__all__ = [ + "fused_moe_func_rs", + "expert_parallel_gmm_rs", + "moe_gmm_local_rs_nodedup", + "_compute_rs_routing", + "_FUSED_RS_MAX_SAFE_SIZE_M", + "_assert_fused_rs_smem_safe", +] diff --git a/tpu_inference/kernels/experimental/fused_moe/gmm_fused/__init__.py b/tpu_inference/kernels/experimental/fused_moe/gmm_fused/__init__.py new file mode 100644 index 0000000000..2e37030330 --- /dev/null +++ b/tpu_inference/kernels/experimental/fused_moe/gmm_fused/__init__.py @@ -0,0 +1,28 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# 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 +# +# http://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. +"""Per-round-push AG+GMM1 EP MoE kernel (example). + +Provides a push-based all-gather fused with GMM1 + activation +(``gmm_v2_ag_gmm1``), driven by a precomputed per-round send schedule +(``per_round_schedule``), plus the paired GMM2 + ICI reduce-scatter kernels. +""" + +from .gmm_v2_ag_rs import (gmm_v2_ag_gmm1, gmm_v2_scatter_ici_dedup, + gmm_v2_scatter_ici_nodedup) + +__all__ = [ + "gmm_v2_ag_gmm1", + "gmm_v2_scatter_ici_dedup", + "gmm_v2_scatter_ici_nodedup", +] diff --git a/tpu_inference/kernels/experimental/fused_moe/gmm_fused/gmm_v2_ag_rs.py b/tpu_inference/kernels/experimental/fused_moe/gmm_fused/gmm_v2_ag_rs.py new file mode 100644 index 0000000000..cbcfa0fbf7 --- /dev/null +++ b/tpu_inference/kernels/experimental/fused_moe/gmm_fused/gmm_v2_ag_rs.py @@ -0,0 +1,2177 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# 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 +# +# http://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. +"""Dedup-aware push-based AG fused with GMM1 + activation (EP MoE). + +Each (sender s, target c, global_token_id g) tuple crosses the ICI link at +most once. Each receiver maintains a persistent HBM cache of shape +(num_tokens_total, K1/NL, NL) indexed directly by `global_token_id`. The +first time chip B needs token T originating from chip A, A pushes T into +B's cache slot T. Subsequent rounds that need T on B read straight from +the cache (zero ICI). + +The schedule is precomputed in JAX (compute_dedup_send_schedule). The +kernel itself just executes: + + Bootstrap: push round 0 NEW tokens to all targets' cache slots. + Per round r in [0, max_num_gm): + - Prefetch: push round (r+1) NEW tokens (slot (r+1) % 2 send_sem). + - Drain send_sem[(r+1) % 2, *] from its prior use, two rounds back. + - (r < my_num_gm) Wait recv_sem[r % 2] for recv_count_new[r] bytes. + - (r < my_num_gm) Per-row indexed gather from cache_ref to VMEM via + dma_gather_gm_start (using lhs_indices; identical to the original + inner_kernel_gather pattern). + - (r < my_num_gm) Wait gather, then run inner_kernel matmul. + +The receiver's gather works whether the token landed in this round or in +an earlier round — the cache is the canonical source. +""" + +import dataclasses +import functools + +import jax +import jax.experimental.pallas as pl +from jax import lax +from jax import numpy as jnp +from jax.experimental.pallas import tpu as pltpu + +# isort: off +# yapf: disable +from ..gmm_v2_gather_scatter import ( + FusedWeightsRef, MetadataRef, WeightsRef, align_to, apply_act_fn, + calculate_tiling, dma_gather_gm_start, dma_gather_gm_wait, + dma_scatter_gm_start, dma_scatter_gm_wait, fill_metadata, + generate_block_specs, get_cost_estimate, get_scope_name, inner_kernel, + make_gmm_configs, zero_out_end, zero_out_end_3d, zero_out_start, + zero_out_start_3d) +# yapf: enable +# isort: on +from .per_round_schedule import compute_dedup_send_schedule + +# ============================================================================= +# Push helpers (NEW tokens only, target = cache_ref[global_id]) +# ============================================================================= + + +def _push_round_to_target( + *, + hidden_states_ref, # 3D HBM: (chunk_size, padded_K1/NL, NL) + target_cache_ref, # 3D HBM: (num_tokens_total, padded_K1/NL, NL) + send_count_for_target, # int32 scalar (SMEM) + send_local_off_for_target, # int32[tile_m] (SMEM) + send_global_id_for_target, # int32[tile_m] (SMEM) — destination cache slot + target_id, # int32 scalar + is_local_target, # bool scalar + send_sem, + recv_sem, + ep_axis_name: str, +): + """Push send_count_for_target NEW tokens to the given target chip. + + Source row = local hidden_states[send_local_off_for_target[i]]. + Destination = target's cache_ref[send_global_id_for_target[i]]. + + Local self-target uses make_async_copy with recv_sem so the receiver- + side wait drains uniformly across local + remote contributions. + Remote uses make_async_remote_copy with both send_sem and recv_sem. + """ + + def _send_row(i, _): + local_off = send_local_off_for_target[i] + cache_slot = send_global_id_for_target[i] + + @pl.when(is_local_target) + def _local(): + pltpu.make_async_copy( + src_ref=hidden_states_ref.at[pl.ds(local_off, 1), :, :], + dst_ref=target_cache_ref.at[pl.ds(cache_slot, 1), :, :], + sem=recv_sem, + ).start() + + @pl.when(jnp.logical_not(is_local_target)) + def _remote(): + pltpu.make_async_remote_copy( + src_ref=hidden_states_ref.at[pl.ds(local_off, 1), :, :], + dst_ref=target_cache_ref.at[pl.ds(cache_slot, 1), :, :], + send_sem=send_sem, + recv_sem=recv_sem, + device_id={ + ep_axis_name: target_id + }, + device_id_type=pltpu.DeviceIdType.MESH, + ).start() + + return _ + + lax.fori_loop(0, send_count_for_target, _send_row, 0) + + +def _push_round( + *, + round_id, + slot, # send_sem slot index (round % 2) + my_id, + hidden_states_ref, + cache_ref, + send_count_ref, + send_local_off_ref, + send_global_id_ref, + send_sem_ref, + recv_sem_ref, + ep_size: int, + ep_axis_name: str, +): + """Issue all outgoing DMAs for one round, to all targets.""" + for c in range(ep_size): + target_id = jnp.int32(c) + is_local_target = my_id == target_id + + send_count_for_target = send_count_ref[c, round_id] + send_local_off_for_target = send_local_off_ref.at[c, round_id, :] + send_global_id_for_target = send_global_id_ref.at[c, round_id, :] + + _push_round_to_target( + hidden_states_ref=hidden_states_ref, + target_cache_ref=cache_ref, + send_count_for_target=send_count_for_target, + send_local_off_for_target=send_local_off_for_target, + send_global_id_for_target=send_global_id_for_target, + target_id=target_id, + is_local_target=is_local_target, + send_sem=send_sem_ref.at[slot, c], + recv_sem=recv_sem_ref.at[slot], + ep_axis_name=ep_axis_name, + ) + + +def _drain_send_sem_for_round( + *, + slot, + round_id, + my_id, + cache_ref, # any HBM buffer used as a byte-count proxy + send_sem_ref, + send_count_ref, + ep_size: int, +): + """Drain the per-(slot, target) send_sem for `round_id` so the (slot, + target) pair can be reused two rounds later. + + The drain byte-count equals `send_count[c, round_id] * row_bytes`. + Skips: + * (slot, my_id) — local self-target writes use recv_sem. + * (slot, c) with send_count[c, round_id] == 0 — no DMA was issued. + """ + for c in range(ep_size): + target_id = jnp.int32(c) + is_local_target = my_id == target_id + n = send_count_ref[c, round_id] + + @pl.when(jnp.logical_and(jnp.logical_not(is_local_target), n > 0)) + def _drain_one(): + # Use cache_ref[0, :n] as a byte-count proxy. The drain only + # cares about TOTAL bytes accumulated, not which physical + # bytes — same trick as `dma_gather_gm_wait`. + pltpu.make_async_copy( + src_ref=cache_ref.at[pl.ds(0, n), :, :], + dst_ref=cache_ref.at[pl.ds(0, n), :, :], + sem=send_sem_ref.at[slot, c], + ).wait() + + +def _drain_recv_sem_for_round( + *, + slot, + round_id, + cache_ref, + recv_sem_ref, + recv_count_new_ref, +): + """Drain recv_sem[slot] by recv_count_new[round_id] * row_bytes. + + The new arrivals land at scattered cache slots indexed by global_id, so + we cannot bound them with a contiguous slice of staging. We drain by + total byte count using any contiguous region of cache_ref of the right + size as a byte-count proxy. + """ + n = recv_count_new_ref[round_id] + + @pl.when(n > 0) + def _drain(): + pltpu.make_async_copy( + src_ref=cache_ref.at[pl.ds(0, n), :, :], + dst_ref=cache_ref.at[pl.ds(0, n), :, :], + sem=recv_sem_ref.at[slot], + ).wait() + + +# ============================================================================= +# Kernel: dedup push + cache + per-row gather + GMM1 + Activation +# ============================================================================= + + +def kernel_main_ag_gmm1( + # Scalar prefetch (8) + lhs_group_sizes_ref, + group_offset_ref, + lhs_indices_ref, + send_count_ref, # SMEM (ep_size, max_num_gm) + send_local_off_ref, # SMEM (ep_size, max_num_gm, tile_m) + send_global_id_ref, # SMEM (ep_size, max_num_gm, tile_m) — cache slot + recv_count_new_ref, # SMEM (max_num_gm,) — new arrivals per round + my_num_gm_ref, # SMEM (1,) + # In (2) + hidden_states_ref, # 3D HBM: local shard (chunk_size, padded_K1/NL, NL) + rhs_ref, # HBM: W1 weights via WeightsRef + # Out (2) + out_ref, # HBM: (size_m, out_N) + cache_ref, # HBM: (num_tokens_total, padded_K1/NL, NL) + # Scratch + partial_out_ref, + acc_ref, + metadata_ref, + zero_ref, + semaphore_ref, + gather_sem_ref, # DMA(2,) — per-row gather from cache to VMEM + gathered_lhs_2x_ref, # VMEM: (2, tile_m, padded_K1/NL, NL) + recv_sem_ref, # DMA(2,) — incoming push completion (per slot) + send_sem_ref, # DMA(2, ep_size) — outgoing push completion + *, + cfgs, + ep_size: int, + ep_axis_name: str, + max_num_gm: int, +): + """Dedup push-based AG fused with GMM1 + activation.""" + my_id = lax.axis_index(ep_axis_name) + num_lanes = pltpu.get_tpu_info().num_lanes + tile_k = cfgs.tiles.tile_k + num_k = pl.cdiv(cfgs.dims.size_k, tile_k) + num_n = pl.cdiv(cfgs.out_size_n, cfgs.tiles.tile_n) + + if cfgs.rhs_cfgs.quant_dtype is not None: + rhs_weight = rhs_ref.weight + rhs_weight = rhs_weight.bitcast(jnp.uint32) + rhs_ref = dataclasses.replace(rhs_ref, weight=rhs_weight) + + num_gm_real = fill_metadata(lhs_group_sizes_ref, + group_offset_ref, + metadata_ref, + cfgs=cfgs) + + # Pad metadata for indices >= num_gm_real so emit_pipeline's RHS index + # map returns a valid group id on padding rounds. + last_valid_offset = metadata_ref.gm_id_to_m_offset[num_gm_real] + for _i in range(max_num_gm): + i = jnp.int32(_i) + + @pl.when(i >= num_gm_real) + def _pad(): + metadata_ref.gm_id_to_group_id[i] = jnp.int32(0) + metadata_ref.gm_id_to_m_offset[i] = last_valid_offset + metadata_ref.gm_id_to_m_offset[i + 1] = last_valid_offset + + in_specs, out_specs = generate_block_specs(metadata_ref, cfgs) + if cfgs.fuse_act is not None: + rhs_up_ref = jax.tree.map(lambda x: x.at[..., cfgs.out_size_n:], + rhs_ref) + rhs_ref = FusedWeightsRef(gate=rhs_ref, up=rhs_up_ref) + _, rhs_spec_orig = in_specs + in_specs = ( + in_specs[0], + FusedWeightsRef(gate=rhs_spec_orig, up=rhs_spec_orig), + ) + _, rhs_in_spec = in_specs + + if cfgs.zero_init: + zero_size = zero_out_start( + out_ref, + zero_ref, + semaphore_ref, + metadata_ref, + num_gm_real, + dims=cfgs.dims, + ) + + my_num_gm = my_num_gm_ref[0] + + # ---- Bootstrap ---- + # 1) Push round 0 NEW tokens to all targets' caches. + # 2) Wait recv_sem[0] for round 0's incoming arrivals (no-op if 0). + # 3) Pre-start gather for tile 0 into VMEM[0]. This kicks off the + # round-0 per-row indexed gather DMAs concurrently with the inner + # loop's prefetch of round 1. + with jax.named_scope("dedup_push_bootstrap_r0"): + _push_round( + round_id=jnp.int32(0), + slot=jnp.int32(0), + my_id=my_id, + hidden_states_ref=hidden_states_ref, + cache_ref=cache_ref, + send_count_ref=send_count_ref, + send_local_off_ref=send_local_off_ref, + send_global_id_ref=send_global_id_ref, + send_sem_ref=send_sem_ref, + recv_sem_ref=recv_sem_ref, + ep_size=ep_size, + ep_axis_name=ep_axis_name, + ) + + _drain_recv_sem_for_round( + slot=jnp.int32(0), + round_id=jnp.int32(0), + cache_ref=cache_ref, + recv_sem_ref=recv_sem_ref, + recv_count_new_ref=recv_count_new_ref, + ) + + # Pre-start gather for tile 0 (no-op if my_num_gm == 0; metadata + # padding makes m_end == m_start so the row-loop is empty). + @pl.when(my_num_gm > 0) + def _bootstrap_gather(): + dma_gather_gm_start( + cache_ref, + gathered_lhs_2x_ref.at[0], + lhs_indices_ref, + gather_sem_ref.at[0], + jnp.int32(0), + metadata_ref, + ) + + # ---- Per-round inner ---- + # At gm_id=r is_first_kn: + # (a) Prefetch push round r+1 (with send_sem drain from 2 rounds back). + # (b) Wait recv_sem[(r+1)%2] for round r+1 arrivals. + # (c) Start gather for tile r+1 into VMEM[(r+1)%2] — async; runs + # concurrently with this round's matmul. + # (d) Wait gather for current tile r (sem_id) — bootstrap + # pre-started r=0; previous iterations pre-started r>=1. + # The matmul body runs gated on r < my_num_gm. + def inner_per_round( + tiled_rhs_ref, + tiled_out_ref, + partial_out_ref_in, + acc_ref_in, + metadata_ref_in, + ): + gm_id = pl.program_id(0) + n_id = pl.program_id(1) + k_id = pl.program_id(2) + num_gm_grid = pl.num_programs(0) + + sem_id = gm_id % 2 + is_first_kn = jnp.logical_and(k_id == 0, n_id == 0) + + # (a) Prefetch sends for round (gm_id + 1). + @pl.when(jnp.logical_and(is_first_kn, gm_id + 1 < num_gm_grid)) + def _prefetch_next(): + next_round = gm_id + 1 + next_slot = 1 - sem_id + + @pl.when(gm_id + 1 >= 2) + def _drain(): + _drain_send_sem_for_round( + slot=next_slot, + round_id=gm_id - 1, + my_id=my_id, + cache_ref=cache_ref, + send_sem_ref=send_sem_ref, + send_count_ref=send_count_ref, + ep_size=ep_size, + ) + + _push_round( + round_id=next_round, + slot=next_slot, + my_id=my_id, + hidden_states_ref=hidden_states_ref, + cache_ref=cache_ref, + send_count_ref=send_count_ref, + send_local_off_ref=send_local_off_ref, + send_global_id_ref=send_global_id_ref, + send_sem_ref=send_sem_ref, + recv_sem_ref=recv_sem_ref, + ep_size=ep_size, + ep_axis_name=ep_axis_name, + ) + + # (b)+(c) Pre-start gather for round r+1: wait its recv_sem, then + # start the per-row gather into the OTHER VMEM slot. This + # overlaps with current matmul. + @pl.when( + jnp.logical_and( + is_first_kn, + jnp.logical_and(gm_id + 1 < num_gm_grid, gm_id + 1 + < my_num_gm), + )) + def _prefetch_gather_next(): + next_slot = 1 - sem_id + _drain_recv_sem_for_round( + slot=next_slot, + round_id=gm_id + 1, + cache_ref=cache_ref, + recv_sem_ref=recv_sem_ref, + recv_count_new_ref=recv_count_new_ref, + ) + dma_gather_gm_start( + cache_ref, + gathered_lhs_2x_ref.at[next_slot], + lhs_indices_ref, + gather_sem_ref.at[next_slot], + gm_id + 1, + metadata_ref_in, + ) + + # (d) Wait gather for current tile r before consuming VMEM[sem_id]. + @pl.when(jnp.logical_and(is_first_kn, gm_id < my_num_gm)) + def _wait_current_gather(): + dma_gather_gm_wait( + gathered_lhs_2x_ref.at[sem_id], + gather_sem_ref.at[sem_id], + gm_id, + metadata_ref_in, + ) + + # Matmul body — only run for r < my_num_gm. + @pl.when(gm_id < my_num_gm) + def _matmul(): + gathered_lhs_k_slice = gathered_lhs_2x_ref.at[ + sem_id, :, + pl.ds(k_id * (tile_k // num_lanes), tile_k // num_lanes), :] + gathered_lhs_data = gathered_lhs_k_slice[...].reshape( + -1, cfgs.dims.size_lhs_sublane, tile_k) + inner_kernel( + gathered_lhs_data, + tiled_rhs_ref, + tiled_out_ref, + partial_out_ref_in, + acc_ref_in, + metadata_ref_in, + cfgs=cfgs, + ) + + pipeline_fn = pltpu.emit_pipeline( + inner_per_round, + grid=(max_num_gm, num_n, num_k), + in_specs=[rhs_in_spec], + out_specs=out_specs, + ) + out_in = out_ref.reshape(-1, cfgs.dims.size_lhs_sublane, out_ref.shape[-1]) + scratches = [partial_out_ref, acc_ref, metadata_ref] + pipeline_fn(rhs_ref, out_in, scratches=scratches) + + # ---- Epilogue: drain in-flight send sems for the last two rounds ---- + with jax.named_scope("dedup_epilogue_drain"): + if max_num_gm >= 2: + last_round_a = jnp.int32(max_num_gm - 1) + last_round_b = jnp.int32(max_num_gm - 2) + slot_a = (max_num_gm - 1) % 2 + slot_b = (max_num_gm - 2) % 2 + _drain_send_sem_for_round( + slot=jnp.int32(slot_a), + round_id=last_round_a, + my_id=my_id, + cache_ref=cache_ref, + send_sem_ref=send_sem_ref, + send_count_ref=send_count_ref, + ep_size=ep_size, + ) + _drain_send_sem_for_round( + slot=jnp.int32(slot_b), + round_id=last_round_b, + my_id=my_id, + cache_ref=cache_ref, + send_sem_ref=send_sem_ref, + send_count_ref=send_count_ref, + ep_size=ep_size, + ) + elif max_num_gm == 1: + _drain_send_sem_for_round( + slot=jnp.int32(0), + round_id=jnp.int32(0), + my_id=my_id, + cache_ref=cache_ref, + send_sem_ref=send_sem_ref, + send_count_ref=send_count_ref, + ep_size=ep_size, + ) + + if cfgs.zero_init: + zero_out_end(out_ref, semaphore_ref, zero_size, dims=cfgs.dims) + + +@jax.jit(static_argnames=[ + "fuse_act", + "ep_size", + "ep_axis_name", +], ) +def gmm_v2_ag_gmm1( + hidden_states_shard: jax.Array, # (chunk_size, size_k1) — local shard + w1: jax.Array, # (size_group, size_k1, size_n1) + group_sizes: jax.Array, # int32[size_lhs_group] + lhs_indices: jax.Array, # int32[size_m] + *, + w1_scale: jax.Array | None = None, + w1_bias: jax.Array | None = None, + group_offset: jax.Array | None = None, + fuse_act: str | None = "silu", + ep_size: int, + ep_axis_name: str, +) -> jax.Array: + """Dedup push-based AG fused with GMM1 + activation.""" + chunk_size = hidden_states_shard.shape[0] + size_k1 = hidden_states_shard.shape[1] + num_tokens_total = chunk_size * ep_size + + if group_offset is None: + group_offset = jnp.array([0], dtype=jnp.int32) + + num_lanes = pltpu.get_tpu_info().num_lanes + sls = pltpu.get_tpu_info().get_sublane_tiling(hidden_states_shard.dtype) + + tile_k_unit = num_lanes * sls + padded_k1 = align_to(size_k1, tile_k_unit) + if padded_k1 != size_k1: + k_pad = padded_k1 - size_k1 + hidden_states_shard = jnp.pad(hidden_states_shard, + ((0, 0), (0, k_pad))) + + size_m = lhs_indices.shape[0] + lhs_for_config = jax.ShapeDtypeStruct((size_m, w1.shape[1]), + hidden_states_shard.dtype) + + _fallback_fuse_act = None + if fuse_act is not None: + _nl = pltpu.get_tpu_info().num_lanes + if w1.shape[2] % (2 * _nl) != 0: + _fallback_fuse_act = fuse_act + fuse_act = None + + vmem_limit_bytes = int(pltpu.get_tpu_info().vmem_capacity_bytes) + + cfgs = make_gmm_configs( + lhs_for_config, + w1, + w1_scale, + w1_bias, + group_sizes, + group_offset, + tile_info=calculate_tiling, + vmem_limit_bytes=vmem_limit_bytes, + out_dtype=None, + acc_dtype=None, + maybe_quantize_lhs=True, + zero_initialize=True, + lhs_indices=lhs_indices, + original_k=size_k1 if padded_k1 != size_k1 else None, + fuse_act=fuse_act, + ) + dims = cfgs.dims + tiles = cfgs.tiles + + if padded_k1 != size_k1: + if tiles.tile_k % tile_k_unit != 0: + aligned_tile_k = (tiles.tile_k // tile_k_unit) * tile_k_unit + if aligned_tile_k == 0: + aligned_tile_k = tile_k_unit + tiles = dataclasses.replace(tiles, tile_k=aligned_tile_k) + cfgs = dataclasses.replace(cfgs, tiles=tiles) + + _out_sls = pltpu.get_tpu_info().get_sublane_tiling(cfgs.out_dtype) + out_n = (align_to(dims.size_n, num_lanes * _out_sls) if cfgs.fuse_act + is None else align_to(cfgs.out_size_n, num_lanes * _out_sls)) + + rhs_scale_spec = rhs_bias_spec = None + if w1_scale is not None: + w1_scale = w1_scale.astype(jnp.float32) + rhs_scale_spec = pl.BlockSpec(memory_space=pltpu.HBM) + if w1_bias is not None: + w1_bias = w1_bias.astype(jnp.float32) + rhs_bias_spec = pl.BlockSpec(memory_space=pltpu.HBM) + rhs_weights = WeightsRef(weight=w1, scale=w1_scale, bias=w1_bias) + rhs_in_spec = WeightsRef( + weight=pl.BlockSpec(memory_space=pltpu.HBM), + scale=rhs_scale_spec, + bias=rhs_bias_spec, + ) + + hidden_3d = hidden_states_shard.reshape(chunk_size, padded_k1 // num_lanes, + num_lanes) + + # Dedup send schedule (each (sender, target, token) sent at most once). + ( + send_count_new, + send_local_off_new, + send_global_id_new, + recv_count_new, + my_num_gm, + max_num_gm_static, + ) = compute_dedup_send_schedule( + lhs_indices, + group_sizes, + group_offset, + ep_axis_name=ep_axis_name, + ep_size=ep_size, + chunk_size=chunk_size, + tile_m=tiles.tile_m, + size_group=dims.size_group, + size_lhs_sublane=dims.size_lhs_sublane, + ) + + max_num_gm = max_num_gm_static + num_n = pl.cdiv(cfgs.out_size_n, tiles.tile_n) + partial_out_n = num_n * tiles.tile_n + acc_cols = 2 * tiles.tile_n if cfgs.fuse_act is not None else tiles.tile_n + + target_zero_ref_bytes = 2 * 1024 * 1024 + out_bytes = jnp.dtype(cfgs.out_dtype).itemsize + tile_zero_m = target_zero_ref_bytes // num_lanes // out_bytes + tile_zero_m = min(tile_zero_m, dims.size_m) + + scratch_shapes = [ + pltpu.VMEM((dims.size_lhs_sublane, partial_out_n), cfgs.out_dtype), + pltpu.VMEM((tiles.tile_m, acc_cols), cfgs.acc_dtype), + MetadataRef( + gm_id_to_group_id=pltpu.SMEM((max_num_gm, ), jnp.int32), + gm_id_to_m_offset=pltpu.SMEM((max_num_gm + 1, ), jnp.int32), + ), + pltpu.VMEM((tile_zero_m, num_lanes), cfgs.out_dtype), + pltpu.SemaphoreType.DMA((1, )), # zero semaphore + pltpu.SemaphoreType.DMA((2, )), # gather_sem + pltpu.VMEM( + (2, tiles.tile_m, padded_k1 // num_lanes, num_lanes), + cfgs.out_dtype, + ), # gathered_lhs_2x + pltpu.SemaphoreType.DMA((2, )), # recv_sem (per slot) + pltpu.SemaphoreType.DMA((2, ep_size)), # send_sem (slot x target) + ] + + out_init = [ + jax.ShapeDtypeStruct((dims.size_m, out_n), cfgs.out_dtype), + # Persistent HBM cache: indexed by global_token_id. + jax.ShapeDtypeStruct( + (num_tokens_total, padded_k1 // num_lanes, num_lanes), + hidden_states_shard.dtype, + ), + ] + + compiler_params = pltpu.CompilerParams( + vmem_limit_bytes=vmem_limit_bytes, + disable_bounds_checks=True, + ) + + result = pl.pallas_call( + functools.partial( + kernel_main_ag_gmm1, + cfgs=cfgs, + ep_size=ep_size, + ep_axis_name=ep_axis_name, + max_num_gm=max_num_gm, + ), + out_shape=out_init, + grid_spec=pltpu.PrefetchScalarGridSpec( + num_scalar_prefetch=8, + in_specs=[ + pl.BlockSpec(memory_space=pltpu.HBM), + rhs_in_spec, + ], + out_specs=[ + pl.BlockSpec(memory_space=pltpu.HBM), + pl.BlockSpec(memory_space=pltpu.HBM), # cache_ref + ], + scratch_shapes=scratch_shapes, + ), + compiler_params=compiler_params, + name=get_scope_name(dims, tiles) + "-ag_gmm1_dedup", + cost_estimate=get_cost_estimate(cfgs), + )( + group_sizes, + group_offset, + lhs_indices, + send_count_new, + send_local_off_new, + send_global_id_new, + recv_count_new, + my_num_gm, + hidden_3d, + rhs_weights, + )[0] + + result = result[:, :cfgs.out_size_n] + + if _fallback_fuse_act is not None: + from ..gmm_v2_gather_scatter import apply_act_fn + + result = apply_act_fn(result, _fallback_fuse_act) + + return result + + +# ============================================================================= +# GMM with ICI direct-write scatter (single GMM2 kernel) +# Manual fori_loop approach — drain at top of loop for overlap. +# ============================================================================= + + +def kernel_main_scatter_ici_dedup( + # Scalar prefetch (6) + lhs_group_sizes_ref, + group_offset_ref, + output_indices_ref, + topk_weights_ref, # (M,) float32 — per-row weights for weighted accumulation + is_last_ref, # (M,) int32 — 1 iff last occurrence of token on this chip + total_recv_count_ref, # (1,) int32 — total remote rows this chip receives + # In (2) + lhs_ref, # HBM: (size_m, K//NL, NL) — 3D for DMA + rhs_ref, # HBM: WeightsRef (size_group, size_k, size_n) + # Out (2) + out_buf_ref, # 3D HBM: (chunk_size * ep_size, N//NL, NL) — ICI target + accumulator_ref, # 3D HBM: (num_tokens, N//NL, NL) — running weighted sum + # Scratch + metadata_ref, + fused_metadata_ref, # single-slot MetadataRef for inner_kernel + tiled_out_2x_ref, # (2, tile_m, aligned_n) — double-buffered matmul output + scatter_staging_ref, # (tile_m, n_cols, num_lanes) — single staging buffer + partial_out_ref, # (sls, tile_n) — dummy for scatter_mode inner_kernel + acc_ref, # (tile_m, tile_n) — matmul accumulator + count_ref, # SMEM(3,): [0]=gm_id, [1]=send_count, [2]=local_count + semaphore_ref, # DMA(1,) for out_buf zero_out + acc_zero_sem_ref, # DMA(1,) for accumulator zero_out + send_sem_ref, # DMA(1,) for ICI remote sends + local_write_sem_ref, # DMA(1,) for local out_buf writes + recv_sem_ref, # DMA(1,) for incoming remote writes + staging_sem_ref, # DMA(1,) for accumulator scatter-back + output_sem_ref, # DMA(1,) for accumulator gather + w_buf_ref, # VMEM: (num_w_bufs, packed_k, tile_n) — weight buffer + w_scale_buf_ref, # VMEM: (num_w_bufs, nqb, 1, tile_n) — scale buffer + w_bias_buf_ref, # VMEM: (1, tile_n) — bias buffer + w_sem_ref, # DMA(num_w_bufs,) — weight DMA semaphores + gathered_lhs_2x_ref, # VMEM: (2, tile_m, K//NL, NL) — double-buffered LHS + gather_sem_ref, # DMA(2,) — LHS gather semaphores + acc_gather_2x_ref, # VMEM: (2, tile_m, N//NL, NL) — double-buffered acc gather + *, + cfgs, + ep_size, + chunk_size, + ep_axis_name, + num_w_bufs: int, +): + """GMM kernel with ICI direct-write scatter + dedup (is_last). + + Uses weighted accumulation: acc[token] += gmm_result * topk_weight. + Only ICI-sends when is_last_for_token == 1 (fully reduced result). + This reduces ICI sends from ~112/tile to ~14/tile (1/top_k). + """ + num_lanes = pltpu.get_tpu_info().num_lanes + my_id = lax.axis_index(ep_axis_name) + out_dtype = cfgs.out_dtype + tile_m = cfgs.tiles.tile_m + tile_k = cfgs.tiles.tile_k + tile_n = cfgs.tiles.tile_n + + num_k = pl.cdiv(cfgs.dims.size_k, tile_k) + num_n = pl.cdiv(cfgs.out_size_n, tile_n) + + # Pack along K for quantized rhs. + if cfgs.rhs_cfgs.quant_dtype is not None: + rhs_weight = rhs_ref.weight.bitcast(jnp.uint32) + rhs_ref = dataclasses.replace(rhs_ref, weight=rhs_weight) + + packing = cfgs.rhs_cfgs.packing + pk = tile_k // packing + has_scale = cfgs.rhs_cfgs.has_scale + has_bias = cfgs.rhs_cfgs.has_bias + nqb = cfgs.num_quant_blocks_per_tile_k if has_scale else 0 + w_dma_n = tile_n * 2 if cfgs.fuse_act else tile_n + + rhs_packed = rhs_ref.weight + + # --- Weight DMA helpers (same pattern as nodedup) --- + @jax.named_scope("start_w_dma") + def start_w_dma(buf_id, expert_id, n_id, k_id=0): + pltpu.make_async_copy( + src_ref=rhs_packed.at[expert_id, + pl.ds(k_id * pk, pk), + pl.ds(n_id * w_dma_n, w_dma_n)], + dst_ref=w_buf_ref.at[buf_id], + sem=w_sem_ref.at[buf_id], + ).start() + if has_scale: + pltpu.make_async_copy( + src_ref=rhs_ref.scale.at[expert_id, + pl.ds(k_id * nqb, nqb), :, + pl.ds(n_id * w_dma_n, w_dma_n)], + dst_ref=w_scale_buf_ref.at[buf_id], + sem=w_sem_ref.at[buf_id], + ).start() + if has_bias: + pltpu.make_async_copy( + src_ref=rhs_ref.bias.at[expert_id, :, + pl.ds(n_id * w_dma_n, w_dma_n)], + dst_ref=w_bias_buf_ref, + sem=w_sem_ref.at[buf_id], + ).start() + + @jax.named_scope("wait_w_dma") + def wait_w_dma(buf_id): + pltpu.make_async_copy( + src_ref=w_buf_ref.at[buf_id], + dst_ref=w_buf_ref.at[buf_id], + sem=w_sem_ref.at[buf_id], + ).wait() + if has_scale: + pltpu.make_async_copy( + src_ref=w_scale_buf_ref.at[buf_id], + dst_ref=w_scale_buf_ref.at[buf_id], + sem=w_sem_ref.at[buf_id], + ).wait() + if has_bias: + pltpu.make_async_copy( + src_ref=w_bias_buf_ref, + dst_ref=w_bias_buf_ref, + sem=w_sem_ref.at[buf_id], + ).wait() + + @jax.named_scope("compute_tile") + def compute_tile(buf_id, n_id, k_id, gm_id): + sem_id = gm_id % 2 + # Read LHS k-slice from VMEM gathered_lhs_2x_ref (pre-loaded via DMA). + k_cols = tile_k // num_lanes + k_offset = k_id * k_cols + sls_check = pltpu.get_tpu_info().get_sublane_tiling(cfgs.out_dtype) + if k_cols >= sls_check and k_cols % sls_check == 0: + # Direct ref slice — efficient, no redundant read. + lhs_tile = gathered_lhs_2x_ref[sem_id, :, + pl.ds(k_offset, k_cols), :].reshape( + -1, cfgs.dims.size_lhs_sublane, + tile_k) + else: + # Fallback for tiny K: read full buffer, slice at JAX level. + lhs_full = gathered_lhs_2x_ref.at[sem_id][...] + lhs_k_slice = lhs_full[:, k_offset:k_offset + k_cols, :] + lhs_tile = lhs_k_slice.reshape(-1, cfgs.dims.size_lhs_sublane, + tile_k) + + w_tile = w_buf_ref.at[buf_id] + w_sc = w_scale_buf_ref.at[buf_id] if has_scale else None + w_bias_tile = w_bias_buf_ref if has_bias else None + w_weights = WeightsRef(weight=w_tile, scale=w_sc, bias=w_bias_tile) + if cfgs.fuse_act is not None: + w_up = WeightsRef( + weight=w_tile.at[:, tile_n:], + scale=w_sc.at[:, :, tile_n:] if w_sc is not None else None, + bias=w_bias_tile.at[:, pl.ds(tile_n, tile_n)] + if has_bias else None, + ) + w_gate = WeightsRef( + weight=w_tile.at[:, :tile_n], + scale=w_sc.at[:, :, :tile_n] if w_sc is not None else None, + bias=w_bias_tile.at[:, pl.ds(0, tile_n)] if has_bias else None, + ) + w_weights = FusedWeightsRef(gate=w_gate, up=w_up) + inner_kernel( + lhs_tile, + w_weights, + tiled_out_2x_ref.at[sem_id, :, + pl.ds(n_id * tile_n, tile_n)], + partial_out_ref, + acc_ref, + fused_metadata_ref, + cfgs=cfgs, + scatter_mode=True, + _k_id=k_id, + _num_k=num_k, + _gm_id=0, + _n_id=n_id, + ) + + # --- Sync barrier --- + @jax.named_scope("sync_barrier") + def sync_barrier(): + barrier_sem = pltpu.get_barrier_semaphore() + for i in range(ep_size): + pltpu.semaphore_signal( + barrier_sem, + device_id={ep_axis_name: jnp.int32(i)}, + device_id_type=pltpu.DeviceIdType.MESH, + ) + pltpu.semaphore_wait(barrier_sem, ep_size) + + sync_barrier() + + # Fill metadata. + num_gm = fill_metadata( + lhs_group_sizes_ref, + group_offset_ref, + metadata_ref, + cfgs=cfgs, + ) + + # Zero-init async start (both out_buf and accumulator). + if cfgs.zero_init: + with jax.named_scope("zero_init_start"): + zero_size_out = zero_out_start_3d( + out_buf_ref, + scatter_staging_ref, + semaphore_ref, + ) + zero_size_acc = zero_out_start_3d( + accumulator_ref, + scatter_staging_ref, + acc_zero_sem_ref, + ) + + # Initialize send/local counts. + count_ref[0] = jnp.int32(0) # gm_id + count_ref[1] = jnp.int32(0) # send_count + count_ref[2] = jnp.int32(0) # local_count + + total_w_steps = num_n * num_k + can_cache_w = num_w_bufs >= total_w_steps + + # --- Main gm loop (nodedup pattern) --- + @jax.named_scope("gm_loop_body") + def gm_loop_body(gm_id, _): + sem_id = gm_id % 2 + + # 1. Wait for previous accumulator scatter before reusing acc_gather_2x_ref. + @pl.when(gm_id > 0) + def _(): + with jax.named_scope("acc_scatter_wait_prev"): + dma_scatter_gm_wait( + acc_gather_2x_ref.at[1 - sem_id], + staging_sem_ref.at[0], + gm_id - 1, + metadata_ref, + ) + + # 2. Setup metadata for this gm tile. + fused_metadata_ref.gm_id_to_group_id[ + 0] = metadata_ref.gm_id_to_group_id[gm_id] + fused_metadata_ref.gm_id_to_m_offset[ + 0] = metadata_ref.gm_id_to_m_offset[gm_id] + fused_metadata_ref.gm_id_to_m_offset[ + 1] = metadata_ref.gm_id_to_m_offset[gm_id + 1] + expert_id = fused_metadata_ref.gm_id_to_group_id[0] + + # 2b. LHS gather: contiguous DMA from HBM to VMEM. + # Bootstrap: start gather for tile 0 on first iteration. + @pl.when(gm_id == 0) + def _(): + with jax.named_scope("lhs_gather_bootstrap"): + m_start_0 = metadata_ref.gm_id_to_m_offset[0] + sls_0 = cfgs.dims.size_lhs_sublane + m_aligned_0 = m_start_0 - m_start_0 % sls_0 + pltpu.make_async_copy( + src_ref=lhs_ref.at[pl.ds(m_aligned_0, tile_m), :, :], + dst_ref=gathered_lhs_2x_ref.at[0], + sem=gather_sem_ref.at[0], + ).start() + + # Prefetch next tile's LHS. + @pl.when(gm_id + 1 < num_gm) + def _(): + with jax.named_scope("lhs_gather_prefetch"): + m_start_next = metadata_ref.gm_id_to_m_offset[gm_id + 1] + sls_n = cfgs.dims.size_lhs_sublane + m_aligned_next = m_start_next - m_start_next % sls_n + pltpu.make_async_copy( + src_ref=lhs_ref.at[pl.ds(m_aligned_next, tile_m), :, :], + dst_ref=gathered_lhs_2x_ref.at[1 - sem_id], + sem=gather_sem_ref.at[1 - sem_id], + ).start() + + # Wait for current tile's LHS gather. + with jax.named_scope("lhs_gather_wait"): + pltpu.make_async_copy( + src_ref=gathered_lhs_2x_ref.at[sem_id], + dst_ref=gathered_lhs_2x_ref.at[sem_id], + sem=gather_sem_ref.at[sem_id], + ).wait() + + # 3. Weight DMA with same-expert caching. + prev_gm_clamped = jnp.maximum(gm_id - 1, 0) + prev_expert = metadata_ref.gm_id_to_group_id[prev_gm_clamped] + is_new_expert = jnp.logical_or(gm_id == 0, prev_expert != expert_id) + is_same_w = jnp.logical_and(jnp.logical_not(is_new_expert), + jnp.bool_(can_cache_w)) + + @pl.when(gm_id == 0) + def _(): + for _i in range(min(num_w_bufs, total_w_steps)): + start_w_dma(_i, expert_id, _i // num_k, _i % num_k) + + # 4. Matmul loop over (N, K) tiles with weight DMA pipelining. + for step in range(total_w_steps): + _n = step // num_k + _k = step % num_k + buf_id = step % num_w_bufs + + @pl.when(jnp.logical_not(is_same_w)) + def _(): + wait_w_dma(buf_id) + + if step + num_w_bufs < total_w_steps: + ns = step + num_w_bufs + start_w_dma(ns % num_w_bufs, expert_id, ns // num_k, + ns % num_k) + + compute_tile(buf_id, _n, _k, gm_id) + + # Cross-gm weight prefetch after first step. + if step == 0: + + @pl.when(gm_id + 1 < num_gm) + def _(): + next_e = metadata_ref.gm_id_to_group_id[gm_id + 1] + next_same = jnp.logical_and(next_e == expert_id, + jnp.bool_(can_cache_w)) + + @pl.when(jnp.logical_not(next_same)) + def _(): + start_w_dma(0, next_e, 0, 0) + + # Cross-gm prefetch remaining buffers. + @pl.when(gm_id + 1 < num_gm) + def _(): + next_e = metadata_ref.gm_id_to_group_id[gm_id + 1] + next_same = jnp.logical_and(next_e == expert_id, + jnp.bool_(can_cache_w)) + + @pl.when(jnp.logical_not(next_same)) + def _(): + for _i in range(1, min(num_w_bufs, total_w_steps)): + start_w_dma(_i, next_e, _i // num_k, _i % num_k) + + # 5. Deferred zero-init wait on first tile (both out_buf and accumulator). + if cfgs.zero_init: + + @pl.when(gm_id == 0) + def _(): + with jax.named_scope("zero_out_end_deferred"): + zero_out_end_3d(out_buf_ref, semaphore_ref, zero_size_out) + zero_out_end_3d(accumulator_ref, acc_zero_sem_ref, + zero_size_acc) + + # 6. Accumulator gather: load current acc values for this tile's rows. + with jax.named_scope("acc_gather_start"): + dma_gather_gm_start( + accumulator_ref, + acc_gather_2x_ref.at[sem_id], + output_indices_ref, + output_sem_ref.at[0], + gm_id, + metadata_ref, + ) + + # 7. Reshape GMM output to staging, then do weighted RMW + conditional ICI send. + m_st = metadata_ref.gm_id_to_m_offset[gm_id] + m_en = metadata_ref.gm_id_to_m_offset[gm_id + 1] + _sls = pltpu.get_tpu_info().get_sublane_tiling(out_dtype) + _ml = m_st % _sls + num_valid = m_en - m_st + + with jax.named_scope("reshape_to_staging"): + scatter_staging_ref[...] = tiled_out_2x_ref[sem_id][...].reshape( + scatter_staging_ref.shape) + + # Wait for accumulator gather to complete. + with jax.named_scope("acc_gather_wait"): + dma_gather_gm_wait( + acc_gather_2x_ref.at[sem_id], + output_sem_ref.at[0], + gm_id, + metadata_ref, + ) + + # Per-row weighted RMW + conditional ICI send (dedup pattern). + acc_slot = acc_gather_2x_ref.at[sem_id] + + @jax.named_scope("weighted_add_and_send") + def _do_weighted_add_and_send(): + + def _row_fn(i, carry): + send_sz, local_sz = carry + row_idx = _ml + i + w = topk_weights_ref[m_st + i].astype(jnp.float32) + token_g = output_indices_ref[m_st + i] + is_last = is_last_ref[m_st + i] + dest_chip = token_g // chunk_size + local_row = token_g % chunk_size + write_pos = local_row * ep_size + my_id + is_local = dest_chip == my_id + + # Weighted RMW: acc[token] += staging[row] * weight + old_r = acc_slot.at[pl.ds(row_idx, 1), :, :] + new_r = scatter_staging_ref.at[pl.ds(row_idx, 1), :, :] + result = (old_r[...].astype(jnp.float32) + + new_r[...].astype(jnp.float32) * w) + acc_slot.at[pl.ds(row_idx, 1), :, :][...] = result.astype( + out_dtype) + + # Conditional ICI send (only when is_last == 1). + @pl.when(jnp.logical_and(is_last == 1, ~is_local)) + def _(): + pltpu.make_async_remote_copy( + src_ref=acc_slot.at[pl.ds(row_idx, 1), :, :], + dst_ref=out_buf_ref.at[pl.ds(write_pos, 1), :, :], + send_sem=send_sem_ref.at[0], + recv_sem=recv_sem_ref.at[0], + device_id={ + ep_axis_name: dest_chip + }, + device_id_type=pltpu.DeviceIdType.MESH, + ).start() + + @pl.when(jnp.logical_and(is_last == 1, is_local)) + def _(): + pltpu.make_async_copy( + src_ref=acc_slot.at[pl.ds(row_idx, 1), :, :], + dst_ref=out_buf_ref.at[pl.ds(write_pos, 1), :, :], + sem=local_write_sem_ref.at[0], + ).start() + + is_last_remote = jnp.logical_and(is_last == 1, ~is_local) + is_last_local = jnp.logical_and(is_last == 1, is_local) + return ( + send_sz + + lax.select(is_last_remote, jnp.int32(1), jnp.int32(0)), + local_sz + + lax.select(is_last_local, jnp.int32(1), jnp.int32(0)), + ) + + return lax.fori_loop(0, num_valid, _row_fn, + (jnp.int32(0), jnp.int32(0))) + + sz_send, sz_local = _do_weighted_add_and_send() + count_ref[1] = count_ref[1] + sz_send + count_ref[2] = count_ref[2] + sz_local + + # Scatter running sum back to accumulator HBM. + with jax.named_scope("acc_scatter_start"): + dma_scatter_gm_start( + acc_gather_2x_ref.at[sem_id], + accumulator_ref, + output_indices_ref, + staging_sem_ref.at[0], + gm_id, + metadata_ref, + ) + + return _ + + lax.fori_loop(0, num_gm, gm_loop_body, None) + + # Handle num_gm == 0 edge case for deferred zero-init. + if cfgs.zero_init: + + @pl.when(num_gm == 0) + def _(): + zero_out_end_3d(out_buf_ref, semaphore_ref, zero_size_out) + zero_out_end_3d(accumulator_ref, acc_zero_sem_ref, zero_size_acc) + + # Epilogue: wait for last accumulator scatter. + @pl.when(num_gm > 0) + def _(): + with jax.named_scope("final_acc_scatter_wait"): + last_sem_id = (num_gm - 1) % 2 + dma_scatter_gm_wait( + acc_gather_2x_ref.at[last_sem_id], + staging_sem_ref.at[0], + num_gm - 1, + metadata_ref, + ) + + # Final barrier. + sync_barrier() + + # Drain send_sem (sparse — only is_last rows sent). + with jax.named_scope("drain_send_sem"): + total_send = count_ref[1] + pltpu.make_async_copy( + src_ref=out_buf_ref.at[pl.ds(0, total_send), :, :], + dst_ref=out_buf_ref.at[pl.ds(0, total_send), :, :], + sem=send_sem_ref.at[0], + ).wait() + + # Drain local_write_sem. + with jax.named_scope("drain_local_write_sem"): + total_local = count_ref[2] + pltpu.make_async_copy( + src_ref=out_buf_ref.at[pl.ds(0, total_local), :, :], + dst_ref=out_buf_ref.at[pl.ds(0, total_local), :, :], + sem=local_write_sem_ref.at[0], + ).wait() + + # Drain recv_sem. + with jax.named_scope("drain_recv_sem"): + total_recv = total_recv_count_ref[0] + pltpu.make_async_copy( + src_ref=out_buf_ref.at[pl.ds(0, total_recv), :, :], + dst_ref=out_buf_ref.at[pl.ds(0, total_recv), :, :], + sem=recv_sem_ref.at[0], + ).wait() + + +@jax.jit(static_argnames=[ + "fuse_act", + "ep_size", + "ep_axis_name", + "chunk_size", + "num_tokens", + "vmem_limit_bytes", + "maybe_quantize_lhs", +], ) +def gmm_v2_scatter_ici_dedup( + lhs: jax.Array, # [size_m, size_k] + rhs: jax.Array, # [size_group, size_k, size_n] + group_sizes: jax.Array, # int32[size_lhs_group] + output_indices: jax.Array, # int32[size_m] — global scatter positions + total_recv_count: jax.Array, # int32[1] — remote rows this chip receives + topk_weights: jax. + Array, # float32[size_m] — per-row weights for accumulation + is_last_for_token: jax. + Array, # int32[size_m] — 1 iff last occurrence on this chip + *, + rhs_scale: jax.Array | None = None, + rhs_bias: jax.Array | None = None, + group_offset: jax.Array | None = None, + fuse_act: str | None = None, + ep_size: int, + ep_axis_name: str, + chunk_size: int, + num_tokens: int, # total tokens across all chips (for accumulator sizing) + vmem_limit_bytes: int | None = None, + maybe_quantize_lhs: bool = True, +) -> jax.Array: + """GMM with ICI direct-write scatter + dedup (is_last). + + Uses weighted accumulation: acc[token] += gmm_result * topk_weight. + Only ICI-sends when is_last_for_token == 1 (fully reduced result). + Output: (chunk_size * ep_size, size_n) — each token has ep_size slots. + Post-kernel: sum(reshape(chunk_size, ep_size, N), axis=1). + + Must be called inside shard_map with ep_axis_name mesh axis. + + Args: + lhs: LHS matrix [size_m, size_k]. + rhs: RHS matrix [size_group, size_k, size_n]. + group_sizes: Group sizes [size_lhs_group]. + output_indices: int32[size_m] — global scatter positions. + dest_chip = oi // chunk_size, local_row = oi % chunk_size. + total_recv_count: int32[1] — number of remote rows this chip receives. + rhs_scale: Optional per-block scale. + rhs_bias: Optional bias. + group_offset: Optional group offset. + fuse_act: Optional activation fusion. + ep_size: Number of expert-parallel chips. + ep_axis_name: Mesh axis name for expert parallelism. + chunk_size: Tokens per chip (output_size // ep_size). + vmem_limit_bytes: Optional VMEM limit. + maybe_quantize_lhs: Whether to quantize LHS. + + Returns: + Output of shape [chunk_size, size_n] (this chip's portion). + """ + if group_offset is None: + group_offset = jnp.array([0], dtype=jnp.int32) + + if vmem_limit_bytes is None: + vmem_limit_bytes = int(pltpu.get_tpu_info().vmem_capacity_bytes) + + num_lanes = pltpu.get_tpu_info().num_lanes + + # Check fuse_act alignment. + _fallback_fuse_act = None + if fuse_act is not None: + if rhs.shape[2] % (2 * num_lanes) != 0: + _fallback_fuse_act = fuse_act + fuse_act = None + + cfgs = make_gmm_configs( + lhs, + rhs, + rhs_scale, + rhs_bias, + group_sizes, + group_offset, + tile_info=calculate_tiling, + vmem_limit_bytes=vmem_limit_bytes, + out_dtype=None, + acc_dtype=None, + maybe_quantize_lhs=maybe_quantize_lhs, + zero_initialize=True, + output_indices=output_indices, + fuse_act=fuse_act, + ) + dims = cfgs.dims + tiles = cfgs.tiles + + # Align N for DMA scatter. + _out_sls = pltpu.get_tpu_info().get_sublane_tiling(cfgs.out_dtype) + aligned_n = align_to(dims.size_n, num_lanes * _out_sls) + + # Ensure tile_n divides aligned_n. + if aligned_n % tiles.tile_n != 0: + _mxu_size = pltpu.get_tpu_info().mxu_column_size + _num_n = pl.cdiv(aligned_n, tiles.tile_n) + _adj_tile_n = (aligned_n // _num_n // _mxu_size) * _mxu_size + if _adj_tile_n > 0: + tiles = dataclasses.replace(tiles, tile_n=_adj_tile_n) + cfgs = dataclasses.replace(cfgs, tiles=tiles) + + n_cols = aligned_n // num_lanes + + # Weight buffer config (same as nodedup pattern). + is_quantized = cfgs.rhs_cfgs.quant_dtype is not None + packing = cfgs.rhs_cfgs.packing + pk = tiles.tile_k // packing + has_scale = cfgs.rhs_cfgs.has_scale + nqb = cfgs.num_quant_blocks_per_tile_k if has_scale else 1 + w_dma_n = tiles.tile_n * 2 if cfgs.fuse_act else tiles.tile_n + num_w_bufs = 2 # double-buffered weights + + # Scratch shapes matching kernel_main_scatter_ici params. + max_num_gm = dims.size_group + pl.cdiv(dims.size_m, tiles.tile_m) - 1 + acc_cols = 2 * tiles.tile_n if cfgs.fuse_act is not None else tiles.tile_n + + scratch_shapes = [ + # metadata_ref + MetadataRef( + gm_id_to_group_id=pltpu.SMEM((max_num_gm, ), jnp.int32), + gm_id_to_m_offset=pltpu.SMEM((max_num_gm + 1, ), jnp.int32), + ), + # fused_metadata_ref (single-slot for inner_kernel) + MetadataRef( + gm_id_to_group_id=pltpu.SMEM((1, ), jnp.int32), + gm_id_to_m_offset=pltpu.SMEM((2, ), jnp.int32), + ), + # tiled_out_2x_ref — double-buffered compute buffer + pltpu.VMEM((2, tiles.tile_m, aligned_n), cfgs.out_dtype), + # scatter_staging_ref — single buffer (no triple-buffering needed with dedup) + pltpu.VMEM((tiles.tile_m, n_cols, num_lanes), cfgs.out_dtype), + # partial_out_ref (dummy for scatter_mode inner_kernel) + pltpu.VMEM((dims.size_lhs_sublane, tiles.tile_n), cfgs.out_dtype), + # acc_ref + pltpu.VMEM((tiles.tile_m, acc_cols), cfgs.acc_dtype), + # count_ref — SMEM: [0]=gm_id, [1]=send_count, [2]=local_count + pltpu.SMEM((3, ), jnp.int32), + # semaphore_ref for out_buf zero_out + pltpu.SemaphoreType.DMA((1, )), + # acc_zero_sem_ref for accumulator zero_out + pltpu.SemaphoreType.DMA((1, )), + # send_sem_ref — single (sparse sends) + pltpu.SemaphoreType.DMA((1, )), + # local_write_sem_ref — single + pltpu.SemaphoreType.DMA((1, )), + # recv_sem_ref — for incoming remote writes + pltpu.SemaphoreType.DMA((1, )), + # staging_sem_ref — for accumulator scatter-back + pltpu.SemaphoreType.DMA((1, )), + # output_sem_ref — for accumulator gather + pltpu.SemaphoreType.DMA((1, )), + # w_buf_ref — weight buffer + pltpu.VMEM( + (num_w_bufs, pk, w_dma_n), + jnp.uint32 if is_quantized else rhs.dtype, + ), + # w_scale_buf_ref + pltpu.VMEM( + (num_w_bufs, nqb, 1, w_dma_n), + jnp.float32, + ), + # w_bias_buf_ref + pltpu.VMEM((1, w_dma_n), jnp.float32), + # w_sem_ref — weight DMA semaphores + pltpu.SemaphoreType.DMA((num_w_bufs, )), + # gathered_lhs_2x_ref — double-buffered LHS only (K//NL cols) + pltpu.VMEM( + (2, tiles.tile_m, dims.size_k // num_lanes, num_lanes), + cfgs.out_dtype, + ), + # gather_sem_ref — LHS gather semaphores + pltpu.SemaphoreType.DMA((2, )), + # acc_gather_2x_ref — double-buffered accumulator gather (N//NL cols) + pltpu.VMEM( + (2, tiles.tile_m, n_cols, num_lanes), + cfgs.out_dtype, + ), + ] + + # Output shapes: out_buf (ICI target) + accumulator (running sum). + out_buf_init = jax.ShapeDtypeStruct( + (chunk_size * ep_size, n_cols, num_lanes), cfgs.out_dtype) + accumulator_init = jax.ShapeDtypeStruct((num_tokens, n_cols, num_lanes), + cfgs.out_dtype) + + # Prepare RHS specs. + rhs_scale_spec = rhs_bias_spec = None + if rhs_scale is not None: + rhs_scale = rhs_scale.astype(jnp.float32) + rhs_scale_spec = pl.BlockSpec(memory_space=pltpu.HBM) + if rhs_bias is not None: + rhs_bias = rhs_bias.astype(jnp.float32) + rhs_bias_spec = pl.BlockSpec(memory_space=pltpu.HBM) + + rhs_weights = WeightsRef(weight=rhs, scale=rhs_scale, bias=rhs_bias) + rhs_in_spec = WeightsRef( + weight=pl.BlockSpec(memory_space=pltpu.HBM), + scale=rhs_scale_spec, + bias=rhs_bias_spec, + ) + + compiler_params = pltpu.CompilerParams( + vmem_limit_bytes=vmem_limit_bytes, + disable_bounds_checks=True, + collective_id=0, + ) + + total_recv_arr = total_recv_count.astype(jnp.int32).reshape((1, )) + topk_w_flat = topk_weights.astype(jnp.float32).flatten() + + # Reshape LHS to 3D for contiguous DMA inside kernel. + lhs_3d = lhs.reshape(dims.size_m, dims.size_k // num_lanes, num_lanes) + + out_buf, _accumulator = pl.pallas_call( + functools.partial( + kernel_main_scatter_ici_dedup, + cfgs=cfgs, + ep_size=ep_size, + chunk_size=chunk_size, + ep_axis_name=ep_axis_name, + num_w_bufs=num_w_bufs, + ), + out_shape=[out_buf_init, accumulator_init], + grid_spec=pltpu.PrefetchScalarGridSpec( + num_scalar_prefetch=6, + in_specs=[ + pl.BlockSpec(memory_space=pltpu.HBM), # lhs_3d + rhs_in_spec, + ], + out_specs=[ + pl.BlockSpec(memory_space=pltpu.HBM), # out_buf + pl.BlockSpec(memory_space=pltpu.HBM), # accumulator + ], + scratch_shapes=scratch_shapes, + ), + compiler_params=compiler_params, + name=get_scope_name(dims, tiles) + "-scatter_ici_dedup", + cost_estimate=get_cost_estimate(cfgs), + )( + group_sizes, + group_offset, + output_indices, + topk_w_flat, + is_last_for_token, + total_recv_arr, + lhs_3d, + rhs_weights, + ) + + result = out_buf.reshape(chunk_size * ep_size, + aligned_n)[:, :cfgs.out_size_n] + + if _fallback_fuse_act is not None: + result = apply_act_fn(result, _fallback_fuse_act) + + return result + + +# ============================================================================= +# GMM with ICI direct-write scatter — NODEDUP version (ALL rows sent) +# Simple triple-buffered staging, no accumulator, no is_last gating. +# ============================================================================= + + +def kernel_main_scatter_ici_nodedup( + # Scalar prefetch (4) + lhs_group_sizes_ref, + group_offset_ref, + output_indices_ref, + total_recv_count_ref, # (1,) int32 — total remote rows this chip receives + # In (2) + lhs_ref, # HBM: (size_m, K//NL, NL) — 3D for DMA + rhs_ref, # HBM: WeightsRef (size_group, size_k, size_n) + # Out (1) + out_buf_ref, # 3D HBM: (chunk_size, N//NL, NL) — ICI target + # Scratch + metadata_ref, + fused_metadata_ref, # single-slot MetadataRef for inner_kernel + tiled_out_2x_ref, # (2, tile_m, aligned_n) — double-buffered matmul output + scatter_staging_3x_ref, # (3, tile_m, n_cols, num_lanes) — triple-buffered staging + partial_out_ref, # (sls, tile_n) — dummy for scatter_mode inner_kernel + acc_ref, # (tile_m, tile_n) — matmul accumulator + count_ref, # SMEM(7,): [0]=gm_id, [1..3]=send counts, [4..6]=local counts + semaphore_ref, # DMA(1,) for out_buf zero_out + send_sems_ref, # DMA(3,) per-staging-slot send sems + local_write_sems_ref, # DMA(3,) per-staging-slot local write sems + recv_sem_ref, # DMA(1,) for incoming remote writes + w_buf_ref, # VMEM: (num_w_bufs, packed_k, tile_n) — weight buffer + w_scale_buf_ref, # VMEM: (num_w_bufs, nqb, 1, tile_n) — scale buffer + w_bias_buf_ref, # VMEM: (1, tile_n) — bias buffer + w_sem_ref, # DMA(num_w_bufs,) — weight DMA semaphores + gathered_lhs_2x_ref, # VMEM: (2, tile_m, K//NL, NL) — double-buffered LHS + gather_sem_ref, # DMA(2,) — LHS gather semaphores + *, + cfgs, + ep_size, + chunk_size, + ep_axis_name, + num_w_bufs: int, +): + """GMM kernel with ICI direct-write scatter (nodedup — ALL rows sent). + + Every GMM output row is ICI-sent to the destination chip. No accumulator, + no is_last gating, no topk_weights. Simple triple-buffered staging with + drain at top of loop. + """ + num_lanes = pltpu.get_tpu_info().num_lanes + my_id = lax.axis_index(ep_axis_name) + out_dtype = cfgs.out_dtype + tile_m = cfgs.tiles.tile_m + tile_k = cfgs.tiles.tile_k + tile_n = cfgs.tiles.tile_n + + num_k = pl.cdiv(cfgs.dims.size_k, tile_k) + num_n = pl.cdiv(cfgs.out_size_n, tile_n) + + # Pack along K for quantized rhs. + if cfgs.rhs_cfgs.quant_dtype is not None: + rhs_weight = rhs_ref.weight.bitcast(jnp.uint32) + rhs_ref = dataclasses.replace(rhs_ref, weight=rhs_weight) + + packing = cfgs.rhs_cfgs.packing + pk = tile_k // packing + has_scale = cfgs.rhs_cfgs.has_scale + has_bias = cfgs.rhs_cfgs.has_bias + nqb = cfgs.num_quant_blocks_per_tile_k if has_scale else 0 + w_dma_n = tile_n * 2 if cfgs.fuse_act else tile_n + + rhs_packed = rhs_ref.weight + + # --- Weight DMA helpers --- + @jax.named_scope("start_w_dma") + def start_w_dma(buf_id, expert_id, n_id, k_id=0): + pltpu.make_async_copy( + src_ref=rhs_packed.at[expert_id, + pl.ds(k_id * pk, pk), + pl.ds(n_id * w_dma_n, w_dma_n)], + dst_ref=w_buf_ref.at[buf_id], + sem=w_sem_ref.at[buf_id], + ).start() + if has_scale: + pltpu.make_async_copy( + src_ref=rhs_ref.scale.at[expert_id, + pl.ds(k_id * nqb, nqb), :, + pl.ds(n_id * w_dma_n, w_dma_n)], + dst_ref=w_scale_buf_ref.at[buf_id], + sem=w_sem_ref.at[buf_id], + ).start() + if has_bias: + pltpu.make_async_copy( + src_ref=rhs_ref.bias.at[expert_id, :, + pl.ds(n_id * w_dma_n, w_dma_n)], + dst_ref=w_bias_buf_ref, + sem=w_sem_ref.at[buf_id], + ).start() + + @jax.named_scope("wait_w_dma") + def wait_w_dma(buf_id): + pltpu.make_async_copy( + src_ref=w_buf_ref.at[buf_id], + dst_ref=w_buf_ref.at[buf_id], + sem=w_sem_ref.at[buf_id], + ).wait() + if has_scale: + pltpu.make_async_copy( + src_ref=w_scale_buf_ref.at[buf_id], + dst_ref=w_scale_buf_ref.at[buf_id], + sem=w_sem_ref.at[buf_id], + ).wait() + if has_bias: + pltpu.make_async_copy( + src_ref=w_bias_buf_ref, + dst_ref=w_bias_buf_ref, + sem=w_sem_ref.at[buf_id], + ).wait() + + @jax.named_scope("compute_tile") + def compute_tile(buf_id, n_id, k_id, gm_id): + sem_id = gm_id % 2 + # Read LHS k-slice from VMEM gathered_lhs_2x_ref (pre-loaded via DMA). + k_cols = tile_k // num_lanes + k_offset = k_id * k_cols + sls_check = pltpu.get_tpu_info().get_sublane_tiling(cfgs.out_dtype) + if k_cols >= sls_check and k_cols % sls_check == 0: + # Direct ref slice — efficient, no redundant read. + lhs_tile = gathered_lhs_2x_ref[sem_id, :, + pl.ds(k_offset, k_cols), :].reshape( + -1, cfgs.dims.size_lhs_sublane, + tile_k) + else: + # Fallback for tiny K: read full buffer, slice at JAX level. + lhs_full = gathered_lhs_2x_ref.at[sem_id][...] + lhs_k_slice = lhs_full[:, k_offset:k_offset + k_cols, :] + lhs_tile = lhs_k_slice.reshape(-1, cfgs.dims.size_lhs_sublane, + tile_k) + + w_tile = w_buf_ref.at[buf_id] + w_sc = w_scale_buf_ref.at[buf_id] if has_scale else None + w_bias_tile = w_bias_buf_ref if has_bias else None + w_weights = WeightsRef(weight=w_tile, scale=w_sc, bias=w_bias_tile) + if cfgs.fuse_act is not None: + w_up = WeightsRef( + weight=w_tile.at[:, tile_n:], + scale=w_sc.at[:, :, tile_n:] if w_sc is not None else None, + bias=w_bias_tile.at[:, pl.ds(tile_n, tile_n)] + if has_bias else None, + ) + w_gate = WeightsRef( + weight=w_tile.at[:, :tile_n], + scale=w_sc.at[:, :, :tile_n] if w_sc is not None else None, + bias=w_bias_tile.at[:, pl.ds(0, tile_n)] if has_bias else None, + ) + w_weights = FusedWeightsRef(gate=w_gate, up=w_up) + inner_kernel( + lhs_tile, + w_weights, + tiled_out_2x_ref.at[sem_id, :, + pl.ds(n_id * tile_n, tile_n)], + partial_out_ref, + acc_ref, + fused_metadata_ref, + cfgs=cfgs, + scatter_mode=True, + _k_id=k_id, + _num_k=num_k, + _gm_id=0, + _n_id=n_id, + ) + + # --- Sync barrier --- + @jax.named_scope("sync_barrier") + def sync_barrier(): + barrier_sem = pltpu.get_barrier_semaphore() + for i in range(ep_size): + pltpu.semaphore_signal( + barrier_sem, + device_id={ep_axis_name: jnp.int32(i)}, + device_id_type=pltpu.DeviceIdType.MESH, + ) + pltpu.semaphore_wait(barrier_sem, ep_size) + + sync_barrier() + + # Fill metadata. + num_gm = fill_metadata( + lhs_group_sizes_ref, + group_offset_ref, + metadata_ref, + cfgs=cfgs, + ) + + # Zero-init async start. + if cfgs.zero_init: + with jax.named_scope("zero_init_start"): + zero_size_out = zero_out_start_3d( + out_buf_ref, + scatter_staging_3x_ref.at[0], + semaphore_ref, + ) + + # Initialize per-slot DMA counts to 0. + for _s in range(3): + count_ref[1 + _s] = jnp.int32(0) # send counts + count_ref[4 + _s] = jnp.int32(0) # local write counts + + total_w_steps = num_n * num_k + can_cache_w = num_w_bufs >= total_w_steps + + # --- Main gm loop --- + @jax.named_scope("gm_loop_body") + def gm_loop_body(gm_id, _): + sem_id = gm_id % 2 + stg_id = gm_id % 3 + + # Drain the staging slot we're about to reuse. With triple buffering, + # this slot was last used 3 iterations ago — giving 2 full iterations + # of compute overlap for its DMAs to complete. + @jax.named_scope("drain_prev_dmas") + @pl.when(gm_id >= 3) + def _(): + prev_send = count_ref[1 + stg_id] + pltpu.make_async_copy( + src_ref=scatter_staging_3x_ref.at[stg_id, + pl.ds(0, prev_send), :, :], + dst_ref=scatter_staging_3x_ref.at[stg_id, + pl.ds(0, prev_send), :, :], + sem=send_sems_ref.at[stg_id], + ).wait() + prev_local = count_ref[4 + stg_id] + pltpu.make_async_copy( + src_ref=out_buf_ref.at[pl.ds(0, prev_local), :, :], + dst_ref=out_buf_ref.at[pl.ds(0, prev_local), :, :], + sem=local_write_sems_ref.at[stg_id], + ).wait() + # Reset counts for this slot's new use. + count_ref[1 + stg_id] = jnp.int32(0) + count_ref[4 + stg_id] = jnp.int32(0) + + # 1. Setup metadata for this gm tile. + fused_metadata_ref.gm_id_to_group_id[ + 0] = metadata_ref.gm_id_to_group_id[gm_id] + fused_metadata_ref.gm_id_to_m_offset[ + 0] = metadata_ref.gm_id_to_m_offset[gm_id] + fused_metadata_ref.gm_id_to_m_offset[ + 1] = metadata_ref.gm_id_to_m_offset[gm_id + 1] + expert_id = fused_metadata_ref.gm_id_to_group_id[0] + + # 2. LHS gather: contiguous DMA from HBM to VMEM. + # Bootstrap: start gather for tile 0 on first iteration. + @pl.when(gm_id == 0) + def _(): + with jax.named_scope("lhs_gather_bootstrap"): + m_start_0 = metadata_ref.gm_id_to_m_offset[0] + sls_0 = cfgs.dims.size_lhs_sublane + m_aligned_0 = m_start_0 - m_start_0 % sls_0 + pltpu.make_async_copy( + src_ref=lhs_ref.at[pl.ds(m_aligned_0, tile_m), :, :], + dst_ref=gathered_lhs_2x_ref.at[0], + sem=gather_sem_ref.at[0], + ).start() + + # Prefetch next tile's LHS. + @pl.when(gm_id + 1 < num_gm) + def _(): + with jax.named_scope("lhs_gather_prefetch"): + m_start_next = metadata_ref.gm_id_to_m_offset[gm_id + 1] + sls_n = cfgs.dims.size_lhs_sublane + m_aligned_next = m_start_next - m_start_next % sls_n + pltpu.make_async_copy( + src_ref=lhs_ref.at[pl.ds(m_aligned_next, tile_m), :, :], + dst_ref=gathered_lhs_2x_ref.at[1 - sem_id], + sem=gather_sem_ref.at[1 - sem_id], + ).start() + + # Wait for current tile's LHS gather. + with jax.named_scope("lhs_gather_wait"): + pltpu.make_async_copy( + src_ref=gathered_lhs_2x_ref.at[sem_id], + dst_ref=gathered_lhs_2x_ref.at[sem_id], + sem=gather_sem_ref.at[sem_id], + ).wait() + + # 3. Weight DMA with same-expert caching. + prev_gm_clamped = jnp.maximum(gm_id - 1, 0) + prev_expert = metadata_ref.gm_id_to_group_id[prev_gm_clamped] + is_new_expert = jnp.logical_or(gm_id == 0, prev_expert != expert_id) + is_same_w = jnp.logical_and(jnp.logical_not(is_new_expert), + jnp.bool_(can_cache_w)) + + @pl.when(gm_id == 0) + def _(): + for _i in range(min(num_w_bufs, total_w_steps)): + start_w_dma(_i, expert_id, _i // num_k, _i % num_k) + + # 4. Matmul loop over (N, K) tiles with weight DMA pipelining. + for step in range(total_w_steps): + _n = step // num_k + _k = step % num_k + buf_id = step % num_w_bufs + + @pl.when(jnp.logical_not(is_same_w)) + def _(): + wait_w_dma(buf_id) + + if step + num_w_bufs < total_w_steps: + ns = step + num_w_bufs + start_w_dma(ns % num_w_bufs, expert_id, ns // num_k, + ns % num_k) + + compute_tile(buf_id, _n, _k, gm_id) + + # Cross-gm weight prefetch after first step. + if step == 0: + + @pl.when(gm_id + 1 < num_gm) + def _(): + next_e = metadata_ref.gm_id_to_group_id[gm_id + 1] + next_same = jnp.logical_and(next_e == expert_id, + jnp.bool_(can_cache_w)) + + @pl.when(jnp.logical_not(next_same)) + def _(): + start_w_dma(0, next_e, 0, 0) + + # Cross-gm prefetch remaining buffers. + @pl.when(gm_id + 1 < num_gm) + def _(): + next_e = metadata_ref.gm_id_to_group_id[gm_id + 1] + next_same = jnp.logical_and(next_e == expert_id, + jnp.bool_(can_cache_w)) + + @pl.when(jnp.logical_not(next_same)) + def _(): + for _i in range(1, min(num_w_bufs, total_w_steps)): + start_w_dma(_i, next_e, _i // num_k, _i % num_k) + + # 5. Deferred zero-init wait on first tile. + if cfgs.zero_init: + + @pl.when(gm_id == 0) + def _(): + with jax.named_scope("zero_out_end_deferred"): + zero_out_end_3d(out_buf_ref, semaphore_ref, zero_size_out) + + # 6. Reshape GMM output to staging. + m_st = metadata_ref.gm_id_to_m_offset[gm_id] + m_en = metadata_ref.gm_id_to_m_offset[gm_id + 1] + _sls = pltpu.get_tpu_info().get_sublane_tiling(out_dtype) + _ml = m_st % _sls + num_valid = m_en - m_st + + with jax.named_scope("reshape_to_staging"): + scatter_staging_3x_ref[stg_id] = tiled_out_2x_ref[sem_id][ + ...].reshape( + scatter_staging_3x_ref.shape[1], + scatter_staging_3x_ref.shape[2], + scatter_staging_3x_ref.shape[3], + ) + + # 7. Per-row ICI send — ALL rows (no is_last gating). + @jax.named_scope("direct_write_rows") + def _do_direct_write(): + + def _write_row(i, carry): + send_sz, local_sz = carry + row_idx = _ml + i + token_g = output_indices_ref[m_st + i] + dest_chip = token_g // chunk_size + local_row = token_g % chunk_size + write_pos = local_row + is_local = dest_chip == my_id + + @pl.when(~is_local) + def _(): + pltpu.make_async_remote_copy( + src_ref=scatter_staging_3x_ref.at[ + stg_id, pl.ds(row_idx, 1), :, :], + dst_ref=out_buf_ref.at[pl.ds(write_pos, 1), :, :], + send_sem=send_sems_ref.at[stg_id], + recv_sem=recv_sem_ref.at[0], + device_id={ + ep_axis_name: dest_chip + }, + device_id_type=pltpu.DeviceIdType.MESH, + ).start() + + @pl.when(is_local) + def _(): + pltpu.make_async_copy( + src_ref=scatter_staging_3x_ref.at[ + stg_id, pl.ds(row_idx, 1), :, :], + dst_ref=out_buf_ref.at[pl.ds(write_pos, 1), :, :], + sem=local_write_sems_ref.at[stg_id], + ).start() + + return ( + send_sz + + lax.select(~is_local, jnp.int32(1), jnp.int32(0)), + local_sz + + lax.select(is_local, jnp.int32(1), jnp.int32(0)), + ) + + return lax.fori_loop(0, num_valid, _write_row, + (jnp.int32(0), jnp.int32(0))) + + send_sz, local_sz = _do_direct_write() + count_ref[1 + stg_id] = count_ref[1 + stg_id] + send_sz + count_ref[4 + stg_id] = count_ref[4 + stg_id] + local_sz + + return _ + + lax.fori_loop(0, num_gm, gm_loop_body, None) + + # Handle num_gm == 0 edge case for deferred zero-init. + if cfgs.zero_init: + + @pl.when(num_gm == 0) + def _(): + zero_out_end_3d(out_buf_ref, semaphore_ref, zero_size_out) + + # Epilogue: drain remaining per-slot DMA counts. + @jax.named_scope("epilogue_drain") + @pl.when(num_gm > 0) + def _(): + for _slot in range(3): + + @pl.when(num_gm > _slot) + def _(): + remaining_send = count_ref[1 + _slot] + pltpu.make_async_copy( + src_ref=scatter_staging_3x_ref.at[ + _slot, pl.ds(0, remaining_send), :, :], + dst_ref=scatter_staging_3x_ref.at[ + _slot, pl.ds(0, remaining_send), :, :], + sem=send_sems_ref.at[_slot], + ).wait() + remaining_local = count_ref[4 + _slot] + pltpu.make_async_copy( + src_ref=out_buf_ref.at[pl.ds(0, remaining_local), :, :], + dst_ref=out_buf_ref.at[pl.ds(0, remaining_local), :, :], + sem=local_write_sems_ref.at[_slot], + ).wait() + + # Final barrier. + sync_barrier() + + # Drain recv_sem. + with jax.named_scope("drain_recv_sem"): + total_recv = total_recv_count_ref[0] + pltpu.make_async_copy( + src_ref=out_buf_ref.at[pl.ds(0, total_recv), :, :], + dst_ref=out_buf_ref.at[pl.ds(0, total_recv), :, :], + sem=recv_sem_ref.at[0], + ).wait() + + +@jax.jit(static_argnames=[ + "fuse_act", + "ep_size", + "ep_axis_name", + "chunk_size", + "vmem_limit_bytes", + "maybe_quantize_lhs", +], ) +def gmm_v2_scatter_ici_nodedup( + lhs: jax.Array, # [size_m, size_k] + rhs: jax.Array, # [size_group, size_k, size_n] + group_sizes: jax.Array, # int32[size_lhs_group] + output_indices: jax.Array, # int32[size_m] — global scatter positions + total_recv_count: jax.Array, # int32[1] — remote rows this chip receives + *, + rhs_scale: jax.Array | None = None, + rhs_bias: jax.Array | None = None, + group_offset: jax.Array | None = None, + fuse_act: str | None = None, + ep_size: int, + ep_axis_name: str, + chunk_size: int, + vmem_limit_bytes: int | None = None, + maybe_quantize_lhs: bool = True, +) -> jax.Array: + """GMM with ICI direct-write scatter (nodedup — ALL rows sent). + + Every GMM output row is ICI-sent to the destination chip. No accumulator, + no is_last gating, no topk_weights. + Output: (chunk_size, size_n) — this chip's portion, no post-kernel reduction. + + Must be called inside shard_map with ep_axis_name mesh axis. + + Args: + lhs: LHS matrix [size_m, size_k]. + rhs: RHS matrix [size_group, size_k, size_n]. + group_sizes: Group sizes [size_lhs_group]. + output_indices: int32[size_m] — global scatter positions. + dest_chip = oi // chunk_size, local_row = oi % chunk_size. + total_recv_count: int32[1] — number of remote rows this chip receives. + rhs_scale: Optional per-block scale. + rhs_bias: Optional bias. + group_offset: Optional group offset. + fuse_act: Optional activation fusion. + ep_size: Number of expert-parallel chips. + ep_axis_name: Mesh axis name for expert parallelism. + chunk_size: Tokens per chip (output_size // ep_size). + vmem_limit_bytes: Optional VMEM limit. + maybe_quantize_lhs: Whether to quantize LHS. + + Returns: + Output of shape [chunk_size, size_n] (this chip's portion). + """ + if group_offset is None: + group_offset = jnp.array([0], dtype=jnp.int32) + + if vmem_limit_bytes is None: + vmem_limit_bytes = int(pltpu.get_tpu_info().vmem_capacity_bytes) + + num_lanes = pltpu.get_tpu_info().num_lanes + + # Check fuse_act alignment. + _fallback_fuse_act = None + if fuse_act is not None: + if rhs.shape[2] % (2 * num_lanes) != 0: + _fallback_fuse_act = fuse_act + fuse_act = None + + cfgs = make_gmm_configs( + lhs, + rhs, + rhs_scale, + rhs_bias, + group_sizes, + group_offset, + tile_info=calculate_tiling, + vmem_limit_bytes=vmem_limit_bytes, + out_dtype=None, + acc_dtype=None, + maybe_quantize_lhs=maybe_quantize_lhs, + zero_initialize=True, + output_indices=output_indices, + fuse_act=fuse_act, + ) + dims = cfgs.dims + tiles = cfgs.tiles + + # Align N for DMA scatter. + _out_sls = pltpu.get_tpu_info().get_sublane_tiling(cfgs.out_dtype) + aligned_n = align_to(dims.size_n, num_lanes * _out_sls) + + # Ensure tile_n divides aligned_n. + if aligned_n % tiles.tile_n != 0: + _mxu_size = pltpu.get_tpu_info().mxu_column_size + _num_n = pl.cdiv(aligned_n, tiles.tile_n) + _adj_tile_n = (aligned_n // _num_n // _mxu_size) * _mxu_size + if _adj_tile_n > 0: + tiles = dataclasses.replace(tiles, tile_n=_adj_tile_n) + cfgs = dataclasses.replace(cfgs, tiles=tiles) + + n_cols = aligned_n // num_lanes + + # Weight buffer config. + is_quantized = cfgs.rhs_cfgs.quant_dtype is not None + packing = cfgs.rhs_cfgs.packing + pk = tiles.tile_k // packing + has_scale = cfgs.rhs_cfgs.has_scale + nqb = cfgs.num_quant_blocks_per_tile_k if has_scale else 1 + w_dma_n = tiles.tile_n * 2 if cfgs.fuse_act else tiles.tile_n + num_w_bufs = 2 # double-buffered weights + + # Scratch shapes. + max_num_gm = dims.size_group + pl.cdiv(dims.size_m, tiles.tile_m) - 1 + acc_cols = 2 * tiles.tile_n if cfgs.fuse_act is not None else tiles.tile_n + + scratch_shapes = [ + # metadata_ref + MetadataRef( + gm_id_to_group_id=pltpu.SMEM((max_num_gm, ), jnp.int32), + gm_id_to_m_offset=pltpu.SMEM((max_num_gm + 1, ), jnp.int32), + ), + # fused_metadata_ref (single-slot for inner_kernel) + MetadataRef( + gm_id_to_group_id=pltpu.SMEM((1, ), jnp.int32), + gm_id_to_m_offset=pltpu.SMEM((2, ), jnp.int32), + ), + # tiled_out_2x_ref — double-buffered compute buffer + pltpu.VMEM((2, tiles.tile_m, aligned_n), cfgs.out_dtype), + # scatter_staging_3x_ref — triple-buffered staging + pltpu.VMEM((3, tiles.tile_m, n_cols, num_lanes), cfgs.out_dtype), + # partial_out_ref (dummy for scatter_mode inner_kernel) + pltpu.VMEM((dims.size_lhs_sublane, tiles.tile_n), cfgs.out_dtype), + # acc_ref + pltpu.VMEM((tiles.tile_m, acc_cols), cfgs.acc_dtype), + # count_ref — SMEM: [0]=gm_id, [1..3]=send counts, [4..6]=local counts + pltpu.SMEM((7, ), jnp.int32), + # semaphore_ref for out_buf zero_out + pltpu.SemaphoreType.DMA((1, )), + # send_sems_ref — triple-buffered + pltpu.SemaphoreType.DMA((3, )), + # local_write_sems_ref — triple-buffered + pltpu.SemaphoreType.DMA((3, )), + # recv_sem_ref — for incoming remote writes + pltpu.SemaphoreType.DMA((1, )), + # w_buf_ref — weight buffer + pltpu.VMEM( + (num_w_bufs, pk, w_dma_n), + jnp.uint32 if is_quantized else rhs.dtype, + ), + # w_scale_buf_ref + pltpu.VMEM( + (num_w_bufs, nqb, 1, w_dma_n), + jnp.float32, + ), + # w_bias_buf_ref + pltpu.VMEM((1, w_dma_n), jnp.float32), + # w_sem_ref — weight DMA semaphores + pltpu.SemaphoreType.DMA((num_w_bufs, )), + # gathered_lhs_2x_ref — double-buffered LHS only (K//NL cols) + pltpu.VMEM( + (2, tiles.tile_m, dims.size_k // num_lanes, num_lanes), + cfgs.out_dtype, + ), + # gather_sem_ref — LHS gather semaphores + pltpu.SemaphoreType.DMA((2, )), + ] + + # Output shape: out_buf only (no accumulator). + out_buf_init = jax.ShapeDtypeStruct((chunk_size, n_cols, num_lanes), + cfgs.out_dtype) + + # Prepare RHS specs. + rhs_scale_spec = rhs_bias_spec = None + if rhs_scale is not None: + rhs_scale = rhs_scale.astype(jnp.float32) + rhs_scale_spec = pl.BlockSpec(memory_space=pltpu.HBM) + if rhs_bias is not None: + rhs_bias = rhs_bias.astype(jnp.float32) + rhs_bias_spec = pl.BlockSpec(memory_space=pltpu.HBM) + + rhs_weights = WeightsRef(weight=rhs, scale=rhs_scale, bias=rhs_bias) + rhs_in_spec = WeightsRef( + weight=pl.BlockSpec(memory_space=pltpu.HBM), + scale=rhs_scale_spec, + bias=rhs_bias_spec, + ) + + compiler_params = pltpu.CompilerParams( + vmem_limit_bytes=vmem_limit_bytes, + disable_bounds_checks=True, + collective_id=0, + ) + + total_recv_arr = total_recv_count.astype(jnp.int32).reshape((1, )) + + # Reshape LHS to 3D for contiguous DMA inside kernel. + lhs_3d = lhs.reshape(dims.size_m, dims.size_k // num_lanes, num_lanes) + + out_buf = pl.pallas_call( + functools.partial( + kernel_main_scatter_ici_nodedup, + cfgs=cfgs, + ep_size=ep_size, + chunk_size=chunk_size, + ep_axis_name=ep_axis_name, + num_w_bufs=num_w_bufs, + ), + out_shape=out_buf_init, + grid_spec=pltpu.PrefetchScalarGridSpec( + num_scalar_prefetch=4, + in_specs=[ + pl.BlockSpec(memory_space=pltpu.HBM), # lhs_3d + rhs_in_spec, + ], + out_specs=pl.BlockSpec(memory_space=pltpu.HBM), # out_buf + scratch_shapes=scratch_shapes, + ), + compiler_params=compiler_params, + name=get_scope_name(dims, tiles) + "-scatter_ici_nodedup", + cost_estimate=get_cost_estimate(cfgs), + )(group_sizes, group_offset, output_indices, total_recv_arr, lhs_3d, + rhs_weights) + + result = out_buf.reshape(chunk_size, aligned_n)[:, :cfgs.out_size_n] + + if _fallback_fuse_act is not None: + result = apply_act_fn(result, _fallback_fuse_act) + + return result diff --git a/tpu_inference/kernels/experimental/fused_moe/gmm_fused/per_round_schedule.py b/tpu_inference/kernels/experimental/fused_moe/gmm_fused/per_round_schedule.py new file mode 100644 index 0000000000..ff58425c32 --- /dev/null +++ b/tpu_inference/kernels/experimental/fused_moe/gmm_fused/per_round_schedule.py @@ -0,0 +1,351 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# 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 +# +# http://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. +"""Per-round push schedule for AG+GMM1 EP MoE kernel. + +Computes, for each chip's local view, a static send schedule that drives a +push-based all-to-all interleaved with the GMM1 compute pipeline. + +Each round `r` corresponds to one gm tile on the receiver. In round `r`, every +sender pushes a (possibly variable) number of rows to each receiver. The +invariant is that the total number of rows arriving at any single receiver in +round `r` equals exactly that receiver's gm tile size (≤ tile_m). + +The schedule is computed via a one-time `jax.lax.all_gather` of `lhs_indices`, +`group_sizes`, and `group_offset` (all tiny int32 arrays). +""" + +from __future__ import annotations + +import jax +import jax.numpy as jnp +from jax import lax + +# ============================================================================= +# JAX reimplementation of fill_metadata (per chip). +# ============================================================================= +# Mirrors the SMEM-mutating loop in gmm_v2_gather_scatter.py but produces JAX +# arrays so we can reason about every chip's gm tiling from outside the kernel. + + +def _fill_metadata_one_chip( + group_sizes: jax.Array, # int32[size_lhs_group] + group_offset: jax.Array, # int32 scalar — may differ across chips + *, + size_group: int, # static + tile_m: int, # static + size_lhs_sublane: int, # static + max_num_gm: int, # static upper bound +): + """Compute (gm_to_m_offset, gm_to_group_id, num_gm) for one chip.""" + max_num_group = group_offset + size_group + init_off = jnp.zeros((max_num_gm + 1, ), dtype=jnp.int32) + init_gid = jnp.zeros((max_num_gm, ), dtype=jnp.int32) + + def outer_group_loop(lhs_group_id, carry): + num_gm, start_m_offset, off_arr, gid_arr = carry + group_id = lhs_group_id - group_offset + group_size = group_sizes[lhs_group_id] + end_m_offset = start_m_offset + group_size + local_offset_outer = start_m_offset % size_lhs_sublane + aligned_group_size = group_size + local_offset_outer + curr_num_gm = (aligned_group_size + tile_m - 1) // tile_m + should_process = jnp.logical_and(group_size > 0, group_id >= 0) + curr_num_gm = jnp.where(should_process, curr_num_gm, 0) + next_num_gm = num_gm + curr_num_gm + + def inner_tm_loop(tm_id, inner_carry): + curr_m_offset, off_in, gid_in = inner_carry + local_off = curr_m_offset % size_lhs_sublane + tm_size = jnp.minimum(tile_m - local_off, + end_m_offset - curr_m_offset) + gid_in = gid_in.at[tm_id].set(group_id) + next_m_offset = curr_m_offset + tm_size + off_in = off_in.at[tm_id].set(curr_m_offset) + off_in = off_in.at[tm_id + 1].set(next_m_offset) + return (next_m_offset, off_in, gid_in) + + _, off_arr, gid_arr = lax.fori_loop( + num_gm, + next_num_gm, + inner_tm_loop, + (start_m_offset, off_arr, gid_arr), + ) + return (next_num_gm, end_m_offset, off_arr, gid_arr) + + num_gm, _, off_arr, gid_arr = lax.fori_loop( + 0, + max_num_group, + outer_group_loop, + (jnp.int32(0), jnp.int32(0), init_off, init_gid), + ) + return off_arr, gid_arr, num_gm + + +def _fill_metadata_all_chips( + all_group_sizes: jax.Array, # int32[ep_size, size_lhs_group] + all_group_offsets: jax.Array, # int32[ep_size] + *, + size_group: int, + tile_m: int, + size_lhs_sublane: int, + max_num_gm: int, +): + """Vectorize _fill_metadata_one_chip over the EP axis.""" + return jax.vmap(lambda gs, go: _fill_metadata_one_chip( + gs, + go, + size_group=size_group, + tile_m=tile_m, + size_lhs_sublane=size_lhs_sublane, + max_num_gm=max_num_gm, + ))(all_group_sizes, all_group_offsets) + + +# ============================================================================= +# Per-round send schedule. +# ============================================================================= + + +def compute_max_num_gm(size_m: int, size_group: int, tile_m: int) -> int: + """Static upper bound on num_gm — same formula used everywhere.""" + return size_group + (size_m + tile_m - 1) // tile_m - 1 + + +def compute_per_round_send_schedule( + lhs_indices: jax.Array, # int32[size_m] — local chip's indices + group_sizes: jax. + Array, # int32[size_lhs_group] — local chip's group sizes + group_offset: jax. + Array, # int32[1] — local chip's group offset (may differ per chip) + *, + ep_axis_name: str, + ep_size: int, + chunk_size: int, + tile_m: int, + size_group: int, # static — same on every chip + size_lhs_sublane: int, # static +): + """Build the per-round send schedule for a push-based AG+GMM1 pipeline. + + Returns: + send_count[c, r]: int32 — rows this chip pushes to chip c in round r. + send_local_off[c, r, k]: int32 — local row offset (in [0, chunk_size)). + send_dest_pos[c, r, k]: int32 — destination row position within chip c's + round-r staging slot (= m_start_local + p). + my_num_gm: int32[1] — actual num_gm for this chip. + max_num_gm_static: Python int — static upper bound. + """ + size_m = int(lhs_indices.shape[0]) + max_num_gm = compute_max_num_gm(size_m, size_group, tile_m) + + my_id = lax.axis_index(ep_axis_name) + + # All-gather routing tables across the EP axis (small int32 arrays). + all_lhs_indices = lax.all_gather(lhs_indices, + ep_axis_name) # (ep_size, size_m) + all_group_sizes = lax.all_gather(group_sizes, + ep_axis_name) # (ep_size, size_lhs_group) + all_group_offsets_2d = lax.all_gather(group_offset, + ep_axis_name) # (ep_size, 1) + all_group_offsets = all_group_offsets_2d[:, 0] # (ep_size,) + + # Per-chip gm tiling. + gm_to_m_offset, _gm_to_group_id, num_gm_per_chip = _fill_metadata_all_chips( + all_group_sizes, + all_group_offsets, + size_group=size_group, + tile_m=tile_m, + size_lhs_sublane=size_lhs_sublane, + max_num_gm=max_num_gm, + ) + + p_arange = jnp.arange(tile_m, dtype=jnp.int32) + starts = gm_to_m_offset[:, :max_num_gm] # (ep_size, max_num_gm) + ends = gm_to_m_offset[:, 1:max_num_gm + 1] # (ep_size, max_num_gm) + m_offsets_3d = starts[..., None] + p_arange[None, None, :] + valid_mask = m_offsets_3d < ends[..., + None] # (ep_size, max_num_gm, tile_m) + + safe_m = jnp.minimum(m_offsets_3d, size_m - 1) + global_ids = jax.vmap(lambda idx_row, sm: idx_row[sm])(all_lhs_indices, + safe_m) + + source_chip = global_ids // chunk_size + local_offset = global_ids % chunk_size + is_mine = jnp.logical_and(source_chip == my_id, valid_mask) + + # Stable partition: place "mine" entries first. + sort_key = jnp.where( + is_mine, + p_arange[None, None, :], + tile_m + p_arange[None, None, :], + ) + sort_perm = jnp.argsort(sort_key, axis=-1) # (ep_size, max_num_gm, tile_m) + + # Receiver staging convention mirrors `dma_gather_gm_start` + # (gmm_v2_gather_scatter.py:840): tile rows occupy + # [m_start_local, m_start_local + expected_rows) within the slot. + m_start_local = (starts % size_lhs_sublane).astype(jnp.int32) + p_full = jnp.broadcast_to(p_arange[None, None, :], + (ep_size, max_num_gm, tile_m)) + dest_pos_unsorted = m_start_local[..., None] + p_full + + send_local_off = jnp.take_along_axis(local_offset, sort_perm, axis=-1) + send_dest_pos = jnp.take_along_axis(dest_pos_unsorted, sort_perm, axis=-1) + send_count = jnp.sum(is_mine.astype(jnp.int32), axis=-1) + + my_num_gm = num_gm_per_chip[my_id].reshape((1, )).astype(jnp.int32) + + return ( + send_count.astype(jnp.int32), + send_local_off.astype(jnp.int32), + send_dest_pos.astype(jnp.int32), + my_num_gm, + max_num_gm, + ) + + +# ============================================================================= +# Per-(sender, target, token) dedup schedule with persistent HBM cache. +# ============================================================================= + + +def compute_dedup_send_schedule( + lhs_indices: jax.Array, # int32[size_m] — local chip's indices + group_sizes: jax.Array, # int32[size_lhs_group] — local chip's group sizes + group_offset: jax.Array, # int32[1] — local chip's group offset + *, + ep_axis_name: str, + ep_size: int, + chunk_size: int, + tile_m: int, + size_group: int, + size_lhs_sublane: int, +): + """Build a dedup-aware send schedule. + + For each (sender s, target c, global_token_id g), `g` is sent over ICI + AT MOST ONCE — at the earliest round r where c needs g. The receiver + keeps a persistent HBM cache of shape `(num_tokens_total, K1/NL, NL)` + indexed directly by `global_token_id`. Subsequent rounds that need the + same token re-read from the cache (zero ICI traffic). + + Returns: + send_count_new[c, r]: int32 — NEW (first-use) tokens this chip + pushes to chip c in round r. + send_local_off_new[c, r, k]: int32 — local row to read (in + [0, chunk_size)). + send_global_id_new[c, r, k]: int32 — destination cache slot in chip + c's HBM cache (= global token id). + recv_count_new[r]: int32 — total NEW token arrivals at + THIS chip in round r (used for the + receiver-side recv-sem wait). + my_num_gm: int32[1] — actual num_gm for this chip. + max_num_gm_static: Python int — static upper bound. + """ + size_m = int(lhs_indices.shape[0]) + num_tokens_total = chunk_size * ep_size + max_num_gm = compute_max_num_gm(size_m, size_group, tile_m) + + my_id = lax.axis_index(ep_axis_name) + + all_lhs_indices = lax.all_gather(lhs_indices, ep_axis_name) + all_group_sizes = lax.all_gather(group_sizes, ep_axis_name) + all_group_offsets_2d = lax.all_gather(group_offset, ep_axis_name) + all_group_offsets = all_group_offsets_2d[:, 0] + + gm_to_m_offset, _gm_to_group_id, num_gm_per_chip = _fill_metadata_all_chips( + all_group_sizes, + all_group_offsets, + size_group=size_group, + tile_m=tile_m, + size_lhs_sublane=size_lhs_sublane, + max_num_gm=max_num_gm, + ) + + p_arange = jnp.arange(tile_m, dtype=jnp.int32) + r_arange = jnp.arange(max_num_gm, dtype=jnp.int32) + starts = gm_to_m_offset[:, :max_num_gm] + ends = gm_to_m_offset[:, 1:max_num_gm + 1] + m_offsets_3d = starts[..., None] + p_arange[None, None, :] + valid_mask = m_offsets_3d < ends[..., + None] # (ep_size, max_num_gm, tile_m) + + safe_m = jnp.minimum(m_offsets_3d, size_m - 1) + global_ids = jax.vmap(lambda idx_row, sm: idx_row[sm])(all_lhs_indices, + safe_m) + # (ep_size, max_num_gm, tile_m) + + source_chip = global_ids // chunk_size + + # ---- First-use determination ---- + # Encode (r, p) as a single ordered rank = r * tile_m + p. For each + # (c, global_id), the first-use is the smallest rank over all valid + # (r, p) with global_ids[c, r, p] == global_id. Computed by + # per-chip scatter-min into a (num_tokens_total,) array. + sentinel = jnp.iinfo(jnp.int32).max + rank_3d = (r_arange[None, :, None] * tile_m + + p_arange[None, None, :]).astype(jnp.int32) + rank_3d = jnp.broadcast_to(rank_3d, (ep_size, max_num_gm, tile_m)) + masked_rank = jnp.where(valid_mask, rank_3d, sentinel) + + flat_global = global_ids.reshape(ep_size, max_num_gm * tile_m) + flat_rank = masked_rank.reshape(ep_size, max_num_gm * tile_m) + + init_min = jnp.full((ep_size, num_tokens_total), sentinel, dtype=jnp.int32) + + def _scatter_min_one_chip(init_row, idx, vals): + # Use .min() to keep the smallest rank per token id. + return init_row.at[idx].min(vals) + + min_rank_per_token = jax.vmap(_scatter_min_one_chip)( + init_min, flat_global, flat_rank) # (ep_size, num_tokens_total) + + # Look up min_rank for each (c, r, p)'s global_id and compare to its rank. + gathered_min = jax.vmap(lambda mr, g: mr[g])( + min_rank_per_token, + global_ids.reshape(ep_size, -1)).reshape(ep_size, max_num_gm, tile_m) + is_first_use = jnp.logical_and(masked_rank == gathered_min, valid_mask) + # (ep_size, max_num_gm, tile_m) + + # ---- Sender side: my chip pushes only first-use entries it owns ---- + is_my_send = jnp.logical_and(is_first_use, source_chip == my_id) + + sort_key = jnp.where( + is_my_send, + p_arange[None, None, :], + tile_m + p_arange[None, None, :], + ) + sort_perm = jnp.argsort(sort_key, axis=-1) + + local_offset = global_ids % chunk_size + send_local_off_new = jnp.take_along_axis(local_offset, sort_perm, axis=-1) + send_global_id_new = jnp.take_along_axis(global_ids, sort_perm, axis=-1) + send_count_new = jnp.sum(is_my_send.astype(jnp.int32), axis=-1) + + # ---- Receiver side: total new arrivals per round on this chip ---- + my_first_use = is_first_use[my_id] # (max_num_gm, tile_m) + recv_count_new = jnp.sum(my_first_use.astype(jnp.int32), axis=-1) + # (max_num_gm,) + + my_num_gm = num_gm_per_chip[my_id].reshape((1, )).astype(jnp.int32) + + return ( + send_count_new.astype(jnp.int32), + send_local_off_new.astype(jnp.int32), + send_global_id_new.astype(jnp.int32), + recv_count_new.astype(jnp.int32), + my_num_gm, + max_num_gm, + ) diff --git a/tpu_inference/kernels/experimental/fused_moe/gmm_fused_rs_nodedup.py b/tpu_inference/kernels/experimental/fused_moe/gmm_fused_rs_nodedup.py new file mode 100644 index 0000000000..a82c9151d8 --- /dev/null +++ b/tpu_inference/kernels/experimental/fused_moe/gmm_fused_rs_nodedup.py @@ -0,0 +1,2044 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# 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 +# +# http://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. +"""Fused GMM kernel with ICI direct-write reduce-scatter. + +Each chip ICI-sends its GMM2 output directly to the correct (token, topk_index) +position in the destination chip's output buffer. No recv buffer, no scatter-add, +no topk weighting inside the kernel. The reduction (topk_weight × sum) happens +as a trivial post-kernel JAX operation. + +Key design principles: +1. All writes via DMA engine (direct to final position) +2. Static Python loop unrolling for routing +3. All routing pre-computed at JAX level +4. No recv buffer, no scatter-add, no periodic barriers +5. Post-kernel weighted reduction in JAX +""" + +import dataclasses +import functools + +import jax +import jax.experimental.pallas as pl +from jax import lax +from jax import numpy as jnp +from jax.experimental.pallas import tpu as pltpu + +# isort: off +# yapf: disable +from .gmm_v2_gather_scatter import ( + Dimensions, FusedDims, FusedWeightsRef, GmmConfigs, InputConfigs, + MetadataRef, TileSizes, WeightsRef, _recover_quant_block_size, align_to, + calculate_tiling, dma_gather_gm_start, dma_gather_gm_wait, fill_metadata, + get_maybe_quantize_lhs, inner_kernel, zero_out_end_3d, zero_out_start_3d) +# yapf: enable +# isort: on + + +def get_fused_rs_tuned_block_sizes( + m, + k1, + n1, + k2, + n2, + num_current_groups, + lhs_dtype, + rhs_dtype, + rhs_quant_block_size, + default_block_sizes, + fuse_act=None, + fp8_direct_write=False, +): + # Device-specific tuned tables were removed; return the default (with the + # fp8 direct-write tile_m clamp for VMEM safety). + result = default_block_sizes + if fp8_direct_write and result[0] > 64: + result = (64, ) + tuple(result[1:]) + return result + + +# ============================================================================= +# Phase 1: Pre-kernel metadata computation (JAX level) +# ============================================================================= + + +def compute_num_gm(group_sizes, tile_m, size_lhs_sublane): + """Compute number of gm tiles from group sizes (JAX level). + + Mirrors fill_metadata logic including sublane alignment offset. + """ + + def _per_group(carry, group_size): + num_gm, m_offset = carry + local_offset = m_offset % size_lhs_sublane + aligned = group_size + local_offset + curr_gm = jnp.where(group_size > 0, pl.cdiv(aligned, tile_m), 0) + return (num_gm + curr_gm, m_offset + group_size), None + + (total_gm, _), _ = lax.scan(_per_group, (jnp.int32(0), jnp.int32(0)), + group_sizes) + return total_gm + + +def compute_send_routing(output_indices, chunk_size): + """Pre-compute per-row routing: dest_chip and local_row for every slot. + + Args: + output_indices: int32[size_m] — scatter indices (global token IDs) + chunk_size: int — num_tokens // ep_size (tokens per chip) + + Returns: + send_dest_chips: int32[size_m] — which chip each row goes to + send_local_rows: int32[size_m] — token position within dest chip's shard + """ + send_dest_chips = output_indices // chunk_size + send_local_rows = output_indices % chunk_size + return send_dest_chips, send_local_rows + + +@dataclasses.dataclass(frozen=True) +class FusedRsBlockSizes: + """Immutable set of block sizes + post-pad dimensions for fused_rs. + + All fields are static Python ints picked before any JAX tracing; they + are safe to use as ``@jax.jit`` static args or in shape computations. + """ + + tile_m: int + tile_k1: int + tile_n1: int + tile_k2: int + tile_n2: int + num_w1_bufs: int + num_w2_bufs: int + # Padded / aligned dimensions (needed by both sites for shape math). + padded_k1: int + aligned_n1: int # new size_n1 after 2*num_lanes alignment + padded_k2: int # new size_k2 after aligned_n1 bump (== aligned_n1 // 2) + aligned_n2: int # new size_n2 after num_lanes*out_sls alignment + + +def _select_fused_rs_block_sizes( + *, + size_m: int, + size_k1: int, + size_n1: int, + size_k2: int, + size_n2: int, + size_group: int, + size_lhs_group: int, + ep_size: int, + out_dtype: jnp.dtype, + w1_dtype: jnp.dtype, + w2_dtype: jnp.dtype, + is_quantized: bool, + quant_block_size: int | None, + act_fn: str | None, + vmem_limit_bytes: int | None = None, + fp8_direct_write: bool = False, +) -> FusedRsBlockSizes: + """Deterministic selection of all tile/block sizes for fused_rs. + + Pure Python — no JAX ops, no random state. Same inputs always return + the same ``FusedRsBlockSizes``. Safe to call from multiple entry + points (``run_gmm_fused_rs`` and ``gmm_v2_fused_rs``) for the same + shape/dtype combination. + + Pipeline: + 1. Pad K1 / N1 / N2 to DMA and MXU alignment, updating K2 when N1 + changes (since GMM2's K is GMM1's intermediate size). + 2. VMEM budget accounting → ``fused_vmem`` for ``calculate_tiling``. + 3. ``calculate_tiling`` for GMM1 and GMM2 to produce *default* tiles. + 4. ``get_fused_rs_tuned_block_sizes`` dict lookup (overrides default). + 5. Alignment post-processing of tile_k1, tile_n1, tile_k2, tile_n2 + so they exactly divide the padded / original dims the kernel + iterates over. + 6. Final divisibility assertions. + """ + if vmem_limit_bytes is None: + vmem_limit_bytes = int(pltpu.get_tpu_info().vmem_capacity_bytes * 0.95) + del ep_size # Reserved for future per-chip heuristics; unused here. + # --- Step 1: Pad dimensions --- + num_lanes = pltpu.get_tpu_info().num_lanes + sls = pltpu.get_tpu_info().get_sublane_tiling(out_dtype) + + tile_k_unit = num_lanes * sls + padded_k1 = align_to(size_k1, tile_k_unit) + + n1_unit = 2 * num_lanes + aligned_n1 = align_to(size_n1, n1_unit) + if aligned_n1 != size_n1: + new_k2 = aligned_n1 // 2 + else: + new_k2 = size_k2 + padded_k2 = new_k2 + + out_sls = pltpu.get_tpu_info().get_sublane_tiling(out_dtype) + aligned_n2 = align_to(size_n2, num_lanes * out_sls) + + # --- Step 2: VMEM budget --- + fixed_vmem = ( + 2 * 256 * size_k1 * jax.dtypes.itemsize_bits(out_dtype) // 8 + + 2 * 256 * aligned_n2 * jax.dtypes.itemsize_bits(out_dtype) // 8 + 256 * + (aligned_n1 // 2 if act_fn else aligned_n1) * 4 + + 256 * padded_k2 * jax.dtypes.itemsize_bits(out_dtype) // 8 + 256 * + (aligned_n2 // num_lanes) * num_lanes * + jax.dtypes.itemsize_bits(out_dtype) // 8) + spill_budget = 6 * 1024 * 1024 + available_vmem = vmem_limit_bytes - fixed_vmem - spill_budget + fused_vmem = max(available_vmem // 2, 1024 * 1024) + + size_lhs_sublane = min(sls, size_m) + + # --- Step 3: calculate_tiling for defaults --- + dims1 = Dimensions( + size_m=size_m, + size_k=size_k1, + size_n=aligned_n1, + size_group=size_group, + size_lhs_group=size_lhs_group, + size_lhs_sublane=size_lhs_sublane, + ) + rhs1_cfgs = InputConfigs( + quant_dtype=w1_dtype if is_quantized else None, + quant_block_size=quant_block_size if quant_block_size else size_k1, + dtype=w1_dtype, + ) + tiles1 = calculate_tiling( + dims1, + InputConfigs(quant_dtype=None, + quant_block_size=size_k1, + dtype=out_dtype), + rhs1_cfgs, + fused_vmem, + ) + + dims2 = Dimensions( + size_m=size_m, + size_k=padded_k2, + size_n=aligned_n2, + size_group=size_group, + size_lhs_group=size_lhs_group, + size_lhs_sublane=size_lhs_sublane, + ) + rhs2_cfgs = InputConfigs( + quant_dtype=w2_dtype if is_quantized else None, + quant_block_size=quant_block_size if quant_block_size else padded_k2, + dtype=w2_dtype, + ) + tiles2 = calculate_tiling( + dims2, + InputConfigs(quant_dtype=None, + quant_block_size=padded_k2, + dtype=out_dtype), + rhs2_cfgs, + fused_vmem, + ) + + default_tiles = ( + tiles1.tile_m, + tiles1.tile_k, + tiles1.tile_n, + tiles2.tile_k, + tiles2.tile_n, + 2, + 2, + ) + + # --- Step 4: Tuned lookup --- + tile_m, tile_k1, tile_n1, tile_k2, tile_n2, num_w1_bufs, num_w2_bufs = ( + get_fused_rs_tuned_block_sizes( + size_m, + size_k1, + aligned_n1, + padded_k2, + aligned_n2, + size_group, + out_dtype, + w1_dtype, + quant_block_size if quant_block_size else None, + default_tiles, + fuse_act=act_fn, + fp8_direct_write=fp8_direct_write, + )) + + # --- Step 5: Alignment post-processing --- + k_align1 = num_lanes + if tile_k1 % k_align1 != 0: + tile_k1 = (tile_k1 // k_align1) * k_align1 or k_align1 + if size_k1 % tile_k1 != 0: + tile_k1 = size_k1 + + out_n1 = aligned_n1 // 2 if act_fn else aligned_n1 + if out_n1 % tile_n1 != 0: + while tile_n1 > num_lanes and out_n1 % tile_n1 != 0: + tile_n1 -= num_lanes + + # Adjust tile_n2 to divide size_n2 (not aligned_n2). The kernel + # iterates num_n2 = original_n2 // tile_n2 real N-tiles only. + if size_n2 % tile_n2 != 0: + mxu_cols = pltpu.get_tpu_info().mxu_column_size + nn2 = pl.cdiv(size_n2, tile_n2) + adj_tn2 = (size_n2 // nn2 // mxu_cols) * mxu_cols + if adj_tn2 > 0: + tile_n2 = adj_tn2 + if size_n2 % tile_n2 != 0: + while tile_n2 > num_lanes and size_n2 % tile_n2 != 0: + tile_n2 -= num_lanes + + if tile_k2 % num_lanes != 0: + tile_k2 = (tile_k2 // num_lanes) * num_lanes or num_lanes + if padded_k2 % tile_k2 != 0: + tk2_adj = tile_k2 + while tk2_adj > num_lanes and padded_k2 % tk2_adj != 0: + tk2_adj -= num_lanes + tile_k2 = tk2_adj if (tk2_adj > 0 + and padded_k2 % tk2_adj == 0) else padded_k2 + + # NOTE: fp8 direct-write tile_m sizing is now handled in Step 4 + # (get_fused_rs_tuned_block_sizes, called with fp8_direct_write): dedicated + # fp8-comm entries carry VMEM-safe tiles (e.g. tile_m=96 full-N), and shapes + # without one fall back to the bf16 tile clamped to tile_m<=64. Both the host + # max-gm calc and the kernel call this helper, so their loop bounds stay + # aligned. (Replaces the former unconditional ``tile_m = min(tile_m, 64)``.) + + # --- Step 6: Assertions --- + assert size_k1 % tile_k1 == 0, f"tile_k1={tile_k1} must divide size_k1={size_k1}" + assert ( + padded_k2 % + tile_k2 == 0), f"tile_k2={tile_k2} must divide padded_k2={padded_k2}" + assert ( + aligned_n1 % + tile_n1 == 0), f"tile_n1={tile_n1} must divide aligned_n1={aligned_n1}" + # tile_n2 must divide original size_n2 (kernel only iterates over real N-tiles). + assert size_n2 % tile_n2 == 0, f"tile_n2={tile_n2} must divide size_n2={size_n2}" + + return FusedRsBlockSizes( + tile_m=tile_m, + tile_k1=tile_k1, + tile_n1=tile_n1, + tile_k2=tile_k2, + tile_n2=tile_n2, + num_w1_bufs=num_w1_bufs, + num_w2_bufs=num_w2_bufs, + padded_k1=padded_k1, + aligned_n1=aligned_n1, + padded_k2=padded_k2, + aligned_n2=aligned_n2, + ) + + +# ============================================================================= +# Phase 2: kernel_main_fused_rs — Core Kernel (Direct-Write ICI) +# ============================================================================= + + +def kernel_main_fused_rs( + # Scalar prefetch (9) — `output_indices` was dropped (== lhs_indices for EP). + # `topk_indices_ref` is included but conditionally used: + # - pack_indices=False: lhs_indices_ref holds raw lhs_idx values, and + # topk_indices_ref holds raw topk_slot values (both size_m). + # - pack_indices=True: lhs_indices_ref holds packed values + # `combined = lhs_idx * top_k + topk_slot`, and topk_indices_ref is a + # 1-element dummy (saves ~512 KB of SMEM at large prefill). + lhs_group_sizes_ref, + group_offset_ref, + lhs_indices_ref, + topk_indices_ref, + max_num_gm_ref, # (1,) int32 — max gm tiles across all chips + total_recv_count_ref, # (1,) int32 — total remote rows this chip receives + w1_gs_gate_ref, # (E,) GMM1 gate global_scale + w1_gs_up_ref, # (E,) GMM1 up global_scale + w2_gs_ref, # (E,) GMM2 global_scale + # In (7) + hidden_states_ref, + w1_ref, + w2_ref, + w1_scale_ref, + w2_scale_ref, + w1_bias_ref, + w2_bias_ref, + # Out (1) + out_buf_ref, # HBM: (chunk_size * top_k, N2//NL, NL) — output & ICI DMA target + # Scratch — compute pipeline (same as kernel_main_fused) + metadata_ref, + fused_metadata_ref, + gathered_lhs_2x_ref, + gmm1_out_ref, + intermediate_ref, + tiled_out_2x_ref, + scatter_staging_3x_ref, # VMEM: (3, tile_m, N2//NL, NL) triple-buffered, first-dim indexed + partial_out1_ref, + partial_out2_ref, + shared_acc_ref, + gather_sem_ref, + staging_sem_ref, + zero_sem_ref, + gm_id_ref, + w1_buf_ref, + w2_buf_ref, + w1_scale_buf_ref, + w2_scale_buf_ref, + w1_bias_buf_ref, + w2_bias_buf_ref, + w1_sem_ref, + w2_sem_ref, + output_sem_ref, + # Scratch — ICI direct-write + send_sems_ref, # DMA(3,) per-staging-slot send sems + local_write_sems_ref, # DMA(3,) per-staging-slot local write sems + recv_sem_ref, # DMA(1,) for incoming remote writes + scale_send_sems_ref=None, # Optional DMA(3,) for FP8 row-scale sends. + scale_local_write_sems_ref=None, # Optional DMA(3,) for FP8 row-scale local writes. + scale_recv_sem_ref=None, # Optional DMA(1,) for incoming FP8 row-scale writes. + scatter_fp8_staging_3x_ref=None, # Optional VMEM staging for FP8 payload rows. + scatter_scale_3x_ref=None, # Optional VMEM staging for per-row FP8 scales. + out_scale_ref=None, # Optional HBM output for per-row FP8 activation scales. + *, + fused_dims: FusedDims, + tile_m: int, + tile_k1: int, + tile_k2: int, + tile_n1: int, + tile_n2: int, + num_w1_bufs: int, + num_w2_bufs: int, + act_fn: str, + out_dtype: jnp.dtype, + cfgs1: GmmConfigs, + cfgs2: GmmConfigs, + ep_size: int, + chunk_size: int, + ep_axis_name: str, + top_k: int, + pack_indices: bool = True, + fp8_direct_write: bool = False, +): + """Fused gather + GMM1 + act + GMM2 + ICI direct-write kernel. + + `pack_indices`: when True (default), `lhs_indices_ref` holds packed values + `combined = lhs_idx * top_k + topk_slot`; `topk_indices_ref` is a dummy. + When False, the two refs hold raw values separately. Packing saves + ~512 KB SMEM at the cost of ~80 us extra scalar-pipe work at small N. + + Pipeline per gm tile: + Steps 1-9: gather → GMM1 → act → GMM2 + Step 10: Per-row direct-write to out_buf on dest chip + """ + dims = fused_dims + num_lanes = pltpu.get_tpu_info().num_lanes + my_id = lax.axis_index(ep_axis_name) + max_num_gm = max_num_gm_ref[0] + + @jax.named_scope("sync_barrier") + def sync_barrier(): + barrier_sem = pltpu.get_barrier_semaphore() + for i in range(ep_size): + pltpu.semaphore_signal( + barrier_sem, + device_id={ep_axis_name: jnp.int32(i)}, + device_id_type=pltpu.DeviceIdType.MESH, + ) + pltpu.semaphore_wait(barrier_sem, ep_size) + + sync_barrier() + + # Build metadata. + meta_dims = Dimensions( + size_m=dims.size_m, + size_k=dims.size_k1, + size_n=dims.size_n1, + size_group=dims.size_group, + size_lhs_group=dims.size_lhs_group, + size_lhs_sublane=dims.size_lhs_sublane, + ) + meta_cfgs = GmmConfigs( + tiles=TileSizes(tile_m=tile_m, tile_k=tile_k1, tile_n=dims.size_n1), + dims=meta_dims, + lhs_cfgs=InputConfigs(quant_dtype=None, + quant_block_size=dims.size_k1, + dtype=out_dtype), + rhs_cfgs=InputConfigs(quant_dtype=None, + quant_block_size=dims.size_k1, + dtype=out_dtype), + out_dtype=out_dtype, + acc_dtype=jnp.float32, + zero_init=False, + ) + local_num_gm = fill_metadata( + lhs_group_sizes_ref, + group_offset_ref, + metadata_ref, + cfgs=meta_cfgs, + ) + + # Zero-init the output buffer (also serves as ICI DMA target). + zero_src = (scatter_fp8_staging_3x_ref.at[0] + if fp8_direct_write else scatter_staging_3x_ref.at[0]) + zero_size = zero_out_start_3d(out_buf_ref, zero_src, zero_sem_ref) + + if cfgs1.rhs_cfgs.quant_dtype is not None: + w1_packed = w1_ref.bitcast(jnp.uint32) + else: + w1_packed = w1_ref + if cfgs2.rhs_cfgs.quant_dtype is not None: + w2_packed = w2_ref.bitcast(jnp.uint32) + else: + w2_packed = w2_ref + + num_k1 = dims.size_k1 // tile_k1 + num_k2 = dims.size_k2 // tile_k2 + out_n1 = dims.size_n1 // 2 if act_fn else dims.size_n1 + num_n1 = out_n1 // tile_n1 + # Use original (unpadded) N for GMM2 loop count, matching the DMA + # gmm_v2 path which uses cfgs.out_size_n (not aligned_n). + # This avoids iterating over padding N-tiles which can cause + # incorrect results with fp8 weights. + num_n2 = dims.original_n2 // tile_n2 + packing = cfgs1.rhs_cfgs.packing + pk1 = tile_k1 // packing + pk2 = tile_k2 // packing + has_scale = cfgs1.rhs_cfgs.has_scale + has_bias = cfgs1.rhs_cfgs.has_bias + nqb1 = cfgs1.num_quant_blocks_per_tile_k if has_scale else 0 + nqb2 = cfgs2.num_quant_blocks_per_tile_k if has_scale else 0 + w1_dma_n = tile_n1 * 2 if act_fn else tile_n1 + + # --- Weight DMA helpers --- + @jax.named_scope("start_w1_dma") + def start_w1_dma(buf_id, expert_id, n_id, k_id=0): + if act_fn and num_n1 > 1: + # Fix: with fuse_act, w1 is [all_gate | all_up] along N. + # Load tile_n1 from gate + tile_n1 from up into buf as [gate|up]. + _gate_offset = n_id * tile_n1 + _up_offset = out_n1 + n_id * tile_n1 + # Gate half -> first tile_n1 cols of buf + pltpu.make_async_copy( + src_ref=w1_packed.at[expert_id, + pl.ds(k_id * pk1, pk1), + pl.ds(_gate_offset, tile_n1)], + dst_ref=w1_buf_ref.at[buf_id, :, :tile_n1], + sem=w1_sem_ref.at[buf_id], + ).start() + # Up half -> last tile_n1 cols of buf + pltpu.make_async_copy( + src_ref=w1_packed.at[expert_id, + pl.ds(k_id * pk1, pk1), + pl.ds(_up_offset, tile_n1)], + dst_ref=w1_buf_ref.at[buf_id, :, tile_n1:], + sem=w1_sem_ref.at[buf_id], + ).start() + else: + pltpu.make_async_copy( + src_ref=w1_packed.at[expert_id, + pl.ds(k_id * pk1, pk1), + pl.ds(n_id * w1_dma_n, w1_dma_n)], + dst_ref=w1_buf_ref.at[buf_id], + sem=w1_sem_ref.at[buf_id], + ).start() + if has_scale: + if act_fn and num_n1 > 1: + _gate_offset_s = n_id * tile_n1 + _up_offset_s = out_n1 + n_id * tile_n1 + pltpu.make_async_copy( + src_ref=w1_scale_ref.at[ + expert_id, + pl.ds(k_id * nqb1, nqb1), + :, + pl.ds(_gate_offset_s, tile_n1), + ], + dst_ref=w1_scale_buf_ref.at[buf_id, :, :, :tile_n1], + sem=w1_sem_ref.at[buf_id], + ).start() + pltpu.make_async_copy( + src_ref=w1_scale_ref.at[ + expert_id, + pl.ds(k_id * nqb1, nqb1), + :, + pl.ds(_up_offset_s, tile_n1), + ], + dst_ref=w1_scale_buf_ref.at[buf_id, :, :, tile_n1:], + sem=w1_sem_ref.at[buf_id], + ).start() + else: + pltpu.make_async_copy( + src_ref=w1_scale_ref.at[ + expert_id, + pl.ds(k_id * nqb1, nqb1), + :, + pl.ds(n_id * w1_dma_n, w1_dma_n), + ], + dst_ref=w1_scale_buf_ref.at[buf_id], + sem=w1_sem_ref.at[buf_id], + ).start() + if has_bias: + pltpu.make_async_copy( + src_ref=w1_bias_ref.at[expert_id, :, + pl.ds(n_id * w1_dma_n, w1_dma_n)], + dst_ref=w1_bias_buf_ref, + sem=w1_sem_ref.at[buf_id], + ).start() + + @jax.named_scope("start_w2_dma") + def start_w2_dma(buf_id, expert_id, n_id, k_id=0): + pltpu.make_async_copy( + src_ref=w2_packed.at[expert_id, + pl.ds(k_id * pk2, pk2), + pl.ds(n_id * tile_n2, tile_n2)], + dst_ref=w2_buf_ref.at[buf_id], + sem=w2_sem_ref.at[buf_id], + ).start() + if has_scale: + pltpu.make_async_copy( + src_ref=w2_scale_ref.at[ + expert_id, + pl.ds(k_id * nqb2, nqb2), + :, + pl.ds(n_id * tile_n2, tile_n2), + ], + dst_ref=w2_scale_buf_ref.at[buf_id], + sem=w2_sem_ref.at[buf_id], + ).start() + if has_bias: + pltpu.make_async_copy( + src_ref=w2_bias_ref.at[expert_id, :, + pl.ds(n_id * tile_n2, tile_n2)], + dst_ref=w2_bias_buf_ref, + sem=w2_sem_ref.at[buf_id], + ).start() + + @jax.named_scope("wait_w1_dma") + def wait_w1_dma(buf_id): + pltpu.make_async_copy( + src_ref=w1_buf_ref.at[buf_id], + dst_ref=w1_buf_ref.at[buf_id], + sem=w1_sem_ref.at[buf_id], + ).wait() + if has_scale: + pltpu.make_async_copy( + src_ref=w1_scale_buf_ref.at[buf_id], + dst_ref=w1_scale_buf_ref.at[buf_id], + sem=w1_sem_ref.at[buf_id], + ).wait() + if has_bias: + pltpu.make_async_copy( + src_ref=w1_bias_buf_ref, + dst_ref=w1_bias_buf_ref, + sem=w1_sem_ref.at[buf_id], + ).wait() + + @jax.named_scope("wait_w2_dma") + def wait_w2_dma(buf_id): + pltpu.make_async_copy( + src_ref=w2_buf_ref.at[buf_id], + dst_ref=w2_buf_ref.at[buf_id], + sem=w2_sem_ref.at[buf_id], + ).wait() + if has_scale: + pltpu.make_async_copy( + src_ref=w2_scale_buf_ref.at[buf_id], + dst_ref=w2_scale_buf_ref.at[buf_id], + sem=w2_sem_ref.at[buf_id], + ).wait() + if has_bias: + pltpu.make_async_copy( + src_ref=w2_bias_buf_ref, + dst_ref=w2_bias_buf_ref, + sem=w2_sem_ref.at[buf_id], + ).wait() + + # --- GMM compute helpers --- + @jax.named_scope("compute_gmm1_tile") + def compute_gmm1_tile(buf_id, n_id, k_id, gm_id): + sem_id = gm_id % 2 + k_cols = tile_k1 // num_lanes + k_offset = k_id * k_cols + lhs_data = gathered_lhs_2x_ref[sem_id, :, + pl.ds(k_offset, k_cols), :].reshape( + -1, dims.size_lhs_sublane, tile_k1) + w1_tile = w1_buf_ref.at[buf_id] + w1_sc = w1_scale_buf_ref.at[buf_id] if has_scale else None + w1_bias_tile = w1_bias_buf_ref if has_bias else None + w1_weights = WeightsRef(weight=w1_tile, scale=w1_sc, bias=w1_bias_tile) + if cfgs1.fuse_act is not None: + w1_up = WeightsRef( + weight=w1_tile.at[:, tile_n1:], + scale=w1_sc.at[:, :, tile_n1:] if w1_sc is not None else None, + bias=w1_bias_tile.at[:, pl.ds(tile_n1, tile_n1)] + if has_bias else None, + ) + w1_gate = WeightsRef( + weight=w1_tile.at[:, :tile_n1], + scale=w1_sc.at[:, :, :tile_n1] if w1_sc is not None else None, + bias=w1_bias_tile.at[:, + pl.ds(0, tile_n1)] if has_bias else None, + ) + w1_weights = FusedWeightsRef(gate=w1_gate, up=w1_up) + inner_kernel( + lhs_data, + w1_weights, + gmm1_out_ref.at[:, pl.ds(n_id * tile_n1, tile_n1)], + partial_out1_ref, + shared_acc_ref.at[:, + pl.ds(0, tile_n1 * 2 if act_fn else tile_n1)], + fused_metadata_ref, + cfgs=cfgs1, + gs_gate_ref=w1_gs_gate_ref, + gs_up_ref=w1_gs_up_ref, + scatter_mode=True, + _k_id=k_id, + _num_k=num_k1, + _gm_id=0, + _n_id=n_id, + ) + + @jax.named_scope("compute_gmm2_tile") + def compute_gmm2_tile(buf_id, n_id, k_id, gm_id): + sem_id = gm_id % 2 + lhs_data = intermediate_ref[:, :, pl.ds(k_id * tile_k2, tile_k2)] + w2_tile = w2_buf_ref.at[buf_id] + w2_sc = w2_scale_buf_ref.at[buf_id] if has_scale else None + inner_kernel( + lhs_data, + WeightsRef(weight=w2_tile, + scale=w2_sc, + bias=w2_bias_buf_ref if has_bias else None), + tiled_out_2x_ref.at[sem_id, :, + pl.ds(n_id * tile_n2, tile_n2)], + partial_out2_ref, + shared_acc_ref.at[:, pl.ds(0, tile_n2)], + fused_metadata_ref, + cfgs=cfgs2, + gs_gate_ref=w2_gs_ref, + gs_up_ref=w2_gs_ref, + scatter_mode=True, + _k_id=k_id, + _num_k=num_k2, + _gm_id=0, + _n_id=n_id, + ) + + # Initialize per-slot DMA counts to 0. + for _s in range(3): + gm_id_ref[1 + _s] = jnp.int32(0) # send counts + gm_id_ref[4 + _s] = jnp.int32(0) # local write counts + + # --- Main padded gm loop --- + @jax.named_scope("gm_loop_body") + def gm_loop_body(gm_id, _): + is_active = gm_id < local_num_gm + sem_id = gm_id % 2 + stg_id = gm_id % 3 + + # Drain the staging slot we're about to reuse. With triple buffering, + # this slot was last used 3 iterations ago — giving 2 full iterations + # of compute overlap for its DMAs to complete. + @jax.named_scope("drain_prev_dmas") + @pl.when(gm_id >= 3) + def _(): + prev_send = gm_id_ref[1 + stg_id] + send_wait_ref = (scatter_fp8_staging_3x_ref + if fp8_direct_write else scatter_staging_3x_ref) + pltpu.make_async_copy( + src_ref=send_wait_ref.at[stg_id, + pl.ds(0, prev_send), :, :], + dst_ref=send_wait_ref.at[stg_id, + pl.ds(0, prev_send), :, :], + sem=send_sems_ref.at[stg_id], + ).wait() + prev_local = gm_id_ref[4 + stg_id] + pltpu.make_async_copy( + src_ref=out_buf_ref.at[pl.ds(0, prev_local), :, :], + dst_ref=out_buf_ref.at[pl.ds(0, prev_local), :, :], + sem=local_write_sems_ref.at[stg_id], + ).wait() + if fp8_direct_write: + pltpu.make_async_copy( + src_ref=scatter_scale_3x_ref.at[stg_id, + pl.ds(0, prev_send), :], + dst_ref=scatter_scale_3x_ref.at[stg_id, + pl.ds(0, prev_send), :], + sem=scale_send_sems_ref.at[stg_id], + ).wait() + pltpu.make_async_copy( + src_ref=out_scale_ref.at[pl.ds(0, prev_local), :], + dst_ref=out_scale_ref.at[pl.ds(0, prev_local), :], + sem=scale_local_write_sems_ref.at[stg_id], + ).wait() + # Reset counts for this slot's new use. + gm_id_ref[1 + stg_id] = jnp.int32(0) + gm_id_ref[4 + stg_id] = jnp.int32(0) + + # Steps 1-9: Only when this chip still has compute work. + @pl.when(is_active) + def _(): + gm_id_ref[0] = gm_id + fused_metadata_ref.gm_id_to_group_id[ + 0] = metadata_ref.gm_id_to_group_id[gm_id] + fused_metadata_ref.gm_id_to_m_offset[ + 0] = metadata_ref.gm_id_to_m_offset[gm_id] + fused_metadata_ref.gm_id_to_m_offset[ + 1] = metadata_ref.gm_id_to_m_offset[gm_id + 1] + expert_id = fused_metadata_ref.gm_id_to_group_id[0] + + # DMA gather. When indices are packed, set divisor=top_k so the + # gather extracts the actual lhs_idx via integer division. + _gather_divisor = top_k if pack_indices else 1 + + @jax.named_scope("dma_gather_start") + @pl.when(gm_id == 0) + def _(): + dma_gather_gm_start( + hidden_states_ref, + gathered_lhs_2x_ref.at[sem_id], + lhs_indices_ref, + gather_sem_ref.at[sem_id], + 0, + metadata_ref, + divisor=_gather_divisor, + ) + + @jax.named_scope("dma_gather_prefetch") + @pl.when(gm_id + 1 < local_num_gm) + def _(): + dma_gather_gm_start( + hidden_states_ref, + gathered_lhs_2x_ref.at[1 - sem_id], + lhs_indices_ref, + gather_sem_ref.at[1 - sem_id], + gm_id + 1, + metadata_ref, + divisor=_gather_divisor, + ) + + # --- Weight DMA (overlapped with gather) --- + # w1 prologue: load first num_w1_bufs tiles (gm==0 only; + # for gm>0, tiles were prefetched during previous gm's GMM2). + total_w1_steps = num_n1 * num_k1 + total_w2_steps = num_n2 * num_k2 + + # ---- Same-expert weight caching (Idea 1) ---- + # When num_w*_bufs >= total_w*_steps, ALL of the current expert's + # W1/W2 weight tiles fit simultaneously in VMEM. If the next gm + # tile uses the same expert, we can SKIP the cross-gm prefetch + # entirely — the buffers already hold the right data. This is + # very common: at M=65536, E=32, tile_m=128 each expert spans + # ~16 consecutive gm tiles → up to 16x fewer weight DMAs. + # + # When buffers can't hold all tiles (bf16 case where total_steps>num_bufs), + # the buffer rotation already overwrites earlier tiles within a + # single gm — same-expert caching is unsafe there. The static + # checks below disable the optimization in that case. + can_cache_w1 = num_w1_bufs >= total_w1_steps + can_cache_w2 = num_w2_bufs >= total_w2_steps + + # Read previous expert id (clamp to gm_id=0 case, masked below). + prev_gm_clamped = jnp.maximum(gm_id - 1, 0) + prev_expert = metadata_ref.gm_id_to_group_id[prev_gm_clamped] + # is_new_expert is True at gm_id=0 OR when expert changes. + is_new_expert = jnp.logical_or(gm_id == 0, prev_expert + != expert_id) + # is_same_expert as previous gm. Used to skip prefetch+wait. + is_same_w1 = jnp.logical_and(jnp.logical_not(is_new_expert), + jnp.bool_(can_cache_w1)) + is_same_w2 = jnp.logical_and(jnp.logical_not(is_new_expert), + jnp.bool_(can_cache_w2)) + + @pl.when(gm_id == 0) + def _(): + for _i in range(min(num_w1_bufs, total_w1_steps)): + start_w1_dma(_i, expert_id, _i // num_k1, _i % num_k1) + + # Early w2 prefetch — skip if cache hit (same expert as prev gm). + @pl.when(jnp.logical_not(is_same_w2)) + def _(): + for _i in range(min(num_w2_bufs, total_w2_steps)): + start_w2_dma(_i, expert_id, _i // num_k2, _i % num_k2) + + # Wait for gather (weight DMAs running in parallel). + dma_gather_gm_wait( + gathered_lhs_2x_ref.at[sem_id], + gather_sem_ref.at[sem_id], + gm_id, + metadata_ref, + ) + + # GMM1 loop. + for step in range(total_w1_steps): + _n1 = step // num_k1 + _k1 = step % num_k1 + buf_id = step % num_w1_bufs + + # Skip wait if W1 was cached from previous gm tile + # (same expert + all tiles fit in buffers). + @pl.when(jnp.logical_not(is_same_w1)) + def _(): + wait_w1_dma(buf_id) + + if step + num_w1_bufs < total_w1_steps: + ns = step + num_w1_bufs + start_w1_dma(ns % num_w1_bufs, expert_id, ns // num_k1, + ns % num_k1) + compute_gmm1_tile(buf_id, _n1, _k1, gm_id) + if step == 0: + + @pl.when(gm_id + 1 < local_num_gm) + def _(): + next_e = metadata_ref.gm_id_to_group_id[gm_id + 1] + # Skip cross-prefetch when next gm reuses same expert + # (caching only makes sense when buffers fit all tiles). + next_same = jnp.logical_and(next_e == expert_id, + jnp.bool_(can_cache_w1)) + + @pl.when(jnp.logical_not(next_same)) + def _(): + start_w1_dma(0, next_e, 0, 0) + + @pl.when(gm_id + 1 < local_num_gm) + def _(): + next_e = metadata_ref.gm_id_to_group_id[gm_id + 1] + next_same = jnp.logical_and(next_e == expert_id, + jnp.bool_(can_cache_w1)) + + @pl.when(jnp.logical_not(next_same)) + def _(): + for _i in range(1, min(num_w1_bufs, total_w1_steps)): + start_w1_dma(_i, next_e, _i // num_k1, _i % num_k1) + + with jax.named_scope("interlude"): + # Copy gmm1_out (f32) to intermediate_ref (bf16, scatter layout). + # Activation was already applied inside inner_kernel via fuse_act. + gmm1_result = gmm1_out_ref[...] + k2_pad = dims.size_k2 - gmm1_result.shape[1] + if k2_pad > 0: + gmm1_result = jnp.concatenate( + [ + gmm1_result, + jnp.zeros( + (tile_m, k2_pad), dtype=gmm1_result.dtype), + ], + axis=1, + ) + intermediate_ref[...] = gmm1_result.astype(out_dtype).reshape( + intermediate_ref.shape) + + # GMM2 loop. + for step in range(total_w2_steps): + _n2 = step // num_k2 + _k2 = step % num_k2 + buf_id = step % num_w2_bufs + + # Skip wait if W2 was cached (same expert as prev + buffers fit all). + @pl.when(jnp.logical_not(is_same_w2)) + def _(): + wait_w2_dma(buf_id) + + if step + num_w2_bufs < total_w2_steps: + ns = step + num_w2_bufs + start_w2_dma(ns % num_w2_bufs, expert_id, ns // num_k2, + ns % num_k2) + compute_gmm2_tile(buf_id, _n2, _k2, gm_id) + + # Finish zero-init on first tile. + @jax.named_scope("zero_out_end") + @pl.when(gm_id == 0) + def _(): + zero_out_end_3d(out_buf_ref, zero_sem_ref, zero_size) + + m_st = metadata_ref.gm_id_to_m_offset[gm_id] + m_en = metadata_ref.gm_id_to_m_offset[gm_id + 1] + _sls = pltpu.get_tpu_info().get_sublane_tiling(out_dtype) + _ml = m_st % _sls + + with jax.named_scope("reshape_gmm2_output"): + scatter_staging_3x_ref[stg_id] = tiled_out_2x_ref[sem_id][ + ...].reshape( + tile_m, + scatter_staging_3x_ref.shape[2], + scatter_staging_3x_ref.shape[3], + ) + + # Step 10: Direct-write — single pass over valid rows only. + # Each row either ICI-sends (remote) or DMA-copies (local) to + # direct_write_buf[local_row * top_k + topk_idx]. + # Loop bound is m_en - m_st (valid rows), not tile_m, eliminating + # wasted iterations on padding rows. + num_valid = m_en - m_st + + @jax.named_scope("direct_write_rows") + def _do_direct_write(): + + def _write_row(i, carry): + send_sz, local_sz = carry + row_idx = _ml + i + # Read indices. Either packed (one ref) or separate. + if pack_indices: + # combined = lhs_idx * top_k + topk_slot + # = (dest_chip * chunk_size + local_row) * top_k + topk_slot + # So: + # write_pos = local_row * top_k + topk_slot + # = combined % (chunk_size * top_k) + # dest_chip = combined // (chunk_size * top_k) + # When chunk_size * top_k is a power of 2 (typical), + # XLA lowers these to a shift + mask — essentially free. + # No need to compute oi/topk_idx/local_row separately. + combined = lhs_indices_ref[m_st + i] + cs_tk = chunk_size * top_k + write_pos = combined % cs_tk + dest_chip = combined // cs_tk + else: + oi = lhs_indices_ref[m_st + i] + topk_idx = topk_indices_ref[m_st + i] + dest_chip = oi // chunk_size + local_row = oi % chunk_size + write_pos = local_row * top_k + topk_idx + is_local = dest_chip == my_id + + if fp8_direct_write: + # Quantize each completed expert row before the direct + # write. Remote ICI then moves FP8 payload plus one + # fp32 row scale instead of the full bf16 activation row. + row_f32 = scatter_staging_3x_ref[stg_id, + row_idx, :, :].astype( + jnp.float32) + fp8_max = jnp.array(jnp.finfo(jnp.float8_e4m3fn).max, + dtype=jnp.float32) + row_scale = (jnp.maximum( + jnp.max(jnp.abs(row_f32)), + jnp.array(1e-6, dtype=jnp.float32), + ) / fp8_max) + scatter_scale_3x_ref[stg_id, row_idx, :] = ( + row_scale + jnp.zeros((128, ), dtype=jnp.float32)) + scatter_fp8_staging_3x_ref[stg_id, + row_idx, :, :] = jnp.clip( + row_f32 / row_scale, + -fp8_max, + fp8_max).astype( + jnp.float8_e4m3fn) + + @pl.when(~is_local) + def _(): + pltpu.make_async_remote_copy( + src_ref=scatter_fp8_staging_3x_ref.at[ + stg_id, pl.ds(row_idx, 1), :, :], + dst_ref=out_buf_ref.at[ + pl.ds(write_pos, 1), :, :], + send_sem=send_sems_ref.at[stg_id], + recv_sem=recv_sem_ref.at[0], + device_id={ + ep_axis_name: dest_chip + }, + device_id_type=pltpu.DeviceIdType.MESH, + ).start() + pltpu.make_async_remote_copy( + src_ref=scatter_scale_3x_ref.at[ + stg_id, pl.ds(row_idx, 1), :], + dst_ref=out_scale_ref.at[ + pl.ds(write_pos, 1), :], + send_sem=scale_send_sems_ref.at[stg_id], + recv_sem=scale_recv_sem_ref.at[0], + device_id={ + ep_axis_name: dest_chip + }, + device_id_type=pltpu.DeviceIdType.MESH, + ).start() + + @pl.when(is_local) + def _(): + pltpu.make_async_copy( + src_ref=scatter_fp8_staging_3x_ref.at[ + stg_id, pl.ds(row_idx, 1), :, :], + dst_ref=out_buf_ref.at[ + pl.ds(write_pos, 1), :, :], + sem=local_write_sems_ref.at[stg_id], + ).start() + pltpu.make_async_copy( + src_ref=scatter_scale_3x_ref.at[ + stg_id, pl.ds(row_idx, 1), :], + dst_ref=out_scale_ref.at[ + pl.ds(write_pos, 1), :], + sem=scale_local_write_sems_ref.at[stg_id], + ).start() + + else: + + @pl.when(~is_local) + def _(): + pltpu.make_async_remote_copy( + src_ref=scatter_staging_3x_ref.at[ + stg_id, pl.ds(row_idx, 1), :, :], + dst_ref=out_buf_ref.at[ + pl.ds(write_pos, 1), :, :], + send_sem=send_sems_ref.at[stg_id], + recv_sem=recv_sem_ref.at[0], + device_id={ + ep_axis_name: dest_chip + }, + device_id_type=pltpu.DeviceIdType.MESH, + ).start() + + @pl.when(is_local) + def _(): + pltpu.make_async_copy( + src_ref=scatter_staging_3x_ref.at[ + stg_id, pl.ds(row_idx, 1), :, :], + dst_ref=out_buf_ref.at[ + pl.ds(write_pos, 1), :, :], + sem=local_write_sems_ref.at[stg_id], + ).start() + + return ( + send_sz + + lax.select(~is_local, jnp.int32(1), jnp.int32(0)), + local_sz + + lax.select(is_local, jnp.int32(1), jnp.int32(0)), + ) + + return lax.fori_loop(0, num_valid, _write_row, + (jnp.int32(0), jnp.int32(0))) + + send_sz, local_sz = _do_direct_write() + gm_id_ref[1 + stg_id] = gm_id_ref[1 + stg_id] + send_sz + gm_id_ref[4 + stg_id] = gm_id_ref[4 + stg_id] + local_sz + + return _ + + lax.fori_loop(0, max_num_gm, gm_loop_body, None) + + # If max_num_gm == 0 OR local_num_gm == 0 (this chip has no work but + # the loop ran padded iterations), zero_out_end was never called inside + # the loop body (it's guarded by is_active AND gm_id==0). Drain it here. + @pl.when(jnp.logical_or(max_num_gm == 0, local_num_gm == 0)) + def _(): + zero_out_end_3d(out_buf_ref, zero_sem_ref, zero_size) + + # Epilogue: drain remaining per-slot DMA counts. + @jax.named_scope("epilogue_drain") + @pl.when(max_num_gm > 0) + def _(): + for _slot in range(3): + + @pl.when(max_num_gm > _slot) + def _(): + remaining_send = gm_id_ref[1 + _slot] + send_wait_ref = (scatter_fp8_staging_3x_ref if fp8_direct_write + else scatter_staging_3x_ref) + pltpu.make_async_copy( + src_ref=send_wait_ref.at[_slot, + pl.ds(0, remaining_send), :, :], + dst_ref=send_wait_ref.at[_slot, + pl.ds(0, remaining_send), :, :], + sem=send_sems_ref.at[_slot], + ).wait() + remaining_local = gm_id_ref[4 + _slot] + pltpu.make_async_copy( + src_ref=out_buf_ref.at[pl.ds(0, remaining_local), :, :], + dst_ref=out_buf_ref.at[pl.ds(0, remaining_local), :, :], + sem=local_write_sems_ref.at[_slot], + ).wait() + if fp8_direct_write: + pltpu.make_async_copy( + src_ref=scatter_scale_3x_ref.at[ + _slot, pl.ds(0, remaining_send), :], + dst_ref=scatter_scale_3x_ref.at[ + _slot, pl.ds(0, remaining_send), :], + sem=scale_send_sems_ref.at[_slot], + ).wait() + pltpu.make_async_copy( + src_ref=out_scale_ref.at[pl.ds(0, remaining_local), :], + dst_ref=out_scale_ref.at[pl.ds(0, remaining_local), :], + sem=scale_local_write_sems_ref.at[_slot], + ).wait() + + # Final barrier: ensures all chips have finished their sends before + # any chip returns. This guarantees all remote DMAs have landed. + sync_barrier() + + # Drain recv_sem — each incoming remote DMA incremented it by 1. + # Must be zero before kernel exit (Mosaic requirement). + with jax.named_scope("drain_recv_sem"): + total_recv = total_recv_count_ref[0] + pltpu.make_async_copy( + src_ref=out_buf_ref.at[pl.ds(0, total_recv), :, :], + dst_ref=out_buf_ref.at[pl.ds(0, total_recv), :, :], + sem=recv_sem_ref.at[0], + ).wait() + if fp8_direct_write: + pltpu.make_async_copy( + src_ref=out_scale_ref.at[pl.ds(0, total_recv), :], + dst_ref=out_scale_ref.at[pl.ds(0, total_recv), :], + sem=scale_recv_sem_ref.at[0], + ).wait() + + # All ICI DMA writes go directly to out_buf_ref (the pallas_call output). + # No input_output_aliases needed — avoids XLA tiling copy on weight inputs. + + +def kernel_main_fused_rs_fp8( + lhs_group_sizes_ref, + group_offset_ref, + lhs_indices_ref, + topk_indices_ref, + max_num_gm_ref, + total_recv_count_ref, + w1_gs_gate_ref, + w1_gs_up_ref, + w2_gs_ref, + hidden_states_ref, + w1_ref, + w2_ref, + w1_scale_ref, + w2_scale_ref, + w1_bias_ref, + w2_bias_ref, + out_buf_ref, + out_scale_ref, + metadata_ref, + fused_metadata_ref, + gathered_lhs_2x_ref, + gmm1_out_ref, + intermediate_ref, + tiled_out_2x_ref, + scatter_staging_3x_ref, + partial_out1_ref, + partial_out2_ref, + shared_acc_ref, + gather_sem_ref, + staging_sem_ref, + zero_sem_ref, + gm_id_ref, + w1_buf_ref, + w2_buf_ref, + w1_scale_buf_ref, + w2_scale_buf_ref, + w1_bias_buf_ref, + w2_bias_buf_ref, + w1_sem_ref, + w2_sem_ref, + output_sem_ref, + send_sems_ref, + local_write_sems_ref, + recv_sem_ref, + scale_send_sems_ref, + scale_local_write_sems_ref, + scale_recv_sem_ref, + scatter_fp8_staging_3x_ref, + scatter_scale_3x_ref, + *, + fused_dims: FusedDims, + tile_m: int, + tile_k1: int, + tile_k2: int, + tile_n1: int, + tile_n2: int, + num_w1_bufs: int, + num_w2_bufs: int, + act_fn: str, + out_dtype: jnp.dtype, + cfgs1: GmmConfigs, + cfgs2: GmmConfigs, + ep_size: int, + chunk_size: int, + ep_axis_name: str, + top_k: int, + pack_indices: bool = True, +): + """FP8 direct-write wrapper with output refs ordered for pallas_call.""" + return kernel_main_fused_rs( + lhs_group_sizes_ref, + group_offset_ref, + lhs_indices_ref, + topk_indices_ref, + max_num_gm_ref, + total_recv_count_ref, + w1_gs_gate_ref, + w1_gs_up_ref, + w2_gs_ref, + hidden_states_ref, + w1_ref, + w2_ref, + w1_scale_ref, + w2_scale_ref, + w1_bias_ref, + w2_bias_ref, + out_buf_ref, + metadata_ref, + fused_metadata_ref, + gathered_lhs_2x_ref, + gmm1_out_ref, + intermediate_ref, + tiled_out_2x_ref, + scatter_staging_3x_ref, + partial_out1_ref, + partial_out2_ref, + shared_acc_ref, + gather_sem_ref, + staging_sem_ref, + zero_sem_ref, + gm_id_ref, + w1_buf_ref, + w2_buf_ref, + w1_scale_buf_ref, + w2_scale_buf_ref, + w1_bias_buf_ref, + w2_bias_buf_ref, + w1_sem_ref, + w2_sem_ref, + output_sem_ref, + send_sems_ref, + local_write_sems_ref, + recv_sem_ref, + scale_send_sems_ref, + scale_local_write_sems_ref, + scale_recv_sem_ref, + scatter_fp8_staging_3x_ref, + scatter_scale_3x_ref, + out_scale_ref, + fused_dims=fused_dims, + tile_m=tile_m, + tile_k1=tile_k1, + tile_k2=tile_k2, + tile_n1=tile_n1, + tile_n2=tile_n2, + num_w1_bufs=num_w1_bufs, + num_w2_bufs=num_w2_bufs, + act_fn=act_fn, + out_dtype=out_dtype, + cfgs1=cfgs1, + cfgs2=cfgs2, + ep_size=ep_size, + chunk_size=chunk_size, + ep_axis_name=ep_axis_name, + top_k=top_k, + pack_indices=pack_indices, + fp8_direct_write=True, + ) + + +# ============================================================================= +# Phase 3: gmm_v2_fused_rs — Public API +# ============================================================================= + + +@jax.jit(static_argnames=[ + "act_fn", + "output_size", + "vmem_limit_bytes", + "ep_size", + "ep_axis_name", + "top_k", + "fp8_direct_write", +], ) +def gmm_v2_fused_rs( + hidden_states: jax.Array, + w1: jax.Array, + w2: jax.Array, + group_sizes: jax.Array, + lhs_indices: jax.Array, + output_indices: jax.Array, + *, + w1_scale: jax.Array | None = None, + w2_scale: jax.Array | None = None, + w1_global_scale: jax.Array | None = None, # (E, 2) gate/up per-expert + w2_global_scale: jax.Array | None = None, # (E,) per-expert + w1_bias: jax.Array | None = None, + w2_bias: jax.Array | None = None, + act_fn: str = "silu", + group_offset: jax.Array | None = None, + output_size: int, + vmem_limit_bytes: int | None = None, + topk_indices: jax.Array, + ep_size: int, + ep_axis_name: str, + max_num_gm: jax.Array, + total_recv_count: jax.Array, + top_k: int, + fp8_direct_write: bool = False, +) -> jax.Array: + """Fused gather + GMM1 + act + GMM2 + ICI direct-write. + + Output: (chunk_size * top_k, size_n2) per chip — raw expert contributions. + Post-kernel reduction applies topk_weights and sums over top_k dim. + """ + size_group, size_k1, size_n1 = w1.shape + _, size_k2, size_n2 = w2.shape + size_m = lhs_indices.shape[0] + chunk_size = output_size // ep_size + + is_quantized = w1_scale is not None + if is_quantized: + assert w2_scale is not None + rhs1_quant_block_size = _recover_quant_block_size( + size_k1, w1_scale.shape[1]) + rhs2_quant_block_size = _recover_quant_block_size( + size_k2, w2_scale.shape[1]) + quant_block_size = rhs1_quant_block_size + w1_scale = w1_scale.astype(jnp.float32) + w2_scale = w2_scale.astype(jnp.float32) + else: + rhs1_quant_block_size = size_k1 + rhs2_quant_block_size = size_k2 + quant_block_size = None + # Per-expert global_scale as scalar prefetch (SMEM): one scalar per expert, + # indexed inside inner_kernel via metadata_ref.gm_id_to_group_id. + if w1_global_scale is not None: + if w1_global_scale.ndim == 2: + _w1_gs_gate = w1_global_scale[:, 0].astype(jnp.float32) + _w1_gs_up = w1_global_scale[:, 1].astype(jnp.float32) + else: + _w1_gs_gate = w1_global_scale.astype(jnp.float32) + _w1_gs_up = _w1_gs_gate + else: + _w1_gs_gate = jnp.ones((size_group, ), dtype=jnp.float32) + _w1_gs_up = _w1_gs_gate + if w2_global_scale is not None: + _w2_gs = w2_global_scale.astype(jnp.float32) + else: + _w2_gs = jnp.ones((size_group, ), dtype=jnp.float32) + + if group_offset is None: + group_offset = jnp.array([0], dtype=jnp.int32) + if vmem_limit_bytes is None: + vmem_limit_bytes = int(pltpu.get_tpu_info().vmem_capacity_bytes * 0.95) + + num_lanes = pltpu.get_tpu_info().num_lanes + sls = pltpu.get_tpu_info().get_sublane_tiling(hidden_states.dtype) + out_dtype = hidden_states.dtype + size_lhs_sublane = min(sls, size_m) + intermediate_size = size_n1 // 2 + + # --- Single source of truth for all block/tile sizes --- + # This helper is pure-Python and deterministic; it is also called from + # ``run_gmm_fused_rs`` with the same inputs so that ``compute_num_gm`` + # / ``max_num_gm`` agree with the kernel's own ``local_num_gm``. + block_sizes = _select_fused_rs_block_sizes( + size_m=size_m, + size_k1=size_k1, + size_n1=size_n1, + size_k2=size_k2, + size_n2=size_n2, + size_group=size_group, + size_lhs_group=group_sizes.shape[0], + ep_size=ep_size, + out_dtype=out_dtype, + w1_dtype=w1.dtype, + w2_dtype=w2.dtype, + is_quantized=is_quantized, + quant_block_size=quant_block_size, + act_fn=act_fn, + vmem_limit_bytes=vmem_limit_bytes, + fp8_direct_write=fp8_direct_write, + ) + tile_m = block_sizes.tile_m + tile_k1 = block_sizes.tile_k1 + tile_n1 = block_sizes.tile_n1 + tile_k2 = block_sizes.tile_k2 + tile_n2 = block_sizes.tile_n2 + num_w1_bufs = block_sizes.num_w1_bufs + num_w2_bufs = block_sizes.num_w2_bufs + padded_k1 = block_sizes.padded_k1 + aligned_n1 = block_sizes.aligned_n1 + aligned_n2 = block_sizes.aligned_n2 + + # --- Apply the padding the helper chose --- + # Pad K1 for DMA gather alignment. + if padded_k1 != size_k1: + k_pad = padded_k1 - size_k1 + hidden_states = jnp.pad(hidden_states, ((0, 0), (0, k_pad))) + + # Pad N1 (and propagate to K2 since GMM2's K is GMM1's intermediate size). + if aligned_n1 != size_n1: + n1_pad = aligned_n1 - size_n1 + w1 = jnp.pad(w1, ((0, 0), (0, 0), (0, n1_pad))) + if is_quantized: + w1_scale = jnp.pad(w1_scale, ((0, 0), (0, 0), (0, 0), (0, n1_pad))) + new_k2 = aligned_n1 // 2 + if new_k2 != size_k2: + w2 = jnp.pad(w2, ((0, 0), (0, new_k2 - size_k2), (0, 0))) + size_k2 = new_k2 + size_n1 = aligned_n1 + + # N2 alignment note: aligned_n2 is used for OUTPUT buffer sizing (DMA + # scatter requires 3D refs at aligned width), but w2 weights do NOT need + # padding — the kernel iterates num_n2 = original_n2 // tile_n2 tiles, + # so DMA only accesses w2 columns [0, original_n2). Removing w2 pad + # eliminates a large copy+pad (~33% of w2) on every forward pass. + + # Pad scales. + if is_quantized: + nqb1 = -(-tile_k1 // rhs1_quant_block_size) + total_scale_blocks1 = (size_k1 // tile_k1) * nqb1 + if w1_scale.shape[1] < total_scale_blocks1: + w1_scale = jnp.concatenate( + [ + w1_scale, + jnp.repeat( + w1_scale[:, -1:, :, :], + total_scale_blocks1 - w1_scale.shape[1], + axis=1, + ), + ], + axis=1, + ) + nqb2 = -(-tile_k2 // rhs2_quant_block_size) + total_scale_blocks2 = (size_k2 // tile_k2) * nqb2 + if w2_scale.shape[1] < total_scale_blocks2: + w2_scale = jnp.concatenate( + [ + w2_scale, + jnp.repeat( + w2_scale[:, -1:, :, :], + total_scale_blocks2 - w2_scale.shape[1], + axis=1, + ), + ], + axis=1, + ) + else: + total_scale_blocks1 = 0 + total_scale_blocks2 = 0 + + has_bias = w1_bias is not None + if has_bias: + assert w2_bias is not None + if w1_bias.shape[-1] != size_n1: + w1_bias = jnp.pad(w1_bias, ((0, 0), (0, 0), + (0, size_n1 - w1_bias.shape[-1]))) + if w2_bias.shape[-1] != size_n2: + w2_bias = jnp.pad(w2_bias, ((0, 0), (0, 0), + (0, size_n2 - w2_bias.shape[-1]))) + w1_bias = w1_bias.astype(jnp.float32) + w2_bias = w2_bias.astype(jnp.float32) + + hidden_3d = hidden_states.reshape(hidden_states.shape[0], + padded_k1 // num_lanes, num_lanes) + + fused_dims = FusedDims( + size_m=size_m, + size_k1=size_k1, + padded_k1=padded_k1, + size_n1=size_n1, + size_k2=size_k2, + size_n2=aligned_n2, + original_n2=size_n2, + size_group=size_group, + size_lhs_group=group_sizes.shape[0], + size_lhs_sublane=size_lhs_sublane, + intermediate_size=intermediate_size, + has_bias=has_bias, + quant_block_size=quant_block_size, + num_scale_blocks1=total_scale_blocks1, + num_scale_blocks2=total_scale_blocks2, + ) + + rhs1_packing = (32 // + jax.dtypes.itemsize_bits(w1.dtype)) if is_quantized else 1 + rhs2_packing = (32 // + jax.dtypes.itemsize_bits(w2.dtype)) if is_quantized else 1 + lhs1_quant_block_size = min(256 if rhs1_quant_block_size < 512 else 512, + size_k1) + lhs2_quant_block_size = min(256 if rhs2_quant_block_size < 512 else 512, + size_k2) + + cfgs1 = GmmConfigs( + dims=Dimensions( + size_m=size_m, + size_k=size_k1, + size_n=size_n1, + size_group=size_group, + size_lhs_group=group_sizes.shape[0], + size_lhs_sublane=size_lhs_sublane, + ), + tiles=TileSizes(tile_m=tile_m, tile_k=tile_k1, tile_n=tile_n1), + lhs_cfgs=InputConfigs( + quant_dtype=( + jnp.float8_e4m3fn.dtype if is_quantized + and get_maybe_quantize_lhs(w1.dtype, rhs1_quant_block_size, + lhs1_quant_block_size) else None), + quant_block_size=lhs1_quant_block_size, + dtype=out_dtype, + ), + rhs_cfgs=InputConfigs( + quant_dtype=w1.dtype if is_quantized else None, + quant_block_size=rhs1_quant_block_size, + dtype=w1.dtype, + has_bias=has_bias, + has_scale=is_quantized, + packing=rhs1_packing, + num_quant_blocks=total_scale_blocks1 if is_quantized else 1, + ), + out_dtype=out_dtype, + acc_dtype=jnp.float32, + zero_init=False, + fuse_act=act_fn, + ) + + cfgs2 = GmmConfigs( + dims=Dimensions( + size_m=size_m, + size_k=size_k2, + size_n=aligned_n2, + size_group=size_group, + size_lhs_group=group_sizes.shape[0], + size_lhs_sublane=size_lhs_sublane, + ), + tiles=TileSizes(tile_m=tile_m, tile_k=tile_k2, tile_n=tile_n2), + lhs_cfgs=InputConfigs( + quant_dtype=( + jnp.float8_e4m3fn.dtype if is_quantized + and get_maybe_quantize_lhs(w2.dtype, rhs2_quant_block_size, + lhs2_quant_block_size) else None), + quant_block_size=lhs2_quant_block_size, + dtype=out_dtype, + ), + rhs_cfgs=InputConfigs( + quant_dtype=w2.dtype if is_quantized else None, + quant_block_size=rhs2_quant_block_size, + dtype=w2.dtype, + has_bias=has_bias, + has_scale=is_quantized, + packing=rhs2_packing, + num_quant_blocks=total_scale_blocks2 if is_quantized else 1, + ), + out_dtype=out_dtype, + acc_dtype=jnp.float32, + zero_init=False, + ) + + # Scratch shapes. + max_num_gm_static = size_group + pl.cdiv(size_m, tile_m) - 1 + n2_cols = aligned_n2 // num_lanes + + scratch_shapes = [ + # metadata_ref + MetadataRef( + gm_id_to_group_id=pltpu.SMEM((max_num_gm_static, ), jnp.int32), + gm_id_to_m_offset=pltpu.SMEM((max_num_gm_static + 1, ), jnp.int32), + ), + # fused_metadata_ref + MetadataRef( + gm_id_to_group_id=pltpu.SMEM((1, ), jnp.int32), + gm_id_to_m_offset=pltpu.SMEM((2, ), jnp.int32), + ), + # gathered_lhs_2x_ref + pltpu.VMEM((2, tile_m, padded_k1 // num_lanes, num_lanes), out_dtype), + # gmm1_out_ref + pltpu.VMEM((tile_m, size_n1 // 2 if act_fn else size_n1), jnp.float32), + # intermediate_ref + pltpu.VMEM((tile_m // size_lhs_sublane, size_lhs_sublane, size_k2), + out_dtype), + # tiled_out_2x_ref + pltpu.VMEM((2, tile_m, aligned_n2), out_dtype), + # scatter_staging_3x_ref — 4D: first-dim selects triple-buffer slot, + # avoids dynamic offset that triggers Mosaic VMEM tiling alignment error. + pltpu.VMEM((3, tile_m, n2_cols, num_lanes), out_dtype), + # partial_out1_ref + pltpu.VMEM((size_lhs_sublane, tile_n1 * 2 if act_fn else tile_n1), + jnp.float32), + # partial_out2_ref + pltpu.VMEM((size_lhs_sublane, tile_n2), jnp.float32), + # shared_acc_ref + pltpu.VMEM((tile_m, max(tile_n1 * 2 if act_fn else tile_n1, tile_n2)), + jnp.float32), + # gather_sem_ref + pltpu.SemaphoreType.DMA((2, )), + # staging_sem_ref + pltpu.SemaphoreType.DMA((1, )), + # zero_sem_ref + pltpu.SemaphoreType.DMA((1, )), + # gm_id_ref: [0]=gm_id, [1..3]=per-slot send counts, [4..6]=per-slot local counts + pltpu.SMEM((7, ), jnp.int32), + # w1_buf_ref + pltpu.VMEM( + ( + num_w1_bufs, + (tile_k1 // (32 // jax.dtypes.itemsize_bits(w1.dtype)) + if is_quantized else tile_k1), + tile_n1 * 2 if act_fn else tile_n1, + ), + jnp.uint32 if is_quantized else w1.dtype, + ), + # w2_buf_ref + pltpu.VMEM( + ( + num_w2_bufs, + (tile_k2 // (32 // jax.dtypes.itemsize_bits(w2.dtype)) + if is_quantized else tile_k2), + tile_n2, + ), + jnp.uint32 if is_quantized else w2.dtype, + ), + # w1_scale_buf_ref + pltpu.VMEM( + ( + num_w1_bufs, + (-(-tile_k1 // rhs1_quant_block_size) if is_quantized else 1), + 1, + tile_n1 * 2 if act_fn else tile_n1, + ), + jnp.float32, + ), + # w2_scale_buf_ref + pltpu.VMEM( + ( + num_w2_bufs, + (-(-tile_k2 // rhs2_quant_block_size) if is_quantized else 1), + 1, + tile_n2, + ), + jnp.float32, + ), + # w1_bias_buf_ref + pltpu.VMEM((1, tile_n1 * 2 if act_fn else tile_n1), jnp.float32), + # w2_bias_buf_ref + pltpu.VMEM((1, tile_n2), jnp.float32), + # w1_sem_ref + pltpu.SemaphoreType.DMA((num_w1_bufs, )), + # w2_sem_ref + pltpu.SemaphoreType.DMA((num_w2_bufs, )), + # output_sem_ref + pltpu.SemaphoreType.DMA((1, )), + # --- ICI direct-write buffers --- + # send_sems_ref (per-staging-slot, 3 slots) + pltpu.SemaphoreType.DMA((3, )), + # local_write_sems_ref (per-staging-slot, 3 slots) + pltpu.SemaphoreType.DMA((3, )), + # recv_sem_ref (for incoming remote writes) + pltpu.SemaphoreType.DMA((1, )), + ] + + compiler_params = pltpu.CompilerParams( + vmem_limit_bytes=vmem_limit_bytes, + disable_bounds_checks=True, + collective_id=0, + ) + + w1_scale_input = w1_scale if is_quantized else jnp.zeros( + (1, 1, 1, 1), jnp.float32) + w2_scale_input = w2_scale if is_quantized else jnp.zeros( + (1, 1, 1, 1), jnp.float32) + w1_bias_input = w1_bias if has_bias else jnp.zeros((1, 1, 1), jnp.float32) + w2_bias_input = w2_bias if has_bias else jnp.zeros((1, 1, 1), jnp.float32) + max_num_gm_arr = jnp.array([max_num_gm], dtype=jnp.int32) + + # Output buffer — also serves as ICI DMA target. In the FP8 direct-write + # path this stores quantized payload rows; row scales are returned as a + # second output and dequantized immediately after the kernel. + payload_dtype = jnp.float8_e4m3fn if fp8_direct_write else out_dtype + out_buf_init = jax.ShapeDtypeStruct( + (chunk_size * top_k, n2_cols, num_lanes), payload_dtype) + # TPU VMEM vector slices must be tile-aligned. Store the scalar row scale + # as a padded 128-wide row and use column 0 when dequantizing. + out_scale_init = jax.ShapeDtypeStruct((chunk_size * top_k, 128), + jnp.float32) + + # Decide whether to pack lhs_indices + topk_slot_indices into a single + # SMEM scalar prefetch. Packing's per-row unpack (1 mod + 1 div via the + # math identity) is cheap (~bit ops for power-of-2 chunk_size*top_k) but + # NOT free — costs ~130-200 us at size_m=65K. So we only pack when needed + # to fit SMEM budget (1 MB). + # + # Two int32[size_m] arrays + internal scratch (~10-20 KB) need to fit in + # 1 MB SMEM. Pack when 2 separate refs would exceed ~960 KB (2 * size_m * 4). + # That's size_m > 120K (e.g., prefill_16384 with top_k=8 has size_m=131K). + pack_indices = size_m > 120_000 + if pack_indices: + # combined = lhs_idx * top_k + topk_slot + primary_idx_ref = lhs_indices.astype( + jnp.int32) * top_k + topk_indices.astype(jnp.int32) + secondary_idx_ref = jnp.zeros((1, ), dtype=jnp.int32) # dummy + else: + primary_idx_ref = lhs_indices + secondary_idx_ref = topk_indices + + pallas_inputs = ( + group_sizes, + group_offset, + primary_idx_ref, + secondary_idx_ref, + max_num_gm_arr, + total_recv_count, + _w1_gs_gate, + _w1_gs_up, + _w2_gs, + hidden_3d, + w1, + w2, + w1_scale_input, + w2_scale_input, + w1_bias_input, + w2_bias_input, + ) + pallas_name = (f"gmm_v2_fused_rs-E_{size_group}-M_{size_m}" + f"-K1_{size_k1}-N1_{size_n1}-K2_{size_k2}-N2_{size_n2}" + f"-EP_{ep_size}-TK_{top_k}" + f"{'-packed' if pack_indices else ''}") + kernel_kwargs = dict( + fused_dims=fused_dims, + tile_m=tile_m, + tile_k1=tile_k1, + tile_k2=tile_k2, + tile_n1=tile_n1, + tile_n2=tile_n2, + num_w1_bufs=num_w1_bufs, + num_w2_bufs=num_w2_bufs, + act_fn=act_fn, + out_dtype=out_dtype, + cfgs1=cfgs1, + cfgs2=cfgs2, + ep_size=ep_size, + chunk_size=chunk_size, + ep_axis_name=ep_axis_name, + top_k=top_k, + pack_indices=pack_indices, + ) + in_specs = [ + pl.BlockSpec(memory_space=pltpu.HBM), # hidden_states_3d + pl.BlockSpec(memory_space=pltpu.HBM), # w1 + pl.BlockSpec(memory_space=pltpu.HBM), # w2 + pl.BlockSpec(memory_space=pltpu.HBM), # w1_scale + pl.BlockSpec(memory_space=pltpu.HBM), # w2_scale + pl.BlockSpec(memory_space=pltpu.HBM), # w1_bias + pl.BlockSpec(memory_space=pltpu.HBM), # w2_bias + ] + + if fp8_direct_write: + fp8_scratch_shapes = scratch_shapes + [ + pltpu.SemaphoreType.DMA((3, )), + pltpu.SemaphoreType.DMA((3, )), + pltpu.SemaphoreType.DMA((1, )), + pltpu.VMEM((3, tile_m, n2_cols, num_lanes), jnp.float8_e4m3fn), + pltpu.VMEM((3, tile_m, 128), jnp.float32), + ] + payload, row_scales = pl.pallas_call( + functools.partial(kernel_main_fused_rs_fp8, **kernel_kwargs), + out_shape=(out_buf_init, out_scale_init), + grid_spec=pltpu.PrefetchScalarGridSpec( + num_scalar_prefetch=9, + in_specs=in_specs, + out_specs=[ + pl.BlockSpec(memory_space=pltpu.HBM), + pl.BlockSpec(memory_space=pltpu.HBM), + ], + scratch_shapes=fp8_scratch_shapes, + ), + compiler_params=compiler_params, + name=f"{pallas_name}-fp8-direct-write", + )(*pallas_inputs) + payload_2d = payload.reshape(chunk_size * top_k, + aligned_n2).astype(jnp.float32) + return (payload_2d * row_scales[:, :1]).astype(out_dtype)[:, :size_n2] + + result = pl.pallas_call( + functools.partial(kernel_main_fused_rs, **kernel_kwargs), + out_shape=out_buf_init, + grid_spec=pltpu.PrefetchScalarGridSpec( + num_scalar_prefetch=9, + in_specs=in_specs, + out_specs=pl.BlockSpec(memory_space=pltpu.HBM), + scratch_shapes=scratch_shapes, + ), + compiler_params=compiler_params, + name=pallas_name, + )(*pallas_inputs) + + return result.reshape(chunk_size * top_k, aligned_n2)[:, :size_n2] + + +# ============================================================================= +# Phase 4: run_gmm_fused_rs — High-level shard_map entry point +# ============================================================================= + +EXPERT = "model" +MLP_DATA = "data" + + +@functools.partial( + jax.jit, + static_argnames=( + "act_fn", + "mesh", + "top_k", + "ep_axis_name", + "tile_m", + "ep_size", + ), + compiler_options={ + # "xla_enable_transpose_trace": True, + }, +) +def run_gmm_fused_rs( + hidden_states: jax.Array, + w1: jax.Array, + w2: jax.Array, + group_sizes: jax.Array, + lhs_indices: jax.Array, + output_indices: jax.Array, + topk_weights: jax.Array, + topk_indices: jax.Array, + ep_size: int, + mesh: jax.sharding.Mesh, + act_fn: str = "silu", + top_k: int = 1, + ep_axis_name: str = EXPERT, + tile_m: int = 128, +) -> jax.Array: + """Run the fused GMM + ICI direct-write kernel via shard_map. + + Includes post-kernel topk_weights reduction and all_gather. + + Args: + hidden_states: bf16[num_tokens, hidden_size] — ungathered input. + w1: bf16[num_experts, hidden_size, intermediate_size*2] — gate+up weights. + w2: bf16[num_experts, intermediate_size, hidden_size] — down weights. + group_sizes: int32[num_experts] — rows per expert. + lhs_indices: int32[size_m] — gather indices into hidden_states. + output_indices: int32[size_m] — scatter indices (global token IDs). + topk_weights: float32[num_tokens, top_k] — topk weights for post-kernel reduction. + topk_indices: int32[size_m] — per-row topk slot (0..top_k-1). + ep_size: number of expert-parallel chips. + mesh: JAX Mesh. + act_fn: activation function ("silu", "gelu", etc.). + top_k: number of topk selections per token. + ep_axis_name: mesh axis name for expert parallelism. + tile_m: tile size for M dimension. + + Returns: + (num_tokens, hidden_size) — post-reduction output with all_gather. + """ + del tile_m # ignored; see docstring. See _select_fused_rs_block_sizes. + from jax.sharding import PartitionSpec as P + + size_m = lhs_indices.shape[0] + num_experts = w1.shape[0] + num_experts_per_shard = num_experts // ep_size + hidden_size = w2.shape[-1] + sls = min(pltpu.get_tpu_info().get_sublane_tiling(hidden_states.dtype), + size_m) + + # Single source of truth for tile_m. Using the same helper the kernel + # calls internally guarantees that ``compute_num_gm`` / ``max_num_gm`` + # (outer loop bound) matches the kernel's ``fill_metadata``-derived + # ``local_num_gm``. Mismatch here causes the kernel's final prefetch + # gather DMA to be unawaited and leaks ``gather_sem`` on kernel exit. + block_sizes = _select_fused_rs_block_sizes( + size_m=size_m, + size_k1=w1.shape[1], + size_n1=w1.shape[2], + size_k2=w2.shape[1], + size_n2=w2.shape[2], + size_group=num_experts_per_shard, + size_lhs_group=group_sizes.shape[0], + ep_size=ep_size, + out_dtype=hidden_states.dtype, + w1_dtype=w1.dtype, + w2_dtype=w2.dtype, + # This entry point does not pipe scales through to gmm_v2_fused_rs, + # so gmm_v2_fused_rs sees ``is_quantized=False`` internally. + is_quantized=False, + quant_block_size=None, + act_fn=act_fn, + ) + kernel_tile_m = block_sizes.tile_m + + group_offset = jnp.arange(0, num_experts, num_experts_per_shard) + + def _run(h, w1l, w2l, gs, go, li, oi, tw, ti): + my_id = lax.axis_index(ep_axis_name) + num_tokens = h.shape[0] + chunk_size = num_tokens // ep_size + + num_local_experts = w1l.shape[0] + local_group_sizes = lax.dynamic_slice(gs, (go[0], ), + (num_local_experts, )) + local_num_gm = compute_num_gm(local_group_sizes, kernel_tile_m, sls) + send_dest_chips = li // chunk_size + max_num_gm = lax.pmax(local_num_gm, axis_name=ep_axis_name) + + go_val = go[0] + gs_cumsum = jnp.cumsum(gs) + local_start = jnp.where(go_val > 0, gs_cumsum[go_val - 1], 0) + local_end = gs_cumsum[go_val + num_local_experts - 1] + # Single-pass recv_count: sum(dest == my_id AND NOT row_is_mine). + rows_arr = jnp.arange(li.shape[0], dtype=jnp.int32) + row_is_mine = jnp.logical_and(rows_arr >= local_start, rows_arr + < local_end) + to_me_remote = jnp.logical_and(send_dest_chips == my_id, + jnp.logical_not(row_is_mine)) + my_recv_count = jnp.sum(jnp.where(to_me_remote, 1, 0)) + total_recv_count = jnp.array([my_recv_count], dtype=jnp.int32) + + out_buf = gmm_v2_fused_rs( + h, + w1l, + w2l, + gs, + li, + oi, + act_fn=act_fn, + output_size=num_tokens, + group_offset=go, + topk_indices=ti, + ep_size=ep_size, + ep_axis_name=ep_axis_name, + max_num_gm=max_num_gm, + total_recv_count=total_recv_count, + top_k=top_k, + ) + + # Post-kernel reduction: apply topk_weights and sum over top_k. + # topk_weights arrives pre-sharded [chunk_size, top_k] via shard_map in_specs. + my_weights = tw + out_3d = out_buf.reshape(chunk_size, top_k, + hidden_size).astype(jnp.float32) + token_hidden = jnp.sum(out_3d * + my_weights.astype(jnp.float32)[:, :, None], + axis=1).astype(h.dtype) + + return lax.all_gather(token_hidden, + axis_name=ep_axis_name, + axis=0, + tiled=True) + + ep_p_spec = P(ep_axis_name) + replicated = P() + + return jax.shard_map( + _run, + mesh=mesh, + in_specs=( + replicated, # hidden_states (replicated on all chips) + ep_p_spec, # w1 (expert-sharded) + ep_p_spec, # w2 (expert-sharded) + replicated, # group_sizes (replicated) + ep_p_spec, # group_offset (expert-sharded) + replicated, # lhs_indices (replicated) + replicated, # output_indices (replicated) + P(ep_axis_name, + None), # topk_weights (EP-sharded, local chunk only) + replicated, # topk_indices (replicated) + ), + out_specs=replicated, + check_vma=False, + )( + hidden_states, + w1, + w2, + group_sizes, + group_offset, + lhs_indices, + output_indices, + topk_weights, + topk_indices, + ) diff --git a/tpu_inference/kernels/experimental/fused_moe/gmm_v2_gather_scatter.py b/tpu_inference/kernels/experimental/fused_moe/gmm_v2_gather_scatter.py new file mode 100644 index 0000000000..8d8944c36c --- /dev/null +++ b/tpu_inference/kernels/experimental/fused_moe/gmm_v2_gather_scatter.py @@ -0,0 +1,1304 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# 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 +# +# http://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. + +import dataclasses +import functools +from typing import Any, Callable, Tuple + +import jax +import jax.numpy as jnp +from jax import lax +from jax.experimental import pallas as pl +from jax.experimental.pallas import tpu as pltpu + +# --- Base dataclass + block-size helpers (kept here so the package is small) --- + + +@dataclasses.dataclass(frozen=True) +class TileSizes: + tile_m: int + tile_k: int + tile_n: int + + +def _recover_quant_block_size(size_k, num_blocks): + """Round the ceil(size_k / num_blocks) estimate up to the next power of two.""" + if num_blocks <= 1: + return size_k + approx = -(-size_k // num_blocks) + return 1 << (approx - 1).bit_length() + + +def get_tuned_block_sizes_v2( + m, + k, + n, + num_current_groups, + lhs_dtype, + rhs_dtype, + maybe_quantize_lhs, + rhs_quant_block_size, + default_block_sizes, + lhs_indices=None, + output_indices=None, + fuse_act=None, + fused=False, +): + # Device-specific tuned tables were removed; return the default. + return default_block_sizes + + +get_tuned_block_sizes = get_tuned_block_sizes_v2 + + +def get_maybe_quantize_lhs( + rhs_dtype=None, + rhs_quant_block_size: int | None = None, + lhs_quant_block_size: int | None = None, +) -> bool: + """Return whether to online-quantize LHS activations. + + FP4 weights use W4A16 when RHS scale blocks are smaller than the LHS + online-quant block; otherwise they use W4A8. Other quantized RHS dtypes keep + the existing online-quantized LHS path. + """ + if rhs_dtype is not None and jnp.dtype(rhs_dtype) == jnp.float4_e2m1fn: + if rhs_quant_block_size is None or lhs_quant_block_size is None: + return False + return rhs_quant_block_size >= lhs_quant_block_size + return True + + +# Util. +def align_to(x, a): + return pl.cdiv(x, a) * a + + +# Define data classes. +@jax.tree_util.register_dataclass +@dataclasses.dataclass(frozen=True) +class MetadataRef: + gm_id_to_group_id: jax.Array + gm_id_to_m_offset: jax.Array + + +@jax.tree_util.register_dataclass +@dataclasses.dataclass(frozen=True) +class WeightsRef: + weight: Any + scale: Any | None + bias: Any | None + + def get_weight(self): + return self.weight[...] + + def get_scale(self): + return self.scale[...] + + def get_bias(self): + return self.bias[...] + + +@jax.tree_util.register_dataclass +@dataclasses.dataclass(frozen=True) +class FusedWeightsRef: + """Wraps gate and up WeightsRef for fused activation.""" + + gate: WeightsRef + up: WeightsRef + + def get_weight(self): + return jnp.concatenate([self.gate.weight[...], self.up.weight[...]], + axis=-1) + + def get_scale(self): + return jnp.concatenate([self.gate.scale[...], self.up.scale[...]], + axis=-1) + + def get_bias(self): + return jnp.concatenate([self.gate.bias[...], self.up.bias[...]], + axis=-1) + + +@dataclasses.dataclass(frozen=True) +class Dimensions: + size_m: int + size_k: int + size_n: int + size_group: int + size_lhs_group: int + size_lhs_sublane: int + + +@dataclasses.dataclass(frozen=True) +class InputConfigs: + quant_dtype: jnp.dtype | None + quant_block_size: int | None + dtype: jnp.dtype + has_bias: bool = False + has_scale: bool = False + packing: int = 1 + num_quant_blocks: int = 1 + + +@dataclasses.dataclass(frozen=True) +class GmmConfigs: + tiles: TileSizes + dims: Dimensions + lhs_cfgs: InputConfigs + rhs_cfgs: InputConfigs + out_dtype: jnp.dtype + acc_dtype: jnp.dtype + zero_init: bool + fuse_act: str | None = None + has_post_norm: bool = False + + @property + def num_quant_blocks_per_tile_k(self) -> int: + return pl.cdiv(self.tiles.tile_k, self.rhs_cfgs.quant_block_size) + + @property + def out_size_n(self) -> int: + if self.fuse_act is None: + return self.dims.size_n + return self.dims.size_n // 2 + + +TileFn = Callable[[Dimensions, InputConfigs, InputConfigs, int, str | None], + TileSizes] + + +def apply_act_fn(acc: jax.Array, fuse_act: str | None) -> jax.Array: + """Apply fused activation (split gate/up, activate, multiply).""" + if fuse_act is None: + return acc + acc_gate, acc_up = jnp.split(acc, 2, -1) + match fuse_act: + case "silu": + return jax.nn.silu(acc_gate) * acc_up + case "gelu": + return jax.nn.gelu(acc_gate) * acc_up + case "swigluoai": + limit = 7.0 + alpha = 1.702 + acc_gate = jnp.clip(acc_gate, max=limit) + acc_up = jnp.clip(acc_up, min=-limit, max=limit) + return (acc_gate * jax.nn.sigmoid(alpha * acc_gate)) * (acc_up + 1) + case _: + raise NotImplementedError(f"Unsupported fuse_act: {fuse_act}") + + +class IndexMaps: + """Index maps for GMM kernel.""" + + def __init__(self, metadata_ref: MetadataRef, cfgs: GmmConfigs): + self.metadata_ref = metadata_ref + self.cfgs = cfgs + + def lhs_index_map(self, gm_id: jax.Array, _: jax.Array, k_id: jax.Array): + m_start = self.metadata_ref.gm_id_to_m_offset[gm_id] + m_end = self.metadata_ref.gm_id_to_m_offset[gm_id + 1] + + row_start = m_start // self.cfgs.dims.size_lhs_sublane + row_end = pl.cdiv(m_end, self.cfgs.dims.size_lhs_sublane) + row_size = row_end - row_start + + return (pl.ds(row_start, row_size), 0, k_id) + + def rhs_weight_index_map(self, gm_id: jax.Array, n_id: jax.Array, + k_id: jax.Array): + group_id = self.metadata_ref.gm_id_to_group_id[gm_id] + return (group_id, k_id, n_id) + + def rhs_bias_index_map(self, gm_id: jax.Array, n_id: jax.Array, + _: jax.Array): + group_id = self.metadata_ref.gm_id_to_group_id[gm_id] + return (group_id, 0, n_id) + + def rhs_scale_index_map(self, gm_id: jax.Array, n_id: jax.Array, + k_id: jax.Array): + group_id = self.metadata_ref.gm_id_to_group_id[gm_id] + b_id = (k_id * + self.cfgs.tiles.tile_k) // self.cfgs.rhs_cfgs.quant_block_size + b_tile_id = b_id // self.cfgs.num_quant_blocks_per_tile_k + return (group_id, b_tile_id, 0, n_id) + + def out_index_map(self, gm_id: jax.Array, n_id: jax.Array, _: jax.Array): + is_last_gm = gm_id == (pl.num_programs(0) - 1) + m_start = self.metadata_ref.gm_id_to_m_offset[gm_id] + m_end = self.metadata_ref.gm_id_to_m_offset[gm_id + 1] + + row_start = m_start // self.cfgs.dims.size_lhs_sublane + capped_row_end = m_end // self.cfgs.dims.size_lhs_sublane + last_row_end = pl.cdiv(m_end, self.cfgs.dims.size_lhs_sublane) + row_end = jnp.where(is_last_gm, last_row_end, capped_row_end) + row_size = row_end - row_start + + return (pl.ds(row_start, row_size), 0, n_id) + + +def generate_block_specs( + metadata_ref: MetadataRef, cfgs: GmmConfigs +) -> Tuple[Tuple[pl.BlockSpec, WeightsRef], pl.BlockSpec]: + """Generates block specs for the given lhs, rhs, and out refs.""" + + index_map = IndexMaps(metadata_ref, cfgs) + bounded_slice_gm = pl.BoundedSlice(cfgs.tiles.tile_m // + cfgs.dims.size_lhs_sublane) + + lhs_block_spec = pl.BlockSpec( + (bounded_slice_gm, cfgs.dims.size_lhs_sublane, cfgs.tiles.tile_k), + index_map.lhs_index_map, + ) + + rhs_weight_spec = pl.BlockSpec( + (None, cfgs.tiles.tile_k // cfgs.rhs_cfgs.packing, cfgs.tiles.tile_n), + index_map.rhs_weight_index_map, + pipeline_mode=pl.Buffered(buffer_count=3), + ) + rhs_scale_block_spec = rhs_bias_block_spec = None + if cfgs.rhs_cfgs.has_bias: + rhs_bias_block_spec = pl.BlockSpec( + (None, 1, cfgs.tiles.tile_n), + index_map.rhs_bias_index_map, + ) + if cfgs.rhs_cfgs.has_scale: + rhs_scale_block_spec = pl.BlockSpec( + (None, cfgs.num_quant_blocks_per_tile_k, 1, cfgs.tiles.tile_n), + index_map.rhs_scale_index_map, + ) + + rhs_block_spec = WeightsRef( + weight=rhs_weight_spec, + scale=rhs_scale_block_spec, + bias=rhs_bias_block_spec, + ) + + out_block_spec = pl.BlockSpec( + (bounded_slice_gm, cfgs.dims.size_lhs_sublane, cfgs.tiles.tile_n), + index_map.out_index_map, + ) + + return (lhs_block_spec, rhs_block_spec), out_block_spec + + +# Define kernels. +def inner_kernel( + # In + tiled_lhs_ref: jax.Array, + # [tile_m // size_lhs_sublane, size_lhs_sublane, tile_k] + tiled_rhs_ref: WeightsRef, # [tile_k, tile_n] + # Out + tiled_out_ref: jax.Array, + # [tile_m // size_lhs_sublane, size_lhs_sublane, tile_n] or [tile_m, tile_n] in scatter_mode + # Scratch + partial_out_ref: jax.Array, # [size_lhs_sublane, partial_out_n] + acc_ref: jax.Array, # [tile_m, tile_n] + metadata_ref: MetadataRef, + *, + cfgs: GmmConfigs, + gs_gate_ref=None, + gs_up_ref=None, + scatter_mode: bool = False, + _k_id: int | None = None, + _num_k: int | None = None, + _gm_id: int | None = None, + _n_id: int | None = None, +): + """Inner kernel invoked by emit_pipeline to perform matmul. + + tiled_lhs_ref and tiled_out_ref points to rows [m_start:m_end] of lhs and out. + Additionally, m_start and m_end does not have to align with tile boundaries + [m_offset:m_offset+tile_m]. Therefore, rows [m_offset:m_start] and + [m_end:m_offset+tile_m] of tiled_lhs_ref and tiled_out_ref will contain + invalid data and needs to be masked out. + + Args: + tiled_lhs_ref: Contains value lhs[m_start:m_end, k_start:k_end] + tiled_rhs_ref: Contains value rhs[g_id, k_start:k_end, n_start:n_end]. where + g_id is the group associated with lhs[m_start:m_end, :] + tiled_out_ref: Contains value out[m_start:m_end, n_start:n_end] + partial_out_ref: Contains last size_lhs_sublane rows of the previous output. + Will be initialized to zero if this is first tile for grid[n_id, :, :]. + acc_ref: Reference to the accumulator. + metadata_ref: Reference to the metadata. + cfgs: GmmConfigs. + """ + # Do not feed an unquantized LHS directly into MXU with a quantized RHS: + # that mixed-dtype matmul has produced numerically wrong TPU results + # (~70% cos-sim vs. truth). The quantized RHS path below avoids that in + # both supported cases: W4A8/W8A8 online-quantizes LHS before matmul, + # while W4A16 promotes/dequantizes RHS to bf16 before matmul. + # Small FP4 RHS scale blocks are folded into RHS and requantized to the + # LHS quant dtype before matmul so their per-block scales are not dropped. + _resolved_gm_id = pl.program_id(0) if _gm_id is None else _gm_id + + def _matmul(is_first_k_step: bool, is_last_k_step: bool): + tiled_lhs = tiled_lhs_ref.reshape(-1, cfgs.tiles.tile_k)[...] + tiled_rhs = tiled_rhs_ref.get_weight() + + num_quant_blocks_per_tile_k = cfgs.num_quant_blocks_per_tile_k + + # 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.packing > 1: + tiled_rhs = pltpu.bitcast(tiled_rhs, cfgs.rhs_cfgs.quant_dtype) + + def _valid_m_mask(shape): + gm_id = _resolved_gm_id + m_start = metadata_ref.gm_id_to_m_offset[gm_id] + m_end = metadata_ref.gm_id_to_m_offset[gm_id + 1] + m_offset = m_start - m_start % cfgs.dims.size_lhs_sublane + m_start_local = m_start - m_offset + m_end_local = m_end - m_offset + iota = lax.broadcasted_iota(jnp.int32, shape, 0) + return jnp.logical_and(m_start_local <= iota, iota < m_end_local) + + def _get_scale_slice(b_id, start_n, end_n): + rhs_scale = tiled_rhs_ref.get_scale() + return rhs_scale[..., b_id, :, start_n:end_n] + + def _online_quantize(block, axis, quant_dtype): + if jnp.issubdtype(quant_dtype, jnp.floating): + dtype_max = float(jnp.finfo(quant_dtype).max) + else: + dtype_max = float(jnp.iinfo(quant_dtype).max) + + block_f32 = block.astype(jnp.float32) + amax = jnp.max(jnp.abs(block_f32), axis=axis, keepdims=True) + amax_safe = jnp.where(amax == 0, jnp.ones_like(amax), amax) + block_q = (block_f32 / amax_safe * dtype_max).astype(quant_dtype) + block_scale = amax / dtype_max + return block_q, block_scale + + valid_k = cfgs.dims.size_k % cfgs.tiles.tile_k + if is_last_k_step and valid_k != 0: + mask_rhs = lax.broadcasted_iota(jnp.int32, tiled_rhs.shape, + 0) < valid_k + tiled_rhs = jnp.where(mask_rhs, tiled_rhs, 0) + mask_lhs = lax.broadcasted_iota(jnp.int32, tiled_lhs.shape, + 1) < valid_k + tiled_lhs = jnp.where(mask_lhs, tiled_lhs, 0) + + if cfgs.rhs_cfgs.quant_dtype is None: + # Unquantized RHS matmul path. + acc_list = [] + mxu_size = pltpu.get_tpu_info().mxu_column_size + rhs_qbs = cfgs.rhs_cfgs.quant_block_size + rhs_tile_n = tiled_rhs.shape[-1] + 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 b_id in range(num_quant_blocks_per_tile_k): + k_start = b_id * rhs_qbs + k_end = (b_id + 1) * rhs_qbs + partial_result = jnp.matmul( + tiled_lhs[:, k_start:k_end], + tiled_rhs[k_start:k_end, start_n:end_n], + preferred_element_type=jnp.float32, + ) + if cfgs.rhs_cfgs.has_scale: + rhs_scale_slice = _get_scale_slice( + b_id, start_n, end_n) + partial_result *= rhs_scale_slice + acc_n = acc_n + partial_result + acc_list.append(acc_n.astype(acc_ref.dtype)) + acc = jnp.concatenate(acc_list, axis=1) + else: + # Quantized RHS path. When one stored RHS scale covers a full LHS + # quant block, apply it after matmul. Otherwise fold the per-slice + # RHS scales into RHS before matmul so small scale blocks are not + # dropped. Online quantization avoids reciprocal scales; zero + # blocks use amax_safe=1 and scale=0. + q_block_size: int | None = cfgs.lhs_cfgs.quant_block_size + fp8_activation_quant = cfgs.lhs_cfgs.quant_dtype is not None + rhs_qbs = cfgs.rhs_cfgs.quant_block_size + rhs_scale_matches_lhs_block = (rhs_qbs >= q_block_size + and rhs_qbs % q_block_size == 0) + apply_rhs_scale_after_matmul = cfgs.rhs_cfgs.has_scale and ( + rhs_scale_matches_lhs_block) + + # 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]. + acc_list = [] + mxu_size = pltpu.get_tpu_info().mxu_column_size + rhs_tile_n = tiled_rhs.shape[-1] + 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): + end_k = min(cfgs.tiles.tile_k, start_k + q_block_size) + + block_lhs = tiled_lhs[:, start_k:end_k] + block_rhs = tiled_rhs[start_k:end_k, start_n:end_n] + + if cfgs.rhs_cfgs.has_scale and not apply_rhs_scale_after_matmul: + scaled_rhs_slices = [] + rhs_k_start = start_k + while rhs_k_start < end_k: + b_id = rhs_k_start // rhs_qbs + rhs_k_end = min((b_id + 1) * rhs_qbs, end_k) + local_k_start = rhs_k_start - start_k + local_k_end = rhs_k_end - start_k + rhs_scale_slice = _get_scale_slice( + b_id, start_n, end_n) + block_rhs_slice = block_rhs[ + local_k_start:local_k_end] + scaled_rhs_slice = block_rhs_slice.astype( + jnp.bfloat16) * rhs_scale_slice.astype( + jnp.bfloat16) + scaled_rhs_slices.append(scaled_rhs_slice) + rhs_k_start = rhs_k_end + block_rhs_for_quant = jnp.concatenate( + scaled_rhs_slices, axis=0) + else: + block_rhs_for_quant = block_rhs.astype(jnp.bfloat16) + + if fp8_activation_quant: + block_lhs_for_matmul, lhs_scale = _online_quantize( + block_lhs, + axis=1, + quant_dtype=cfgs.lhs_cfgs.quant_dtype) + if cfgs.rhs_cfgs.has_scale and not apply_rhs_scale_after_matmul: + # Fold small RHS scale blocks into RHS first, then + # requantize to the activation quant dtype so + # matmul uses fp8 x fp8. + block_rhs_for_matmul, rhs_online_scale = _online_quantize( + block_rhs_for_quant, + axis=0, + quant_dtype=cfgs.lhs_cfgs.quant_dtype, + ) + else: + block_rhs_for_matmul = block_rhs + rhs_online_scale = None + else: + block_lhs_for_matmul = block_lhs.astype(jnp.bfloat16) + block_rhs_for_matmul = block_rhs_for_quant + + block_acc = jnp.matmul( + block_lhs_for_matmul, + block_rhs_for_matmul, + preferred_element_type=jnp.float32, + ).astype(acc_ref.dtype) + + if fp8_activation_quant: + block_acc *= lhs_scale.astype(acc_ref.dtype) + if rhs_online_scale is not None: + block_acc *= rhs_online_scale.astype(acc_ref.dtype) + + if apply_rhs_scale_after_matmul: + b_id = start_k // rhs_qbs + rhs_scale_slice = _get_scale_slice( + b_id, start_n, end_n) + block_acc *= rhs_scale_slice.astype(acc_ref.dtype) + + acc_n += block_acc + + acc_list.append(acc_n) + acc = jnp.concatenate(acc_list, axis=1) + + if not is_first_k_step: + acc += acc_ref[...] + + if is_last_k_step: + if cfgs.rhs_cfgs.has_scale and gs_gate_ref is not None: + _gid = metadata_ref.gm_id_to_group_id[_resolved_gm_id] + if cfgs.fuse_act is not None: + n_iota = lax.broadcasted_iota(jnp.int32, acc.shape, 1) + global_scale = jnp.where( + n_iota >= acc.shape[1] // 2, + gs_up_ref[_gid], + gs_gate_ref[_gid], + ) + else: + global_scale = gs_gate_ref[_gid] + acc *= global_scale.astype(acc.dtype) + + if cfgs.rhs_cfgs.has_bias: + acc += tiled_rhs_ref.get_bias().astype(acc.dtype) + + acc = apply_act_fn(acc, cfgs.fuse_act) + + gm_id = _resolved_gm_id + # Mask out rows that does not belong to the current group. + m_start = metadata_ref.gm_id_to_m_offset[gm_id] + m_end = metadata_ref.gm_id_to_m_offset[gm_id + 1] + m_offset = m_start - m_start % cfgs.dims.size_lhs_sublane + m_start_local = m_start - m_offset + m_end_local = m_end - m_offset + + iota = lax.broadcasted_iota(jnp.int32, acc.shape, 0) + mask = jnp.logical_and(m_start_local <= iota, iota < m_end_local) + + if scatter_mode: + # In scatter mode, write masked acc directly to 2D output. + # No 3D reshape needed — avoids VMEM tiling mismatch with DMA. + # No partial_out needed — each scatter tile is independent. + acc_masked = jnp.where(mask, acc, 0) + tiled_out_ref[...] = acc_masked.astype(tiled_out_ref.dtype) + else: + acc_masked = jnp.where(mask, acc, + 0).reshape(tiled_out_ref.shape) + + # Write the final output to the output ref. + tiled_out_ref[...] = acc_masked.astype(tiled_out_ref.dtype) + + # partial_out is n-aware: use n_id to index into the correct + # n-tile slice of partial_out_ref (shape: sls × total_n). + n_id = pl.program_id(1) if _n_id is None else _n_id + n_offset = n_id * cfgs.tiles.tile_n + tile_n = cfgs.tiles.tile_n + sls = cfgs.dims.size_lhs_sublane + partial_out_zeros = jnp.zeros((sls, tile_n), + dtype=partial_out_ref.dtype) + + # Accumulate the partial output from the previous gm step + # at the same n position. + tiled_out_ref[0] += jnp.where( + gm_id == 0, + partial_out_zeros, + partial_out_ref[:, pl.ds(n_offset, tile_n)], + ) + + # Consider following case where size_lhs_sublane = 4, number denotes group + # id and | denotes boundaries between sublanes: + # | 0 0 1 2 | 2 2 2 2 | 3 3 4 4 | + # + # Assuming group id of current step is 1, current step will not completely + # fill size_lhs_sublane rows and will be revisited at the next step. By + # storing the partial rows into the partial_out_ref, the next step can + # read them and accumulate to them. Additionally, for group id of 2, + # since it completely fills the size_lhs_sublane rows, we need to zero out + # partial_out_ref to avoid numeric error for group 3. + last_row = m_end_local // cfgs.dims.size_lhs_sublane + partial_out_ref[:, pl.ds(n_offset, tile_n)] = jnp.where( + m_end_local % cfgs.dims.size_lhs_sublane == 0, + partial_out_zeros, + tiled_out_ref[last_row], + ) + else: + acc_ref[...] = acc + + # Define matmul wrapper functions. + @jax.named_scope("matmul_first_last") + def matmul_first_last(): + _matmul(is_first_k_step=True, is_last_k_step=True) + + @jax.named_scope("matmul_first") + def matmul_first(): + _matmul(is_first_k_step=True, is_last_k_step=False) + + @jax.named_scope("matmul") + def matmul(): + _matmul(is_first_k_step=False, is_last_k_step=False) + + @jax.named_scope("matmul_last") + def matmul_last(): + _matmul(is_first_k_step=False, is_last_k_step=True) + + # Select and execute matmul function based on the current step. + num_k = pl.num_programs(2) if _num_k is None else _num_k + k_id = pl.program_id(2) if _k_id is None else _k_id + + is_first_k_step = k_id == 0 + is_last_k_step = k_id == (num_k - 1) + + lax.cond( + is_first_k_step, + lambda: lax.cond( + is_last_k_step, + matmul_first_last, + matmul_first, + ), + lambda: lax.cond( + is_last_k_step, + matmul_last, + matmul, + ), + ) + + +@jax.named_scope("fill_metadata") +def fill_metadata( + lhs_group_sizes_ref: jax.Array, # int32[size_lhs_group] + group_offset_ref: jax.Array, # int32[1] + metadata_ref: MetadataRef, + *, + cfgs: GmmConfigs, +) -> jax.Array: + """Fills the metadata for the given lhs group sizes and group offset. + + Iterates over the lhs group sizes and if the group id is valid, determines + the number of gm tiles that are needed to process the current group. Then, + it fills starting and ending offset (gm_id_to_m_offset), and the group id + (gm_id_to_group_id) for each gm tile. + + Args: + lhs_group_sizes_ref: The group sizes of lhs. + group_offset_ref: Offset of the first group to process. + metadata_ref: Metadata that is used to determine the group id and m offsets + for each gmm tile. + cfgs: GmmConfigs. + + Returns: + The number of gm tiles to process lhs with given group offset. + """ + + group_offset = group_offset_ref[0] + max_num_group = group_offset + cfgs.dims.size_group + metadata_ref.gm_id_to_m_offset[0] = 0 + + @jax.named_scope("inner_tm_loop") + def inner_tm_loop(tm_id, curr_m_offset, *, end_m_offset, group_id): + local_offset = curr_m_offset % cfgs.dims.size_lhs_sublane + tm_size = jnp.minimum(cfgs.tiles.tile_m - local_offset, + end_m_offset - curr_m_offset) + + metadata_ref.gm_id_to_group_id[tm_id] = group_id + + next_m_offset = curr_m_offset + tm_size + metadata_ref.gm_id_to_m_offset[tm_id] = curr_m_offset + metadata_ref.gm_id_to_m_offset[tm_id + 1] = next_m_offset + + return next_m_offset + + @jax.named_scope("outer_group_loop") + def outer_group_loop(lhs_group_id, carry): + num_gm, start_m_offset = carry + + group_id = lhs_group_id - group_offset + group_size = lhs_group_sizes_ref[lhs_group_id] + end_m_offset = start_m_offset + group_size + + # Assume following arguments: + # - size_lhs_sublane & tile_m = 4 + # - group_size = 3 + # - start_m_offset = 7 + # + # If we visualize it, it will look like this where: + # - |: denotes boundaries between sublanes + # - 0: denotes values for other groups + # - 1: denotes values for the current group + # | 0 0 0 0 | 0 0 0 1 | 1 1 0 0 | + # + # In this example, we see that we require processing 2 m tiles. + # But, performing a naive cdiv(group_size, tile_m) will return 1. + # Instead, adding local_offset will give us the correct value. + local_offset = start_m_offset % cfgs.dims.size_lhs_sublane + aligned_group_size = group_size + local_offset + curr_num_gm = pl.cdiv(aligned_group_size, cfgs.tiles.tile_m) + + # We need to handle cases where we should not process the group. + # 1. Even if group_size is 0, if local_offset is not 0, cdiv will return 1. + # 2. If group comes before the group_offset, we should not process it. + should_process = jnp.logical_and(group_size > 0, group_id >= 0) + curr_num_gm = jnp.where(should_process, curr_num_gm, 0) + next_num_gm = num_gm + curr_num_gm + + tm_loop_fn = functools.partial( + inner_tm_loop, + end_m_offset=end_m_offset, + group_id=group_id, + ) + lax.fori_loop(num_gm, next_num_gm, tm_loop_fn, start_m_offset) + + return next_num_gm, end_m_offset + + num_gm, _ = lax.fori_loop(0, max_num_group, outer_group_loop, (0, 0)) + return num_gm + + +@jax.named_scope("zero_out_start") +def zero_out_start( + out_ref: jax.Array, # [size_m, size_n] + zero_ref: jax.Array, # [tile_zero_m, num_lanes] + semaphore_ref: jax.Array, # [1] + metadata_ref: MetadataRef, + num_gm: jax.Array, + *, + dims: Dimensions, +): + """Zero out output rows that are not used in the computation.""" + + num_lanes = pltpu.get_tpu_info().num_lanes + assert num_lanes == zero_ref.shape[-1] + zero_ref[...] = jnp.zeros_like(zero_ref) + + zero_dma = zero_ref.reshape(-1, dims.size_lhs_sublane, num_lanes) + out_dma = out_ref.reshape(-1, dims.size_lhs_sublane, out_ref.shape[-1]) + row_size = zero_dma.shape[0] + + compute_start = metadata_ref.gm_id_to_m_offset[0] + compute_end = metadata_ref.gm_id_to_m_offset[num_gm] + + left_zero_start = 0 + left_zero_end = compute_start // dims.size_lhs_sublane + left_zero_size = left_zero_end - left_zero_start + left_num_loops = pl.cdiv(left_zero_size, row_size) + + right_zero_start = pl.cdiv(compute_end, dims.size_lhs_sublane) + right_zero_end = out_dma.shape[0] + right_zero_size = right_zero_end - right_zero_start + right_num_loops = pl.cdiv(right_zero_size, row_size) + + def fill_zero(i, zero_size, *, start, end): + dma_start = start + i * row_size + dma_end = jnp.minimum(dma_start + row_size, end) + dma_size = dma_end - dma_start + + # Static loop. Will be unrolled during compile time. + for n_start in range(0, out_ref.shape[-1], num_lanes): + n_end = n_start + num_lanes + pltpu.make_async_copy( + src_ref=zero_dma.at[pl.ds(0, dma_size)], + dst_ref=out_dma.at[pl.ds(dma_start, dma_size), :, + n_start:n_end], + sem=semaphore_ref.at[0], + ).start(priority=1) + + return zero_size + dma_size + + @jax.named_scope("left_fill_zero") + def left_fill_zero(i, zero_size): + return fill_zero(i, + zero_size, + start=left_zero_start, + end=left_zero_end) + + @jax.named_scope("right_fill_zero") + def right_fill_zero(i, zero_size): + return fill_zero(i, + zero_size, + start=right_zero_start, + end=right_zero_end) + + zero_size = lax.fori_loop(0, left_num_loops, left_fill_zero, 0) + zero_size = lax.fori_loop(0, right_num_loops, right_fill_zero, zero_size) + return zero_size + + +@jax.named_scope("zero_out_end") +def zero_out_end( + out_ref: jax.Array, # [size_m, size_n] + semaphore_ref: jax.Array, # [1] + zero_size: jax.Array, + *, + dims: Dimensions, +): + out_dma = out_ref.reshape(-1, dims.size_lhs_sublane, out_ref.shape[-1]) + pltpu.make_async_copy( + src_ref=out_dma.at[pl.ds(0, zero_size)], + dst_ref=out_dma.at[pl.ds(0, zero_size)], + sem=semaphore_ref.at[0], + ).wait() + + +@jax.named_scope("zero_out_start_3d") +def zero_out_start_3d( + out_ref: jax.Array, # [size_m, size_n // num_lanes, num_lanes] + zero_src_ref: jax.Array, # [tile_rows, size_n // num_lanes, num_lanes] + semaphore_ref: jax.Array, # [1] +): + """Zero out ALL output rows via DMA. Required for DMA scatter where output + positions are non-contiguous (scattered). + + Reuses an existing 3D VMEM scratch (e.g. scatter_staging_ref) as the zero + source — no dedicated zero_ref allocation needed. The caller must ensure + zero_src_ref is not in use when this function runs. + """ + + zero_src_ref[...] = jnp.zeros_like(zero_src_ref) + row_size = zero_src_ref.shape[0] + + total_rows = out_ref.shape[0] + num_loops = pl.cdiv(total_rows, row_size) + + def fill_zero(i, zero_size): + dma_start = i * row_size + dma_end = jnp.minimum(dma_start + row_size, total_rows) + dma_size = dma_end - dma_start + + pltpu.make_async_copy( + src_ref=zero_src_ref.at[pl.ds(0, dma_size)], + dst_ref=out_ref.at[pl.ds(dma_start, dma_size)], + sem=semaphore_ref.at[0], + ).start(priority=1) + + return zero_size + dma_size + + zero_size = lax.fori_loop(0, num_loops, fill_zero, 0) + return zero_size + + +@jax.named_scope("zero_out_end_3d") +def zero_out_end_3d( + out_ref: jax.Array, # [size_m, size_n // num_lanes, num_lanes] + semaphore_ref: jax.Array, # [1] + zero_size: jax.Array, +): + """Wait for all zero-fill DMAs to complete. Works with 3D output refs.""" + pltpu.make_async_copy( + src_ref=out_ref.at[pl.ds(0, zero_size)], + dst_ref=out_ref.at[pl.ds(0, zero_size)], + sem=semaphore_ref.at[0], + ).wait() + + +@jax.named_scope("dma_gather_gm_start") +def dma_gather_gm_start(src_ref, + dst_ref, + indices_ref, + sem_ref, + gm_id, + metadata_ref, + divisor: int = 1): + """Start gathering rows for a specific gm tile via DMA. + + src_ref and dst_ref must be 3D: (rows, k // num_lanes, num_lanes). + No reshape — reshape on refs breaks dynamic pl.ds offsets. + + `divisor`: optional integer divisor applied to indices_ref values before + use. Set > 1 when indices_ref contains packed values (e.g., + `combined = lhs_idx * divisor + extra_field`) and we need to recover the + actual src_row via integer division. Default 1 = no unpacking. + """ + m_start = metadata_ref.gm_id_to_m_offset[gm_id] + m_end = metadata_ref.gm_id_to_m_offset[gm_id + 1] + sls = pltpu.get_tpu_info().get_sublane_tiling(src_ref.dtype) + m_start_local = m_start % sls + + def _gather_body(i, _): + row = m_start + i + src_row = indices_ref[row] + if divisor != 1: + src_row = src_row // divisor + pltpu.make_async_copy( + src_ref=src_ref.at[pl.ds(src_row, 1), :, :], + dst_ref=dst_ref.at[pl.ds(m_start_local + i, 1), :, :], + sem=sem_ref, + ).start() + return _ + + lax.fori_loop(0, m_end - m_start, _gather_body, 0) + + +@jax.named_scope("dma_gather_gm_wait") +def dma_gather_gm_wait(dst_ref, sem_ref, gm_id, metadata_ref): + """Wait for all gather DMAs for a specific gm tile to complete. + + dst_ref must be 3D: (rows, k // num_lanes, num_lanes). + """ + m_start = metadata_ref.gm_id_to_m_offset[gm_id] + m_end = metadata_ref.gm_id_to_m_offset[gm_id + 1] + num_rows = m_end - m_start + sls = pltpu.get_tpu_info().get_sublane_tiling(dst_ref.dtype) + m_start_local = m_start % sls + pltpu.make_async_copy( + src_ref=dst_ref.at[pl.ds(m_start_local, num_rows), :, :], + dst_ref=dst_ref.at[pl.ds(m_start_local, num_rows), :, :], + sem=sem_ref, + ).wait() + + +@jax.named_scope("dma_scatter_gm_start") +def dma_scatter_gm_start(src_ref, dst_ref, indices_ref, sem_ref, gm_id, + metadata_ref): + """Start scattering rows for a specific gm tile via DMA. + + src_ref and dst_ref must be 3D: (rows, n // num_lanes, num_lanes). + No reshape — reshape on refs breaks dynamic pl.ds offsets. + """ + m_start = metadata_ref.gm_id_to_m_offset[gm_id] + m_end = metadata_ref.gm_id_to_m_offset[gm_id + 1] + sls = pltpu.get_tpu_info().get_sublane_tiling(src_ref.dtype) + m_start_local = m_start % sls + + def _scatter_body(i, _): + row = m_start + i + dst_row = indices_ref[row] + pltpu.make_async_copy( + src_ref=src_ref.at[pl.ds(m_start_local + i, 1), :, :], + dst_ref=dst_ref.at[pl.ds(dst_row, 1), :, :], + sem=sem_ref, + ).start() + return _ + + lax.fori_loop(0, m_end - m_start, _scatter_body, 0) + + +@jax.named_scope("dma_scatter_gm_wait") +def dma_scatter_gm_wait(src_ref, sem_ref, gm_id, metadata_ref): + """Wait for all scatter DMAs for a specific gm tile to complete. + + src_ref must be 3D: (rows, n // num_lanes, num_lanes). + """ + m_start = metadata_ref.gm_id_to_m_offset[gm_id] + m_end = metadata_ref.gm_id_to_m_offset[gm_id + 1] + num_rows = m_end - m_start + sls = pltpu.get_tpu_info().get_sublane_tiling(src_ref.dtype) + m_start_local = m_start % sls + pltpu.make_async_copy( + src_ref=src_ref.at[pl.ds(m_start_local, num_rows), :, :], + dst_ref=src_ref.at[pl.ds(m_start_local, num_rows), :, :], + sem=sem_ref, + ).wait() + + +def calculate_tiling( + dims: Dimensions, + lhs_cfgs: InputConfigs, + rhs_cfgs: InputConfigs, + vmem_limit_bytes: int, + fuse_act: str | None = None, +) -> TileSizes: + """Calculate optimal tile sizes for GMM kernel.""" + + lhs_dtype = lhs_cfgs.quant_dtype or lhs_cfgs.dtype + rhs_dtype = rhs_cfgs.dtype + lhs_bits = jax.dtypes.itemsize_bits(lhs_dtype) + rhs_bits = jax.dtypes.itemsize_bits(rhs_dtype) + + # When using bf16 for lhs and rhs, 128 is the largest tile_m value that is + # safe to use for most scenarios. But if lower bitwidth is used, we need + # to tweak tile_m to account for using faster hardware unit. + # TODO(kyuyeunk): Account for different TPU hardware specs. + bf16_bf16_tile_m = 128 + lhs_mod = min(pl.cdiv(16, lhs_bits), 2) + rhs_mod = min(pl.cdiv(16, rhs_bits), 2) + tile_m = bf16_bf16_tile_m * lhs_mod // rhs_mod + tile_m = min(tile_m, dims.size_m) + + # Subtract non-rhs VMEM overhead before computing per-buffer budget. + # Overhead includes: gathered_lhs_2x (DMA gather), acc/partial_out buffers, + # tiled_out_2x (DMA scatter), and compiler spill headroom. + lhs_bits_item = jax.dtypes.itemsize_bits(lhs_cfgs.dtype) + _overhead = ( + 2 * tile_m * dims.size_k * lhs_bits_item // 8 # gathered_lhs_2x + + 5 * 1024 * 1024 # acc, partial_out, spill headroom + ) + rhs_vmem_budget = max(vmem_limit_bytes - _overhead, vmem_limit_bytes // 2) + + # Calculate vmem limit for a single rhs buffer when using triple buffers. + num_rhs_buffers = 3 + rhs_vmem_target = rhs_vmem_budget // num_rhs_buffers + base_rhs_size_bytes = dims.size_k * dims.size_n * rhs_bits // 8 + + # To avoid stalling MXU, we add some buffer room where tile_n cannot go + # smaller than 2x of mxu_column_size. + tile_n_limit = pltpu.get_tpu_info().mxu_column_size * 2 + tile_n_limit = min(tile_n_limit, dims.size_n) + + # When fuse_act is set, tile_n tiles the output N (= size_n // 2). + # base_rhs_size_bytes still uses size_n to account for loading both halves. + size_n_per_rhs = dims.size_n + if fuse_act is not None: + size_n_per_rhs //= 2 + tile_n_limit = min(tile_n_limit, size_n_per_rhs) + + # Initialize tile_k and tile_n to their maximum valid values. + num_k_tiles = num_n_tiles = 1 + num_lanes = pltpu.get_tpu_info().num_lanes + tile_k = align_to(dims.size_k, num_lanes) + tile_n = align_to(size_n_per_rhs, num_lanes) + + # Multiple k tiles will introduce accumulation overhead. Thus, we first try + # to fit rhs into vmem by only adjusting tile_n. + + # Decrease tile_n until rhs fits in vmem target. + while (pl.cdiv(base_rhs_size_bytes, num_n_tiles) > rhs_vmem_target + and tile_n > tile_n_limit): + num_n_tiles += 1 + tile_n = align_to(size_n_per_rhs, + num_n_tiles * num_lanes) // num_n_tiles + + # If decreasing tile_n is no longer possible, we decrease tile_k instead. + if tile_n < tile_n_limit: + num_n_tiles -= 1 + tile_n = align_to(size_n_per_rhs, + num_n_tiles * num_lanes) // num_n_tiles + + # Decrease tile_k until rhs fits in vmem target. + base_rhs_size_bytes = pl.cdiv(base_rhs_size_bytes, num_n_tiles) + while pl.cdiv(base_rhs_size_bytes, num_k_tiles) > rhs_vmem_target: + num_k_tiles += 1 + tile_k = align_to(dims.size_k, + num_k_tiles * num_lanes) // num_k_tiles + + if tile_n == 0 or tile_k == 0: + raise ValueError( + f"Could not find valid tile sizes for {dims=} and {rhs_vmem_target=}." + ) + + return TileSizes(tile_m=tile_m, tile_k=tile_k, tile_n=tile_n) + + +def validate_inputs( + lhs: jax.Array, + rhs: jax.Array, + rhs_scale: jax.Array | None, + rhs_bias: jax.Array | None, + group_sizes: jax.Array, + group_offset: jax.Array, + fuse_act: str | None = None, + packed_nvfp4: bool = False, +): + """Validates the inputs for the GMM kernel.""" + + size_m = lhs.shape[0] + size_group, size_k_raw, size_n = rhs.shape + size_k = size_k_raw * 2 if packed_nvfp4 else size_k_raw + size_lhs_group = group_sizes.shape[0] + + assert size_group <= size_lhs_group + assert lhs.shape == ( + size_m, + size_k, + ), f"lhs.shape={lhs.shape} expected=({size_m}, {size_k}), rhs.shape={rhs.shape}, group_sizes.shape={group_sizes.shape}" + assert rhs.shape == (size_group, size_k_raw, size_n) + if rhs_bias is not None: + assert rhs_bias.shape == (size_group, 1, size_n) + if rhs_scale is not None: + num_quant_blocks = rhs_scale.shape[1] + assert rhs_scale.shape == (size_group, num_quant_blocks, 1, size_n) + # When K is zero-padded for DMA alignment, size_k may not divide + # num_quant_blocks evenly. The original (unpadded) K always does. + # The inner kernel uses quant_block_size from make_gmm_configs which + # is derived from the original K, so this is safe. + + assert group_offset.shape == (1, ) + + size_lhs_sublane = pltpu.get_tpu_info().get_sublane_tiling(lhs.dtype) + size_lhs_sublane = min(size_lhs_sublane, size_m) + + return Dimensions( + size_m=size_m, + size_k=size_k, + size_n=size_n, + size_group=size_group, + size_lhs_group=size_lhs_group, + size_lhs_sublane=size_lhs_sublane, + ) + + +def get_cost_estimate(cfgs: GmmConfigs): + """Returns the cost estimate for the GMM kernel.""" + + dims = cfgs.dims + lhs_dtype = cfgs.lhs_cfgs.quant_dtype or cfgs.lhs_cfgs.dtype + rhs_dtype = cfgs.rhs_cfgs.dtype + + flops = 2 * dims.size_m * dims.size_k * dims.size_n + + lhs_bytes = dims.size_m * dims.size_k * lhs_dtype.itemsize + + rhs_bytes = (dims.size_group * dims.size_k * dims.size_n * + jax.dtypes.itemsize_bits(rhs_dtype)) // 8 + if cfgs.rhs_cfgs.has_scale: + rhs_bytes += (dims.size_group * cfgs.rhs_cfgs.num_quant_blocks * + dims.size_n * jnp.dtype(jnp.float32).itemsize) + if cfgs.rhs_cfgs.has_bias: + rhs_bytes += dims.size_group * dims.size_n * jnp.dtype( + jnp.float32).itemsize + + out_bytes = dims.size_m * dims.size_n * jnp.dtype(cfgs.out_dtype).itemsize + + total_bytes = lhs_bytes + rhs_bytes + out_bytes + + return pl.CostEstimate( + flops=flops, + bytes_accessed=total_bytes, + transcendentals=0, + ) + + +def get_scope_name(dims: Dimensions, tiles: TileSizes) -> str: + return ( + f"gmm_v2-g_{dims.size_group}-m_{dims.size_m}-k_{dims.size_k}" + f"-n_{dims.size_n}-tm_{tiles.tile_m}-tk_{tiles.tile_k}-tn_{tiles.tile_n}" + ) + + +def make_gmm_configs( + lhs: jax.Array, + rhs: jax.Array, + rhs_scale: jax.Array | None, + rhs_bias: jax.Array | None, + group_sizes: jax.Array, + group_offset: jax.Array, + *, + tile_info: TileSizes | TileFn, + vmem_limit_bytes: int | None, + out_dtype: jnp.dtype | None, + acc_dtype: jnp.dtype | None, + maybe_quantize_lhs: bool, + zero_initialize: bool, + lhs_indices: jax.Array | None = None, + output_indices: jax.Array | None = None, + original_k: int | None = None, + fuse_act: str | None = None, + post_expert_norm_weight: jax.Array | None = None, + original_size_m: int | None = None, + original_size_k: int | None = None, + original_size_n: int | None = None, + packed_nvfp4: bool = False, +): + """Fills the GMM config for the GMM kernel.""" + + dims = validate_inputs( + lhs, + rhs, + rhs_scale, + rhs_bias, + group_sizes, + group_offset, + fuse_act=fuse_act, + packed_nvfp4=packed_nvfp4, + ) + + if rhs_scale is not None: + has_scale = True + rhs_quant_dtype = jnp.float4_e2m1fn.dtype if packed_nvfp4 else rhs.dtype + num_blocks = rhs_scale.shape[1] + # When K was zero-padded for DMA alignment, use the original K + # to compute block_size so quant block boundaries stay correct. + _k_for_blocks = original_k if original_k is not None else dims.size_k + block_size = _recover_quant_block_size(_k_for_blocks, num_blocks) + rhs_packing = 8 if packed_nvfp4 else 32 // jax.dtypes.itemsize_bits( + rhs.dtype) + else: + has_scale = False + rhs_quant_dtype = None + num_blocks = 1 + block_size = dims.size_k + rhs_packing = 1 + + rhs_cfgs = InputConfigs( + quant_dtype=rhs_quant_dtype, + quant_block_size=block_size, + dtype=rhs.dtype, + has_bias=rhs_bias is not None, + has_scale=has_scale, + packing=rhs_packing, + num_quant_blocks=num_blocks, + ) + + lhs_quant_block_size = 256 if rhs_cfgs.quant_block_size < 512 else 512 + + lhs_q_dtype = None + if (maybe_quantize_lhs and get_maybe_quantize_lhs( + rhs_quant_dtype, + rhs_cfgs.quant_block_size, + lhs_quant_block_size, + ) and rhs_quant_dtype is not None): + # Choose lhs quantization dtype based on TPU hardware support. + is_rhs_float = jnp.issubdtype(rhs_quant_dtype, jnp.floating) + tpu_info = pltpu.get_tpu_info() + # Check if there is hardware compute support for rhs dtype group. + if is_rhs_float: + if tpu_info.fp8_ops_per_second > 0: + lhs_q_dtype = jnp.float8_e4m3fn.dtype + else: + if tpu_info.int8_ops_per_second > 0: + lhs_q_dtype = jnp.int8.dtype + + lhs_cfgs = InputConfigs( + quant_dtype=lhs_q_dtype, + # Input quantization involves reading all elements in a block to compute + # 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=lhs_quant_block_size, + dtype=lhs.dtype, + ) + + if out_dtype is None: + out_dtype = lhs.dtype + + if acc_dtype is None: + acc_dtype = acc_dtype = jnp.float32.dtype + + if isinstance(tile_info, TileSizes): + tiles = tile_info + else: + tiles = tile_info(dims, lhs_cfgs, rhs_cfgs, vmem_limit_bytes, fuse_act) + default_block_sizes = (tiles.tile_m, tiles.tile_k, tiles.tile_n) + # Use original sizes (before padding) if provided, to keep lookup keys consistent + # across different kernels + tile_m, tile_k, tile_n = get_tuned_block_sizes( + dims.size_m if original_size_m is None else original_size_m, + dims.size_k if original_size_k is None else original_size_k, + dims.size_n if original_size_n is None else original_size_n, + dims.size_group, + lhs_cfgs.dtype, + rhs_cfgs.dtype, + maybe_quantize_lhs, + rhs_cfgs.quant_block_size, + default_block_sizes, + lhs_indices, + output_indices, + fuse_act, + ) + tiles = TileSizes(tile_m=tile_m, tile_k=tile_k, tile_n=tile_n) + + return GmmConfigs( + dims=dims, + tiles=tiles, + lhs_cfgs=lhs_cfgs, + rhs_cfgs=rhs_cfgs, + out_dtype=out_dtype, + acc_dtype=acc_dtype, + zero_init=zero_initialize, + fuse_act=fuse_act, + has_post_norm=post_expert_norm_weight is not None, + ) + + +# ============================================================================= +# Fused kernel: gather + GMM1 + activation + GMM2 + scatter +# ============================================================================= +# Time-shares the VMEM weight buffer between GMM1 and GMM2, keeping the +# intermediate activation result in VMEM to avoid an HBM round-trip. +# Prototype: bf16 only, synchronous weight reads (no pipelining). +# ============================================================================= +@dataclasses.dataclass(frozen=True) +class FusedDims: + """Dimensions for the fused gather+GMM1+act+GMM2+scatter kernel.""" + + size_m: int # Number of gathered/scattered rows + size_k1: int # K for GMM1 (= hidden_size H, original unpadded) + size_n1: int # N for GMM1 (= 2 * intermediate_size, gate+up) + size_k2: int # K for GMM2 (= intermediate_size I) + size_n2: int # N for GMM2 (= hidden_size H), aligned for DMA + original_n2: int # Original (unpadded) N for GMM2 + size_group: int # Number of experts on this shard + size_lhs_group: int # Total number of expert groups + size_lhs_sublane: int # Sublane tiling for LHS dtype + intermediate_size: int = 0 # Original I (= N1_orig // 2) for activation split + has_bias: bool = False # Whether bias refs are present + quant_block_size: int | None = None # RHS quantization block size + num_scale_blocks1: int = 0 # K1 // quant_block_size (for scale DMA) + num_scale_blocks2: int = 0 # K2 // quant_block_size (for scale DMA) + padded_k1: int = 0 # DMA-aligned K1 (>= size_k1); 0 means no padding