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 struct

TransientScheduler

struct TransientScheduler[tile_shape: UInt32, num_heads: UInt32, flip_prompt_idx: Bool, pair_cta: Bool = False, splitk_partitions: UInt32 = UInt32(1)]

Implements a non-persistent (transient) tile scheduler for the MHA kernel.

Each CTA processes exactly one work tile identified by its block_idx, with optional prompt-index flipping, pair-CTA clustering, or split-K partitioning to widen the launch grid.

Parameters

• ​tile_shape (UInt32): Size of each query tile along the sequence dimension, in tokens.
• ​num_heads (UInt32): Number of attention heads mapped across the grid's y dimension.
• ​flip_prompt_idx (Bool): Whether to reverse the prompt-tile traversal direction so that the last tile is processed first.
• ​pair_cta (Bool): Whether to cluster two CTAs per query tile for a 2-SM launch width (defaults to False).
• ​splitk_partitions (UInt32): Number of split-K partitions used to widen the launch grid along x (defaults to 1).

Implemented traits

AnyType, Copyable, Defaultable, Deinitable, DevicePassable, ImplicitlyCopyable, MHATileScheduler, Movable, RegisterPassable, TrivialRegisterPassable

comptime members

cluster_size

comptime cluster_size = (UInt32(2) if pair_cta else UInt32(1) * splitk_partitions)

device_type

comptime device_type = TransientScheduler[tile_shape, num_heads, flip_prompt_idx, pair_cta, splitk_partitions]

may_advance

comptime may_advance = False

mha_schedule

comptime mha_schedule = MHASchedule.DEFAULT

Methods

__init__

def __init__() -> Self

get_type_name

static def get_type_name() -> String

Returns:

String

get_current_work_info

def get_current_work_info(self, num_prompt_tiles: UInt32) -> WorkInfo

Returns:

WorkInfo

def get_current_work_info[ValidLengthType: OptionalPointer, //](self, ts: MHATileSummary[ValidLengthType], state: MHATileState) -> WorkInfo

Returns:

WorkInfo

advance

def advance[ValidLengthType: OptionalPointer, //, producer: Bool, sync: MHASchedulerSynchronization = MHASchedulerSynchronization.DEFAULT](self, ts: MHATileSummary[ValidLengthType], mut state: MHATileState, pipeline_idx: UInt32) -> OptionalReg[SeqInfo]

Returns:

OptionalReg[SeqInfo]

grid_dim

static def grid_dim(batch_size: UInt32, max_num_prompt_tiles: UInt32) -> Tuple[Int, Int, Int]

Returns:

Tuple[Int, Int, Int]

static def grid_dim(batch_size: UInt32, max_num_prompt_tiles: UInt32, num_partitions: UInt32) -> Tuple[Int, Int, Int]

Returns:

Tuple[Int, Int, Int]

initial_state

def initial_state[ValidLengthType: OptionalPointer, //](self, ptr: Pointer[UInt32, MutAnyOrigin, address_space=AddressSpace.SHARED], tile_summary: MHATileSummary[ValidLengthType]) -> MHATileState

Returns:

MHATileState

unsafe_seq_info

def unsafe_seq_info[ValidLengthType: OptionalPointer, //](self, ts: MHATileSummary[ValidLengthType], state: MHATileState) -> SeqInfo

Returns:

SeqInfo