one-shot limit

Author: dsidler

benchmark/ops/bench_matmul_all_reduce.py

@@ -12,7 +12,7 @@
 from tritonblas.matmul import persistent_matmul_lt
 
 
-_MAX_ONE_SHOT_INBOX_ELEMENTS = 2**32
+_MAX_ONE_SHOT_INBOX_ELEMENTS = 2**31
 
 
 @bench.register
@@ -106,9 +106,9 @@ def matmul_all_reduce(state, ctx):
     world_size = ctx.get_num_ranks()
     rank = ctx.get_rank()
 
-    if variant == "one_shot" and world_size * M * N >= _MAX_ONE_SHOT_INBOX_ELEMENTS:
+    if variant == "one_shot" and world_size * M * N > _MAX_ONE_SHOT_INBOX_ELEMENTS:
         state.skip(
-            "one_shot requires world_size * M * N < 2**32 elements "
+            "one_shot requires world_size * M * N <= 2**31 elements "
             f"(got {world_size * M * N})"
         )
 

benchmark/ops/bench_matmul_all_reduce_copy_engine.py

@@ -16,6 +16,9 @@
 from tritonblas.matmul import _make_matmul_selector
 
 
+_MAX_ONE_SHOT_INBOX_ELEMENTS = 2**31
+
+
 def _make_selector(M: int, N: int, K: int, dtype: torch.dtype, device: torch.device):
     return _make_matmul_selector(
         M,
@@ -44,6 +47,11 @@ def _register_copy_engine(state, ctx, *, variant: str) -> None:
         state.skip(f"N={N} must be divisible by block_size_n={selector.block_n}")
     if K % selector.block_k != 0:
         state.skip(f"K={K} must be divisible by block_size_k={selector.block_k}")
+    if variant == "one_shot" and world_size * M * N > _MAX_ONE_SHOT_INBOX_ELEMENTS:
+        state.skip(
+            "one_shot requires world_size * M * N <= 2**31 elements "
+            f"(got {world_size * M * N})"
+        )
 
     torch.manual_seed(123 + rank)
     A = ctx.randn((M, K), dtype=dtype)

iris/ops/matmul_all_reduce.py

@@ -23,7 +23,7 @@
 
 
 _SUPPORTED_VARIANTS = ("one_shot", "two_shot")
-_MAX_ONE_SHOT_INBOX_ELEMENTS = 2**32
+_MAX_ONE_SHOT_INBOX_ELEMENTS = 2**31
 
 
 @triton.jit()
@@ -371,9 +371,9 @@ def _validate_variant(variant: str):
 
 def _validate_one_shot_inbox_size(M: int, N: int, world_size: int):
     inbox_elements = world_size * M * N
-    if inbox_elements >= _MAX_ONE_SHOT_INBOX_ELEMENTS:
+    if inbox_elements > _MAX_ONE_SHOT_INBOX_ELEMENTS:
         raise ValueError(
-            "matmul_all_reduce one_shot requires world_size * M * N < 2**32 elements "
+            "matmul_all_reduce one_shot requires world_size * M * N <= 2**31 elements "
             "because the current publish/reduce kernels use 32-bit element offsets; "
             f"got world_size={world_size}, M={M}, N={N}, elements={inbox_elements}."
         )

iris/ops/matmul_all_reduce_copy_engine.py

@@ -24,6 +24,9 @@
 from .tritonblas_launch_wave_schedule import build_launch_wave_plan
 
 
+_MAX_ONE_SHOT_INBOX_ELEMENTS = 2**31
+
+
 @triton.jit()
 def _partitioned_xcd_matmul_kernel(
     A,
@@ -423,6 +426,16 @@ def _ceil_div(value: int, divisor: int) -> int:
     return (value + divisor - 1) // divisor
 
 
+def _validate_one_shot_inbox_size(M: int, N: int, world_size: int):
+    inbox_elements = world_size * M * N
+    if inbox_elements > _MAX_ONE_SHOT_INBOX_ELEMENTS:
+        raise ValueError(
+            "matmul_all_reduce_copy_engine one_shot requires world_size * M * N <= 2**31 elements "
+            "because the current local-reduce kernels use 32-bit element offsets; "
+            f"got world_size={world_size}, M={M}, N={N}, elements={inbox_elements}."
+        )
+
+
 def _partitioned_xcd_gemm_num_sms(launch: dict, world_size: int) -> int:
     cached = launch.get("gemm_num_sms")
     if cached is not None:
@@ -916,9 +929,11 @@ def matmul_all_reduce_copy_engine_preamble(
     dtype = A.dtype
     world_size = shmem.get_num_ranks()
 
-    # Validate config
     config.validate(world_size=world_size)
 
+    if config.all_reduce_variant == "one_shot":
+        _validate_one_shot_inbox_size(M, N, world_size)
+
     if config.all_reduce_variant == "two_shot" and M % world_size != 0:
         raise ValueError(
             "matmul_all_reduce_copy_engine two_shot requires M to be divisible by world_size "
@@ -1047,6 +1062,9 @@ def matmul_all_reduce_copy_engine(
 
     config.validate(world_size=world_size)
 
+    if config.all_reduce_variant == "one_shot":
+        _validate_one_shot_inbox_size(M, N, world_size)
+
     # Prepare workspace if needed
     if workspace is None:
         workspace = matmul_all_reduce_copy_engine_preamble(