Mojo struct
QueuedTileScheduler
@register_passable(trivial)
struct QueuedTileScheduler[tile_shape: SIMD[uint32, 1], num_heads: SIMD[uint32, 1], /, decoding: Bool, 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)]
If decoding == False
, then num_heads
is q_num_heads
. If decoding == True
, then num_heads
is kv_num_heads
.
Fields
- gidx_ptr (
UnsafePointer[SIMD[uint32, 1], address_space=AddressSpace(1)]
):
Implemented traits
AnyType
,
Copyable
,
MHATileScheduler
,
Movable
,
UnknownDestructibility
Aliases
may_advance
alias may_advance = True
mha_schedule
alias mha_schedule = schedule
Methods
__init__
__init__(gidx_ptr: UnsafePointer[SIMD[uint32, 1]]) -> Self
get_current_work_info
get_current_work_info(self, ts: MHATileSummary, 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]
The parameter func
must return a Bool
indicating whether the WorkInfo
arg is valid. This function returns whether the current idx corresponds to a valid WorkInfo
. Note that if MHASchedulerSynchronization
is NONE
, then we assume it is only called by thread_idx.x==0
.
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!