From 4da30776aece1ebd5fd6ff4ad896cc889a568a23 Mon Sep 17 00:00:00 2001 From: Muhammad Awad Date: Sat, 2 May 2026 08:35:52 -0700 Subject: [PATCH 1/2] Fix matmul_reduce_scatter wrapper passing unsupported bias argument The OpsNamespace.matmul_reduce_scatter() wrapper was passing `bias` as a positional argument to the underlying function which doesn't accept it. This caused TypeError when calling ctx.ops.matmul_reduce_scatter(). The bias parameter was being passed as async_op, shifting all subsequent arguments by one position. Co-Authored-By: Claude Opus 4.6 --- iris/ops/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/iris/ops/__init__.py b/iris/ops/__init__.py index e0d12ba51..32d3041bb 100644 --- a/iris/ops/__init__.py +++ b/iris/ops/__init__.py @@ -164,7 +164,7 @@ def matmul_reduce_scatter(self, output_tensor, A, B, bias=None, async_op=False, >>> output = shmem.zeros((M, N_local), dtype=torch.float16) >>> shmem.ops.matmul_reduce_scatter(output, A, B) """ - return matmul_reduce_scatter(self._shmem, output_tensor, A, B, bias, async_op, config, workspace) + return matmul_reduce_scatter(self._shmem, output_tensor, A, B, async_op, config, workspace) # Export public API From 6798ff6aa91a29368823b3e52af1eb666e21a07c Mon Sep 17 00:00:00 2001 From: Muhammad Awad Date: Sat, 2 May 2026 09:25:09 -0700 Subject: [PATCH 2/2] Remove unsupported bias parameter from matmul_all_reduce and matmul_reduce_scatter wrappers Neither matmul_all_reduce() nor matmul_reduce_scatter() accepts a bias argument. The wrapper signatures falsely exposed bias=None, silently dropping it. Remove from both signatures and docstrings. all_gather_matmul and matmul_all_gather correctly support bias. Co-Authored-By: Claude Opus 4.6 --- iris/ops/__init__.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/iris/ops/__init__.py b/iris/ops/__init__.py index 32d3041bb..0e177c187 100644 --- a/iris/ops/__init__.py +++ b/iris/ops/__init__.py @@ -65,17 +65,16 @@ def __init__(self, shmem): """ self._shmem = shmem - def matmul_all_reduce(self, output_tensor, A, B, bias=None, async_op=False, config=None, workspace=None): + def matmul_all_reduce(self, output_tensor, A, B, async_op=False, config=None, workspace=None): """ Fused matrix multiplication and all-reduce. - Computes: output = all_reduce(A @ B + bias) + Computes: output = all_reduce(A @ B) Args: output_tensor: Output tensor (M, N) A: Input matrix A (M, K) B: Input matrix B (K, N) - bias: Optional bias vector (M,) or (N,) async_op: If False, performs barrier at end config: Optional FusedConfig for tuning workspace: Optional pre-allocated workspace @@ -141,17 +140,16 @@ def matmul_all_gather(self, output_tensor, A, B, bias=None, async_op=False, conf """ return matmul_all_gather(self._shmem, output_tensor, A, B, bias, async_op, config, workspace) - def matmul_reduce_scatter(self, output_tensor, A, B, bias=None, async_op=False, config=None, workspace=None): + def matmul_reduce_scatter(self, output_tensor, A, B, async_op=False, config=None, workspace=None): """ Fused matrix multiplication and reduce-scatter. - Computes: output = reduce_scatter(A @ B + bias) along N dimension + Computes: output = reduce_scatter(A @ B) along N dimension Args: output_tensor: Output tensor (M, N_local) where N_local = N / world_size A: Input matrix A (M, K) B: Input matrix B (K, N) - bias: Optional bias vector (M,) or (N,) async_op: If False, performs barrier at end config: Optional FusedConfig for tuning workspace: Optional pre-allocated workspace