TokenSpeed-Kernel-AMD

TokenSpeed-Kernel-AMD is a standalone collection of performance-oriented AMD GPU kernels for LLM inference. Its performance kernels are written in Gluon. The package currently covers architecture-specific implementations for:

• Attention: MHA, MLA, DSA, and KDA
• MoE: BF16, MXFP4 weights
• GEMM

Performance Numbers

The following measurements use GPT-OSS 120B workloads on one AMD Instinct MI355X GPU. They were collected at TokenSpeed commit 1492030, with AITER 0.1.13 and ROCm 7.2.1. See the TokenSpeed-Kernel PyTorch blog for the complete methodology and analysis.

Attention

The benchmark uses BF16 Q/K/V, head dimension 64, 64 query heads, 8 KV heads, full causal attention, and attention sinks. It covers sequence lengths 1K, 4K, and 8K with batch sizes from 1 to 16. Bars report throughput in TFLOP/s; higher is better. The Gluon kernel is the fastest evaluated backend. It is 1.4-2.3x faster than the Triton baseline and 1.1-1.3x faster than AITER.

MoE

This benchmark measures full-MoE latency, including routing, both GEMMs, clamped SwiGLU, and combine, for GPT-OSS 120B with 128 experts, top-4 routing, MXFP4 weights, FP8 activations, and D = I = 2880. For small decode batches (M = 1-4), Gluon is 1.7-2.1x faster than Triton and 1.1-1.6x faster than AITER. At M = 8-16, Gluon remains 1.3-1.4x faster than Triton. The prefill results show Gluon substantially ahead of Triton and competitive with AITER across M = 512-8192.

Package Organization

Kernels are organized by AMD architecture and operator family. This keeps architecture-specific tuning local while giving each family a consistent place for its implementations and supporting utilities.

tokenspeed-kernel-amd/
└── python/
    └── tokenspeed_kernel_amd/
        └── ops/
            ├── gfx950/
            │   ├── attention/
            │   │   ├── mha/   # Multi-head attention
            │   │   ├── rmha/  # Relative-bias multi-head attention
            │   │   ├── mla/   # Multi-head latent attention
            │   │   ├── dsa/   # DeepSeek sparse attention
            │   │   └── kda/   # Kimi delta attention
            │   ├── gemm/      # GEMM implementations
            │   ├── moe/       # MoE implementations
            │   └── sampling/  # Sampling implementations
            └── gfx1250/
                ├── attention/
                │   ├── mha/
                │   ├── mla/
                │   └── kda/
                └── moe/

Public entry points currently remain architecture-specific. Consumers should import the implementation matching the target GPU, or use TokenSpeed-Kernel to select a compatible implementation through its registry.

Usage

Install a ROCm-compatible PyTorch build first, then install the package from PyPI:

pip install tokenspeed-kernel-amd

For development from this repository:

pip install -e ./tokenspeed-kernel-amd