Mojo struct
TileScheduler
@register_passable(trivial)
struct TileScheduler[tile_shape: UInt32, num_heads: UInt32, /, num_ctas: UInt32 = GPUInfo("H100", Vendor(2), "cuda", "hopper", 9, "sm_90a", 132, 32, 2048, 233472, 65536, 1024).sm_count, schedule: MHASchedule = MHASchedule(0)]
Implemented traits
AnyType,
Copyable,
Defaultable,
ImplicitlyCopyable,
MHATileScheduler,
Movable,
UnknownDestructibility
Aliases
__copyinit__is_trivial
alias __copyinit__is_trivial = True
__del__is_trivial
alias __del__is_trivial = True
__moveinit__is_trivial
alias __moveinit__is_trivial = True
may_advance
alias may_advance = True
mha_schedule
alias mha_schedule = schedule
Methods
__init__
__init__() -> Self
get_current_work_info
get_current_work_info[ValidLengthType: OptionalPointer, //](self, ts: MHATileSummary[ValidLengthType], state: MHATileState) -> WorkInfo
Returns:
fetch_next_work
fetch_next_work(self, ts: MHATileSummary[ValidLengthType], mut state: MHATileState) -> WorkInfo
Returns:
advance
advance[ValidLengthType: OptionalPointer, //, producer: Bool, sync: MHASchedulerSynchronization = MHASchedulerSynchronization(1)](self, ts: MHATileSummary[ValidLengthType], mut state: MHATileState, pipeline_idx: UInt32) -> OptionalReg[SeqInfo]
Returns:
grid_dim
static grid_dim(batch_size: UInt32, max_num_prompt_tiles: UInt32) -> Tuple[Int, Int, Int]
Returns:
initial_state
initial_state[ValidLengthType: OptionalPointer, //](self, ptr: UnsafePointer[UInt32, address_space=AddressSpace(3)], tile_summary: MHATileSummary[ValidLengthType]) -> MHATileState
Returns:
unsafe_seq_info
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!