Mojo module
grouped_1d1d_tile_scheduler
Work scheduler for grouped 1D-1D block-scaled SM100 matmul.
Provides work iteration using offset-based addressing for the 1D-1D tensor layout. This is a port of the TileScheduler from grouped_matmul_tile_scheduler.mojo to the structured kernels architecture with context manager patterns.
Key characteristics:
- Uses a_offsets tensor for group boundaries (prefix sum of token counts)
- Each iteration returns (m_coord, n_coord, expert_id, expert_scale)
- Supports block swizzling for L2 cache efficiency
- 3-warp specialization (no scheduler warp)
Structsβ
- β
GroupedWorkContext1D1D: Context for current work tile, used with context manager pattern. - β
GroupedWorkInfo1D1D: Work tile information for 1D-1D grouped matmul. - β
GroupedWorkIterator1D1D: Work iterator for 1D-1D grouped block-scaled matmul.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!