Tensor core operations

Parameters:

2-CTA Scaled MMA: When two_ctas=True, the scaled MMA operates across two CTAs in a cluster. Key considerations:

Example: 2-CTA Scaled MMA

    # B data split across CTAs, but B scale is full
    desc_b = tl.make_tensor_descriptor(b_ptr, ..., block_shape=[BLOCK_K, BLOCK_N // 2])
    desc_b_scale = tl.make_tensor_descriptor(b_scale_ptr, ..., block_shape=[BLOCK_N // 128, ...])  # Full scale

    # Load B with CTA offset, B scale without offset
    tlx.async_descriptor_load(desc_b, b_tile[0], [0, cluster_cta_rank * BLOCK_N // 2], bar_b)
    tlx.async_descriptor_load(desc_b_scale, b_scale_tile[0], [0, 0, 0, 0], bar_b_scale)  # Full B scale

    # CTA sync: "Arrive Remote, Wait Local"
    tlx.barrier_arrive(cta_bars[0], 1, remote_cta_rank=0)
    tlx.barrier_wait(cta_bars[0], phase=0, pred=pred_cta0)

    # 2-CTA scaled MMA with mBarriers for completion tracking
    tlx.async_dot_scaled(
        a_tile[0], b_tile[0], c_tile[0],
        a_scale_tile[0], "e4m3",
        b_scale_tile[0], "e4m3",
        use_acc=False,
        two_ctas=True,
        mBarriers=[mma_done_bar],
    )
    tlx.barrier_wait(mma_done_bar, tl.constexpr(0))

Alternative: Using tcgen05_commit for MMA completion

    # Issue MMA without mBarriers
    tlx.async_dot_scaled(..., two_ctas=True)

    # Use tcgen05_commit to track all prior MMA ops
    tlx.tcgen05_commit(mma_done_bar, two_ctas=True)
    tlx.barrier_wait(mma_done_bar, tl.constexpr(0))

TMEM-backed MX Scales:

For scaled MMA operations on Blackwell GPUs, scales can be stored in Tensor Memory (TMEM) for efficient access. TLX provides automatic layout resolution for TMEM scale buffers.

*Allocating TMEM Scale Buffers:*

When allocating TMEM buffers for uint8/int8 types (used for MX scales), TLX uses a placeholder layout (DummyTMEMLayoutAttr) that gets automatically resolved to TensorMemoryScalesEncodingAttr during compilation when the buffer is used with async_dot_scaled.

    # Allocate TMEM buffers for scales (layout is automatically resolved)
    a_scale_tmem = tlx.local_alloc((128, 8), tl.uint8, num=1, storage=tlx.storage_kind.tmem)
    b_scale_tmem = tlx.local_alloc((256, 4), tl.uint8, num=1, storage=tlx.storage_kind.tmem)

*Copying Scales from SMEM to TMEM:*

Use tlx.tmem_copy [sm100] to efficiently transfer scale data from shared memory to tensor memory:

    # Copy scales from SMEM to TMEM (asynchronous, uses tcgen05.cp instruction)
    tlx.tmem_copy(a_scale_smem, a_scale_tmem)
    tlx.tmem_copy(b_scale_smem, b_scale_tmem)

*Using TMEM Scales with Scaled MMA:*

    # TMEM scales are automatically detected and used with the correct layout
    tlx.async_dot_scaled(
        a_smem, b_smem, acc_tmem,
        A_scale=a_scale_tmem, A_format="e4m3",
        B_scale=b_scale_tmem, B_format="e4m3",
        use_acc=True,
        mBarriers=[mma_bar],
    )

*Complete Example: TMEM-backed Scaled GEMM:*

    @triton.jit
    def scaled_gemm_kernel(...):
        # Allocate TMEM for accumulator and scales
        acc = tlx.local_alloc((BLOCK_M, BLOCK_N), tl.float32, num=1, storage=tlx.storage_kind.tmem)
        a_scale_tmem = tlx.local_alloc((BLOCK_M // 128, BLOCK_K // 32), tl.uint8, num=1, storage=tlx.storage_kind.tmem)
        b_scale_tmem = tlx.local_alloc((BLOCK_N // 128, BLOCK_K // 32), tl.uint8, num=1, storage=tlx.storage_kind.tmem)

        # Load scales from global memory to SMEM
        tlx.async_descriptor_load(a_scale_desc, a_scale_smem, [...], barrier=bar)
        tlx.async_descriptor_load(b_scale_desc, b_scale_smem, [...], barrier=bar)
        tlx.barrier_wait(bar, phase)

        # Copy scales from SMEM to TMEM
        tlx.tmem_copy(a_scale_smem[0], a_scale_tmem[0])
        tlx.tmem_copy(b_scale_smem[0], b_scale_tmem[0])

        # Perform scaled MMA with TMEM scales
        tlx.async_dot_scaled(
            a_smem[0], b_smem[0], acc[0],
            A_scale=a_scale_tmem[0], A_format="e4m3",
            B_scale=b_scale_tmem[0], B_format="e4m3",
            use_acc=False,
        )

Note: Multibuffering is automatically cancelled for scale buffers since TMEM scales don't support multibuffering. 3D allocations (1×M×K) are automatically flattened to 2D (M×K).

Shift a mutable tensor-memory buffer in place with Blackwell tcgen05.shift.down instructions. Within each 32-row TMEM group, destination row r receives the previous contents of source row r + 1 for r in 0 through 30; destination row 31 retains its previous contents.

The instruction completes asynchronously. Call tlx.tcgen05_commit with an mbarrier after the shift, then wait on that barrier before loading or otherwise consuming the buffer. The buffer must be a mutable rank-2 TMEM allocation using the standard tensor-memory encoding. It must occupy exactly 128 physical TMEM rows, and its physical column count must be a multiple of 8. The compiler emits one shift instruction per 32-byte column segment.

Parameters:

Example:

    buffers = tlx.local_alloc(
        (128, 16),
        tl.float32,
        tl.constexpr(1),
        tlx.storage_kind.tmem,
    )
    buffer = tlx.local_view(buffers, 0)
    barriers = tlx.alloc_barriers(tl.constexpr(1))
    barrier = tlx.local_view(barriers, 0)

    tlx.local_store(buffer, values)
    tlx.tmem_shift(buffer)
    tlx.tcgen05_commit(barrier)
    tlx.barrier_wait(barrier, tl.constexpr(0))
    shifted = tlx.local_load(buffer)

Wait for completion of prior asynchronous dot operations. The pendings argument indicates the number of in-flight operations not completed.

Example:

    acc = tlx.async_dot(a_smem, b_smem)
    acc = tlx.async_dot_wait(tl.constexpr(0), acc)
    tl.store(C_ptrs, acc)

Explicit MFMA scheduling

amd_scheduled_mfma and amd_mfma_commit support [gfx942, gfx950] native BF16/F16 MFMA layouts. The register-class constraint operations are AMD-only and currently verified on [gfx950].

These are low-level backend-tuning controls. Prefer tl.dot unless a kernel needs explicit native accumulator chains, register-class constraints, or an MFMA completion boundary. Inspect the generated code and benchmark any use; the register allocator still chooses physical registers.

Related operations are tlx.extract_slice, which selects an aligned register fragment without cross-thread movement; tlx.rematerialized_range, which recreates inexpensive coordinates near a use; and tlx.amd_sched_barrier, which prevents selected AMD instruction classes from crossing a source boundary. None of these operations is a workgroup barrier or memory fence.

tlx.amd_register_resident

value = tlx.amd_register_resident(
    value,
    register_class="agpr",
    registers_per_group=1,
)

Returns value unchanged, with the same shape, element type, and distributed layout. At one allocator-visible point, every packed per-thread group must be available in the requested "agpr" or "vgpr" class.

registers_per_group is the number of consecutive 32-bit registers in each group, not the number of groups. It must be one of 1, 2, 4, 8, 16, or 32. A group of width R contains R 32-bit elements or 2 * R 16-bit elements, and the per-thread element count must divide evenly into groups. Registers are consecutive within one group; different groups need not be adjacent.

This constraint does not select physical register numbers, reserve registers for the value's full lifetime, prevent copies or spills, or guarantee a particular occupancy. Wider groups strengthen the tuple constraint and can reduce allocator flexibility. Always use the returned value.

tlx.amd_register_class_anchor

value = tlx.amd_register_class_anchor(value, register_class="vgpr")

Returns value unchanged after passing each packed 32-bit register value through a separate, side-effecting tied "agpr" or "vgpr" constraint. This provides a source-local allocation and scheduling anchor without requiring all tensor groups to meet at one point.

A tied input and output may be coalesced by LLVM. The operation does not guarantee a distinct physical register or live interval, emit a copy, shorten a live range, prevent spills, or improve occupancy. Use amd_register_resident when simultaneous all-group allocation is required, and use the value returned by amd_register_class_anchor for downstream work.

Both register-class operations accept distributed integer or floating-point tensors with 16-bit or 32-bit elements. They constrain allocation at a source point; they do not wait for MFMA completion or synchronize lanes, waves, or memory.

tlx.amd_scheduled_mfma

acc = tlx.amd_scheduled_mfma(
    a,
    b,
    acc,
    accumulator_role="transient",
    resident_operand=None,
    accumulator_register_class=None,
    initialize=False,
)

With initialize=False, the operation computes the native-fragment equivalent of acc + a @ b for the entire accumulator tile passed to the call. That tile may be an explicit subtile of a larger output. Within a call, the operation keeps one SSA chain per output fragment and creates updates in K-major, N-major, M-minor order:

for k in native_k_fragments:
    for n in output_n_fragments:
        for m in output_m_fragments:
            acc[m, n] = mfma(a[m, k], b[k, n], acc[m, n])

This source order round-robins a K slice over independent accumulators before returning to the same dependency chain. LLVM may still reschedule independent instructions on both intrinsic paths.

To schedule output regions independently, slice the corresponding operands and carry a separate accumulator for each region. Each scheduled call updates its entire passed subtile, so independent loads or other work can be placed between those calls. Keep these accumulators separate through the matching amd_mfma_commit boundary.

a and b must be matching BF16 or F16 tensors with dot-operand layouts using kWidth=4 or 8. acc must be F32 with the corresponding unit-tile MFMA layout. All tensors must have matching rank two or three. Rank-three tensors have matching leading batch dimensions distributed over waves, with one batch repetition per wave. Matrix shapes and per-wave native fragments must match.

resident_operand and an explicit accumulator class impose hard class constraints only on the current persistent lowering. The transient intrinsic path leaves physical placement to LLVM; use amd_register_resident separately when a transient source needs an explicit allocation point. Set initialize=True only for the first band of a multi-band accumulation, or the earlier accumulated value will be discarded. When carrying independent output subtiles, apply this initialization rule to each accumulator separately.

LLVM models MFMA latency and hazards through native intrinsics for both roles. The compiler no longer runs a separate Python pass to repair MFMA assembly hazards. Persistent register-class pins use inline assembly without hiding the MFMA arithmetic. When results can reach opaque consumers, the compiler adds target-specific completion waits. Explicit amd_mfma_commit boundaries retain their register constraints and waits, including for adjacent commits. Register pins constrain allocation at a point; they do not guarantee exact machine instruction order or prevent spills.

gfx942 rejects every explicit accumulator_register_class="agpr" request. Because persistent auto also resolves to AGPR, persistent gfx942 calls must select "vgpr" explicitly.

All active lanes of a wave must execute amd_scheduled_mfma uniformly.

tlx.amd_mfma_commit

committed = tlx.amd_mfma_commit(acc)
committed, live = tlx.amd_mfma_commit(acc, preserve=live)

Applies the target-specific MFMA result-readiness boundary and returns every input numerically unchanged. value may be one F32 MFMA-layout tensor or a nonempty tuple of independent results. Optional preserve is one BF16 or F16 dot-operand tensor threaded through the same liveness and allocation boundary; it is not another numerical MFMA operand.

Each F32 chain must reach this boundary without an ordinary SSA fork. The verifier also accepts the narrow runtime-loop form where a block argument has mutually exclusive body and exit uses, because only one use executes on each dynamic iteration. Downstream code must use the returned result. When a preserve dependency remains live after the boundary, its downstream uses must similarly use the returned live_out; the returned dependency may be discarded when it is no longer needed.

Pass independent output subtiles as a tuple to commit them at the same boundary. Concatenate the returned results only after the commit. Concatenation may produce a different distributed layout; use tlx.require_layout on the joined value when its consumer requires a particular layout.

The current lowering constrains a results-only boundary to AGPRs. A boundary with preserve constrains its F32 results to VGPRs and the preserved dot operand to AGPRs; directly committing an AGPR-resident scheduled-MFMA result with preserve is rejected. These placement choices and the exact wait padding are implementation details, not portable synchronization semantics.

amd_mfma_commit is not an MFMA queue instruction, s_waitcnt, hardware memory fence, workgroup barrier, or cross-wave synchronization operation.

Independent output subtiles on gfx950

Assume a has shape [64, 16] and b has shape [16, 128], with matching dot-operand layouts for the MFMA layout below. Each [64, 64] accumulator is one native output fragment per wave. The column slices align to the CTA tile, so extract_slice selects their registers without moving values across waves.

mma: tl.constexpr = tlx.amd_mfma_layout(
    version=4,
    instr_shape=[32, 32, 16],
    transposed=True,
    warps_per_cta=[2, 2],
)
b_left = tlx.extract_slice(b, [16, 64], [0, 0])
b_right = tlx.extract_slice(b, [16, 64], [0, 64])
acc_left = tlx.zeros((64, 64), tl.float32, layout=mma)
acc_right = tlx.zeros((64, 64), tl.float32, layout=mma)

acc_left = tlx.amd_scheduled_mfma(
    a, b_left, acc_left, accumulator_role="persistent", initialize=True,
)
tlx.amd_sched_barrier(0)
# Independent loads or other work may be placed here.
tlx.amd_sched_barrier(0)
acc_right = tlx.amd_scheduled_mfma(
    a, b_right, acc_right, accumulator_role="persistent", initialize=True,
)
acc_left, acc_right = tlx.amd_mfma_commit((acc_left, acc_right))
acc = tlx.require_layout(tl.cat(acc_left, acc_right, dim=1), mma, pin=False)

For accumulation over multiple reduction bands, carry both subtiles through the loop with initialize=False after the first band, then commit their final values together before joining them.

Transient-chain example

Assume a, b, and acc already carry matching dot-operand and MFMA layouts:

b_live = tlx.amd_register_resident(
    b,
    register_class="agpr",
    registers_per_group=4,
)
acc = tlx.amd_scheduled_mfma(
    a,
    b_live,
    acc,
    accumulator_role="transient",
    initialize=True,
)
acc, _ = tlx.amd_mfma_commit(acc, b_live)
tl.store(output_ptrs, acc)

For later accumulation bands, pass initialize=False. Here, amd_register_resident creates a common allocation point for b_live, while preserve keeps it live through the commit boundary. If later code needs that operand, consume the returned value instead of _. resident_operand controls only persistent lowering.

Scaled dot

tlx.dot_scaled(lhs, lhs_scale, lhs_format, rhs, rhs_scale, rhs_format, acc=None, *, fast_math=False, lhs_k_pack=True, rhs_k_pack=True, out_dtype=tl.float32, tiles_per_warp=None) is a thin wrapper around tl.dot_scaled. Without tiles_per_warp it is exactly equivalent to tl.dot_scaled — pass it only when you need the AMD-specific WMMA scheduling hint described below.

tiles_per_warp — what it controls

So tiles_per_warp=[2, 2] does not mean "distribute 4 tiles across 4 warps." It means *each* warp emits a 2×2 block of WMMA instruction tiles, holding the corresponding 2×2 accumulator registers. Concretely, for a tt.dot_scaled lowered to gfx1250 WMMA (instrShape = [16, 16, K]), 4 warps, warpsPerCTA = [2, 2]:

For a 256 × 256 result, [1, 1] repeats the layout 8 × 8 times, [2, 2] repeats it 4 × 4. Larger tiles_per_warp gives each warp more contiguous accumulator state (better register reuse for preshuffled MXFP scales, fewer warp-level reductions), at the cost of more registers per warp.

Together, instrShape, warpsPerCTA, and tiles_per_warp define the M/N extent of one CTA-level WMMA layout period: period[d] = instrShape[d] * warpsPerCTA[d] * tiles_per_warp[d]. If the result tile is larger than this period, the period repeats. The K entry of instrShape is the per-instruction reduction depth and is handled separately from this M/N tiling.

tiles_per_warp is validated by AccelerateAMDMatmul: it must have one entry per result-tile dim, each entry must be positive, and instrShape[d] * warpsPerCTA[d] * tiles_per_warp[d] must fit in the result tile shape.

Example: tiles_per_warp

import triton.language.extra.tlx as tlx

acc = tlx.dot_scaled(
    a, a_scale, "e5m2",
    b, b_scale, "e5m2",
    acc,
    tiles_per_warp=[2, 2],   # pack 2x2 WMMA tiles per warp for preshuffled MXFP
)

Mechanism (for IR-level users)

The wrapper attaches amdg.wmma_tiles_per_warp = array<i32: m, n> on the resulting tt.dot_scaled op. ScaledBlockedToScaledWMMAF8F6F4 reads the attribute and substitutes m, n for the default 1, 1 when building the WMMA encoding. Setting the attribute directly on a tt.dot[_scaled] op in MLIR has the same effect; the wrapper just spares Python kernels from hand-poking attributes.

Currently consumed only by the scaled-WMMA pattern (gfx1250). Regular tt.dot WMMA and the MFMA patterns do not read it.