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
MatmulConfig
struct MatmulConfig[a_type: DType, b_type: DType, c_type: DType, transpose_b: Bool = True]
Static configuration for SM90 (Hopper) GPU warp-specialized matmul.
Encapsulates all compile-time parameters that define a matmul kernel
variant: tile shapes, MMA instruction shape, cluster dimensions, pipeline
depth, consumer warp group count, and scheduling options. Used by the
dispatch layer to select and instantiate HopperMatmulSM90Kernel.
Parameters
- a_type (
DType): ElementDTypeof the A input matrix. - b_type (
DType): ElementDTypeof the B input matrix. - c_type (
DType): ElementDTypeof the output C matrix. - transpose_b (
Bool): WhetherBis stored transposed (K-major layout) (defaults toTrue).
Fields
- block_tile_shape (
IndexList[Int(3)]): - mma_shape (
IndexList[Int(3)]): - cluster_shape (
IndexList[Int(3)]): - num_pipeline_stages (
Int): - num_k_partitions (
Int): - num_consumer (
Int): - partitioned_multicast (
Bool): - k_group_size (
Int):
Implemented traits
AnyType,
Copyable,
Equatable,
Hashable,
ImplicitlyCopyable,
ImplicitlyDeletable,
Movable,
RegisterPassable,
TrivialRegisterPassable,
Writable
Methods
__init__
def __init__(block_tile_shape: IndexList[Int(3)], mma_shape: IndexList[Int(3)], cluster_shape: IndexList[Int(3)], num_pipeline_stages: Int, num_k_partitions: Int, num_consumer: Int, partitioned_multicast: Bool, pdl_level: PDLLevel, k_group_size: Int) -> Self
Initialize MatmulConfig with explicit values for all fields.
Args:
- block_tile_shape (
IndexList[Int(3)]): The(BM, BN, BK)tile shape per CTA in the M, N, and K dimensions. - mma_shape (
IndexList[Int(3)]): The(m, n, k)shape of a single MMA instruction. - cluster_shape (
IndexList[Int(3)]): The thread block cluster dimensions in(M, N, K). - num_pipeline_stages (
Int): Number of pipeline stages for double-buffering loads. - num_k_partitions (
Int): Number of partitions to split the K dimension into. - num_consumer (
Int): Number of consumer warp groups. - partitioned_multicast (
Bool): Whether to use partitioned TMA multicast. - pdl_level (
PDLLevel): Programmatic dependent launch level for grid controls. - k_group_size (
Int): Number of K iterations grouped per pipeline stage.
def __init__(m: Int, n: Int, k: Int, num_k_partitions: Int = Int(1), partitioned_multicast: Bool = False, pdl_level: PDLLevel = PDLLevel.OFF, k_groups: Optional[Int] = None, consumer_groups: Optional[Int] = None, swapAB: Bool = False) -> Self
Initialize MatmulConfig by computing optimal values from M, N, K.
Args:
- m (
Int): The M dimension of the matmul. - n (
Int): The N dimension of the matmul. - k (
Int): The K dimension of the matmul. - num_k_partitions (
Int): Number of K partitions. - partitioned_multicast (
Bool): Whether to use partitioned multicast. - pdl_level (
PDLLevel): PDL level for grid controls. - k_groups (
Optional[Int]): How many pipeline (loads and stores) are grouped together. - consumer_groups (
Optional[Int]): The number of consumer groups. - swapAB (
Bool): Whether to swap A and B.
__eq__
adjust_kgroup_size
static def adjust_kgroup_size(mma_m: Int, mma_n: Int, K: Int, BK: Int, num_pipeline_stages: Int) -> Int
Returns:
pdl_level
to_base_config
def to_base_config(self) -> MatmulConfig[a_type, b_type, c_type, transpose_b]
Convert to base MatmulConfig from utils_gpu.
Returns:
write_to
def write_to(self, mut writer: T)
write_repr_to
def write_repr_to(self, mut writer: T)
__hash__
def __hash__[H: Hasher](self, mut hasher: H)
Updates hasher with the underlying bytes.
Parameters:
- H (
Hasher): The hasher type.
Args:
- hasher (
H): The hasher instance.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!