diff --git a/tests/kernels/ragged_gather_reduce_v2_test.py b/tests/kernels/ragged_gather_reduce_v2_test.py new file mode 100644 index 0000000000..be170ecbe7 --- /dev/null +++ b/tests/kernels/ragged_gather_reduce_v2_test.py @@ -0,0 +1,271 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# 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 functools +import itertools +import time + +import jax +import jax.numpy as jnp +import numpy as np +from absl.testing import absltest, parameterized +from jax._src import test_util as jtu + +from tpu_inference.kernels.sparse_core.ragged_gather_reduce import \ + ragged_gather_reduce as ragged_gather_reduce_v1 +from tpu_inference.kernels.sparse_core.ragged_gather_reduce_v2 import \ + ragged_gather_reduce as ragged_gather_reduce_v2 +from tpu_inference.kernels.sparse_core.ragged_scatter import ragged_scatter + +jax.config.parse_flags_with_absl() + + +def reference_ragged_gather_reduce( + x: jax.Array, + indices: jax.Array, + topk_weights: jax.Array, + valid_rows_mask: jax.Array, + reduce_group_size: int, +) -> jax.Array: + """Reference implementation of ragged gather reduce.""" + out = x[indices] * topk_weights[:, None].astype(jnp.float32) + out = jnp.where(valid_rows_mask[:, None], out, 0) + out = out.reshape(-1, reduce_group_size, out.shape[-1]) + out = jnp.sum(out, axis=1).astype(jnp.bfloat16) + return out + + +@functools.partial(jax.jit, static_argnames="reduce_group_size") +def ragged_scatter_and_reduce( + x: jax.Array, + indices: jax.Array, + topk_weights: jax.Array, + valid_rows_mask: jax.Array, + start: jax.Array, + end: jax.Array, + reduce_group_size: int, +) -> jax.Array: + """Reference implementation of ragged gather reduce.""" + x = ragged_scatter(x, indices, start, end) + out = x.reshape((-1, reduce_group_size, x.shape[-1])) + topk_weights = topk_weights.reshape((-1, reduce_group_size))[..., None] + out = out * topk_weights + out = jnp.where( + valid_rows_mask.reshape((-1, reduce_group_size))[:, :, None], out, 0.0) + out = out.sum(axis=-2) + return out + + +def _time_function(fn, *args, n_repeats=100): + # Warmup + for _ in range(10): + fn(*args).block_until_ready() + + # Asynchronous dispatch to hide Python overhead + start = time.perf_counter() + results = [fn(*args) for _ in range(n_repeats)] + results[-1].block_until_ready() + end = time.perf_counter() + + return (end - start) / n_repeats + + +@jtu.with_config(jax_numpy_dtype_promotion="standard") +class ScatterTest(jtu.JaxTestCase): + _test_cases = [ + dict(out_size=o, + start_end=se, + hidden_size=h, + dtype=d, + reduce_group_size=rg) for o, se, h, d, rg in itertools.chain( + itertools.product( + [400, 840], + [(3, 338), (10, 255)], + [128, 512, 8192], + [jnp.bfloat16, jnp.float32], + [8, 5], + ), + itertools.product( + [16384], + [(99, 1120)], + [7168], + [jnp.bfloat16], + [8], + ), + itertools.product( + [16384], + [(300, 2358)], + [6144], + [jnp.bfloat16], + [8], + ), + itertools.product( + [20480], + [(300, 2850)], + [4096], + [jnp.bfloat16], + [10], + ), + ) + ] + + @parameterized.parameters(*_test_cases) + def test_sc_ragged_gather_reduce(self, out_size, hidden_size, start_end, + dtype, reduce_group_size): + start, end = start_end + start = min(start, out_size) + end = min(end, out_size) + key = jax.random.key(0) + x = jax.random.normal(key, (out_size, hidden_size), jnp.float32) + x = x.astype(dtype) + indices = jax.random.permutation(key, out_size) + topk_weights = jax.random.normal(key, (out_size, ), jnp.bfloat16) + valid_rows_mask = jnp.where( + jnp.logical_and( + jnp.array([start], jnp.int32) <= indices, + indices < jnp.array([end], jnp.int32), + ), + True, + False, + ) + # Correctness check. + desired = reference_ragged_gather_reduce(x, indices, topk_weights, + valid_rows_mask, + reduce_group_size) + for rgr, name in ( + (ragged_gather_reduce_v1, "ragged_gather_reduce_v1"), + (ragged_gather_reduce_v2, "ragged_gather_reduce_v2"), + ): + try: + actual = rgr(x, indices, topk_weights, valid_rows_mask, + reduce_group_size) + np.testing.assert_allclose(actual, + desired, + atol=1e-2, + rtol=1e-2) + except AssertionError: + raise + except Exception as e: # pylint: disable=broad-except + print(f"Skipping {name} correctness check due to error: {e}") + + # The first perf test case approximates the DeepSeekV3, 2k-batch-size, EP=16. + # The second case approximates the Qwen3-Coder-480B, 2k-batch-size, EP=8. + _perf_test_cases = [ + dict( + out_size=o, + start_end=se, + hidden_size=h, + dtype=d, + reduce_group_size=rg, + col_chunk_size=c_sz, + ) for o, se, h, d, rg, c_sz in itertools.chain( + itertools.product( + [16384], + [(99, 1120)], + [7168], + [jnp.bfloat16], + [8], + [3584], + ), + itertools.product( + [16384], + [(300, 2400)], + [6144], + [jnp.bfloat16], + [8], + [2048], + ), + itertools.product( + [65536], + [(100, 8300)], + [6144], + [jnp.bfloat16], + [8], + [2048], + ), + ) + ] + + @parameterized.parameters(*_perf_test_cases) + def test_perf( + self, + out_size, + hidden_size, + start_end, + dtype, + reduce_group_size, + col_chunk_size, + ): + start, end = start_end + start = min(start, out_size) + end = min(end, out_size) + key = jax.random.key(0) + x = jax.random.normal(key, (out_size, hidden_size), jnp.float32) + x = x.astype(dtype) + indices = jax.random.permutation(key, out_size) + topk_weights = jax.random.normal(key, (out_size, ), jnp.bfloat16) + valid_rows_mask = jnp.where( + jnp.logical_and( + jnp.array([start], jnp.int32) <= indices, + indices < jnp.array([end], jnp.int32), + ), + True, + False, + ) + + print(f"\n=== Running shape: out={out_size}," + f" hidden={hidden_size}, start={start}, end={end} ===") + + def run_and_time(name, fn, *args): + try: + t_val = _time_function(fn, *args) + print(f"{name}: {t_val*1000:.3f} ms") + except Exception as e: # pylint: disable=broad-except + print(f"{name} failed: {e}") + + run_and_time( + "ragged_scatter_and_reduce", + ragged_scatter_and_reduce, + x, + indices, + topk_weights, + valid_rows_mask, + start, + end, + reduce_group_size, + ) + + run_and_time( + "ragged_gather_reduce_v1", + ragged_gather_reduce_v1, + x, + indices, + topk_weights, + valid_rows_mask, + reduce_group_size, + ) + + run_and_time( + "ragged_gather_reduce_v2", + ragged_gather_reduce_v2, + x, + indices, + topk_weights, + valid_rows_mask, + reduce_group_size, + ) + + +if __name__ == "__main__": + absltest.main(testLoader=jtu.JaxTestLoader()) diff --git a/tests/kernels/ragged_gather_v2_test.py b/tests/kernels/ragged_gather_v2_test.py new file mode 100644 index 0000000000..aeea46c4ad --- /dev/null +++ b/tests/kernels/ragged_gather_v2_test.py @@ -0,0 +1,167 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# 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 time + +import jax +import jax.numpy as jnp +import numpy as np +from absl.testing import absltest, parameterized +from jax._src import test_util as jtu + +from tpu_inference.kernels.sparse_core.ragged_gather import ragged_gather +from tpu_inference.kernels.sparse_core.ragged_gather_v2 import ragged_gather_v2 + +jax.config.parse_flags_with_absl() + + +@jtu.with_config(jax_numpy_dtype_promotion="standard") +class GatherTest(jtu.JaxTestCase): + + @parameterized.product( + in_out_size=[(512, 400), (512, 1024)], + start_end=[(3, 338), (10, 422)], + hidden_size=[128, 512, 8192], + dtype=[jnp.int4, jnp.int8, jnp.bfloat16, jnp.float32], + kernel_version=[2], + ) + def test_sc_gather(self, in_out_size, hidden_size, start_end, dtype, + kernel_version): + in_size, out_size = in_out_size + start, end = start_end + start = min(start, out_size) + end = min(end, out_size) + key = jax.random.key(0) + x = jax.random.normal(key, (in_size, hidden_size), jnp.float32) + x = x.astype(dtype) + indices = jax.random.randint(key, (out_size, ), 0, in_size, jnp.int32) + + start_arr = jnp.array([start], jnp.int32) + end_arr = jnp.array([end], jnp.int32) + + kernel = ragged_gather if kernel_version == 1 else ragged_gather_v2 + + actual = kernel(x, indices, start_arr, end_arr) + actual.block_until_ready() + + # Correctness check. + actual = actual[start:end] + desired = x[indices][start:end] + + self.assertArraysEqual(actual, desired) + + def test_benchmark(self): + benchmark_shapes = [ + # in_size, out_size, hidden_size + (512, 1024, 8192), + (1024, 2048, 8192), + (2048, 4096, 8192), + ] + dtype = jnp.bfloat16 + + def _time_function(fn, *args, n_repeats=100): + # Warmup + for _ in range(10): + fn(*args).block_until_ready() + + # Asynchronous dispatch to hide Python overhead + start = time.perf_counter() + results = [fn(*args) for _ in range(n_repeats)] + results[-1].block_until_ready() + end = time.perf_counter() + + return (end - start) / n_repeats + + for in_size, out_size, hidden_size in benchmark_shapes: + + print(f"\n=== Running shape: in={in_size}, out={out_size}," + f" hidden={hidden_size} ===") + + key = jax.random.key(0) + x = jax.random.normal(key, (in_size, hidden_size), + jnp.float32).astype(dtype) + indices = jax.random.randint(key, (out_size, ), 0, in_size, + jnp.int32) + + # For benchmark, use full coverage to test peak performance + start = 0 + end = out_size + start_arr = jnp.array([start], jnp.int32) + end_arr = jnp.array([end], jnp.int32) + + @jax.jit + def run_v1(x, indices, start_arr, end_arr): + return ragged_gather(x, indices, start_arr, end_arr) + + @jax.jit + def run_v2(x, indices, start_arr, end_arr): + return ragged_gather_v2(x, indices, start_arr, end_arr) + + @jax.jit + def run_jax(x, indices): + return x[indices] + + # Run JAX baseline first + t_jax = _time_function(run_jax, x, indices) + res_jax = run_jax(x, indices) + res_jax_sliced = res_jax[start:end] + + print(f"JAX: {t_jax*1000:.3f} ms") + + # V1 Kernel + t_v1_str = "FAILED" + err_v1_str = "FAILED" + try: + t_v1 = _time_function(run_v1, x, indices, start_arr, end_arr) + t_v1_str = f"{t_v1*1000:.3f} ms" + res_v1 = run_v1(x, indices, start_arr, end_arr) + res_v1_sliced = res_v1[start:end] + err_v1 = jnp.max(jnp.abs(res_v1_sliced - res_jax_sliced)) + err_v1_str = f"{float(err_v1):.6f}" + np.testing.assert_allclose(res_v1_sliced, + res_jax_sliced, + atol=1e-2, + rtol=1e-2) + except Exception: # pylint: disable=broad-except + print( + f"[Warning] V1 Kernel failed for shape ({in_size}, {out_size}," + f" {hidden_size})") + + # V2 Kernel + t_v2_str = "FAILED" + err_v2_str = "FAILED" + try: + t_v2 = _time_function(run_v2, x, indices, start_arr, end_arr) + t_v2_str = f"{t_v2*1000:.3f} ms" + res_v2 = run_v2(x, indices, start_arr, end_arr) + res_v2_sliced = res_v2[start:end] + err_v2 = jnp.max(jnp.abs(res_v2_sliced - res_jax_sliced)) + err_v2_str = f"{float(err_v2):.6f}" + np.testing.assert_allclose(res_v2_sliced, + res_jax_sliced, + atol=1e-2, + rtol=1e-2) + except Exception: # pylint: disable=broad-except + print( + f"[Warning] V2 Kernel failed for shape ({in_size}, {out_size}," + f" {hidden_size})") + + print(f"V1 Kernel: {t_v1_str}") + print(f"V2 Kernel: {t_v2_str}") + print(f"V1 Kernel vs JAX Max Err: {err_v1_str}") + print(f"V2 Kernel vs JAX Max Err: {err_v2_str}") + + +if __name__ == "__main__": + absltest.main(testLoader=jtu.JaxTestLoader()) diff --git a/tpu_inference/envs.py b/tpu_inference/envs.py index e893fdfc0f..756bba6844 100644 --- a/tpu_inference/envs.py +++ b/tpu_inference/envs.py @@ -45,6 +45,13 @@ JITTED_MM_MODULE_KEYS: list[str] = [] REGISTER_MM_MODULE_CUSTOM_PYTREE_CLASSES: list[str] = [] RAGGED_GATED_DELTA_RULE_IMPL: str = "chunked_jax_pd" + # SparseCore MoE gather kernel version used by fused_moe_gmm. + # "v2" (default) = ragged_gather_v2; "v1" = legacy ragged_gather. + RAGGED_GATHER_VERSION: str = "v2" + # SparseCore MoE gather-reduce (combine) kernel version used by + # fused_moe_gmm. "v2" (default) = ragged_gather_reduce_v2; "v1" = legacy + # ragged_gather_reduce. + RAGGED_GATHER_REDUCE_VERSION: str = "v2" MOE_ALL_GATHER_ACTIVATION_DTYPE: str = "" TPU_OFFLOAD_SKIP_JAX_PRECOMPILE: bool = False TPU_OFFLOAD_DECODE_SAVE: bool = False @@ -332,6 +339,10 @@ def _get_int_list_env() -> list[int]: "chunked_kernel_p_recurrent_kernel_d" ], ), + "RAGGED_GATHER_VERSION": + env_with_choices("RAGGED_GATHER_VERSION", "v2", ["v1", "v2"]), + "RAGGED_GATHER_REDUCE_VERSION": + env_with_choices("RAGGED_GATHER_REDUCE_VERSION", "v2", ["v1", "v2"]), "MOE_ALL_GATHER_ACTIVATION_DTYPE": lambda: os.getenv("MOE_ALL_GATHER_ACTIVATION_DTYPE", ""), # kv offload to dram: skip pre-compiling swap-related jax functions diff --git a/tpu_inference/kernels/sparse_core/ragged_gather_reduce_v2.py b/tpu_inference/kernels/sparse_core/ragged_gather_reduce_v2.py new file mode 100644 index 0000000000..2a5e60ea85 --- /dev/null +++ b/tpu_inference/kernels/sparse_core/ragged_gather_reduce_v2.py @@ -0,0 +1,691 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# 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 + +import jax +import jax.numpy as jnp +from jax.experimental import pallas as pl +from jax.experimental.pallas import tpu as pltpu +from jax.experimental.pallas import tpu_sc as plsc + +from tpu_inference.kernels.sparse_core import core_map_helper + + +@dataclasses.dataclass(frozen=True) +class _Config: + num_row_partitions: int + num_column_partitions: int + reduce_group_size: int + col_size: int + col_chunk_size: int + num_row_subchunks: int + num_simd_lanes: int + topk_dtype: Any + in_dtype: Any + core_axis_name: str + subcore_axis_name: str + + @property + def row_chunk_size(self) -> int: + """Number of rows handled per row-pipeline block.""" + return self.num_simd_lanes * self.num_row_subchunks + + @property + def row_shift(self) -> int: + """log2 of how many source rows pack into one uint32 gather element. + + The SparseCore indirect DMA requires 32-bit elements: bfloat16 packs two + source rows per uint32 (shift 1), float32 is 1:1 (shift 0). + """ + input_packing = 32 // jax.dtypes.itemsize_bits(self.in_dtype) + return input_packing.bit_length() - 1 + + +@jax.tree_util.register_dataclass +@dataclasses.dataclass(frozen=True) +class _Inputs: + num_src_rows_per_row_partition: Any + x: Any + indices: Any + topk_weights: Any + sorted_by_validity: Any + + +@jax.tree_util.register_dataclass +@dataclasses.dataclass(frozen=True) +class _Scratch: + num_rows_per_row_partition_vmem: Any + prev_iter_last_row_vmem: Any + prev_dst_row_smem: Any + sorted_by_validity_vmem: Any + src_indices_vmem: Any + dst_indices_vmem: Any + tw_f32_vmem: Any + dma_src_row_vmem: Any + dma_dst_row_vmem: Any + prev_dst_val_vmem: Any + out_vmem: Any + sem: Any + + def __len__(self) -> int: + return len(dataclasses.fields(self)) + + def __getitem__(self, index: Any): + return getattr(self, dataclasses.fields(self)[index].name) + + +# ceil up to the nearest multiple of b. +def _align_to(a, b): + return pl.cdiv(a, b) * b + + +def _fallback_implementation( + x: jax.Array, + indices: jax.Array, + topk_weights: jax.Array, + valid_rows_mask: jax.Array, + reduce_group_size: int, +) -> jax.Array: + out = x[indices] * topk_weights[:, None].astype(jnp.float32) + out = jnp.where(valid_rows_mask[:, None], out, 0) + out = out.reshape(-1, reduce_group_size, out.shape[-1]) + out = jnp.sum(out, axis=1).astype(jnp.bfloat16) + return out + + +def _calculate_num_column_partitions(hidden_size: int, num_cores: int, + num_lanes: int) -> int: + """Calculates the number of row partitions.""" + # Each column partition should be multiple of 128 (number of lanes) due to + # DMA requirements. + # Prefer to use a large number of column partitions, as long as each + # partition's size is not too small for DMA pipeline efficiency and each + # partition's size can divide the hidden size. + + # Each column partition will do DMA pipelining on col_size. + preferred_num_stages = 4 + num_column_partitions = 1 + while (num_cores % (num_column_partitions * 2) == 0 + and hidden_size % (num_lanes * num_column_partitions * 2) == 0 + and hidden_size // + (num_column_partitions * 2 * num_lanes) >= preferred_num_stages): + num_column_partitions *= 2 + return num_column_partitions + + +def _calculate_col_chunk_size(col_size: int, num_simd_lanes: int) -> int: + """Picks the column chunk size the inner pipeline gathers at a time. + + The chunk is the largest divisor of ``col_size`` whose gather double-buffer + still fits comfortably in SparseCore VMEM. + """ + generation = pltpu.get_tpu_info().generation + match generation: + case 6: + target_bytes = int(256 * 1024 * 0.95) + case 7: + target_bytes = int(512 * 1024 * 0.95) + case _: + target_bytes = int(128 * 1024 * 0.95) + + # uint32 gather buffer, double-buffered by emit_pipeline. + bytes_per_col = num_simd_lanes * 4 * 2 + max_safe_col = (target_bytes // bytes_per_col // 128) * 128 + + start_col = (min(col_size, max_safe_col) // 128) * 128 + for chunk in range(start_col, 127, -128): + if col_size % chunk == 0: + return chunk + return 128 + + +def _preprocess( + valid_rows_mask: jax.Array, + reduce_group_size: int, + num_row_partitions: int, + num_simd_lanes: int, + row_chunk_size: int, +) -> tuple[jax.Array, jax.Array, jax.Array]: + """Sorts valid source rows to the front of each row partition. + + Returns: + sorted_by_validity: original row index of each slot after the stable + sort, flattened across partitions and padded to ``row_chunk_size``. + num_src_rows_per_row_partition: valid row count per partition, padded to + ``num_simd_lanes`` so the kernel can load it as a single vector. + mask: per output group, whether the group has any valid source row. + """ + row_partition_size = valid_rows_mask.shape[0] // num_row_partitions + valid_rows_mask_2d = valid_rows_mask.reshape(num_row_partitions, -1) + + # Stable sort of a boolean key is a stable partition: valid rows keep their + # relative order and move ahead of the invalid ones. + sorted_by_validity = jnp.argsort(~valid_rows_mask_2d, + descending=False, + stable=True, + axis=-1) + sorted_by_validity += (jnp.arange(num_row_partitions)[:, None] * + row_partition_size) + + pad_to = _align_to(row_partition_size, row_chunk_size) + if pad_to > row_partition_size: + sorted_by_validity = jnp.pad( + sorted_by_validity, + ((0, 0), (0, pad_to - row_partition_size)), + constant_values=0, + ) + sorted_by_validity = sorted_by_validity.reshape(-1) + + num_src_rows_per_row_partition = jnp.pad( + jnp.sum(valid_rows_mask_2d, axis=-1).astype(jnp.int32), + (0, max(0, num_simd_lanes - num_row_partitions)), + ) + mask = jnp.any(valid_rows_mask.reshape(-1, reduce_group_size), axis=-1) + return ( + sorted_by_validity.astype(jnp.int32), + num_src_rows_per_row_partition, + mask, + ) + + +def _pack_scalars_to_vector(scalar_list: list[jax.Array], + num_simd_lanes: int) -> jax.Array: + """Builds a lane vector from per-lane scalars. + + SparseCore cannot store individual scalars into VMEM lanes, so the vector is + assembled with masked accumulation before being stored. + """ + idx_vec = jnp.arange(num_simd_lanes) + vec = jnp.zeros((num_simd_lanes, ), jnp.int32) + for i in range(num_simd_lanes): + vec += (idx_vec == i).astype(jnp.int32) * scalar_list[i] + return vec + + +def _row_gather_spec( + sorted_by_validity_vmem: jax.Ref, + sub: int, + *, + num_simd_lanes: int, + row_chunk_size: int, +) -> pl.BlockSpec: + """Indirect BlockSpec gathering sub-chunk ``sub``'s rows of a 1-D input.""" + return pl.BlockSpec( + (pl.Indirect(num_simd_lanes), ), + lambda i, s=sub: (sorted_by_validity_vmem[pl.ds( + i * row_chunk_size + s * num_simd_lanes, num_simd_lanes)], ), + ) + + +def main_kernel( + inputs: _Inputs, + out_hbm_ref: jax.Ref, + scratch: _Scratch, + *, + cfg: _Config, +): + # Step 1: Resolve this core's row/column partition and its column slice. + num_simd_lanes = cfg.num_simd_lanes + col_chunk_size = cfg.col_chunk_size + num_row_subchunks = cfg.num_row_subchunks + row_chunk_size = cfg.row_chunk_size + + num_col_chunks = cfg.col_size // col_chunk_size + + core_id = jax.lax.axis_index((cfg.core_axis_name, cfg.subcore_axis_name)) + row_partition_id = core_id // cfg.num_column_partitions + col_partition_id = core_id % cfg.num_column_partitions + + row_partition_size_padded = (inputs.sorted_by_validity.shape[0] // + cfg.num_row_partitions) + row_start_padded = row_partition_id * row_partition_size_padded + col_start = col_partition_id * cfg.col_size + + # Step 2: Stage this partition's row count and sort permutation into VMEM. + recv_sem = scratch.sem.at[0] + num_rows_dma = pltpu.make_async_copy( + inputs.num_src_rows_per_row_partition.at[pl.ds(0, num_simd_lanes)], + scratch.num_rows_per_row_partition_vmem, + recv_sem, + ) + sorted_dma = pltpu.make_async_copy( + inputs.sorted_by_validity.at[pl.ds(row_start_padded, + row_partition_size_padded)], + scratch.sorted_by_validity_vmem, + recv_sem, + ) + num_rows_dma.start() + sorted_dma.start() + num_rows_dma.wait() + sorted_dma.wait() + + num_rows_per_row_partition = scratch.num_rows_per_row_partition_vmem[...] + num_rows_current_row_partition = jnp.array(0, jnp.int32) + for i in range(cfg.num_row_partitions): + num_rows_current_row_partition = jnp.where( + row_partition_id == i, + num_rows_per_row_partition[i], + num_rows_current_row_partition, + ) + num_row_blocks = pl.cdiv(num_rows_current_row_partition, row_chunk_size) + + # Step 3: Run the gather / weighted segmented-reduce / scatter pipeline. + + # The SparseCore indirect DMA requires 32-bit elements, so x is gathered + # through a uint32 reinterpretation. bfloat16 packs two source rows per + # uint32 row (row index >> 1); float32 is 1:1 (row index unchanged). + in_32b_hbm_ref = inputs.x.bitcast(jnp.uint32) + + # Sentinel for the cross-block reduction carry (no previous group). + scratch.prev_dst_row_smem[0] = -1 + + # One gather per sub-chunk for ``indices``, then the same for + # ``topk_weights``. + row_pipeline_in_specs = (tuple( + _row_gather_spec( + scratch.sorted_by_validity_vmem, + sub, + num_simd_lanes=num_simd_lanes, + row_chunk_size=row_chunk_size, + ) for sub in range(num_row_subchunks)) * 2) + + @functools.partial( + pltpu.emit_pipeline, + grid=(num_row_blocks, ), + in_specs=row_pipeline_in_specs, + out_specs=(), + ) + def row_pipeline(*args): + src_indices_refs = args[:num_row_subchunks] + topk_weights_refs = args[num_row_subchunks:2 * num_row_subchunks] + ( + src_indices_vmem_sc, + dst_indices_vmem_sc, + tw_f32_vmem_sc, + dma_src_row_vmem_sc, + dma_dst_row_vmem_sc, + prev_dst_val_vmem_sc, + out_vmem_sc, + sem_sc, + ) = args[-8:] + + row_block_id = pl.program_id(0) + + # Destination output row of each source row in this block. + dst_indices_list = [ + scratch.sorted_by_validity_vmem[pl.ds( + row_block_id * row_chunk_size + s * num_simd_lanes, + num_simd_lanes, + )] // cfg.reduce_group_size for s in range(num_row_subchunks) + ] + + # Stage the gathered indices/weights and the destinations in VMEM. + for s in range(num_row_subchunks): + sub = pl.ds(s * num_simd_lanes, num_simd_lanes) + src_indices_vmem_sc[sub] = src_indices_refs[s][...] + dst_indices_vmem_sc[sub] = dst_indices_list[s] + + tw = topk_weights_refs[s][...] + if cfg.topk_dtype == jnp.bfloat16: + tw_f32 = plsc.bitcast(jnp.bitwise_left_shift(tw, 16), + jnp.float32) + else: + tw_f32 = plsc.bitcast(tw, jnp.float32) + tw_f32_vmem_sc[sub] = tw_f32 + + # For each sub-chunk, the destination of the row just before it -- the + # seed for the segmented reduction's "same group as previous row" test. + for s in range(num_row_subchunks): + if s == 0: + prev_dst = scratch.prev_dst_row_smem[0] + else: + prev_dst = dst_indices_list[s - 1][num_simd_lanes - 1] + prev_dst_val_vmem_sc[pl.ds(s * num_simd_lanes, + num_simd_lanes)] = (jnp.broadcast_to( + prev_dst, (num_simd_lanes, ))) + + def get_dst_idx(global_idx): + return dst_indices_list[global_idx // + num_simd_lanes][global_idx % + num_simd_lanes] + + # For each source row, find the VMEM row that will hold its group's fully + # reduced value -- the last row of the group within this block. Scanning + # backwards, a row inherits its successor's merge target when they share + # a destination, otherwise it is its own target. + src_row_idx_in_vmem = [] + row_valid_vec = [] + for row_vmem_idx in reversed(range(row_chunk_size)): + global_row_idx = row_block_id * row_chunk_size + row_vmem_idx + row_valid_vec.append( + global_row_idx < num_rows_current_row_partition) + if row_vmem_idx == row_chunk_size - 1: + src_row_idx_in_vmem.append(row_vmem_idx) + else: + same_group_as_next = jnp.logical_and( + row_valid_vec[-2], + get_dst_idx(row_vmem_idx) == get_dst_idx(row_vmem_idx + 1), + ).astype(jnp.int32) + src_row_idx_in_vmem.append( + same_group_as_next * src_row_idx_in_vmem[-1] + + (1 - same_group_as_next) * row_vmem_idx) + src_row_idx_in_vmem.reverse() + row_valid_vec.reverse() + + # Per source row, the (VMEM source row, HBM destination row) of its + # scatter. Rows whose group is not yet fully reduced in this sub-chunk, + # and padding rows, are routed to a throwaway row. + garbage_dst = out_hbm_ref.shape[0] - 1 + dma_src_rows = [] + dma_dst_rows = [] + for s in range(num_row_subchunks): + sub_src = [] + sub_dst = [] + for i in range(num_simd_lanes): + global_idx = s * num_simd_lanes + i + merge_target = src_row_idx_in_vmem[global_idx] + is_final_write = jnp.logical_and( + row_valid_vec[global_idx], + merge_target < (s + 1) * num_simd_lanes, + ) + sub_src.append( + jnp.where(is_final_write, merge_target % num_simd_lanes, + 0)) + sub_dst.append( + jnp.where(is_final_write, dst_indices_list[s][i], + garbage_dst)) + dma_src_rows.append(sub_src) + dma_dst_rows.append(sub_dst) + + for s in range(num_row_subchunks): + sub = pl.ds(s * num_simd_lanes, num_simd_lanes) + dma_src_row_vmem_sc[sub] = _pack_scalars_to_vector( + dma_src_rows[s], num_simd_lanes) + dma_dst_row_vmem_sc[sub] = _pack_scalars_to_vector( + dma_dst_rows[s], num_simd_lanes) + + @functools.partial( + pltpu.emit_pipeline, + grid=(num_row_subchunks, num_col_chunks), + in_specs=pl.BlockSpec( + (pl.Indirect(num_simd_lanes), col_chunk_size), + lambda s, c: ( + jnp.bitwise_right_shift( + src_indices_vmem_sc[pl.ds(s * num_simd_lanes, + num_simd_lanes)], + cfg.row_shift, + ), + col_start // col_chunk_size + c, + ), + ), + out_specs=(), + ) + def col_pipeline(gather_ref, sem_inner): + s = pl.program_id(0) + c = pl.program_id(1) + col_hbm_start = col_start + c * col_chunk_size + send_sem = sem_inner.at[1] + + row_slice = pl.ds(s * num_simd_lanes, num_simd_lanes) + tw_slice = tw_f32_vmem_sc[row_slice] + dst_slice = dst_indices_vmem_sc[row_slice] + src_idx_slice = src_indices_vmem_sc[row_slice] + prev_dst_vals_vec = prev_dst_val_vmem_sc[row_slice] + + def col_loop(col_compute_offset): + col_slice = pl.ds(col_compute_offset, num_simd_lanes) + # Running sum, seeded by the carry from the previous sub-chunk. + previous_accumulated_data = scratch.prev_iter_last_row_vmem[ + c, col_slice] + + for row_src in range(num_simd_lanes): + val_u32 = gather_ref[row_src, col_slice] + if cfg.in_dtype == jnp.bfloat16: + # The two bfloat16 rows packed in one uint32 word sit in the low + # (even row) or high (odd row) 16 bits. Shift the wanted half + # into the float32 sign/exponent position and clear the rest. + shift = jnp.where( + jnp.bitwise_and(src_idx_slice[row_src], 1) == 0, + 16, 0) + shifted = jnp.bitwise_and( + jnp.left_shift(val_u32, shift), + jnp.uint32(0xFFFF0000)) + data_f32 = plsc.bitcast(shifted, jnp.float32) + else: + data_f32 = plsc.bitcast(val_u32, jnp.float32) + data_f32 *= tw_slice[row_src] + + # Reduction: accumulate while the destination group is unchanged, + # restart otherwise. Sorting guarantees rows of one group are + # contiguous. + dst_row_hbm = dst_slice[row_src] + if row_src == 0: + prev_dst = prev_dst_vals_vec[0] + else: + prev_dst = dst_slice[row_src - 1] + accumulated_data = jnp.where( + dst_row_hbm == prev_dst, + previous_accumulated_data + data_f32, + data_f32, + ) + previous_accumulated_data = accumulated_data + + # The output buffer stays float32: a bfloat16 output would be + # (16, 128)-tiled and the per-row scatter below writes a single + # row at an arbitrary, non-tile-aligned destination, which is only + # legal for 32-bit elements. The cast happens in the wrapper. + out_vmem_sc[row_src, col_slice] = accumulated_data + if row_src == num_simd_lanes - 1: + scratch.prev_iter_last_row_vmem[ + c, col_slice] = accumulated_data + + plsc.parallel_loop(0, col_chunk_size, + step=num_simd_lanes)(col_loop) + + # Scatter every source row's reduced value to its output row. Rows + # that share a group write the same value (idempotent); rows routed to + # the garbage destination are harmless. + dma_src_row_slice = dma_src_row_vmem_sc[row_slice] + dma_dst_row_slice = dma_dst_row_vmem_sc[row_slice] + copies = [] + for i in range(num_simd_lanes): + copy = pltpu.make_async_copy( + out_vmem_sc.at[dma_src_row_slice[i], + pl.ds(0, col_chunk_size)], + out_hbm_ref.at[dma_dst_row_slice[i], + pl.ds(col_hbm_start, col_chunk_size)], + send_sem, + ) + copy.start() + copies.append(copy) + for copy in copies: + copy.wait() + + col_pipeline(in_32b_hbm_ref, scratches=(sem_sc, )) + scratch.prev_dst_row_smem[0] = dst_indices_list[-1][num_simd_lanes - 1] + + row_pipeline( + *([inputs.indices] * num_row_subchunks), + *([inputs.topk_weights] * num_row_subchunks), + scratches=( + scratch.src_indices_vmem, + scratch.dst_indices_vmem, + scratch.tw_f32_vmem, + scratch.dma_src_row_vmem, + scratch.dma_dst_row_vmem, + scratch.prev_dst_val_vmem, + scratch.out_vmem, + scratch.sem, + ), + ) + + +@functools.partial(jax.jit, static_argnames=("reduce_group_size", )) +def ragged_gather_reduce( + x: jax.Array, + indices: jax.Array, + topk_weights: jax.Array, + valid_rows_mask: jax.Array, + reduce_group_size: int, +) -> jax.Array: + """Gathers ``x`` by ``indices``, weights and masks, then reduces by group. + + Args: + x: 2-D input features, ``(num_rows, hidden_size)``. + indices: 1-D gather indices, ``(input_size,)``. + topk_weights: 1-D per-row weights, ``(input_size,)``. + valid_rows_mask: 1-D bool mask of valid gathered rows, ``(input_size,)``. + reduce_group_size: number of consecutive rows summed into one output row. + + Returns: + Reduced output, ``(input_size // reduce_group_size, hidden_size)``. + """ + # Step 1: Choose the implementation (TensorCore fallback or SparseCore). + sc_info = pltpu.get_tpu_info().sparse_core + if sc_info is None: + return _fallback_implementation(x, indices, topk_weights, + valid_rows_mask, reduce_group_size) + + # For a small {input + output} both likely fit in TensorCore VMEM, where a + # plain TC gather-reduce beats routing through SparseCore and HBM. This + # also keeps the kernel off configs with num_row_partitions > num_simd_lanes. + dtype_bytes = jax.dtypes.itemsize_bits(x.dtype) // 8 + if (jnp.size(x) * dtype_bytes * 2 + < pltpu.get_tpu_info().vmem_capacity_bytes * 0.6): + return _fallback_implementation(x, indices, topk_weights, + valid_rows_mask, reduce_group_size) + + # Step 2: Derive the kernel configuration (core grid and column tiling). + hidden_size = x.shape[-1] + input_size = indices.size + num_simd_lanes = sc_info.num_lanes + num_lanes = pltpu.get_tpu_info().num_lanes + num_cores = sc_info.num_cores * sc_info.num_subcores + + num_column_partitions = _calculate_num_column_partitions( + hidden_size, num_cores, num_lanes) + num_row_partitions = num_cores // num_column_partitions + assert (num_row_partitions <= num_simd_lanes + ), f"{num_row_partitions=} must be <= {num_simd_lanes=}" + base_block_size = num_simd_lanes * num_row_partitions + num_row_subchunks = max( + 1, + min( + 4, + pl.cdiv(input_size, base_block_size), + ), + ) + row_chunk_size = num_simd_lanes * num_row_subchunks + + aligned_hidden_size = _align_to(hidden_size, 128 * num_column_partitions) + col_size = aligned_hidden_size // num_column_partitions + col_chunk_size = _calculate_col_chunk_size(col_size, num_simd_lanes) + + # Step 3: Pre-process inputs (weights, padding, sort by validity). + # The kernel gathers x through a uint32 reinterpretation; carry the weights + # the same way so they can be bitcast back to float32 on SparseCore. + if topk_weights.dtype == jnp.bfloat16: + topk_weights_u32 = jax.lax.bitcast_convert_type( + topk_weights, jnp.uint16).astype(jnp.uint32) + else: + topk_weights_u32 = jax.lax.bitcast_convert_type( + topk_weights, jnp.uint32) + + # Pad the input so each row partition holds a whole number of reduce + # groups; no group is then split across two physical cores. + padded_input_size = _align_to(input_size, + num_row_partitions * reduce_group_size) + valid_rows_mask = jnp.pad( + valid_rows_mask, + (0, padded_input_size - input_size), + constant_values=False, + ) + + sorted_by_validity, num_src_rows_per_row_partition, mask = _preprocess( + valid_rows_mask, + reduce_group_size, + num_row_partitions, + num_simd_lanes, + row_chunk_size, + ) + + # Step 4: Launch the SparseCore kernel. + vector_mesh = plsc.VectorSubcoreMesh( + num_cores=sc_info.num_cores, + num_subcores=sc_info.num_subcores, + core_axis_name="core", + subcore_axis_name="subcore", + ) + + cfg = _Config( + num_row_partitions=num_row_partitions, + num_column_partitions=num_column_partitions, + reduce_group_size=reduce_group_size, + col_size=col_size, + col_chunk_size=col_chunk_size, + num_row_subchunks=num_row_subchunks, + num_simd_lanes=num_simd_lanes, + topk_dtype=topk_weights.dtype, + in_dtype=x.dtype, + core_axis_name=vector_mesh.core_axis_name, + subcore_axis_name=vector_mesh.subcore_axis_name, + ) + + # The output gets one extra row: the kernel's garbage scatter destination. + out = core_map_helper.kernel( + functools.partial(main_kernel, cfg=cfg), + out_type=jax.ShapeDtypeStruct( + (padded_input_size // reduce_group_size + 1, aligned_hidden_size), + jnp.float32, + ), + compiler_params=pltpu.CompilerParams( + use_tc_tiling_on_sc=True, + disable_bounds_checks=True, + needs_layout_passes=False, + ), + scratch_types=(_Scratch( + num_rows_per_row_partition_vmem=pltpu.VMEM((num_simd_lanes, ), + jnp.int32), + prev_iter_last_row_vmem=pltpu.VMEM( + (col_size // col_chunk_size, col_chunk_size), jnp.float32), + prev_dst_row_smem=pltpu.SMEM((1, ), jnp.int32), + sorted_by_validity_vmem=pltpu.VMEM( + (sorted_by_validity.size // num_row_partitions, ), jnp.int32), + src_indices_vmem=pltpu.VMEM((row_chunk_size, ), jnp.int32), + dst_indices_vmem=pltpu.VMEM((row_chunk_size, ), jnp.int32), + tw_f32_vmem=pltpu.VMEM((row_chunk_size, ), jnp.float32), + dma_src_row_vmem=pltpu.VMEM((row_chunk_size, ), jnp.int32), + dma_dst_row_vmem=pltpu.VMEM((row_chunk_size, ), jnp.int32), + prev_dst_val_vmem=pltpu.VMEM((row_chunk_size, ), jnp.int32), + out_vmem=pltpu.VMEM((num_simd_lanes, col_chunk_size), jnp.float32), + sem=pltpu.SemaphoreType.DMA((2, )), + ), ), + mesh=vector_mesh, + name="sc_ragged_gather_reduce_v2", + )(_Inputs( + num_src_rows_per_row_partition=num_src_rows_per_row_partition, + x=x, + indices=indices, + topk_weights=topk_weights_u32, + sorted_by_validity=sorted_by_validity, + ), ) + + # Step 5: Post-process the output (drop padding, zero empty groups, cast). + out = out[:input_size // reduce_group_size, :hidden_size] + out = jnp.where(mask[:input_size // reduce_group_size, None], out, + jnp.zeros_like(out)) + return out.astype(x.dtype) diff --git a/tpu_inference/kernels/sparse_core/ragged_gather_v2.py b/tpu_inference/kernels/sparse_core/ragged_gather_v2.py new file mode 100644 index 0000000000..8bf08bef68 --- /dev/null +++ b/tpu_inference/kernels/sparse_core/ragged_gather_v2.py @@ -0,0 +1,292 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# 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 functools + +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 +from jax.experimental.pallas import tpu_sc as plsc + +from tpu_inference.kernels.sparse_core import core_map_helper + + +def calculate_col_size(hidden_size: int, packing: int) -> int: + """Calculates the max column size bounded by VMEM limits and hidden_size divisibility.""" + tpu_info = pltpu.get_tpu_info() + sc_info = tpu_info.sparse_core + assert sc_info is not None, "SparseCore info is missing." + lanes = sc_info.num_lanes + + match tpu_info.generation: + case 6: + target_bytes = (256 * 1024) * 0.8 + case 7: + target_bytes = (512 * 1024) * 0.8 + case _: + target_bytes = (128 * 1024) * 0.8 + + # Calculate max safe column size based on VMEM budget and aligne to 128. + num_buffers = 2 + bytes_per_col = (lanes + lanes // packing) * 4 * num_buffers + max_safe_col = int((target_bytes // bytes_per_col) // 128) * 128 + + # Search for the largest divisor of hidden_size bounded by max_safe_col. + # The first divisor found is the maximum. + start_col = (min(hidden_size, max_safe_col) // 128) * 128 + for c in range(start_col, 127, -128): + if hidden_size % c == 0: + return c + return max_safe_col + + +def main_kernel_v2( + start_ref: jax.Ref, + end_ref: jax.Ref, + in_hbm_ref: jax.Ref, + indices_hbm_ref: jax.Ref, + out_hbm_ref: jax.Ref, + start_vmem_ref: jax.Ref, + end_vmem_ref: jax.Ref, + sem_ref: jax.Ref, + *, + core_axis_name: str, + subcore_axis_name: str, + num_row_subchunks: int, +): + tpu_info = pltpu.get_tpu_info() + sc_info = tpu_info.sparse_core + assert sc_info is not None + num_simd_lanes = sc_info.num_lanes + hidden_size = in_hbm_ref.shape[-1] + dtype_bits = jax.dtypes.itemsize_bits(out_hbm_ref.dtype) + packing = 32 // dtype_bits + col_size = calculate_col_size(hidden_size, packing) + + assert isinstance(hidden_size, + int), f"hidden_size must be int, got {type(hidden_size)}" + num_cores = jax.lax.axis_size((core_axis_name, subcore_axis_name)) + row_subchunk_size = num_simd_lanes + row_chunk_size = row_subchunk_size * num_row_subchunks + block_size = row_chunk_size * num_cores + + recv_sem = sem_ref.at[0] + + copy_start = pltpu.make_async_copy(start_ref, start_vmem_ref.at[:1], + recv_sem) + copy_end = pltpu.make_async_copy(end_ref, end_vmem_ref.at[:1], recv_sem) + copy_start.start() + copy_end.start() + copy_start.wait() + copy_end.wait() + + start = start_vmem_ref[...][0] + end = end_vmem_ref[...][0] + + block_start = start // block_size + block_end = pl.cdiv(end, block_size) + num_blocks = block_end - block_start + num_blocks = jnp.where(end <= start, 0, num_blocks) + + num_cols = pl.cdiv(hidden_size, col_size) + + dtype = out_hbm_ref.dtype + dtype_bits = jax.dtypes.itemsize_bits(dtype) + packing = 32 // dtype_bits + + core_index = lax.axis_index((core_axis_name, subcore_axis_name)) + + # SparseCore `.bitcast()` leverages hardware Row-Packing for 16-bit -> 32-bit conversion. + # The logical row count halves, while physical column dimensions remain unchanged. + in_hbm_i32 = in_hbm_ref.bitcast(jnp.int32) + out_hbm_i32 = out_hbm_ref.bitcast(jnp.int32) + + num_phys_cols = col_size + + # The outer pipeline runs on the Vector Core, hoisting the integer arithmetic required + # to decode the row-packed indices. This prevents scalar-core instruction starvation + # during the execution of the indirect `in_specs` block lambdas. + + # TODO(guoweij): The nested `emit_pipeline` design creates a pipeline bubble (DMA wait) + # when the inner pipeline empties and restarts with new indices. This is a known + # high-level API limitation, currently amortized by setting a large num_row_subchunks + # and shouldn't be an issue for most use cases. We should still monitor + # the bubble's impact on E2E performance and, if needed, manually reimplement + # this using primitive operations to eliminate the bubble entirely. + + def col_loop(col_base, gather_ref, out_ref, idx_rem, unpack_col_chunk): + col_slice = pl.ds(col_base, unpack_col_chunk) + if packing == 1: + gather_dt = gather_ref.bitcast(dtype) + out_dt = out_ref.bitcast(dtype) + out_dt[:, col_slice] = gather_dt[:, col_slice] + else: + # Manual bitwise extraction and packing for packing >= 2 (bfloat16, int8, int4) + # bf16: 0xFFFF, int8: 0xFF, int4: 0xF + mask = (1 << dtype_bits) - 1 + shift_multiplier = dtype_bits.bit_length() - 1 + for i in range(num_simd_lanes // packing): + packed_row = jnp.zeros((1, unpack_col_chunk), dtype=jnp.int32) + for j in range(packing): + k = i * packing + j + dynamic_shift = jnp.left_shift(idx_rem[k], + shift_multiplier) + val = jnp.bitwise_right_shift( + gather_ref[pl.ds(k, 1), col_slice], dynamic_shift) + val = jnp.bitwise_and(val, mask) + pack_shift = j * dtype_bits + packed_row = jnp.bitwise_or( + packed_row, jnp.left_shift(val, pack_shift)) + out_ref[pl.ds(i, 1), col_slice] = packed_row + + def inner_pipeline(gather_ref, out_ref, idx_ref, unpack_col_chunk): + row_slice = pl.ds( + pl.program_id(0) * row_subchunk_size, row_subchunk_size) + subchunk_idxs = idx_ref[row_slice] + if packing > 1: + # Equivalent to `subchunk_idxs % packing` + idx_rem = jnp.bitwise_and(subchunk_idxs, packing - 1) + else: + idx_rem = jnp.zeros_like(subchunk_idxs) + + col_loop_fn = functools.partial( + col_loop, + gather_ref=gather_ref, + out_ref=out_ref, + idx_rem=idx_rem, + unpack_col_chunk=unpack_col_chunk, + ) + plsc.parallel_loop(0, num_phys_cols, + step=unpack_col_chunk)(col_loop_fn) + + def outer_pipeline(idx_ref): + b = pl.program_id(0) + b_global = b + block_start + + unpack_col_chunk = 128 + assert num_phys_cols % unpack_col_chunk == 0 + shift_amount = packing.bit_length() - 1 + pltpu.emit_pipeline( + functools.partial(inner_pipeline, + idx_ref=idx_ref, + unpack_col_chunk=unpack_col_chunk), + grid=(num_row_subchunks, num_cols), + in_specs=pl.BlockSpec( + (pl.Indirect(row_subchunk_size), num_phys_cols), + lambda r, col_id: ( + jnp.bitwise_right_shift( + idx_ref[pl.ds(r * row_subchunk_size, row_subchunk_size) + ], + shift_amount, + ), + col_id, + ), + ), + out_specs=pl.BlockSpec( + (row_subchunk_size // packing, num_phys_cols), + lambda r, col_id: ( + (b_global * num_cores + core_index) * num_row_subchunks + + r, + col_id, + ), + ), + )(in_hbm_i32, out_hbm_i32) + + pltpu.emit_pipeline( + outer_pipeline, + grid=(num_blocks, ), + in_specs=pl.BlockSpec( + (row_chunk_size, ), + lambda b: ((b + block_start) * num_cores + core_index, ), + ), + )(indices_hbm_ref) + + +@jax.jit +def ragged_gather_v2(x: jax.Array, indices: jax.Array, start: jax.Array, + end: jax.Array) -> jax.Array: + """Perform gather on indices within dynamic array start and end using BlockSpec.""" + + assert x.ndim == 2, "Ragged gather only supports 2d inputs." + assert indices.ndim == 1, "Ragged gather only supports 1d indices." + + if jnp.isscalar(start): + start = start[None] + if jnp.isscalar(end): + end = end[None] + + dtype = x.dtype + if dtype not in (jnp.bfloat16, jnp.float32, jnp.int8, jnp.int4): + raise ValueError( + f"dtype must be f32, bf16, int8, or int4, but got {dtype}") + + sc_info = pltpu.get_tpu_info().sparse_core + if sc_info is None: + return x[indices] + + hidden_size = x.shape[-1] + out_size = indices.size + + dtype_bits = jax.dtypes.itemsize_bits(dtype) + packing = 32 // dtype_bits + col_size = calculate_col_size(hidden_size, packing) + + aligned_hidden_size = ((hidden_size + col_size - 1) // col_size) * col_size + + num_simd_lanes = sc_info.num_lanes + num_cores = sc_info.num_cores * sc_info.num_subcores + base_block_size = num_simd_lanes * num_cores + + # Calculate ideal num_row_subchunks to avoid too much padding overhead. + num_row_subchunks = max( + 1, min(4, (out_size + base_block_size - 1) // base_block_size)) + + row_subchunk_size = num_simd_lanes + row_chunk_size = row_subchunk_size * num_row_subchunks + block_size = row_chunk_size * num_cores + + out_pad_size = ( + (out_size + block_size - 1) // block_size) * block_size - out_size + indices = jnp.pad(indices, ((0, out_pad_size))) + + vector_mesh = plsc.VectorSubcoreMesh( + num_cores=sc_info.num_cores, + num_subcores=sc_info.num_subcores, + core_axis_name="core", + subcore_axis_name="subcore", + ) + return core_map_helper.kernel( + functools.partial( + main_kernel_v2, + core_axis_name=vector_mesh.core_axis_name, + subcore_axis_name=vector_mesh.subcore_axis_name, + num_row_subchunks=num_row_subchunks, + ), + out_type=jax.ShapeDtypeStruct( + (out_size + out_pad_size, aligned_hidden_size), dtype), + compiler_params=pltpu.CompilerParams( + use_tc_tiling_on_sc=True, + needs_layout_passes=True, + disable_bounds_checks=True, + ), + scratch_types=[ + pltpu.VMEM((16, ), jnp.int32), + pltpu.VMEM((16, ), jnp.int32), + pltpu.SemaphoreType.DMA((1, )), + ], + mesh=vector_mesh, + name="sc_ragged_gather_v2", + )(start, end, x, indices)[:out_size, :hidden_size] diff --git a/tpu_inference/layers/common/fused_moe_gmm.py b/tpu_inference/layers/common/fused_moe_gmm.py index 52cd05fe4a..171ab53f5b 100644 --- a/tpu_inference/layers/common/fused_moe_gmm.py +++ b/tpu_inference/layers/common/fused_moe_gmm.py @@ -24,9 +24,13 @@ from tpu_inference.kernels.collectives import \ hierarchical_reduce_scatter as hier_rs from tpu_inference.kernels.megablox.gmm_v2 import gmm_v2 -from tpu_inference.kernels.sparse_core.ragged_gather import ragged_gather +from tpu_inference.kernels.sparse_core.ragged_gather import \ + ragged_gather as ragged_gather_v1 from tpu_inference.kernels.sparse_core.ragged_gather_reduce import \ - ragged_gather_reduce + ragged_gather_reduce as ragged_gather_reduce_v1 +from tpu_inference.kernels.sparse_core.ragged_gather_reduce_v2 import \ + ragged_gather_reduce as ragged_gather_reduce_v2 +from tpu_inference.kernels.sparse_core.ragged_gather_v2 import ragged_gather_v2 from tpu_inference.layers.common.quantization import quantize_tensor from tpu_inference.layers.common.sharding import ShardingAxisName from tpu_inference.logger import init_logger @@ -34,6 +38,22 @@ logger = init_logger(__name__) +# Select the SparseCore MoE gather and gather-reduce kernels independently at +# import time, driven by the RAGGED_GATHER_VERSION / RAGGED_GATHER_REDUCE_VERSION +# env vars (set in the server process before boot, so they are fixed for the +# server's lifetime). Both default to "v2" (the new kernels); set either to "v1" +# to fall back to the legacy kernel. Call sites below use the bound names. +if envs.RAGGED_GATHER_VERSION == "v1": + ragged_gather = ragged_gather_v1 +else: + ragged_gather = ragged_gather_v2 +if envs.RAGGED_GATHER_REDUCE_VERSION == "v1": + ragged_gather_reduce = ragged_gather_reduce_v1 +else: + ragged_gather_reduce = ragged_gather_reduce_v2 +logger.info("fused_moe_gmm SparseCore kernels: gather=%s gather_reduce=%s", + envs.RAGGED_GATHER_VERSION, envs.RAGGED_GATHER_REDUCE_VERSION) + # Target chunk size of 2048 slots was found empirically to be optimal # for MoE workloads (e.g., Qwen) to hide ICI/DMA latency during AllReduce. TARGET_SLOT_CHUNK_SIZE = 2048