Sari la conținut

Română · beta. Traducerea a fost făcută automat și este încă în curs de revizuire. Citește pagina în engleză

perspectiva aiqbenchProduse AI

analiza aiqbench

This article details a significant performance breakthrough in deploying MXFP4 quantization for smaller LLMs by fusing online rotation kernels on AMD's MI355X hardware. The optimization reduces the latency overhead of maintaining accuracy during inference from 5–10% down to under 1%, making aggressive low-bit formats viable for production without sacrificing model fidelity.
Sursă care acoperăAMD
Ce este aceastaLectura unui model asupra articolului și a acoperirii conexe, scrisă pentru aiqbench. Nu sunt cuvintele publicației.

Production-Ready MXFP4 Online Rotation with Fused Kernels on AMD Instinct™ MI355X

AMD7 min de cititArticolul și perspectiva sunt în engleză.

Serving large language models affordably increasingly depends on low-bit quantization, and MXFP4 is one of the most aggressive options — but the smaller models that need it most rely on online rotation to stay accurate, and that rotation has historically carried a steep latency tax. In this post you will learn how a single fused Gluon (Triton) kernel on AMD Instinct™ MI355X (CDNA4) removes that tax: we walk through the kernel-fusion design, the RS=64 optimization, the GEAK + Hyperloom tuning workflow, and end-to-end measurements showing online-rotation overhead falling from a prohibitive +5–10% to just +0.3–0.8% on Dense models — with no measurable accuracy loss. By the end, you will understand how to make accuracy-preserving MXFP4 online rotation cheap enough to ship by default.

Serving large language models affordably increasingly depends on low-bit quantization, and MXFP4 is one of the most aggressive options—but the smaller models that need it most rely on online rotation to stay accurate, and that rotation has historically carried a steep latency tax. In this post you will learn how a single fused Gluon (Triton) kernel on AMD Instinct™ MI355X (CDNA4) removes that tax: we walk through the kernel-fusion design, the RS=64 optimization, the GEAK + Hyperloom tuning workflow, and end-to-end measurements showing online-rotation overhead falling from a prohibitive +5–10% to just +0.3–0.8% on Dense models—with no measurable accuracy loss. By the end, you will understand how to make accuracy-preserving MXFP4 online rotation cheap enough to ship by default.

Large language models keep growing, and serving them affordably increasingly depends on low-bit quantization. MXFP4—a 4-bit microscaling format—is one of the most aggressive options, shrinking memory and bandwidth roughly 4× versus bf16. But aggressive compression has a catch: the very largest models absorb MXFP4 with little degradation, while smaller models such as the Qwen3 8B/14B/32B Dense family lose noticeable accuracy under plain round-to-nearest (RTN) quantization. Rotation is the technique that recovers that loss—applying an orthogonal transform to activations and weights to suppress the outliers that break low-bit quantization. The catch, in turn, is that its most accurate form—online rotation, applied at every layer on every decode step—brings back a real cost in latency. Eliminating that latency cost is what this post is about.

In our previous post Advanced MXFP4 Quantization: Combining Fine-Tuned Rotations with SmoothQuant for Near-Lossless Compression we demonstrated that learned rotation significantly improves MXFP4 quantization accuracy for Qwen3 Dense models (8B/14B/32B) on AMD Instinct MI355X (gfx950). The training learns two rotations: R₂ is merged offline into v_proj output / o_proj input weights and incurs zero inference cost. R₁ is applied online at every layer to inputs of qkv_proj and gate_up_proj. R₁ can also be fused offline (the QuaRot/SpinQuant approach of absorbing R₁ into the preceding RMSNorm), but the previous post showed online R₁ yields measurably better MXFP4 accuracy than the offline-fused variant—so online R₁ is our design choice for accuracy. Together, the trained R₁ + R₂ recover most of the perplexity gap between RTN (no rotation training) and the bf16 baseline. For example, Qwen3-8B Wikitext PPL (Quark token-level methodology, same as the previous post) improves from RTN = 10.72 to Sep = 10.17, against a BF16 reference of 9.72.

This online R₁ at every layer per decode step is the cost our fused kernel eliminates. One important structural property: R₁ is block-diagonal, not a full K×K matrix. For a hidden dimension of K=4096 with RS=128, R₁ is applied as 32 independent 128×128 Hadamard rotations to non-overlapping channel blocks, as illustrated in our previous post. This means per-rotation FLOPs scale as K×RS rather than K², and halving RS (from 128 to 64) halves the per-rotation compute cost while preserving the full accuracy benefit. The figure below shows this block-diagonal structure for K=4096 with RS=128.

The separated implementation achieves this accuracy improvement but comes at a cost. For Dense models (8B/14B/32B), each quantization layer goes from 1 kernel (RTN: just MXFP4 quant) to 2 kernels (separated: rotation matmul + MXFP4 quant), with one intermediate bf16 global-memory round-trip in between. For MoE models (30B-A3B), where expert sort is already a separate kernel, the pipeline grows from 2 kernels (RTN: quant + sort) to 3 kernels (separated: rotation + quant + sort), with two intermediate GMEM round-trips. Every decode step now carries this overhead across all layers.

This post is about closing that performance gap—and the result is what makes online rotation production-ready. We implemented a Gluon (Triton) fused kernel on AMD CDNA4 that performs all operations in a single launch, keeping intermediate values in registers. The result:

✅ Further reduction with RS=64: Dense model overhead drops to just +0.3–0.8% vs RTN ✅ Further reduction with RS=64: Dense model overhead drops to just +0.3–0.8% vs RTN

The bottom line: fusion turns online rotation from a feature you pay a steep latency tax for into one that is effectively free. The separated implementation's +5–10% TPOT overhead—high enough to discourage enabling rotation in production—drops to a tolerable +0.3–0.8% on Dense models (and is nearly halved on MoE at RS=128, from +10% to +5.6%). With that overhead reduced to near-zero, accuracy-preserving MXFP4 online rotation becomes practical to ship by default on AMD Instinct MI355X.

As the figure above shows, adding online rotation means an extra kernel launch and one bf16 GMEM round-trip per layer per decode step.

For MoE models, expert sort is already a separate kernel even in the RTN path. Adding online rotation grows the pipeline from 2 kernels to 3:

As the figure above shows, adding online rotation means an extra kernel launch plus an additional bf16 GMEM round-trip per MoE layer per decode step.

We implemented the fused kernel using Gluon (Triton): the f32 accumulator values produced by the rotation matmul are fed directly into the FP4 quantization instruction in registers, eliminating the intermediate bf16 GMEM round-trip. For MoE models, the expert sort scatter is also fused into the same kernel launch. The pipeline comparison across all three variants (RTN / Separated / Fused) is shown above in the Dense and MoE pipeline figures.

1. Shared memory layout for rotation matrix tiles The rotation matrix R (shape [K, K]) is tiled into shared memory (LDS) so each wave reads its tile once and reuses it across MFMA iterations. We use ds_read_tr16_b64—a CDNA4 hardware instruction that reads and transposes a 16-element tile from LDS in a single operation—to feed the MFMA instruction in the correct layout without an explicit transpose step.

2. MFMA tile selection v_mfma_f32_16x16x32_bf16 was chosen over larger tiles (32×32×16) because the rotation matmul at M=1 is severely compute-underutilized; the 16×16 tile allows more waves to run concurrently, improving occupancy. The larger tile looked promising on paper but ran into lane mapping complexity that GEAK explored and did not pursue further.

3. In-register FP4 quantization The f32 accumulators from MFMA are fed directly into v_cvt_scalef32_pk_fp4_f32 without any register spill. We pre-compute the per-group scale in a reduction across the wave before the conversion, keeping everything in VGPRs.

4. Gluon (Triton) auto-tuning Gluon uses Triton JIT compilation to automatically select well-suited tile configurations across different batch sizes, and integrates cleanly with CUDAGraph's fixed schedule replay.

Before writing the fused kernel, we profiled the separated implementation using rocprofv3 to understand where time was actually going. The breakdown per MoE layer at M=1 (single-token decode) is shown in the figure below:

The rotation matmul dominates (56%), but it is unavoidable—it is the computation that delivers the accuracy benefit.

The rotation matmul dominates (56%), but it is unavoidable—it is the computation that delivers the accuracy benefit.

The GMEM writes between kernels (18% combined: bf16 write 10% + uint8 write 8%) are pure waste—they exist only because the three operations are in separate kernels that cannot share register state. The corresponding read costs are hidden in the subsequent kernels' launch and execution time, so the true waste exceeds 18%.

The GMEM writes between kernels (18% combined: bf16 write 10% + uint8 write 8%) are pure waste—they exist only because the three operations are in separate kernels that cannot share register state. The corresponding read costs

Text integral sub licența deschisă a publicației · 1209 cuvinte

Subiecte conexe