Mojo trait
MHATileScheduler
Implemented traits
AnyType,
UnknownDestructibility
Methods
get_current_work_info
get_current_work_info(self: _Self, ts: MHATileSummary, state: MHATileState) -> WorkInfo
Returns the current WorkInfo.
advance
advance[ragged: Bool, producer: Bool, sync: MHASchedulerSynchronization = MHASchedulerSynchronization(__init__[__mlir_type.!pop.int_literal](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.
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.
initial_state
initial_state(self: _Self, ptr: UnsafePointer[SIMD[uint32, 1], address_space=AddressSpace(3)], tile_summary: MHATileSummary) -> MHATileState
Create the initial state object.
unsafe_seq_info
unsafe_seq_info[ragged: Bool](self: _Self, ts: MHATileSummary, state: MHATileState) -> SeqInfo
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!