在大型模型推理的战场上，算子的性能是构建一切的基础。随着Transformer模型的参数规模突破千亿级别，并且MoE架构逐渐成为行业标配，单一的Kernel实现早已无法满足不同形状、不同精度以及不同硬件代际的多样化需求。

• AITER（面向ROCm的AI张量引擎）是AMD推出的一款高性能AI算子库，专为基于ROCm的推理与训练负载提供经过深度优化的GPU核心运算程序。该库整合了多种可直接投入生产环境的算子，框架开发者能够将其无缝集成到自身的架构体系中。
• https://github.com/ROCm/aiter
• https://rocm.github.io/aiter/

AMD推出的AITER（AI Tensor Engine for ROCm）正是为了解决这一痛点——它并非一个简单的算子集合，而是一个融合了汇编手写内核（ASM）、Composable Kernel（CK）代码生成、Triton JIT以及自研FlyDSL四种后端的统一调度框架。

AITER作为AMD的集中式代码库，为AI工作负载加速提供了各类高性能算子，同时也是统一接收所有自定义算子级需求的优质平台，能够满足多样化的定制需求。开发者可以专注于算子开发，而客户则能将这个算子集合集成到自己的私有、公有或其他各类框架中。

更为关键的是，它内部构建了一个从“未调优形状发现”到“在线自动Tuning”，再到“CSV配置热加载”的完整闭环，确保每一次GEMM和MoE调用都能在运行时命中性能最优的内核。

本文将从架构设计、JIT编译基础设施、Fused MoE核心流水线以及GEMM自动调优四个维度，带你深入AITER的源码核心区域。

该框架的集成详情如下表：

关键算子的性能表现如下，更多详情可参考ATOM/benchmark-dashboard[1]：

当前支持的AMD GPU型号如下：

unsetunset本文目录unsetunset

• 快速上手
• 一、架构总览与设计哲学
  • 1.1 一个入口，四种后端
  • 1.2 核心数据流
• 二、JIT编译基础设施：从源码到.so的全自动流水线
  • 2.1 compile_ops装饰器：声明即编译
  • 2.2 build_module：多进程安全的编译引擎
  • 2.3 torch_compile_guard：与torch.compile无缝对接
• 三、Fused MoE：Token排序与多阶段GEMM的精密编排
  • 3.1 Token排序：为GPU并行铺路
  • 3.2 一阶段 vs 两阶段：动态决策引擎
  • 3.3 Block Size自适应：榨干每一个CU
  • 3.4 K-Split：低Token数场景下的CU利用率救星
• 四、GEMM自动调优：从形状发现到最优内核的闭环
  • 4.1 多后端路由表
  • 4.2 Skinny GEMM：小M场景的秘密武器
  • 4.3 在线调优与离线调优的双轨机制
• 五、量化生态：从FP4到INT8的全精度覆盖
  • 5.1 量化类型的精细分级
  • 5.2 MXFP4量化的融合排序
• 六、有趣的地方
  • 6.1 LRU缓存的大量使用
  • 6.2 FlyDSL回退机制
  • 6.3 ctypes FFI：绕过PyTorch的无Torch调用路径
• 结语

unsetunset快速上手unsetunset

AITER 的安装过程极为简洁，仅需执行以下三条命令即可完成：

git clone –recursive https://github.com/ROCm/aiter.git
cd aiter
python3 setup.py develop

若需要启用 FlyDSL 混合精度 MoE 支持，则需额外执行一条安装指令：

pip install –pre flydsl

关于更多安装选项（例如 Triton 通信库、Iris 库等）的详细信息，请参考项目根目录下的 README.md 文件。安装完成后，开发者便可以在 Python 环境中通过 import aiter 命令来调用所有算子。例如，运行内置测试脚本以验证安装是否成功：

python3 op_tests/test_layernorm2d.py

unsetunset一、架构总览与设计哲学unsetunset

1.1 单一入口，四重后端

AITER 最核心的设计理念可以浓缩为一句话：借助 Python 层的统一 API 来屏蔽底层内核的异构性，并通过配置驱动的方式在运行时动态选择最优后端执行。

从 aiter/__init__.py 文件的导入结构中，我们可以清晰地看到这种分层设计的实现方式：

来源：aiter/init.py

from .ops.gemm_op_a8w8 import * # INT8/FP8 量化 GEMM
from .ops.gemm_op_a16w16 import * # BF16/FP16 GEMM
from .ops.gemm_op_a4w4 import * # FP4 量化 GEMM
from .ops.moe_op import * # MoE 算子
from .ops.attention import * # MHA/MLA/PA
from .ops.activation import * # SiLU/GeLU 等激活函数
from .ops.rope import * # RoPE 位置编码

每个 ops 模块对外暴露统一的函数签名，而在内部，系统会根据硬件架构（如 gfx942/gfx950）、矩阵形状（M/N/K）、量化类型（如 per_Token/per_1x128/per_1x32）以及预调优配置（存储在 CSV 文件中） 等多个维度，将请求动态路由到四种不同的后端实现上。这四种后端分别是：

• ASM 汇编内核：采用手写的 GCN/RDNA 汇编，性能极致但受限于特定形状。
• CK（Composable Kernel）：基于 AMD 开源的模板化 C++ 代码生成框架。
• Triton：基于 OpenAI Triton 开发的 Python DSL 内核。
• FlyDSL：AMD 自研的领域特定语言，专为混合精度 MoE 场景设计。

1.2 核心数据流

一次典型的 Fused MoE 推理调用的数据流如下所示：

用户调用 fused_moe()
→ 量化类型推断 & 形状对齐
→ get_2stage_cfgs() 查询 CSV 配置
→ moe_sorting() Token 按专家排序
→ 1-stage 或 2-stage GEMM 执行
→ 输出 moe_buf

该流水线中的每一个环节都包含了精巧的工程决策，接下来我们将逐一进行深入剖析。

unsetunset二、JIT 编译基础设施：从源码到 .so 的全自动流水线unsetunset

2.1 compile_ops 装饰器：声明即编译

AITER 的 JIT 系统构成了整个项目的骨架。

开发者只需使用 @compile_ops 装饰一个仅包含类型签名的空函数，框架便会自动完成整个流程：查找预编译的 .so 文件、若未找到则触发 JIT 编译、加载生成的模块、最终调用 C++ 算子。

来源：aiter/ops/gemm_op_a8w8.py

@compile_ops(
“module_gemm_a8w8″, fc_name=”gemm_a8w8”,
gen_fake=gen_gemm_a8w8_ck_fake_tensors
)
def gemm_a8w8_ck(
XQ: torch.Tensor, WQ: torch.Tensor,
x_scale: torch.Tensor, w_scale: torch.Tensor,
Out: torch.Tensor, bias: Optional[torch.Tensor] = None,
splitK: int = 0,
) -> torch.Tensor: …

请注意，函数体是 ...（Ellipsis）——它根本不需要 Python 实现。compile_ops 内部的 wrapper 函数会执行以下步骤：

1. 尝试通过 get_module(md_name) 加载已编译的模块。
2. 如果抛出 ModuleNotFoundError 异常，则调用 build_module() 启动编译过程。
3. 编译完成后，通过 getattr(module, loadName) 获取 C++ 函数的指针。

2.2 build_module：多进程安全的编译引擎

build_module 是 JIT 系统的重型核心，它需要解决几个关键难题。

多进程编译锁：在多 GPU 训练场景中，多个进程可能会同时触发同一模块的编译。AITER 采用文件锁（FileBaton）机制来确保只有一个进程执行编译操作：

来源：aiter/jit/core.py

def mp_lock(lockPath, MainFunc, FinalFunc=None, WaitFunc=None):
baton = FileBaton(lockPath)
if baton.try_acquire():
try:
ret = MainFunc()
finally:
if FinalFunc is not None:
FinalFunc()
baton.release()
else:
baton.wait() # 其他进程等待

好的，请查收根据您的要求深度重写后的文章片段。

编译标志自适应

该框架具备自动检测HIP编译器版本的能力，并能据此逐一添加优化标志。具体来说，当运行在ROCm 6.2及以上版本时，它会启用 amdgpu-early-inline-all 和 amdgpu-function-calls=false 这两个标志。此外，针对gfx950架构，框架还会自动加入对FP4数据类型的支持标志，其实现逻辑如下：

# 代码出处：aiter/jit/core.py（位于 build_module 内部）
if hip_version > Version("6.2.41132"):
flags_hip += [
"-mllvm -amdgpu-early-inline-all=true",
"-mllvm -amdgpu-function-calls=false",
]
if get_gfx() == "gfx950" and int(os.getenv("AITER_FP4x2", "1")) > 0:
flags_hip += ["-D__Float4_e2m1fn_x2"]

2.3 torch_compile_guard：与 torch.compile 无缝集成

在AITER中，每一个算子都通过 torch_compile_guard 装饰器被注册为 torch.ops.aiter.* 下的自定义算子。这一机制确保了它们能够被 torch.compile 的计算图所捕获并进行优化。该装饰器能够自动推导出算子的 schema，生成用于符号化追踪（symbolic tracing）的 FakeTensor，并妥善处理CUDA与CPU双端的分发逻辑：

# 代码出处：aiter/jit/utils/torch_guard.py
aiter_lib.define(op_schema, tags=tags)
aiter_lib.impl(f"aiter::{loadName}", custom_func, dispatch_key="CUDA")
aiter_lib.impl(f"aiter::{loadName}", custom_func, dispatch_key="CPU")
aiter_lib._register_fake(f"{loadName}", fake_func)

得益于此，用户可以直接在 torch.compile(model) 的上下文中使用AITER算子。编译器将能正确地处理形状推导与内存分配等底层细节。

三、Fused MoE：Token排序与多阶段GEMM的精密编排

Fused MoE是AITER中复杂度最高且最为关键的算子，其核心实现文件 fused_moe.py 的代码量接近2000行。它的核心目标是将“门控路由选择 → Token重排序 → 专家GEMM计算 → 加权聚合”这一完整的MoE流水线，融合成尽可能少的kernel启动次数。

3.1 Token排序：为GPU并行计算铺平道路

MoE推理的第一步是对Token进行重排序，依据是它们被分配到的专家ID。此举旨在确保由同一专家处理的Token在内存中是连续存放的，从而使后续的GEMM操作能够以高效的密集矩阵乘法形式执行。

# 代码出处：aiter/fused_moe.py
def _moe_sorting_impl(...):
M, topk = topk_ids.shape
max_num_tokens_padded = topk_ids.numel() + num_experts * block_size - topk
max_num_m_blocks = (max_num_tokens_padded + block_size - 1) // block_size

sorted_ids = torch.empty(max_num_tokens_padded, dtype=dtypes.i32, device=device)
sorted_weights = torch.empty(max_num_tokens_padded, dtype=dtypes.fp32, device=device)
sorted_expert_ids = torch.empty(max_num_m_blocks, dtype=dtypes.i32, device=device)

fwd_fn = aiter.moe_sorting_opus_fwd if use_opus else aiter.moe_sorting_fwd
fwd_fn(topk_ids, topk_weights, sorted_ids, sorted_weights, ...)

此处运用了一个精妙的 padding 策略：max_num_tokens_padded 不仅包含了所有Token乘以TopK后的展开数量，还为每一位专家额外预留了 block_size 大小的对齐空间。这样设计的好处是，可以确保后续基于Tile的GEMM操作在处理尾部Token时，不会因为其数量不足一个Block而需要编写特殊的边界处理逻辑。

3.2 一阶段 vs 两阶段：动态决策引擎

AITER为MoE提供了两种不同的执行路径：

• 1-stage：在一个kernel内同时完成Gate+Up投影、激活函数计算以及Down投影。
• 2-stage：将计算拆分为两个kernel——第一阶段执行Gate+Up+激活，第二阶段执行Down投影。

具体选择哪条路径，由 get_2stage_cfgs() 函数决定，这个函数堪称整个MoE调度逻辑的“大脑”。其决策依据主要包含以下几个方面：

# 代码出处：aiter/fused_moe.py（位于 get_2stage_cfgs 内部）
# 一阶段路径的启用条件因量化类型而异
if q_type == QuantType.per_1x128:
run_1stage = token > 32 and (inter_dim % 128 == 0)
elif q_type == QuantType.per_Token and q_dtype_w == dtypes.fp8:
run_1stage = token > 16 or inter_dim % 128 != 0
elif q_type != QuantType.per_1x32:
run_1stage = token < 256

这背后的设计直觉是：

• 当Token数量较少（例如decode场景）时，采用1-stage模式可以有效减少kernel启动带来的开销。
• 当Token数量较大（例如prefill场景）时，2-stage模式允许中间结果在两次GEMM操作之间进行更精细的量化处理（例如MXFP4动态量化），从而获得更高的计算密度。

3.3 Block Size 自适应：榨干每一个计算单元

get_block_size_M 函数展现的是一种极致的硬件感知调度策略：

# 来源：aiter/fused_moe.py  
@functools.lru_cache(maxsize=2048)  
def get_block_size_M(token, topk, expert, inter_dim):  
cu_num = get_cu_num()          # GPU Compute Unit 数量  
tileN = 128  
tgN = (inter_dim + tileN - 1) // tileN  
support_list = [32, 64, 128]  

tmp = []  
for el in support_list:  
max_num_tokens = token * topk + expert * el - topk  
tg_num = tgN * (max_num_tokens + el - 1) // el  # 总 Tile Group 数  
rnd = (tg_num + cu_num - 1) // cu_num            # 每 CU 需执行的轮数  
empty = cu_num - tg_num % cu_num                  # 空闲 CU 数  
tmp.append((rnd, empty, el))  
return sorted(tmp, key=lambda x: x[:2])[0][-1]  

该函数遍历所有候选的 Block Size（32、64、128），针对每种配置分别计算总 Tile 数量 → CU 利用率 → 空闲 CU 数量，最终挑选出“执行轮数最少、空闲 CU 最少”的最优方案。这相当于在 Python 层执行了一次精简的 GPU occupancy 分析。

3.4 K-Split：低 Token 数场景下的 CU 利用率救星

在 decode 场景下（Token 数量极少，例如 1-8），MoE 的 GEMM 矩阵 M 维度非常小，导致 Tile 总数不足以充分利用所有 CU。get_ksplit 通过在 K 维度上拆分计算来人为创造更多并行度：

# 来源：aiter/fused_moe.py  
def get_ksplit(token, topk, expert, inter_dim, model_dim):  
if token * topk > expert:  
return 0  # Token 足够多，不需要 split  
cu_num = get_cu_num()  
tg_num = tgN * tgM  
if tg_num >= cu_num:  
return 0  # CU 已经满了  
split_max = (cu_num + tg_num - 1) // tg_num  
for i in reversed(range(2, split_max + 1)):  
if (model_dim % i == 0) and ((model_dim // i) % 256 == 0):  
return i  
return 0  

这就好比将一个大蛋糕（K 维度的累加运算）切成多份，让更多人（CU）同时处理，最后再将结果合并——这是一种经典的 split-K 并行策略。

四、GEMM 自动调优：从形状发现到最优内核的闭环

4.1 多后端路由表

AITER 的 BF16/FP16 GEMM 支持五种后端，通过一个简洁的路由表进行切换：

# 来源：aiter/tuned_gemm.py  
solMap = {  
"torch":     torch_gemm,     # PyTorch F.linear 兜底  
"hipblaslt": hipb_gemm,      # hipBLASLt 库  
"skinny":    skinny_gemm,    # 超窄矩阵专用汇编  
"asm":       asm_gemm,       # 手写汇编 kernel  
"triton":    triton_gemm,    # Triton JIT kernel  
}  

路由决策完全由 CSV 配置文件驱动。get_GEMM_A16W16_config 函数会使用 (cu_num, M, N, K, dtype, otype, ...) 作为索引进行查表：

# 来源：aiter/tuned_gemm.py  
config = cfg.get(  
(cu_num, padded_M, N, K, bias, str(dtype), str(otype), scaleAB, bpreshuffle),  
None,  
)  

如果精确形状未能命中，系统还会尝试两种 padding 策略（get_padded_m 的 gl=0 和 gl=1）进行模糊匹配。

4.2 Skinny GEMM：小 M 场景的秘密武器

当 M 极小（1-16）且 N 不超过 CU 数量的 1-2 倍时，AITER 会绕过通用 GEMM 库，直接使用手写的 Split-K 汇编内核：

# 来源：aiter/tuned_gemm.py（默认路由逻辑）  
if (  
((M == 1 and N <= 2 * cu_num) or (M > 1 and M <= 4 and N <= cu_num))  
and K <= 9216  
or (M > 4 and M <= 8 and N <= cu_num) and K <= 5120  
):  
default_config["libtype"] = "skinny"  
default_config["solidx"] = 2  

这些内核专门针对 MI300X 的 304 个 CU 进行了优化，在 decode 场景（M=1）下通常比通用库快 2-3 倍。

4.3 在线调优与离线调优的双轨机制

AITER 的 MoE 层支持在线调优：当运行时遇到 CSV 中尚未记录的形状组合，且环境变量 AITER_ONLINE_TUNE=1 已设置时，框架会自动触发一次调优流程：

# 来源：aiter/fused_moe.py（get_2stage_cfgs 内部）  
if cfg is None and os.environ.get("AITER_ONLINE_TUNE", "0") == "1":  
lock_path = os.path.join(bd_dir, f"lock_fmoe_tune_{keys}")  
mp_lock(lock_path, MainFunc=MainFunc, FinalFunc=FinalFunc)  

调优脚本 gemm_moe_tune.py 会被直接调用，并将结果写回 CSV。这意味着第一次推理时可能会有几分钟的调优延迟，但后续所有推理都能享受到最优配置。

在离线调优场景下，gradlib提供了一套完整流程：首先通过AITER_TUNE_GEMM=1来捕捉矩阵形状，接着使用gemm_tuner.py执行批量调优，最后将生成的CSV结果放回aiter/configs/目录。

五、量化生态：从FP4到INT8的全精度覆盖

5.1 量化类型的精细分级

AITER中MoE支持的量化粒度从粗到细排列如下：

这套体系在fused_moe_1stage_dict中以查找表形式固化，每种(activation, quant_type, dtype, q_dtype_a, q_dtype_w, isG1U1)组合都映射到一个专用kernel：

# 来源：aiter/fused_moe.py  
fused_moe_1stage_dict = {  
"gfx942": {  
(ActivationType.Silu, QuantType.No,       dtypes.bf16, dtypes.bf16, dtypes.bf16, False, False): aiter.fmoe,  
(ActivationType.Silu, QuantType.per_1x128, dtypes.bf16, dtypes.fp8,  dtypes.fp8,  True,  False): aiter.fmoe_g1u1,  
# ... 共11种组合  
},  
"gfx950": {  
# ... 7种组合，包括gfx950特有的fp4x2和blockscale  
}  
}  

5.2 MXFP4量化的融合排序

针对MXFP4（per_1x32）量化，AITER实现了一项巧妙的优化——将动态量化与Token排序融合进同一个kernel：

# 来源：aiter/fused_moe.py（fused_moe_2stages内部）  
a1, a1_scale = fused_dynamic_mxfp4_quant_moe_sort(  
hidden_states,  
sorted_ids=sorted_ids,  
num_valid_ids=num_valid_ids,  
token_num=token_num,  
topk=topk,  
block_size=block_size_M,  
num_rows=num_local_tokens,  
)  

这种方法避免了“先量化、再排序”带来的两次全局内存读写，对于显存带宽受限的场景意义重大。

六、有趣的地方

6.1 LRU缓存的大量使用

AITER在性能关键路径上大量运用了@functools.lru_cache，涉及get_inter_dim、get_block_size_M、get_2stage_cfgs、get_GEMM_A16W16_config等函数。由于推理阶段的矩阵形状有限且重复，这些缓存几乎能确保每种形状只做一次决策计算，后续全部走O(1)查表。

6.2 FlyDSL回退机制

当调优配置指定了FlyDSL kernel但运行环境中未安装FlyDSL时，AITER不会直接报错，而是查找CSV中标记为flydsl_fallback的备选配置：

# 来源：aiter/fused_moe.py  
if cfg is not None and not is_flydsl_available():  
kn1 = str(cfg.get("kernelName1", ""))  
if kn1.startswith("flydsl_"):  
fallback = fallback_cfgs.get(keys)  
if fallback is not None:  
cfg = fallback  

这种优雅降级设计让AITER在不同部署环境中都能稳定工作。

6.3 ctypes FFI：绕过PyTorch的无Torch调用路径

对于某些对延迟极度敏感的算子（如ASM GEMM），AITER提供了ffi_type="ctypes"选项，通过ctypes直接调用C ABI的.so，完全绕过PyTorch的pybind11开销。_ctypes_call函数会根据Python类型标注自动生成ctypes参数映射，并在每次调用时将torch.Tensor转换为轻量的aiter_tensor_t结构体。

结语

AITER展示了一个高性能AI算子库应有的样子：它不追求单点极致，而是构建了一套配置驱动、多后端协同、自动调优的系统化工程。

在大模型推理日益成为产业刚需的今天，这种“让每块GPU都跑在最优点”的基础设施能力，正是AMD ROCm生态最需要的底层支撑。对于从事GPU算子开发或大模型推理优化的工程师而言，AITER的JIT编译框架设计、MoE多阶段调度策略和自动调优闭环，都是极具参考价值的工程实践。

参考资料[1]

ATOM/benchmark-dashboard: https://rocm.github.io/ATOM/benchmark-dashboard/