Fused Communication and Compute Operations

Fused Communication and Compute Operations#

Operations that fuse communication with compute in a single kernel, accessible via the ops property on the Iris instance (e.g. ctx.ops.matmul_all_reduce(...)). The compute side is currently GEMM, supplied by tritonBLAS; the communication side is an Iris collective. Fusing them means the collective overlaps the matmul instead of running as a separate kernel after it.

matmul_all_reduce#

OpsNamespace.matmul_all_reduce(output_tensor, A, B, bias=None, async_op=False, config=None, workspace=None)[source]#

Fused matrix multiplication and all-reduce.

Computes: output = all_reduce(A @ B + bias)

Parameters:

  • 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

Returns:

Updated workspace object

Return type:

workspace

Example

>>> output = shmem.zeros((M, N), dtype=torch.float16)
>>> shmem.ops.matmul_all_reduce(output, A, B)

all_gather_matmul#

OpsNamespace.all_gather_matmul(output_tensor, A_sharded, B, bias=None, async_op=False, config=None, workspace=None)[source]#

Fused all-gather and matrix multiplication.

Computes: output = all_gather(A_sharded) @ B + bias

Parameters:

  • output_tensor – Output tensor (M, N)

  • A_sharded – Sharded input matrix (M, K_local)

  • B – Input matrix B (K, N) where K = K_local * world_size

  • 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

Returns:

Updated workspace object

Return type:

workspace

Example

>>> K_local = K // world_size
>>> A_sharded = shmem.randn((M, K_local), dtype=torch.float16)
>>> output = shmem.zeros((M, N), dtype=torch.float16)
>>> shmem.ops.all_gather_matmul(output, A_sharded, B)

matmul_all_gather#

OpsNamespace.matmul_all_gather(output_tensor, A, B, bias=None, async_op=False, config=None, workspace=None)[source]#

Fused matrix multiplication and all-gather.

Computes: output = all_gather(A @ B + bias) along M dimension

Parameters:

  • output_tensor – Output tensor (M*world_size, 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

Returns:

Updated workspace object

Return type:

workspace

Example

>>> M_local = M // world_size
>>> A = shmem.randn((M_local, K), dtype=torch.float16)
>>> output = shmem.zeros((M, N), dtype=torch.float16)
>>> shmem.ops.matmul_all_gather(output, A, B)

matmul_reduce_scatter#

OpsNamespace.matmul_reduce_scatter(output_tensor, A, B, bias=None, async_op=False, config=None, workspace=None)[source]#

Fused matrix multiplication and reduce-scatter.

Computes: output = reduce_scatter(A @ B + bias) along N dimension

Parameters:

  • 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

Returns:

Updated workspace object

Return type:

workspace

Example

>>> N_local = N // world_size
>>> output = shmem.zeros((M, N_local), dtype=torch.float16)
>>> shmem.ops.matmul_reduce_scatter(output, A, B)

atomic_sub#

atomic_sub(pointer, val, from_rank, to_rank, heap_bases, mask=None, sem=None, scope=None, hint=None)[source]#

Atomically subtracts data from the specified rank’s memory location.

This function performs an atomic subtraction operation by translating the pointer from the from_rank’s address space to the to_rank’s address space and atomically subtracting the provided data from the to_rank memory location. If the from_rank and to_rank are the same, this function performs a local atomic subtraction operation.

Parameters:

  • pointer (triton.PointerType, or block of dtype=triton.PointerType) – Pointer in the from_rank’s address space that will be translated to the to_rank’s address space. Must be the current rank where the pointer is local.

  • val (Block) – The tensor of elements to be subtracted atomically.

  • from_rank (int) – The rank ID from which the pointer originates. Must be the current rank where the pointer is local.

  • to_rank (int) – The rank ID to which the atomic operation will be performed.

  • heap_bases (triton.PointerType) – Array containing the heap base addresses for all ranks.

  • mask (Block of triton.int1, optional) – If mask[idx] is false, do not perform the atomic operation at address pointer[idx]. Defaults to None.

  • sem (str, optional) – Specifies the memory semantics for the operation. Acceptable values are “acquire”, “release”, “acq_rel” (stands for “ACQUIRE_RELEASE”), and “relaxed”. Defaults to “acq_rel”.

  • scope (str, optional) – Defines the scope of threads that observe the synchronizing effect of the atomic operation. Acceptable values are “gpu” (default), “cta” (cooperative thread array, thread block), or “sys” (stands for “SYSTEM”). Defaults to “gpu”.

  • hint (int or tuple, optional) – Vectorization hint passed to tl.multiple_of / tl.max_contiguous on the translated pointer. Defaults to None (no hint).

Returns:

The value at the memory location before the atomic subtraction.

Return type:

Block

Example

>>> @triton.jit
>>> def kernel(ptr, heap_bases):
>>>     # Atomically subtract 3 from rank 2's memory from rank 0
>>>     cur_rank = 0      # Current rank (source)
>>>     remote_rank = 2   # Remote rank (destination)
>>>     decrement = 3
>>>     old_val = iris.atomic_sub(ptr, decrement, cur_rank, remote_rank, heap_bases)

quiet#

quiet(copy_engine_ctx, to_rank)[source]#

Wait for all submitted SDMA operations to complete for the specified destination rank.

Polls the hardware read pointer until it catches up to the committed write pointer, ensuring all previously submitted SDMA packets to the destination rank have been processed.

Parameters:

  • copy_engine_ctx (tensor) – Copy engine context tensor containing queue metadata

  • to_rank – The destination rank whose SDMA queue should be drained

Example

>>> @triton.jit
>>> def kernel(src, dst, heap_bases, copy_engine_ctx):
>>>     # Submit SDMA operations
>>>     iris.put(
>>>         src, dst, 0, 1, heap_bases, mask=mask,
>>>         copy_engine_ctx=copy_engine_ctx, use_copy_engine=True, contiguous_copy=True
>>>     )
>>>     # Wait for completion
>>>     iris.quiet(copy_engine_ctx, 1)
Contents