Skip to main content

Mojo struct

TileScheduler

@register_passable(trivial) struct TileScheduler[tile_shape: SIMD[uint32, 1], num_heads: SIMD[uint32, 1], /, num_ctas: SIMD[uint32, 1] = SIMD(Info("H100", Vendor(2), "cuda", "hopper", "nvptx-short-ptr=true", 9, "sm_90a", 132, 32, 2048, 32, 64, 2048, 32, 233472, 65536, 256, "warp", 255, 65536, 32, 128, 4, 1024).sm_count), schedule: MHASchedule = MHASchedule(0)]

Implemented traits

AnyType, Copyable, Defaultable, MHATileScheduler, Movable, UnknownDestructibility

Aliases

may_advance

alias may_advance = True

mha_schedule

alias mha_schedule = schedule

Methods

__init__

__init__() -> Self

get_current_work_info

get_current_work_info(self, ts: MHATileSummary, state: MHATileState) -> WorkInfo

fetch_next_work

fetch_next_work(self, ts: MHATileSummary, mut state: MHATileState) -> WorkInfo

advance

advance[ragged: Bool, producer: Bool, sync: MHASchedulerSynchronization = MHASchedulerSynchronization(1)](self, ts: MHATileSummary, mut state: MHATileState, pipeline_idx: SIMD[uint32, 1]) -> OptionalReg[SeqInfo]

grid_dim

static grid_dim(batch_size: SIMD[uint32, 1], max_num_prompt_tiles: SIMD[uint32, 1]) -> Tuple[Int, Int, Int]

initial_state

initial_state(self, ptr: UnsafePointer[SIMD[uint32, 1], address_space=AddressSpace(3)], tile_summary: MHATileSummary) -> MHATileState

unsafe_seq_info

unsafe_seq_info[ragged: Bool](self, ts: MHATileSummary, state: MHATileState) -> SeqInfo

Was this page helpful?