| # SPDX-License-Identifier: MIT | |
| # Copyright (C) 2025, Advanced Micro Devices, Inc. All rights reserved. | |
| """ | |
| Triton-based communication primitives for AITER. | |
| This submodule contains communication operations implemented using Triton, | |
| including Iris-based GPU-initiated communication. | |
| If Iris is not available, importing this module will raise ImportError. | |
| """ | |
| # Import all Iris-based communication primitives | |
| # If Iris is not installed, this import will fail and the entire | |
| # aiter.ops.triton.comms module will be unavailable | |
| from .iris import IrisCommContext, calculate_heap_size | |
| from .reduce_scatter import reduce_scatter | |
| from .all_gather import all_gather | |
| from .fused import reduce_scatter_rmsnorm_quant_all_gather | |
| __all__ = [ | |
| "IrisCommContext", | |
| "calculate_heap_size", | |
| "reduce_scatter", | |
| "all_gather", | |
| "reduce_scatter_rmsnorm_quant_all_gather", | |
| "IRIS_COMM_AVAILABLE", | |
| ] | |
| # If we got here, Iris is available | |
| IRIS_COMM_AVAILABLE = True | |