Posts by Jehandad Khan
Towards Feature Complete Triton Support in JAX-Triton
- 08 July 2026
In this blog we’ll explore recent improvements to JAX-Triton project and learn about new features available at the AMD fork.
OpenXLA and JAX - ROCm Support and the State of CI
- 29 June 2026
The OpenXLA compiler stack — XLA at the foundation, JAX as the front end — now runs upstream on AMD ROCm. XLA gates every pull request on real AMD Instinct silicon through its GitHub Actions workflow, side by side with the CUDA path; JAX runs the same hardware on every ROCm PR through its own workflows, with the merge gate rolling out next. pip install "jax[rocm7-local]" is a first-class entry point. This post documents how that backend is structured, what landed in the last twelve months, and how the CI pipeline that keeps it healthy is wired together. Part 1 covers OpenXLA on AMD — the XLA backend, what landed this year, and CI. Part 2 covers JAX on AMD — the plugin architecture, JAX-side changes, and the four-workflow test matrix.
JAX-AITER: Bringing AMD’s Optimized AI Kernels to JAX on ROCm™
- 24 February 2026
If you’re building large models in JAX on AMD GPUs, you want fast, reliable kernels without spending weeks tuning them yourself. That’s exactly the need that led us to create JAX-AITER.
ROCm MaxText Testing — Decoupled (Offline) and Cloud-Integrated Modes
- 06 January 2026
In this blog, you will learn how to run MaxText unit tests on AMD ROCm GPUs in two complementary modes: offline (decoupled) and fully cloud-integrated. By the end, you will know when to use each mode, how to interpret the results, and how to fold them into your CI and debugging workflows.
ROCm Fork of MaxText: Structure and Strategy
- 06 January 2026
In this blog you will explore how the ROCm fork of MaxText is structured and how that structure supports ROCm and fully offline, decoupled workflows across platforms.