IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /max/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. /max/get-started.md).

Mojo module

softmax

Online softmax for gfx950 attention kernels.

Score, warp_scratch, and output are TileTensors. Warp + fragment geometry is expressed via TileLayout / Coord: WarpLayoutT is the col-major (warp_rows, warp_cols) lane layout, and FragmentLayoutT describes the per-lane MMA fragment shape. All lane decomposition / stride queries go through Layout methods rather than hand-rolled integer arithmetic.

Structs