AMD MI450 Attention Decode: Gluon Kernel Hits 85% HBM Bandwidth

AMD's MI450 and the Memory Wall

Agentic AI pushes LLM inference to a million-token contexts. During decode, each new token attends to all previous tokens in the KV cache. The kernel repeatedly reads past states from HBM. The bottleneck shifts from compute to memory.

AMD's MI450 series targets this regime. It doubles VGPRs per SIMD (1024 vs 512), doubles LDS per WGP (320 KB vs 160 KB), and boosts HBM bandwidth to 19.6 TB/s (from 8.1 TB/s). It also adds two new hardware features: TDM and workgroup clusters.

TDM and Workgroup Clusters

TDM (Tensor Data Movement) is a specialized unit that moves structured tensor data between global memory and LDS. Instead of many small vector loads, the kernel describes the tensor with address, shape, strides, and layout, then issues one bulk transfer asynchronously. MI450 has two TDM units per WGP, compared to one memory unit on MI350.

Workgroup clusters let multiple workgroups coordinate via hardware barriers. They can use multicast loads to share data across WGPs. This avoids separate kernel launches or global-memory synchronization.

Gluon: A Lower-Level Triton

Gluon is a Triton-based DSL that requires explicit tensor layouts. The layout describes how elements distribute across registers, lanes, waves, and workgroups. It gives kernel experts control over generated instructions.

The blog assumes familiarity with Gluon basics. For examples, see the MI450 Gluon Examples.

Attention Decode with MQA

Attention decode processes one query token against many KV tokens. Grouped-query attention (GQA) groups Q heads that share K/V heads. The kernel maps each (batch, KV head) pair to a workgroup. Within a workgroup, multiple Q heads are processed together.

The kernel follows Flash Attention: iterate over KV tiles, compute QK, softmax, then PV. The workload is memory-bound.

Optimization 1: Tensor Layout

Gluon's AMDWMMALayout describes WMMA output layout. For example, instruction shape [16,16,128] with wmma_scaled generates v_wmma_scale_f32_16x16x128_f8f6f4. It consumes two FP8 operands and produces an FP32 tile.

Choose the number of waves carefully. Since softmax reduces along T_kv, distribute waves along the grouped Q dimension to avoid cross-wave communication. For H_q/H_kv=32, use 2 waves.

Layout conversion can be expensive. Match the QK output layout to the PV input layout. Use transpose=True in AMDWMMALayout to transpose the output without extra instructions. Also set the PV input layout's K Width to 8 to match the transposed QK output.

Optimization 2: Data Loading with TDM

TDM can move data directly to LDS or through the cache. For memory-bound workloads, bypass the cache. The direct path requires the innermost dimension to be at least 128 bytes; 256 bytes is recommended.

For FP8 with head dimension 128, reshape K and V to increase the innermost dimension before loading. Then reshape back when moving from LDS to registers.

Optimization 3: Pipelining

The hot loop can stall on memory latency. Pipeline the loop to overlap TDM loads with compute. The blog describes a multi-stage pipeline (though the excerpt truncates the details).

Optimization 4: Split-k Parallelization

Split-k parallelizes over the KV dimension. Multiple workgroups process different KV ranges, then combine partial results. This increases parallelism and can improve bandwidth utilization.

Results and Next Steps

The optimized Gluon kernel achieves 85% of peak HBM bandwidth on MI450. This is an early result.

To apply these techniques: start with the Gluon examples, experiment with layouts, and profile TDM usage. The MI450 is a significant step for memory-bound AI workloads.