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 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'sydimension. - 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 toFalse). - splitk_partitions (
UInt32): Number of split-K partitions used to widen the launch grid alongx(defaults to1).
Implemented traits
AnyType,
Copyable,
Defaultable,
DevicePassable,
ImplicitlyCopyable,
ImplicitlyDeletable,
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
get_current_work_info
def get_current_work_info[ValidLengthType: OptionalPointer, //](self, ts: MHATileSummary[ValidLengthType], state: MHATileState) -> WorkInfo
Returns:
advance
def advance[ValidLengthType: OptionalPointer, //, producer: Bool, sync: MHASchedulerSynchronization = MHASchedulerSynchronization.DEFAULT](self, ts: MHATileSummary[ValidLengthType], mut state: MHATileState, pipeline_idx: UInt32) -> OptionalReg[SeqInfo]
Returns:
grid_dim
static def grid_dim(batch_size: UInt32, max_num_prompt_tiles: UInt32) -> Tuple[Int, Int, Int]
Returns:
static def grid_dim(batch_size: UInt32, max_num_prompt_tiles: UInt32, num_partitions: UInt32) -> Tuple[Int, Int, Int]
Returns:
initial_state
def initial_state[ValidLengthType: OptionalPointer, //](self, ptr: Pointer[UInt32, MutAnyOrigin, address_space=AddressSpace.SHARED, _safe=False], tile_summary: MHATileSummary[ValidLengthType]) -> MHATileState
Returns:
unsafe_seq_info
def unsafe_seq_info[ValidLengthType: OptionalPointer, //](self, ts: MHATileSummary[ValidLengthType], state: MHATileState) -> SeqInfo
Returns:
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!