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
SplitKTileScheduler
struct SplitKTileScheduler[locks_origin: MutOrigin, //, problem_shape_nk: IndexList[Int(2)], tile_shape: IndexList[Int(3)], splits: UInt32, num_consumer: UInt32, num_pipeline_stages: UInt32, cluster_shape: IndexList[Int(2)], raster_order: RasterOrder, reduction_mode: ReductionMode = ReductionMode.Deterministic]
Tile scheduler for split-K GPU matmul reductions on SM90.
Partitions the K dimension into splits independent chunks, each processed
by a separate thread block. Partial results are accumulated into a workspace
buffer and reduced using semaphore-based ordering (Deterministic) or
atomic additions (Nondeterministic). Cluster-aware rasterization keeps
spatially close tiles on the same set of SMs for L2 reuse.
Parameters
- locks_origin (
MutOrigin): Memory origin for the synchronization lock buffer. - problem_shape_nk (
IndexList[Int(2)]): Static (N, K) dimensions of the matmul problem. - tile_shape (
IndexList[Int(3)]): Static (BM, BN, BK) tile sizes. - splits (
UInt32): Number of K-dimension splits. - num_consumer (
UInt32): Number of warp groups consuming each output tile. - num_pipeline_stages (
UInt32): Number of software-pipeline stages. - cluster_shape (
IndexList[Int(2)]): (cluster_m, cluster_n) block cluster dimensions. - raster_order (
RasterOrder): Traversal direction for the output tile grid. - reduction_mode (
ReductionMode): Whether reduction is deterministic or non-deterministic.
Fields
- prob_shape (
IndexList[Int(3)]): - block_id_in_cluster (
IndexList[Int(2)]): - blocks_per_problem (
UInt32): - current_work_linear_idx (
UInt32): - log_cluster_shape_major (
UInt32): - log_cluster_shape_minor (
UInt32): - cluster_blk_major (
UInt32): - locks_ptr (
Pointer[Int32, locks_origin, _safe=False]):
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
ImplicitlyDeletable,
Movable,
RegisterPassable,
TrivialRegisterPassable
comptime members
k_tiles_per_output_tile
comptime k_tiles_per_output_tile = SIMD(ceildiv(problem_shape_nk[Int(1)], tile_shape[Int(2)]))
k_tiles_per_split
comptime k_tiles_per_split = (SIMD(ceildiv(problem_shape_nk[Int(1)], tile_shape[Int(2)])) // splits)
log_cluster_size
comptime log_cluster_size = log2_floor(Int((mul cluster_shape[Int(0)], cluster_shape[Int(1)])))
WorkTileType
comptime WorkTileType[dtype: DType, layout: Layout] = LayoutTensor[dtype, layout, MutAnyOrigin]
Parameters
Methods
__init__
def __init__(prob_shape: IndexList[Int(3)], block_id_in_cluster: IndexList[Int(2)], locks_ptr: Pointer[UInt8, locks_origin, _safe=False]) -> Self
get_sm_num
get_problem_blocks_shape
static def get_problem_blocks_shape(problem_shape: IndexList[Int(3)], dyn_tile_shape: IndexList[Int(3)], dyn_cluster_shape: IndexList[Int(2)]) -> IndexList[Int(2)]
Returns:
initial_work_tile_info
get_current_work_info
get_worktile_m_n_idx
def get_worktile_m_n_idx(mut self, mut work_tile_info: WorkInfo, linear_tile_id: UInt32)
assign_work
def assign_work(mut self, mut work_tile_info: WorkInfo, linear_idx: UInt32)
get_k_start_and_linear_tile_id
def get_k_start_and_linear_tile_id(mut self, mut work_tile_info: WorkInfo, linear_idx: UInt32) -> UInt32
Returns:
fetch_next_work
requires_reduction
advance_to_next_work
def advance_to_next_work(mut self)
is_last_split
get_grid_shape
static def get_grid_shape(dyn_cluster_shape: IndexList[Int(3)], dyn_raster_order: RasterOrder = RasterOrder.AlongN) -> IndexList[Int(3)]
Returns:
get_num_tiles
static def get_num_tiles(problem_shape: IndexList[Int(3)], dyn_tile_shape: IndexList[Int(3)], dyn_cluster_shape: IndexList[Int(2)]) -> Int
Returns:
get_required_locks_buffer_size_bytes
static def get_required_locks_buffer_size_bytes[accum_type: DType, dyn_num_consumer: UInt32](problem_shape: IndexList[Int(3)], dyn_tile_shape: IndexList[Int(3)], dyn_cluster_shape: IndexList[Int(2)]) -> Int
Returns:
get_linear_idx_from_m_and_n
output_tile_index
reduction
def reduction[accum_type: DType, c_reg_layout: Layout, workspace_layout: TensorLayout](self, reduction_workspace: TileTensor[accum_type, workspace_layout], c_reg_tile: LayoutTensor[accum_type, c_reg_layout, MutAnyOrigin, address_space=AddressSpace.LOCAL], work_tile_info: WorkInfo, num_barriers: UInt32, warp_group_local_idx: UInt32)
def reduction[accum_type: DType, c_reg_layout: Layout, workspace_layout: Layout](self, reduction_workspace: LayoutTensor[accum_type, workspace_layout, MutAnyOrigin], c_reg_tile: LayoutTensor[accum_type, c_reg_layout, MutAnyOrigin, address_space=AddressSpace.LOCAL], work_tile_info: WorkInfo, num_barriers: UInt32, warp_group_local_idx: UInt32)
wait_eq
static def wait_eq(lock_ptr: Pointer[Int32, _safe=False], barrier_id: Int32, barrier_group_thread_idx: Int, lock_idx: UInt32, val: UInt32)
wait_lt
static def wait_lt(lock_ptr: Pointer[Int32, _safe=False], barrier_id: Int32, barrier_group_thread_idx: Int, lock_idx: UInt32, count: UInt32)
arrive_set
static def arrive_set(lock_ptr: Pointer[Int32, _safe=False], barrier_id: Int32, barrier_group_thread_idx: Int, lock_idx: UInt32, increment: UInt32)
store_accumulator
def store_accumulator[accum_type: DType, c_reg_layout: Layout, workspace_layout: Layout](self, reduction_workspace: LayoutTensor[accum_type, workspace_layout, MutAnyOrigin], c_reg_tile: LayoutTensor[accum_type, c_reg_layout, MutAnyOrigin, address_space=AddressSpace.LOCAL], reduction_tile_idx: UInt32, warp_group_local_idx: UInt32, warp_group_thread_idx: UInt32)
reduce_add
def reduce_add[accum_type: DType, c_reg_layout: Layout, workspace_layout: Layout, //, *, write_back: Bool](self, reduction_workspace: LayoutTensor[accum_type, workspace_layout, MutAnyOrigin], c_reg_tile: LayoutTensor[accum_type, c_reg_layout, MutAnyOrigin, address_space=AddressSpace.LOCAL], reduction_tile_idx: UInt32, warp_group_local_idx: UInt32, warp_group_thread_idx: UInt32)
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!