AMD ROCm team details GLM-5.2 inference optimizations for MI355X GPUs

AMD's ROCm engineering team published a detailed technical breakdown of optimizations that accelerate GLM-5.2 FP8 inference on eight AMD Instinct MI355X GPUs. The work addresses bottlenecks in long-context Mixture-of-Experts models by co-optimizing parallelization strategy, KV-cache layout, collective communication, kernels, and the SGLang serving framework rather than tuning individual components in isolation.

The team used EvalScope benchmark workloads with 90% shared-prefix configurations at 64K and 128K lengths, plus an extended 1M-token evaluation, to better reflect production serving conditions. Key optimizations include prefill context parallelism (CP8), AITER sparse-attention kernels, PyHIP MoE tuning with reduced communication payloads, fused FP8 query quantization, GEMM tuning tables, and EAGLE multi-token prediction speculative decoding. Throughput improved from 163K to 629K tok/s at 64K context, from 137K to 553K tok/s at 128K, and from 22K to 157K tok/s at 1M tokens with concurrency 1. GSM8K validation confirmed preserved model accuracy across the optimized stack.

End-to-end inference stack optimizations

• Prefill context parallelism (CP8): Partitions query tokens rather than attention heads across ranks, eliminating redundant KV loads and head replication while enabling direct use of MI355X's native gqa:64 assembly-optimized MLA kernel; required SGLang extensions for CP metadata, KV reconstruction, tensor layouts, and numerical edge cases.
• FP8 KV cache and fused quantization: Uses fp8_e4m3 format to reduce memory footprint and enable FP8 attention kernels; replaces multi-operator dynamic quantization with a single fused Triton kernel using a persistent device buffer, cutting kernel launches, temporary allocations, and memory traffic.
• PyHIP MoE communication reduction: Reorders router/top-k and activation quantization before the collective gather, transmitting FP8 hidden states with scales and routing results instead of BF16, reducing per-token communication from 12,288 bytes to 6,408 bytes; falls back to standard path for unsupported shapes.
• AITER kernel integration: Provides persistent MLA implementation for 64-head FP8 execution, block-scale FP8 kernels, and offline-tuned GEMM configurations stored in runtime lookup tables for both A8W8 and BF16 paths.
• EAGLE speculative decoding: Configured with 4 steps, top-k = 1, and 5 draft tokens to address memory-bound decode phases, complementing the compute-focused prefill optimizations.

Why it matters

This work demonstrates that competitive large-model inference on AMD hardware requires open-source stack co-optimization across SGLang, AITER, PyHIP, and custom collectives rather than proprietary black-box solutions. The detailed methodology and integration with open serving frameworks provides a reproducible path for developers deploying MoE architectures on ROCm, though the optimizations remain tightly coupled to MI355X-specific assembly kernels and AMD's ecosystem rather than portable GPU code.