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(__init__[__mlir_type.!kgen.string]("H100"), Vendor(__init__[__mlir_type.!pop.int_literal](2)), __init__[__mlir_type.!kgen.string]("cuda"), __init__[__mlir_type.!kgen.string]("hopper"), __init__[__mlir_type.!kgen.string]("nvptx-short-ptr=true"), __init__[__mlir_type.!pop.float_literal](9), __init__[__mlir_type.!kgen.string]("sm_90a"), 132, 32, 2048, 32, 64, 2048, 32, 233472, 65536, 256, __init__[__mlir_type.!kgen.string]("warp"), 255, 65536, 32, 128, 4, 1024)), schedule: MHASchedule = MHASchedule(__init__[__mlir_type.!pop.int_literal](0))]
Aliases
may_advance = True:mha_schedule = schedule:
Implemented traits
AnyType,
Copyable,
ExplicitlyCopyable,
MHATileScheduler,
Movable,
UnknownDestructibility
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(__init__[__mlir_type.!pop.int_literal](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?
Thank you! We'll create more content like this.
Thank you for helping us improve!