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 module

attention

SM90 (Hopper) multi-head attention producer-side helpers.

Re-exports shared NVIDIA attention primitives from common and provides the SM90-specific implementations of produce (TMA-based Q/K/V tile staging), _apply_mask, _get_position, get_q_head_idx, output_reg_to_smem, and _optional_lt_to_tt consumed by the sm90/mha.mojo kernel.

Functions

• ​get_q_head_idx: Computes the query head indices owned by the current thread.
• ​output_reg_to_smem: Stores the output accumulator registers from local memory into shared memory.
• ​produce: Stages Q, K, and V tiles into shared memory via TMA for the SM90 attention kernel.