IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /get-started.md). For the complete documentation index, see llms.txt.

Skip to main content

For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo package

amd_rdna

TileTensor-native attention kernels for AMD RDNA3+ (gfx11xx/gfx12xx).

Wave32 with 16x16x16 WMMA. 16-element A/B fragments per lane (full K), 8-element C/D fragments per lane. Supports MHA prefill and decode.

Modules

attention: Attention struct for RDNA Wave32 MHA kernels (prefill + decode).
buffers: K, V, Q, P, and Output buffers for RDNA Wave32 attention kernels.
config: RDNA Wave32 attention config.
mha_decode: RDNA Wave32 MHA decode kernel.
mha_prefill: RDNA Wave32 MHA prefill kernel.
mma: RDNA Wave32 WMMA helper for attention kernels.
softmax: Online softmax for RDNA Wave32 attention kernels.
utils: Shared helpers for RDNA Wave32 attention kernels.

Modules