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

KernelContext

struct KernelContext[num_clc_pipeline_stages: Int, cta_group: Int, CLUSTER_M: Int, CLUSTER_N: Int]

Shared kernel state: election vars, CTA coords, multicast masks, pipeline states.

Parameters

num_clc_pipeline_stages (Int): Number of CLC pipeline stages managed by the scheduler work iterator.
cta_group (Int): Number of CTAs that cooperate as a single group, either 1 or 2 for dual-CTA (2SM) mode.
CLUSTER_M (Int): Number of CTAs in the cluster along the M dimension.
CLUSTER_N (Int): Number of CTAs in the cluster along the N dimension.

Fields

elect_one_warp (Bool):
elect_one_thread (Bool):
elect_one_cta (Bool):
is_first_cta_in_cluster (Bool):
warp_id (UInt32):
rank_m (Int):
rank_n (Int):
peer_cta_coord (Tuple[Int, Int, Int]):
a_multicast_mask (UInt16):
b_multicast_mask (UInt16):
mma_complete_mask (Int):
ptr_tmem_addr (Pointer[UInt32, MutUntrackedOrigin, address_space=AddressSpace.SHARED, _safe=False]):

Implemented traits

AnyType, Copyable, ImplicitlyDeletable, Movable

`comptime` members

`TmemAddrArray`

comptime TmemAddrArray = SMemArray[UInt32, Int(1)]

Methods

`init`

def __init__(out self, ptr_tmem_addr: Pointer[UInt32, MutUntrackedOrigin, address_space=AddressSpace.SHARED, _safe=False])

Initialize context from TMEM pointer; computes all derived state.

Args:

ptr_tmem_addr (Pointer[UInt32, MutUntrackedOrigin, address_space=AddressSpace.SHARED, _safe=False]): Shared-memory pointer to the TMEM address used by the epilogue warps.

def __init__(out self, tmem_addr: SMemArray[UInt32, Int(1)])

Initialize context from typed TMEM address array.

Args:

tmem_addr (SMemArray[UInt32, Int(1)]): Typed TMEM address array to initialize the context from.

Parameters​

Fields​

Implemented traits​

comptime members​

TmemAddrArray​

Methods​

__init__​