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

mha_tile_scheduler

Provides tile schedulers for multi-head attention kernels on NVIDIA GPUs.

Defines the MHATileScheduler trait and concrete schedulers (TileScheduler, TransientScheduler, QueuedTileScheduler) that map work tiles to thread blocks, along with supporting state and summary types used by the persistent attention kernel.

Structs

• ​MHASchedule: Enumerates the scheduling strategy for mapping work tiles to thread blocks.
• ​MHASchedulerSynchronization: Enumerates synchronization modes for advancing the MHA scheduler.
• ​MHATileState: Tracks the mutable per-CTA state of the tile scheduler during kernel execution.
• ​MHATileSummary: Summarizes the dimensions and valid-length metadata of the MHA work grid.
• ​QueuedTileScheduler: If decoding == False, then num_heads is q_num_heads. If decoding == True, then num_heads is kv_num_heads.
• ​SeqInfo: Describes a sequence's length and starting offset for a work tile.
• ​TileScheduler: Implements a persistent tile scheduler that cycles CTAs through work tiles.
• ​TransientScheduler: Implements a non-persistent (transient) tile scheduler for the MHA kernel.
• ​WorkInfo: Holds the coordinates and validity of a single work tile.

Traits

• ​MHATileScheduler: Describes a schedule for the persistent MHA kernel.