Packed Sort: speed up MoE routing argsort via bit-packed single-key sort

tests/layers/common/test_sort.py

@@ -0,0 +1,54 @@
+# 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 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.layers.common.sort import can_pack_int32, packed_argsort
+
+
+class CanPackInt32Test(jtu.JaxTestCase):
+
+    def test_fits_exactly(self):
+        # 21 key bits + 10 index bits == 31 usable bits.
+        self.assertTrue(can_pack_int32(n=1024, max_key=2**21 - 1))
+
+    def test_one_bit_too_wide(self):
+        self.assertFalse(can_pack_int32(n=1024, max_key=2**21))
+
+
+class PackedArgsortTest(jtu.JaxTestCase):
+
+    @parameterized.named_parameters(
+        ("random", np.random.default_rng(0).integers(0, 64, size=512), 63),
+        # All keys tied: order can only come from the packed index.
+        ("all_ties", np.zeros(128, dtype=np.int32), 0),
+        # The shape ragged_gather_reduce_v2 passes: 2-D, sorted on axis=-1.
+        ("boolean_2d", np.array([[1, 0, 1, 0], [0, 0, 1, 1]]), 1),
+    )
+    def test_matches_argsort(self, keys, max_key):
+        keys = jnp.asarray(keys)
+        expected = jnp.argsort(keys, axis=-1, stable=True)
+        self.assertArraysEqual(packed_argsort(keys, max_key=max_key), expected)
+
+    def test_raises_when_packing_does_not_fit(self):
+        keys = jnp.zeros(1024, dtype=jnp.int32)
+        with self.assertRaises(ValueError):
+            packed_argsort(keys, max_key=2**21)
+
+
+if __name__ == "__main__":
+    absltest.main(testLoader=jtu.JaxTestLoader())

tpu_inference/kernels/sparse_core/ragged_gather_reduce_v2.py

@@ -25,6 +25,7 @@
 from tpu_inference.kernels.sparse_core import core_map_helper
 from tpu_inference.kernels.sparse_core.ragged_gather_reduce_tuned_params import (
     TunableParams, TuningKey, get_tuned_params)
+from tpu_inference.layers.common.sort import can_pack_int32, packed_argsort
 
 
 @dataclasses.dataclass(frozen=True)
@@ -232,10 +233,20 @@ def _preprocess(
 
     # 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)
+    if can_pack_int32(n=valid_rows_mask_2d.shape[-1], max_key=1):
+        # Sort by the packed [key | index] rather than argsorting, when the
+        # two widths together fit in an int32. Keys need to be an integer
+        # rather than the bool to be packed into a single int32 value.
+        invalid_rows_mask_2d = (~valid_rows_mask_2d).astype(jnp.int32)
+        sorted_by_validity = packed_argsort(
+            invalid_rows_mask_2d,
+            max_key=1,  # boolean 0 or 1
+            axis=-1)
+    else:
+        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)
 

tpu_inference/layers/common/fused_moe_gmm.py

@@ -31,6 +31,7 @@
     ragged_gather_v2 as ragged_gather
 from tpu_inference.layers.common.quantization import quantize_tensor
 from tpu_inference.layers.common.sharding import ShardingAxisName
+from tpu_inference.layers.common.sort import can_pack_int32, packed_argsort
 from tpu_inference.logger import init_logger
 from tpu_inference.utils import get_mesh_shape_product
 
@@ -674,7 +675,14 @@ def fused_moe_func(
     def _process_tokens_locally(hidden_states_local, topk_indices_local):
         num_tokens_local = hidden_states_local.shape[0]
         topk_indices_flat = topk_indices_local.flatten()
-        topk_argsort_indices = jnp.argsort(topk_indices_flat)
+        max_expert_id = global_num_experts - 1
+        if can_pack_int32(n=topk_indices_flat.shape[0], max_key=max_expert_id):
+            # Sort by the packed [key | index] rather than argsorting, when
+            # the two widths together fit in an int32.
+            topk_argsort_indices = packed_argsort(topk_indices_flat,
+                                                  max_key=max_expert_id)
+        else:
+            topk_argsort_indices = jnp.argsort(topk_indices_flat)
         token_indices = jnp.arange(num_tokens_local,
                                    dtype=jnp.int32).repeat(topk)
         token_indices_sorted = token_indices[topk_argsort_indices]

tpu_inference/layers/common/sort.py

@@ -0,0 +1,60 @@
+# 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 jax
+import jax.numpy as jnp
+
+# int32 has 32 bits; reserve the sign bit so packed keys stay non-negative.
+_USABLE_BITS = 31
+
+
+def _index_bits(n: int) -> int:
+    return max(1, (n - 1).bit_length())
+
+
+def _key_bits(max_key: int) -> int:
+    return max(1, max_key.bit_length())
+
+
+def can_pack_int32(n: int, max_key: int) -> bool:
+    # Static in both args, so callers can branch in Python rather than
+    # tracing both the packed and fallback graphs.
+    return _key_bits(max_key) + _index_bits(n) <= _USABLE_BITS
+
+
+def packed_argsort(keys: jax.Array,
+                   max_key: int,
+                   *,
+                   axis: int = -1) -> jax.Array:
+    """Faster drop-in replacement for jnp.argsort when keys pack into int32.
+
+    Packs (key, index) into one int32 and runs a single-key unstable sort.
+    Stability comes from the index in the low bits, not the sort itself.
+    Caller must check can_pack_int32 first since there's no fallback here.
+    Keys must be non-negative and at most `max_key`. Ascending order only.
+    """
+    n = keys.shape[axis]
+    if not can_pack_int32(n, max_key):
+        raise ValueError(
+            f"packing keys up to {max_key} over {n} rows needs "
+            f"{_key_bits(max_key) + _index_bits(n)} bits, "
+            f"{_USABLE_BITS} available; check can_pack_int32() before calling")
+
+    shift = _index_bits(n)
+    index_shape = [1] * keys.ndim
+    index_shape[axis] = n
+    index = jnp.arange(n, dtype=jnp.int32).reshape(index_shape)
+    packed = (keys.astype(jnp.int32) << shift) | index
+    sorted_packed = jnp.sort(packed, axis=axis, stable=False)
+    return sorted_packed & ((1 << shift) - 1)