JAX-AITER: Bringing AMD’s Optimized AI Kernels to JAX on ROCm™

2026-02-24 · Source: AMD ROCm Blogs · Field: Technology & Digital — Artificial Intelligence & Machine Learning, Software Development & Engineering · Depth: Advanced, long

Summary

JAX-AITER is an open-source bridge developed by AMD that integrates AMD's optimized AITER (AI operator) kernels into the JAX ecosystem for ROCm-enabled GPUs. This initiative aims to provide JAX users with high-performance operators, such as multi-head attention (MHA/FMHA), mixture-of-experts (MoE), and matrix multiplication (GEMM), without requiring manual tuning. The architecture consists of a JAX-friendly Python frontend, a C++/FFI bridge for buffer and stream management, and the AITER backend. Benchmarks on AMD Instinct MI350 GPUs demonstrate significant performance improvements, with median speedups of 4.39x and mean speedups of 4.23x for attention workloads compared to pure JAX implementations. While MHA/FMHA kernels are framework-agnostic, some GEMM and custom ops currently rely on PyTorch dependencies, with a roadmap to achieve full framework neutrality.

Key takeaway

For NLP Engineers and AI Scientists developing large models in JAX on AMD Instinct GPUs, integrating JAX-AITER can dramatically accelerate attention-heavy workloads. You should clone the JAX-AITER repository, build it for your ROCm environment, and swap out your existing JAX attention implementations with `flash_attn` to realize median speedups of over 4x, particularly for longer sequences and higher head counts.

Key insights

JAX-AITER bridges JAX with AMD's optimized AITER kernels for significant performance gains on ROCm GPUs.

Principles

• Reuse optimized kernels to avoid reinvention.
• Bridge frameworks with FFI for performance.
• Prioritize framework-agnostic C++ APIs.

Method

JAX-AITER uses `jax.custom_vjp` for autodiff, a C++/FFI layer to map JAX buffers to AITER, and synchronizes HIP streams to ensure correct execution order and memory visibility.

In practice

• Replace JAX attention with `jax_aiter.mha.flash_attn`.
• Benchmark JAX-AITER against pure JAX for your workload.
• Use `block_until_ready()` for accurate timing.

Topics

• JAX-AITER
• ROCm Optimization
• AITER Kernels
• Multi-Head Attention
• AMD Instinct GPUs

Code references

• ROCm/jax-aiter
• ROCm/jax-aiter

Best for: NLP Engineer, AI Scientist, Research Scientist, AI Engineer, Machine Learning Engineer, AI Researcher

Related on AIssential

• Speed-up JAX LLM Training on Intel® Xeon® 6 CPU: Activation Offloading on Heterogeneous Systems — JAX-based activation offloading to Intel Xeon 6 CPU memory improves LLM training throughput by 3-15.5% on heterogeneous systems.
• AMD Instinct™ GPUs MLPerf Inference v6.0 Submission — AMD's MI355X GPU and ROCm stack deliver competitive MLPerf Inference v6.0 results, especially for LLMs and multi-node setups.
• AI giveth and AI taketh CPU — AMD's AI strategy leverages heterogeneous computing and modular chiplet designs for performance, efficiency, and adaptability across diverse workloads.
• Customizing Kernels with hipBLASLt TensileLite GEMM Tuning - Advanced User Guide — TensileLite Tuning generates custom GEMM kernels for AMD GPUs, significantly outperforming existing solutions.
• Taking on CUDA With ROCm: ‘One Step After Another’ — AMD's ROCm software stack is rapidly maturing, leveraging open-source tools and community engagement to challenge Nvidia's CUDA dominance.
• CuTeGen: An LLM-Based Agentic Framework for Generation and Optimization of High-Performance GPU Kernels using CuTe — CuTeGen iteratively refines GPU kernels using CuTe abstraction and delayed profiling, achieving competitive performance.

Open in AIssential →

Editorial summary, takeaway, and curation by AIssential. Original article published by AMD ROCm Blogs.