File size: 962 Bytes
2976eec | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 | # 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
|