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_ws_partial

def bulk_mma_ws_partial[kind: UMMAKind, a_dtype: DType, b_dtype: DType, *, a_BMN: Int, a_BK: Int, a_swizzle: TensorMapSwizzle, a_is_k_major: Bool, b_BMN: Int, b_BK: Int, b_swizzle: TensorMapSwizzle, b_is_k_major: Bool, num_k_mmas: Int, operand_size: Int, tcgen05_mma_type: String, mma_k: Int = Int(16), k_start: Int = Int(0), b_page_dense: Bool = False](idesc: UMMAInsDescriptor[kind], a: MMASmemDescriptorPair, b: MMASmemDescriptorPair, c_tmem: UInt32, c_scale: UInt32, elect: Int32, valid_k_mmas: UInt32)

Issues a partial-K SS warp-specialized contraction for a partially-loaded last KV tile.

Both A and B come from SMEM descriptors; each block's MMA carries a warp-uniform validity guard derived from valid_k_mmas, kept separate from elect.

Parameters:

kind (UMMAKind): UMMAKind selecting the tcgen05.mma.ws instruction variant.
a_dtype (DType): Element dtype of the A operand, used to derive the A SMEM tile layout.
b_dtype (DType): Element dtype of the B operand, used to derive the B SMEM tile layout.
a_BMN (Int): M (or N) dimension of the A operand tile in elements, used to derive the A layout.
a_BK (Int): K dimension of the A operand tile in elements, used to derive the A layout.
a_swizzle (TensorMapSwizzle): SMEM swizzle mode for the A tile.
a_is_k_major (Bool): Whether A is stored k-major (True) or mn-major (False).
b_BMN (Int): M (or N) dimension of the B operand tile in elements, used to derive the B layout.
b_BK (Int): K dimension of the B operand tile in elements, used to derive the B layout.
b_swizzle (TensorMapSwizzle): SMEM swizzle mode for the B tile.
b_is_k_major (Bool): Whether B is stored k-major (True) or mn-major (False).
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.
tcgen05_mma_type (String): tcgen05.mma.ws instruction string prefix, including the CTA-group selector.
mma_k (Int): K-dimension tile size per MMA block, in elements (defaults to 16).
k_start (Int): Absolute K-block index of the first block in this stage (defaults to 0).
b_page_dense (Bool): Whether B uses the row-major page-fold layout (defaults to False).

Args:

idesc (UMMAInsDescriptor[kind]): UMMA instruction descriptor encoding the accumulator and operand dtypes and the output tile shape.
a (MMASmemDescriptorPair): Un-offset (stage-0) SMEM descriptor pair for 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.