Mojo module
attention_rdna
RDNA-specific Attention struct for Wave32 WMMA operations.
This module provides an Attention implementation optimized for AMD RDNA consumer GPUs (Radeon RX 7000/8000 series) using Wave32 WMMA instructions.
Key differences from CDNA Attention:
- Wave size: 32 lanes (vs 64 for CDNA)
- MMA shape: 16x16x16 only (vs multiple shapes for CDNA)
- Fragment sizes: A/B = 16 elements, C/D = 8 elements per lane
- k_group_size = 1 (single MMA per K iteration)
comptime values
RDNA_K_GROUP_SIZE
comptime RDNA_K_GROUP_SIZE = 1
Structs
-
AttentionRDNA: RDNA-specific Attention implementation for Wave32 WMMA.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!