DOCSEARCHbyAlgolia
Submit Docs

ROCm AITER

repository·main·Indexed 19 days ago

https://github.com/rocm/aiter

AMD's high-performance AI operator library providing optimized GPU kernels for inference and training on ROCm. It includes production-ready operators for framework developers, featuring FlyDSL AOT (Ahead-Of-Time) pre-compilation for MoE and GEMM kernels, and the Opus a16w16 GEMM optimized for AMD gfx950 (MI300X class) hardware. The library supports performance tuning for models such as Llama3.3-70B, Llama-405B, and Qwen3-32B.

Tokens
86.6K
Snippets
205
Records
371
Agent score
63%

What's inside rocm-aiter

  1. What is opus and when to use it

    main

    opus is a lightweight, templated C++ DSL designed to sit above hand-written HIP kernels but below highly optimized template libraries like ck/cutlass. It provides essential abstractions without the overhead of a heavy framework.

    Use opus if you need:

    • AMDGPU data type declaration and conversion.
    • Automated vectorized buffer load/store dispatch.
    • Support for various matrix core instructions (MFMA) with minimal code changes.
    • A collection of utility device functions.
    • Simple layout abstractions for index calculations.

    Do NOT use opus if you need:

    • Pre-optimized kernels (e.g., GEMM, attention, reduction) for direct use.
    • Reusable device-side pipelines for GEMM/attention/reduction.
    • A comprehensive layout system for arbitrary tensor transformations. (For these, use ck or aiter kernels instead.)

  2. Overview of AITER

    main

    AITER (AI Tensor Engine for ROCm) is a high-performance AI operator library from AMD that provides optimized GPU kernels for both inference and training workloads on ROCm.

    Key Characteristics:

    • Dual APIs: Accessible via both C++ and Python.
    • Multiple Backends: Utilizes Triton, Composable Kernel (CK), and hand-tuned ASM.
    • Framework Agnostic: Designed to be integrated into frameworks like vLLM, SGLang, or custom stacks.
    • Workload Support: Optimized for inference (serving) and training (including GEMM+communication fused kernels).

  3. Overview of Opus MoE Stage2 Kernels

    main

    Opus MoE provides stage2 kernels and Python bindings specifically optimized for the gfx950 architecture. The current implementation focuses on fused MoE enablement through tuned A8W4 stage2 configurations.

    Key Paths

    • A8W4 Decode Stage2: The active fused MoE path. It uses public algorithm IDs and generates effective inter-dimension specializations. Runtime parameters logical_inter_dim and inter_dim_pad determine the effective-K specialization, while topk, hidden, and experts are handled at runtime.
    • Private BF16 Stage2: A route-reduce kernel currently retained for future use. It is not exposed through the fused_moe or Python user APIs in the current version.

  4. Overview of HK MLA V40 Gen.1 Decode Kernel

    main

    The mi35x_v40_fwd_decode_m16x8_fp8bf16_fp8bf16_gen1 kernel is a V4 MLA (Multi-head Latent Attention) decode kernel designed for the gfx950 (MI350) architecture. It implements causal-free per-token attention for DeepSeek-style MLA layouts, where a single shared latent KV cache is read by multiple query heads.

    Key Characteristics:

    • Gen.1 Design: All inputs (Q and KV) arrive as fp8 (with E8M0 scales) but are converted to bf16 at load time. There is no native fp8 MFMA used; instead, it uses v_mfma_f32_16x16x32_bf16 for both QK and PV operations.
    • Partitions: The kernel operates in two modes based on the total per-workgroup work $W = H \cdot \mathrm{mtp}$:
      • m16x8 ($W=128$): 8 ptiles per workgroup, occupancy 1. Uses the KvManager8to16bitsV2 pipeline. This is the primary focus of the design spec.
      • m16x4 ($W=64$): 4 ptiles per workgroup, occupancy 2. Uses the older KvManager8to16bitsV1 pipeline with kBlockN = 32.
    • Dispatching: The router aiter/mla.py::mla_v40_decode_fwd selects the m16x8 partition when $H \cdot \mathrm{mtp} = 128$ and m16x4 when $H \cdot \mathrm{mtp} = 64$.

  5. Overview of AITER core features

    main

    AITER (AMD Inference and Training Enhanced Repository) is a high-performance AI operator library optimized for AMD GPUs and the ROCm platform. It provides optimized kernels for both inference and training workloads using Triton, Composable Kernel (CK), and hand-written assembly.

    Key operator categories include:

    • Attention Kernels: Multi-Head Attention (MHA), Multi-Latent Attention (MLA) (DeepSeek-style), and Paged Attention for KV-cache management.
    • GEMM Operations: Mixed Precision (FP16, BF16, FP8, INT4), Tuned GEMM for common shapes, and Fused Operations (GEMM with activation fusion).
    • Mixture of Experts (MoE): Fused MoE, multiple routing strategies, and support for quantized experts (FP8 and INT4).
    • Normalization: RMSNorm, LayerNorm, and fused variants.
    • Other Operators: RoPE (Rotary position embeddings), Quantization (BF16/FP16 to FP8/INT4), Element-wise operations, and Communication (AllReduce/collectives via Triton/Iris).

  6. Overview of AITER tutorials

    main

    AITER provides tutorials categorized into three main areas:

    Basic Tutorials

    • basic_usage: Your first AITER program.
    • attention_tutorial: Understanding attention kernels.
    • variable_length: Handling variable-length sequences.

    Advanced Topics

    • add_new_op: Step-by-step guide on how to add a new operator.
    • moe_tutorial: Mixture of Experts (MoE) optimization.
    • custom_kernels: Writing custom ROCm kernels.
    • quantization: INT8 quantization for inference.
    • triton_comms: Triton-based communication primitives.

    Integration Guides

    • vllm_integration: Using AITER with vLLM.
    • pytorch_lightning: PyTorch Lightning integration.
    • deepspeed: DeepSpeed integration.

  7. Understand the Opus GEMM (C++ side) architecture and layout

    main

    The csrc/opus_gemm directory contains the C++ and JIT build inputs for the Opus a16w16 GEMM. The architecture uses a dispatch pattern where a top-level entry point probes the hardware architecture at runtime and routes the call to a specific architecture-optimized implementation.

    Key Components:

    • opus_gemm.cu: Top-level entry points (opus_gemm() / opus_gemm_a16w16_tune()) and arch routers.
    • include/opus_gemm.h & include/opus_gemm_arch.cuh: Cross-arch declarations, the OpusGfxArch enum, and the opus_get_arch_info() probe.
    • include/gfx950/*.cuh: Architecture-specific pipelines (e.g., a16w16 split-barrier, flatmm), traits, and heuristic dispatch logic.
    • gen_instances.py: The JIT codegen driver that bakes tuned CSV data into opus_gemm_lookup.h.
    • opus_gemm_tune.py: The offline tuner CLI.

    Dispatch Flow (e.g., gfx950):

    1. opus_gemm() is called.
    2. opus_get_gfx_arch() identifies the architecture (e.g., OpusGfxArch::Gfx950).
    3. A switch statement routes to the architecture-specific dispatcher (e.g., opus_dispatch_a16w16_gfx950<T>).
    4. The dispatcher uses either a tuned (M,N,K) lookup map or a heuristic dispatch function (opus_a16w16_heuristic_dispatch_gfx950<T>) to select the optimal kernel.

  8. What is Opus?

    main
    Opus is a lightweight, single-header C++ template library (opus.hpp) designed for writing HIP kernels on AMD GPUs. It provides abstractions for vectorized load/store, layouts, and MFMA wrappers. A primary benefit of Opus is its focus on build-time optimization, claiming up to 61x faster builds compared to standard PyTorch extensions.

  9. What is Gluon and how does it differ from Triton?

    main

    Gluon is a GPU programming language used for writing high-performance kernels. It operates at the same level as Triton but provides more explicit control over:

    • Memory layouts
    • Async copy operations
    • MFMA (Matrix Fused Multiply-Add) intrinsics

    Note: Some advanced features, such as scheduling hints like sched_barrier, require the AMD Gluon Extension.

  11. How sytrd (HIP) works

    main

    The sytrd implementation performs Householder symmetric tridiagonalization using a blocked approach. It operates on the lower triangle of the matrix (consistent with rocBLAS/rocSOLVER using rocblas_fill_lower).

    Core Components:

    • larfg: Computes a Householder vector v and scalar tau such that (I - tau v v^T)[alpha; x] = [beta; 0]. The larfg_kernel stores v in-place in the lower part of the column, with beta in E and tau in TAU.
    • sytd2: A non-blocked reduction that applies reflectors one column at a time. It outputs the diagonal D, subdiagonal E, and TAU.
    • latrd: A blocked variant that processes a panel of nb columns. It accumulates coupling with prior columns and uses a symmetric rank-2k trailing update (A <- A - V W^T - W V^T) via GEMM after processing the panel.
    • sytrd: The blocked driver that orchestrates latrd on panels, applies trailing updates, handles remainders with sytd2, and finally extracts D and E.

  12. Winograd 3x3 convolution pipeline

    main

    The Winograd kernels (winograd_f4x3*) implement a 3-stage pipeline to perform convolution efficiently:

    1. Input Transform: Transforms the input $X$ into $V$ using $V = B^T X B$. For $F(4,3)$, this involves reading $6 imes 6$ input patches.
    2. Batched GEMM: Performs a batched matrix multiplication on the transformed data ($36 imes T imes C_{pad}$), reducing over the $C_{pad}$ dimension to produce $M$ ($36 imes T imes K_{out}$).
    3. Output Transform: Transforms the result $M$ back to the output space $Y$ using $Y = A^T M A$.