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

• ​Softmax:

Was this page helpful?

Thank you! We'll create more content like this.

Thank you for helping us improve!

😔 What went wrong?

Some code doesn’t work

It includes inaccurate information

It's missing information I need

It was difficult to understand

Other

Submit