IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /get-started.md). For the complete documentation index, see llms.txt.

Skip to main content

For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo function

mma

def mma[kind: UMMAKind, //, cta_group: Int = Int(1), /, *, c_scale: UInt32 = UInt32(1)](a_desc: MMASmemDescriptor, b_desc: MMASmemDescriptor, c_tmem: UInt32, inst_desc: UMMAInsDescriptor[kind])

Perform a matrix multiply-accumulate operation using the tcgen05.mma instruction.

Parameters:

kind (UMMAKind): Data type of the matrices.
cta_group (Int): Number of ctas used by MMA.
c_scale (UInt32): Scale factor for the C matrix, 0 or 1.

Args:

a_desc (MMASmemDescriptor): The descriptor for the A matrix.
b_desc (MMASmemDescriptor): The descriptor for the B matrix.
c_tmem (UInt32): The address of the C matrix in the tensor memory.
inst_desc (UMMAInsDescriptor[kind]): The descriptor for the MMA instruction.

def mma[kind: UMMAKind, //, cta_group: Int = Int(1), /](a_desc: MMASmemDescriptor, b_desc: MMASmemDescriptor, c_tmem: UInt32, inst_desc: UMMAInsDescriptor[kind], sfa_tmem: UInt32, sfb_tmem: UInt32, c_scale: UInt32)

Perform a matrix multiply-accumulate operation using the tcgen05.mma instruction.

Parameters:

kind (UMMAKind): Data type of the matrices.
cta_group (Int): Number of ctas used by MMA.

Args:

a_desc (MMASmemDescriptor): The descriptor for the A matrix.
b_desc (MMASmemDescriptor): The descriptor for the B matrix.
c_tmem (UInt32): The address of the C matrix in the tensor memory.
inst_desc (UMMAInsDescriptor[kind]): The descriptor for the MMA instruction.
sfa_tmem (UInt32): The address of the block scale factor A in the tensor memory.
sfb_tmem (UInt32): The address of the block scale factor B in the tensor memory.
c_scale (UInt32): Scale factor for the C matrix, 0 or 1.

def mma[kind: UMMAKind, //, cta_group: Int = Int(1), /](a_desc: MMASmemDescriptor, b_desc: MMASmemDescriptor, c_tmem: UInt32, inst_desc: UMMAInsDescriptor[kind], c_scale: UInt32)

Perform a matrix multiply-accumulate operation using the tcgen05.mma instruction.

Parameters:

kind (UMMAKind): Data type of the matrices.
cta_group (Int): Number of ctas used by MMA.

Args:

a_desc (MMASmemDescriptor): The descriptor for the A matrix.
b_desc (MMASmemDescriptor): The descriptor for the B matrix.
c_tmem (UInt32): The address of the C matrix in the tensor memory.
inst_desc (UMMAInsDescriptor[kind]): The descriptor for the MMA instruction.
c_scale (UInt32): Scale factor for the C matrix. Any non-zero value is translated to 1.

def mma[kind: UMMAKind, //, cta_group: Int = Int(1), /](a_desc: UInt32, b_desc: MMASmemDescriptor, c_tmem: UInt32, inst_desc: UMMAInsDescriptor[kind], c_scale: UInt32)

Perform a matrix multiply-accumulate operation using the tcgen05.mma instruction.

Parameters:

kind (UMMAKind): Data type of the matrices.
cta_group (Int): Number of ctas used by MMA.

Args:

a_desc (UInt32): The descriptor for the A matrix.
b_desc (MMASmemDescriptor): The descriptor for the B matrix.
c_tmem (UInt32): The address of the C matrix in the tensor memory.
inst_desc (UMMAInsDescriptor[kind]): The descriptor for the MMA instruction.
c_scale (UInt32): Scale factor for the C matrix. Any non-zero value is interpreted as 1.

def mma[kind: UMMAKind, //, cta_group: Int = Int(1), /, *, c_scale: UInt32 = UInt32(1)](a_desc: UInt32, b_desc: MMASmemDescriptor, c_tmem: UInt32, inst_desc: UMMAInsDescriptor[kind])

Perform a matrix multiply-accumulate operation using the tcgen05.mma instruction.

Parameters:

kind (UMMAKind): Data type of the matrices.
cta_group (Int): Number of ctas used by MMA.
c_scale (UInt32): Scale factor for the C matrix, 0 or 1.

Args:

a_desc (UInt32): The descriptor for the A matrix.
b_desc (MMASmemDescriptor): The descriptor for the B matrix.
c_tmem (UInt32): The address of the C matrix in the tensor memory.
inst_desc (UMMAInsDescriptor[kind]): The descriptor for the MMA instruction.