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[group_offsets_origin: ImmutOrigin, offsets_layout: Layout, //, *, static_MN: Int, tile_shape: IndexList[3], cluster: IndexList[3] = Index[Int, Int, Int](1, 1, 1), cta_group: Int = 1, swizzle: Bool = False, swapAB: Bool = True]
Fields
- num_active_experts (
Int): - group_offsets (
LayoutTensor[DType.uint32, offsets_layout, group_offsets_origin]): - current_iter (
Int32): - current_group_idx (
UInt32): - current_dynamic_dim_cumsum (
UInt32): - block_idx_start (
UInt32):
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
ImplicitlyDeletable,
Movable,
RegisterPassable,
TrivialRegisterPassable
comptime members
cta_group_tile_shape
comptime cta_group_tile_shape = Index[Int, Int]((tile_shape[0] * cta_group), (tile_shape[1] * cta_group))
div_dynamic_block
comptime div_dynamic_block = FastDiv(TileScheduler[static_MN=static_MN, tile_shape=tile_shape, cluster=cluster, cta_group=cta_group, swizzle=swizzle, swapAB=swapAB].cta_group_tile_shape[TileScheduler[static_MN=static_MN, tile_shape=tile_shape, cluster=cluster, cta_group=cta_group, swizzle=swizzle, swapAB=swapAB].dynamic_dim])
dynamic_dim
comptime dynamic_dim = 1 if swapAB else 0
kNum1DBlocksPerGroup
comptime kNum1DBlocksPerGroup = UInt32(16)
num_static_dim_blocks
comptime num_static_dim_blocks = SIMD(ceildiv(static_MN, tile_shape[TileScheduler[static_MN=static_MN, tile_shape=tile_shape, cluster=cluster, cta_group=cta_group, swizzle=swizzle, swapAB=swapAB].static_dim]))
static_dim
comptime static_dim = 0 if swapAB else 1
Methods
__init__
def __init__(num_active_experts: Int, group_offsets: LayoutTensor[DType.uint32, offsets_layout, group_offsets_origin]) -> Self
fetch_next_work
def fetch_next_work(mut self) -> WorkInfo
Returns:
WorkInfo
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!