For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).
Mojo struct
TileScheduler
struct TileScheduler[problem_shape: IndexList[Int(3)], tile_shape: IndexList[Int(3)], grid_shape: IndexList[Int(2)], cluster: IndexList[Int(3)] = Index[Int, Int, Int](Int(1), Int(1), Int(1)), raster_dim: UInt32 = UInt32(1), schedule: MatmulSchedule = MatmulSchedule.TILE2D]
Persistent tile scheduler for GPU matmul output tiles.
Maps thread blocks to (m, n) output tile coordinates using 1D, 2D-wave,
or DeepSeek scheduling strategies. Supports block clusters with multicast
and advances through tiles in a persistent kernel loop via fetch_next_work.
Parameters
- problem_shape (
IndexList[Int(3)]): Static (M, N, K) dimensions of the matmul problem. - tile_shape (
IndexList[Int(3)]): Static (BM, BN, BK) tile sizes. - grid_shape (
IndexList[Int(2)]): (grid_x, grid_y) number of thread blocks in the launch grid. - cluster (
IndexList[Int(3)]): Block cluster shape (default 1×1×1 for no clustering). - raster_dim (
UInt32): Axis along which to rasterize (0 = M, 1 = N). - schedule (
MatmulSchedule): Tile scheduling strategy to use.
Fields
- idx (
UInt32): - prob_shape (
IndexList[Int(3)]): - num_waves_m (
UInt32): - num_waves_n (
UInt32): - log_num_waves_n (
FastDiv[DType.uint32]): - current_iter (
Int): - num_aligned_m_blocks (
UInt32): - num_blocks (
UInt32):
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
ImplicitlyDeletable,
Movable,
RegisterPassable,
TrivialRegisterPassable
comptime members
kNum1DBlocksPerGroup
comptime kNum1DBlocksPerGroup = UInt32(16)
kNumNBlocks
comptime kNumNBlocks = SIMD(ceildiv(problem_shape[Int(1)], tile_shape[Int(1)]))
num_grids
comptime num_grids = SIMD(Int((mul grid_shape[Int(0)], grid_shape[Int(1)])))
wave_shape
comptime wave_shape = Index[Int, Int, dtype=DType.uint32](Int((mul grid_shape[Int(1)], tile_shape[Int(0)])), Int((mul grid_shape[Int(0)], tile_shape[Int(1)])))
Methods
__init__
def __init__(prob_shape: IndexList[Int(3)]) -> Self
get_current_work_info
def get_current_work_info(mut self) -> WorkInfo
Returns:
WorkInfo
advance
def advance(mut self)
fetch_next_work
def fetch_next_work(mut self) -> WorkInfo
Returns:
WorkInfo
num_output_tiles
fetch_next_work_ds
def fetch_next_work_ds(mut self) -> WorkInfo
Returns:
WorkInfo
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!