File size: 2,405 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 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | # SPDX-License-Identifier: MIT
# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
"""AITER Triton kernels for AMD ROCm.
Repackaged from the ``aiter/ops/triton/**`` subtree of the
`ROCm/aiter <https://github.com/ROCm/aiter>`_ project as a self-contained
Hugging Face Hub kernel. Each subpackage under :mod:`aiter_kernels` maps 1:1
to the equivalent upstream module under ``aiter.ops.triton``.
Flash Attention is **not** included here — it lives in
``kernels-community/aiter-flash-attn`` and is synced separately.
"""
from __future__ import annotations
from . import _aiter_compat # noqa: F401 (must be importable before ops)
from . import activation
from . import causal_conv1d
from . import causal_conv1d_update_single_token
from . import fusions
from . import gated_delta_net
from . import gather_kv_b_proj
from . import gemm
from . import gluon
from . import gmm
from . import kv_cache
from . import moe
from . import normalization
from . import quant
from . import rope
# Top-level re-exports for drop-in parity with the standalone ``aiter-rope`` repo.
from .rope import RotateStyle, apply_rotary_transformers
from . import softmax
from . import topk
from . import utils
# ``comms`` pulls in iris and is optional — make it import-safe.
try:
from . import comms
# Re-export communication primitives at this level for convenience
from .comms import (
IrisCommContext,
reduce_scatter,
all_gather,
reduce_scatter_rmsnorm_quant_all_gather,
IRIS_COMM_AVAILABLE,
)
_COMMS_AVAILABLE = True
except ImportError:
_COMMS_AVAILABLE = False
IRIS_COMM_AVAILABLE = False
comms = None # type: ignore[assignment]
__kernel_metadata__ = {
"license": "mit",
}
__all__ = [
"__kernel_metadata__",
"RotateStyle",
"apply_rotary_transformers",
"activation",
"causal_conv1d",
"causal_conv1d_update_single_token",
"comms",
"fusions",
"gated_delta_net",
"gather_kv_b_proj",
"gemm",
"gluon",
"gmm",
"kv_cache",
"moe",
"normalization",
"quant",
"rope",
"softmax",
"topk",
"utils",
]
if _COMMS_AVAILABLE:
__all__.extend(
[
"IrisCommContext",
"reduce_scatter",
"all_gather",
"reduce_scatter_rmsnorm_quant_all_gather",
"IRIS_COMM_AVAILABLE",
]
)
|