Primus-Turbo

September 21, 2026 · View on GitHub

🔍 What's Primus-Turbo?

Primus-Turbo is a high-performance acceleration library dedicated to large-scale model training on AMD GPUs. Built and optimized for the AMD ROCm platform, it covers the full training stack — including core compute operators (GEMM, Attention, GroupedGEMM), communication primitives, low-precision computation (FP8), and compute–communication overlap kernels.

With High Performance, Full-Featured, and Developer-Friendly as its guiding principles, Primus-Turbo is designed to fully unleash the potential of AMD GPUs for large-scale training workloads, offering a robust and complete acceleration foundation for next-generation AI systems.

Note: JAX support is under active development. Optim support is planned but not yet available.

Part of the Primus Ecosystem: Primus-Turbo is the high-performance operator layer of the Primus ecosystem, working together with Primus-LM (training framework) and Primus-SaFE (stability & platform).

🚀 What's New

[2026/09/16] 🔥 Release v0.5.0 — Mega MoE / GroupedMLP MXFP8 maturity, FlyDSL grouped GEMM BF16, and operator auto-tune on MI355X (changes since v0.4.0).
[2026/08/25–09/14] 🔥 Fused GroupedMLP (FP8 / MXFP8) — expert MLP with SwiGLU folded into grouped GEMM; pad-aware paths for GPT-OSS-20B-style MoE (#476, #488, #503).
[2026/08/18] 🔥 Mega MoE (MXFP8) — fused forward and backward in the MXFP8 Mega MoE path: dispatch+FC1 and FC2+combine with intra-node EP (#456).
[2026/07/15] 🔥 Mega MoE (BF16) — FlyDSL fused MoE layer with comm–compute overlap inside the grouped GEMMs (overview) (#412).
[2026/07] 🔥 MXFP4 GEMM & grouped GEMM — FlyDSL and hipBLASLt stacks for dense and expert GEMMs in low-precision training (#424, #483).
[2026/07/30] 🔥 Release v0.4.0 — grouped GEMM autotune, MXFP4 correctness, and Mega MoE stability fixes on gfx950.
[2026/06/10] 🔥 DeepEP — JAX intranode/internode token dispatch (#344); PyTorch DeepEPTokenDispatcher (introduced in #114). The separate rocSHMEM backend integration is tracked in #409 (guide).
[2025/12/16] 🔥 MoE training best practices on AMD GPUs — DeepEP, grouped GEMM, and Primus recipes end to end.

📦 Quick Start

Requirements

Software

ROCm >= 7.0
Python >= 3.10
PyTorch >= 2.6.0 (with ROCm support)
AITER (required for some operators, e.g. FlashAttention / FP8): pip3 install "amd-aiter @ git+https://github.com/ROCm/aiter.git@v0.1.14.post1"
FlyDSL (required; provides the FlyDSL kernel backend for GEMM / GroupedGEMM / Attention / MoE)
rocSHMEM (optional, required for experimental DeepEP). Please refer to our DeepEP Installation Guide for instructions.

Hardware

Architecture	Supported GPUs
GFX942	✅MI300X, ✅MI325X
GFX950	✅MI350X, ✅MI355X

See AMD GPU Architecture to find the architecture for your GPU.

1. Installation

Docker (Recommended)

Use the pre-built AMD ROCm image from Docker Hub:

# PyTorch Ecosystem
docker pull rocm/primus:v26.2

# JAX Ecosystem
docker pull rocm/jax-training:maxtext-v26.2

You can also use the official ROCm PyTorch image from Docker Hub.

Install from Prebuilt Index

Prerequisite: install inside an environment that already has ROCm PyTorch — e.g. the rocm/primus image above, or the official rocm/pytorch image. Primus-Turbo builds against your existing torch and does not install torch for you; in a bare environment pip would otherwise pull a non-ROCm torch.

# PyTorch backend (latest)
pip3 install --no-build-isolation "primus-turbo[pytorch]" \
    --extra-index-url https://amd-agi.github.io/Primus-Turbo/simple/

# Pin a specific version
pip3 install --no-build-isolation "primus-turbo[pytorch]==0.1.0" \
    --extra-index-url https://amd-agi.github.io/Primus-Turbo/simple/

The index currently serves source distributions (sdist), so install compiles HIP kernels locally (needs the ROCm toolchain; supports gfx942 / gfx950). Prebuilt wheels are planned. Keep --no-build-isolation so the build uses your preinstalled torch.

Install from Source

git clone https://github.com/AMD-AGI/Primus-Turbo.git
cd Primus-Turbo

# Install build/runtime dependencies first
pip3 install -r requirements.txt

# Default backend: PyTorch
pip3 install --no-build-isolation ".[pytorch]"

# JAX backend
PRIMUS_TURBO_FRAMEWORK="JAX" pip3 install --no-build-isolation ".[jax]"

Install from GitHub URL (without cloning)

# Install from default branch
pip3 install --no-build-isolation "git+https://github.com/AMD-AGI/Primus-Turbo.git"

# Install from a specific branch
pip3 install --no-build-isolation "git+https://github.com/AMD-AGI/Primus-Turbo.git@main"

Note:

".[pytorch]" / ".[jax]" means install from current local repo with extras.

Extras select Python dependencies. Source compilation target is controlled by PRIMUS_TURBO_FRAMEWORK.

2. Development

For contributors, use editable mode (-e) so that code changes take effect immediately without reinstalling.

git clone https://github.com/AMD-AGI/Primus-Turbo.git
cd Primus-Turbo

pip3 install -r requirements.txt
pip3 install --no-build-isolation -e ".[pytorch]" -v

# (Optional) Set GPU_ARCHS environment variable to specify target AMD GPU architectures.
GPU_ARCHS="gfx942;gfx950" pip3 install --no-build-isolation -e ".[pytorch]" -v

# (Optional) Set PRIMUS_TURBO_FRAMEWORK to compile for a specific framework.
# Supported values: PYTORCH (default), JAX.
# For example, to compile for JAX:
PRIMUS_TURBO_FRAMEWORK="JAX" pip3 install --no-build-isolation -e ".[jax]" -v

# (Optional) ccache/sccache are auto-detected on PATH to speed up incremental rebuilds.
# Just install ccache or sccache and the build will use it automatically.

3. Testing

Option 1: Single-process mode (slow but simple)

pytest tests/pytorch/    # run all PyTorch tests
pytest tests/jax/        # run all JAX tests

Option 2: Multi-process mode (faster)

# PyTorch tests
## single-GPU tests (parallel)
pytest tests/pytorch/ -n 8
## deterministic tests (parallel)
pytest tests/pytorch/ -n 8 --deterministic-only
## multi-GPU tests
pytest tests/pytorch/ --dist-only

# JAX tests
## single-GPU tests (parallel)
pytest tests/jax/ -n 8
## multi-GPU tests
pytest tests/jax/ --dist-only

4. Packaging

pip installation behavior:

Use a compatible wheel (.whl) if available.
Fall back to source distribution (sdist, .tar.gz) when no wheel matches.

Artifact roles:

wheel: prebuilt binary package, fast install, no local C++/HIP build.
sdist: source package, slower install, requires local toolchain, fallback path.

Build artifacts

# Build wheel (binary distribution)
python3 -m build --wheel --no-isolation

# Build sdist (source distribution)
python3 -m build --sdist --no-isolation

Verify wheel install

pip3 install --no-build-isolation ./dist/primus_turbo-XXX.whl

Verify source fallback install

pip3 install --no-build-isolation ./dist/primus_turbo-XXX.tar.gz

Tip: Run import checks outside the source tree (for example under /tmp) to avoid importing local source files by accident.

5. Minimal Example

import torch
import primus_turbo.pytorch as turbo

dtype = torch.bfloat16
device = "cuda:0"

a = torch.randn((128, 256), dtype=dtype, device=device)
b = torch.randn((256, 512), dtype=dtype, device=device)
c = turbo.ops.gemm(a, b)

print(c)
print(c.shape)

┌─────────────────────────────────────────────────────┐
│                   Primus-SaFE                       │
│         (Stability & Platform Layer)                │
│   Cluster Management | Fault Tolerance | Scheduling │
└────────────────────────┬────────────────────────────┘
                         │
┌────────────────────────▼────────────────────────────┐
│                   Primus-LM                         │
│              (Training Framework)                   │
│    Megatron | TorchTitan | Unified CLI | Workflows  │
└────────────────────────┬────────────────────────────┘
                         │
┌────────────────────────▼────────────────────────────┐
│                  Primus-Turbo                       │
│           (High-Performance Operators)              │
│  Attention | GEMM | GroupedGEMM | MoE | DeepEP      │
│  Mega MoE | FP8/MXFP8/MXFP4 | AITER | CK | FlyDSL    │
└─────────────────────────────────────────────────────┘

📦 Component Details

Component	Role	Key Features	Repository
Primus (Primus-LM)	Training framework	Multi-backend training (Megatron, TorchTitan, MaxText), unified CLI, projection and tuning agent, MegaMoE integration in Megatron recipes	Primus
Primus-Turbo	Performance layer	FlashAttention-class kernels, GEMM / GroupedGEMM (BF16, FP8, MXFP8, MXFP4), Mega MoE (BF16 + MXFP8), DeepEP, fused GroupedMLP; backends CK, hipBLASLt, AITER, Triton, FlyDSL	This repo
Primus-SaFE	Platform layer	Cluster sanity checks, topology-aware scheduling, fault tolerance	Primus-SaFE

🔗 How They Work Together

Primus-LM provides the training framework and workflow orchestration.
Primus-Turbo supplies the optimized compute kernels (for example Mega MoE, grouped GEMM, attention, and low-precision GEMMs) that Primus recipes call into.
Primus-SaFE ensures stability and efficient resource utilization at scale.

This separation of concerns allows each component to evolve independently while staying integrated through pinned versions in Primus training images.

🙏 Acknowledgements

Primus-Turbo builds on excellent open-source work from the wider community. We especially thank:

FlyDSL — a Flexible Layout Python DSL and MLIR compiler stack for authoring high-performance AMD GPU kernels. Many of our kernels (GEMM, GroupedGEMM, Attention, MoE) are built with FlyDSL; those files carry FlyDSL attribution and remain under Apache-2.0 — see LICENSE and LICENSE-APACHE. We thank the FlyDSL team for their close collaboration and support.
AITER — AI Tensor Engine for ROCm, providing high-performance operator backends (e.g. FlashAttention, FP8) that Primus-Turbo integrates.
tritonBLAS — high-quality Triton GEMM kernels for AMD GPUs. Our persistent BF16/FP16 and FP8 GEMM kernels are adapted from it.
Triton — our Triton attention kernel is adapted from the AMD performance kernels in python/perf-kernels/flash-attention.py, which implement the FlashAttention v2 algorithm by Tri Dao.
Triton-distributed — a distributed compiler for computation-communication overlapping. Our Mega MoE comm-compute fused kernels reference its overlapping-kernel design.
DeepGEMM — a clean and efficient FP8/BF16 GEMM library. Our Mega MoE barrier and symmetric-heap layout designs reference it.
NVIDIA TransformerEngine — our Triton MoE permute/unpermute kernels are adapted from TransformerEngine. That file remains under Apache-2.0 and carries NVIDIA attribution — see LICENSE and LICENSE-APACHE.
NVIDIA Megatron-LM — parts of our MoE token permutation and dispatch layer are adapted from Megatron-LM. Those files carry NVIDIA attribution and remain under the 3-clause BSD license — see LICENSE.

📜 License

Primus-Turbo is licensed under the MIT License.