tpu-inference 0.11.1.dev202511150811__py3-none-any.whl

tpu_inference/kernels/collectives/all_gather_matmul_tuned_block_sizes.py

@@ -0,0 +1,60 @@
+# SPDX-License-Identifier: Apache-2.0
+"""All-gather matmul kernel's tuned block sizes."""
+
+import jax
+
+# key:
+#    - tpu_version
+#    - m
+#    - n
+#    - k
+#    - dtype
+#    - tp_size
+# value:
+#    - bn
+#    - bk
+TUNED_BLOCK_SIZES = {
+    # go/keep-sorted start
+    (6, 1024, 51200, 5120, 'bfloat16', 8): (6400, 2560),
+    (6, 1024, 57344, 8192, 'bfloat16', 8): (7168, 8192),
+    (6, 2048, 51200, 5120, 'bfloat16', 8): (1280, 5120),
+    (6, 2048, 57344, 8192, 'bfloat16', 8): (1024, 8192),
+    (6, 4096, 51200, 5120, 'bfloat16', 8): (3200, 5120),
+    (6, 8192, 51200, 5120, 'bfloat16', 8): (1280, 5120),
+    # go/keep-sorted end
+}
+
+
+def get_tpu_version() -> int:
+    """Returns the numeric version of the TPU, or -1 if not on TPU."""
+    kind = jax.devices()[0].device_kind
+    if 'TPU' not in kind:
+        return -1
+    if kind.endswith(' lite'):
+        kind = kind[:-len(' lite')]
+    assert kind[:-1] == 'TPU v', kind
+    return int(kind[-1])
+
+
+def get_key(
+    m,
+    n,
+    k,
+    dtype,
+    tp_size,
+):
+    """Returns the key for the given parameters."""
+    return (
+        get_tpu_version(),
+        m,
+        n,
+        k,
+        dtype,
+        tp_size,
+    )
+
+
+def get_tuned_block_sizes(m, n, k, dtype_name, tp_size):
+    """Returns the tuned block sizes for the given parameters."""
+    key = get_key(m, n, k, dtype_name, tp_size)
+    return TUNED_BLOCK_SIZES.get(key, (None, None))

tpu_inference/kernels/collectives/util.py

@@ -0,0 +1,47 @@
+# SPDX-License-Identifier: Apache-2.0
+"""utilities for collective kernels."""
+
+import functools
+
+from jax.experimental import pallas as pl
+from jax.experimental.pallas import tpu as pltpu
+
+
+def local_barrier(left_neighbor, right_neighbor, double_barrier=True):
+    """Performs a barrier with neighbors on the global barrier semaphore.
+
+  Optionally performs a second barrier, which prevents a potential race
+  when reusing the same collective_id across kernel invocations.
+
+  Args:
+    left_neighbor: Left neighbor device id.
+    right_neighbor: Right neighbor device id.
+    double_barrier: Whether to perform a second barrier.
+  """
+    barrier_sem = pltpu.get_barrier_semaphore()
+    for neighbor in [left_neighbor, right_neighbor]:
+        pltpu.semaphore_signal(
+            barrier_sem,
+            inc=1,
+            device_id=(neighbor, ),
+            device_id_type=pltpu.DeviceIdType.MESH,
+        )
+    pltpu.semaphore_wait(barrier_sem, 2)
+    if double_barrier:
+        # The double-barrier prevents a race condition where one neighbor can
+        # re-enter the kernel again on a subsequent call and increment the
+        # barrier semaphore a second time. This would unblock the current device
+        # even if the other neighbor is not ready yet.
+        # To implement a double-barrier, we stack-allocate a second REGULAR
+        # semaphore using run_scoped.
+        @functools.partial(pl.run_scoped,
+                           second_barrier=pltpu.SemaphoreType.REGULAR)
+        def _(second_barrier):
+            for neighbor in [left_neighbor, right_neighbor]:
+                pltpu.semaphore_signal(
+                    second_barrier,
+                    inc=1,
+                    device_id=(neighbor, ),
+                    device_id_type=pltpu.DeviceIdType.MESH,
+                )
+            pltpu.semaphore_wait(second_barrier, 2)