mma.sync

mma.sync#

The warp tile pattern supports the explicit mma.sync.aligned.m16n8k{8,16}.row.col.f32.{f16,bf16}.{f16,bf16}.f32 forms registered by the PTX dialect. Each issue multiplies fixed register fragments and accumulates into four float32 registers per lane. The tile layer does not select K=8 versus K=16, change input type, load shared memory, or transpose an operand implicitly.

Operand coordinates and ownership#

The call takes D,A,B,C. After optional execution-group and batch axes, the canonical shapes are:

Operand

Shape

Storage

D / tensor C

(M_repeat,N_repeat,16,8)

float32

A

(M_repeat,K_repeat,16,k/2)

b32 words containing two adjacent f16/bf16 K values

B

(N_repeat,K_repeat,8,k/2)

The same packed representation

B’s logical atom axes are (N,K/2), even though the mathematical product uses B transposed. For lane l, native register r corresponds to:

Operand

Atom coordinate

Registers per lane

A

(l//4+8*(r%2), l%4+4*(r//2))

k/4

B

(l//4, l%4+4*r)

k/8

D / C

(l//4+8*(r//2), 2*(l%4)+r%2)

4

Local offsets can be strided; lane ownership must match these coordinates. All 32 lanes execute the warp instruction. Warpgroup/CTA scope prepends a warp group axis; that axis is not an additional matrix dimension.

Example: a 16x8 result with two K16 issues#

The complete kernel below loads the exact packed A/B words each lane needs. The host prepares these words from ordinary row-major half matrices.

from tvm.script import tirx as T
from tvm.backend.cuda.axis import laneid
from tvm.tirx.layout import S, TileLayout

la = TileLayout(S[(1, 2, 2, 8, 2, 4) : (8, 4, 1, 4 @ laneid, 2, 1 @ laneid)])
lb = TileLayout(S[(1, 2, 8, 2, 4) : (4, 2, 4 @ laneid, 1, 1 @ laneid)])
ld = TileLayout(S[(1, 1, 2, 8, 4, 2) : (4, 4, 2, 4 @ laneid, 1 @ laneid, 1)])


@T.function
def warp_mma(
    A: T.Tensor((1, 2, 16, 8), "uint32"),
    B: T.Tensor((1, 2, 8, 8), "uint32"),
    Out: T.Tensor((16, 8), "float32"),
):
    T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=32))
    lane = T.cuda.lane_id()
    ar = T.alloc_local((8,), "uint32", layout="default")
    br = T.alloc_local((4,), "uint32", layout="default")
    dr = T.alloc_local((4,), "float32", layout="default")
    av = ar.view(1, 2, 16, 8, layout=la)
    bv = br.view(1, 2, 8, 8, layout=lb)
    dv = dr.view(1, 1, 16, 8, layout=ld)
    for kk, r in T.grid(2, 4):
        ar[kk * 4 + r] = A[0, kk, lane // 4 + 8 * (r % 2), lane % 4 + 4 * (r // 2)]
    for kk, r in T.grid(2, 2):
        br[kk * 2 + r] = B[0, kk, lane // 4, lane % 4 + 4 * r]
    T.cuda.tile.apply(
        T.ptx["mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32"],
        dv,
        av,
        bv,
        T.float32(0),
        scope="warp",
    )
    for r in range(4):
        Out[lane // 4 + 8 * (r // 2), 2 * (lane % 4) + r % 2] = dr[r]

import numpy as np
import tvm

rng = np.random.default_rng(0)
a = (rng.integers(-4, 5, (16, 32)) / 4).astype("float16")
b = (rng.integers(-4, 5, (8, 32)) / 4).astype("float16")
ap = np.ascontiguousarray(a.reshape(1, 16, 2, 16).transpose(0, 2, 1, 3)).view("uint32")
bp = np.ascontiguousarray(b.reshape(1, 8, 2, 16).transpose(0, 2, 1, 3)).view("uint32")
compiled = tvm.tirx.compile(warp_mma, target="cuda", pipeline="tirx")
out = tvm.runtime.tensor(np.zeros((16, 8), dtype="float32"), tvm.cuda())
compiled(tvm.runtime.tensor(ap, tvm.cuda()), tvm.runtime.tensor(bp, tvm.cuda()), out)
np.testing.assert_allclose(
    out.numpy(), a.astype("float32") @ b.astype("float32").T, rtol=2e-4, atol=2e-4
)

M_repeat=N_repeat=1 and K_repeat=2, so each warp issues exactly two K16 instructions. The first uses scalar C=0; the second uses the preceding D registers. No separate zeroing of dr or add instruction is required.

Accumulation, transpose and alpha/beta#

Repeat order is batch, M, N, then K. C can be a matching register tensor or an explicit scalar broadcast. Only the first K issue reads the original C; later issues read D. This is the documented recurrence, rather than an arbitrary traversal inferred from layouts.

To transpose a logical input, prepare a view whose canonical A/B coordinates refer to the intended values while preserving the required lane map. A view alone cannot exchange values between lanes; an incompatible physical owner requires explicit preparation. alpha/beta scaling likewise belongs before or after MMA at the call site, with the intended rounding order.

The tile call inserts no shared loads, ldmatrix, communication, or barrier. Use ldmatrix / stmatrix to prepare fragments when that native shared/register image matches your operands. Sparse and other MMA shapes require their own explicit pattern; this page does not claim those forms.

Implementation: python/tvm/backend/cuda/tile/patterns/mma.py.

Contents