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 function

bulk_mma_partial

def bulk_mma_partial[kind: UMMAKind, //, layout_b: Layout, *, mma_k: Int, num_k_mmas: Int, operand_size: Int, k_start: Int = Int(0), cta_group: Int = Int(1)](idesc: UMMAInsDescriptor[kind], a: UInt32, b: MMASmemDescriptorPair, c_tmem: UInt32, c_scale: UInt32, elect: Int32, valid_k_mmas: UInt32)

Issues a partial-K TS contraction for a partially-loaded last KV tile, non-warp-specialized.

Each block's MMA carries a warp-uniform validity guard derived from valid_k_mmas, kept separate from the elect predicate to preserve identical elect codegen.

Parameters:

kind (UMMAKind): UMMAKind selecting the tcgen05.mma instruction variant.
layout_b (Layout): SMEM layout of the B operand tile, used to compute per-K-block B descriptor offsets.
mma_k (Int): K-dimension tile size per MMA block, in elements.
num_k_mmas (Int): Number of mma_k-sized K-dimension blocks to contract over in this stage.
operand_size (Int): Size in bytes of the A and B operand elements.
k_start (Int): Absolute K-block index of the first block in this stage (defaults to 0).
cta_group (Int): Number of cooperating CTAs, 1 or 2 (defaults to 1).

Args:

idesc (UMMAInsDescriptor[kind]): UMMA instruction descriptor encoding the accumulator and operand dtypes and the output tile shape.
a (UInt32): Un-offset (stage-0) TMEM base address of the A operand.
b (MMASmemDescriptorPair): Un-offset (stage-0) SMEM descriptor pair for the B operand.
c_tmem (UInt32): TMEM base address of the output accumulator C.
c_scale (UInt32): Accumulator init/accumulate scale; nonzero on the first block to initialize the accumulator, zero to accumulate.
elect (Int32): elect() result selecting the single thread that issues the MMA.
valid_k_mmas (UInt32): Count of loaded mma_k-sized blocks; blocks whose absolute index reaches or exceeds this count are predicated off.