Mojo trait
MHATileScheduler
Implemented traits
AnyType,
UnknownDestructibility
Aliases
may_advance
alias may_advance
mha_schedule
alias mha_schedule
The MHATileScheduler trait describes a schedule for the persistent kernel.
Methods
get_current_work_info
get_current_work_info(self: _Self, ts: MHATileSummary, state: MHATileState) -> WorkInfo
Returns the current WorkInfo.
Returns:
advance
advance[ragged: Bool, producer: Bool, sync: MHASchedulerSynchronization = MHASchedulerSynchronization(1)](self: _Self, ts: MHATileSummary, mut state: MHATileState, pipeline_idx: SIMD[uint32, 1]) -> OptionalReg[SeqInfo]
Advance state to the next work item. func must return a Bool indicating whether there is more work. Returns True if there is more work.
Returns:
grid_dim
static grid_dim(batch_size: SIMD[uint32, 1], max_num_prompt_tiles: SIMD[uint32, 1]) -> Tuple[Int, Int, Int]
Return the grid_dim required for the kernel.
Returns:
initial_state
initial_state(self: _Self, ptr: UnsafePointer[SIMD[uint32, 1], address_space=AddressSpace(3)], tile_summary: MHATileSummary) -> MHATileState
Create the initial state object.
Returns:
unsafe_seq_info
unsafe_seq_info[ragged: Bool](self: _Self, ts: MHATileSummary, state: MHATileState) -> SeqInfo
Returns:
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!