tilelang.cuda.language.gemm_op

CUDA dialect of the GEMM operators.

Shadows the common gemm / gemm_blockscaled / gemm_sp with CUDA-specific parameters and hosts the explicit, instruction-pinned CUDA variants (wgmma_gemm, tcgen05_gemm, tcgen05_gemm_blockscaled, mma_gemm_blockscaled) that have no target-neutral meaning.

Functions

gemm(A, B, C[, transpose_A, transpose_B, policy, ...])

TileLang GEMM operator for CUDA.

gemm_blockscaled(A, B, C, SFA, SFB[, transpose_A, ...])

Block-scaled GEMM for CUDA: C (+)= (A * SFA) @ (B * SFB).

gemm_sp(A_sparse, E, B, C[, transpose_A, transpose_E, ...])

Sparse GEMM (2:4 structured sparsity) for CUDA.

wgmma_gemm(A, B, C[, transpose_A, transpose_B, ...])

Explicit Hopper WGMMA GEMM without an implicit wait.

tcgen05_gemm(A, B, C[, transpose_A, transpose_B, ...])

Explicit Blackwell TCGEN05 GEMM without an implicit wait.

tcgen05_gemm_blockscaled(A, B, C, SFA_tmem, SFB_tmem)

Explicit Blackwell TCGEN05 block-scaled GEMM without an implicit wait.

mma_gemm_blockscaled(A, B, C, SFA, SFB[, transpose_A, ...])

Explicit SM120 warp-level block-scaled MMA GEMM.

make_blockscaled_gemm_layout(C, A[, transpose_A])

Build the TMEM store layout for the C accumulator of a block-scaled GEMM.

Module Contents

tilelang.cuda.language.gemm_op.gemm(A, B, C, transpose_A=False, transpose_B=False, policy=GemmWarpPolicy.Square, clear_accum=False, mbar=None, annotations=None)

TileLang GEMM operator for CUDA.

Same semantics as the common tilelang.language.gemm_op.gemm(): the default synchronous GEMM. On Hopper, if the compiler selects WGMMA lowering, TileLang inserts the corresponding wait implicitly. On Blackwell TCGEN5MMA, TileLang inserts the corresponding mbarrier_wait_parity(...) implicitly after issue.

For manual asynchronous scheduling, use T.wgmma_gemm(...) with T.wait_wgmma(...) on Hopper, or T.tcgen05_gemm(...) with T.mbarrier_wait_parity(...) on Blackwell.

Parameters:

  • A (BufferLikeType, i.e. Buffer | BufferLoad | BufferRegion, or Var) – Input buffer A.

  • B (BufferLikeType) – Input buffer B.

  • C (BufferLikeType) – Output buffer C.

  • transpose_A (bool) – Whether to transpose A. Defaults to False.

  • transpose_B (bool) – Whether to transpose B. Defaults to False.

  • policy (GemmWarpPolicy) – GEMM warp partition policy.

  • clear_accum (bool) – Whether to clear the accumulator.

  • mbar (BarrierType, i.e. Buffer | BufferLoad, or Var, optional) – Mbarrier in Blackwell. Required when this GEMM lowers to TCGEN5MMA. Defaults to None.

  • annotations (Optional[dict]) – Additional annotations.

Returns:

A handle to the GEMM operation.

Return type:

tirx.Call

tilelang.cuda.language.gemm_op.gemm_blockscaled(A, B, C, SFA, SFB, transpose_A=False, transpose_B=False, policy=GemmWarpPolicy.Square, clear_accum=False, *, k_start, sf_a_granularity_k, sf_b_granularity_k, mbar=None, use_2cta=False, sf_layout=None, annotations=None)

Block-scaled GEMM for CUDA: C (+)= (A * SFA) @ (B * SFB).

Extends tilelang.language.gemm_op.gemm_blockscaled() with CUDA completion and scale-layout parameters. Like T.gemm, this is the synchronous interface. The compiler selects the block-scaled instruction from the target and operand scopes:

  • Blackwell SM100, C in tensor memory: TCGEN05 kind::mxf8f6f4.block_scale. A/B are FP8/FP6/FP4 in shared memory, and SFA/SFB are E8M0 scale factors already in tensor memory. mbar is required: the MMA posts completion to it and TileLang inserts the matching mbarrier_wait_parity implicitly after issue, so the barrier must flip once per call. use_2cta=True requests cta_group::2 and requires cluster_dims of (2,1,1) or (1,2,1).

  • SM120, C in a fragment: synchronous mma.sync.m16n8k64.kind::mxf4nvf4.block_scale with E2M1 operands, UE4M3 scale factors and FP32 accumulation. A/B and packed scale words reside in shared memory. sf_layout describes the scale words; this path ignores mbar.

Unsupported combinations fail compilation instead of dropping scales. Use T.tcgen05_gemm_blockscaled for explicit asynchronous TCGEN05 scheduling without the implicit wait, or T.mma_gemm_blockscaled to pin the warp-level path.

Parameters:

  • A (tilelang._typing.BufferLikeType) – Left operand tile in shared memory.

  • B (tilelang._typing.BufferLikeType) – Right operand tile in shared memory.

  • C (tilelang._typing.BufferLikeType) – Accumulator in tensor memory on SM100 or a fragment on SM120.

  • SFA (tilelang._typing.BufferLikeType) – Scale factors for A.

  • SFB (tilelang._typing.BufferLikeType) – Scale factors for B.

  • transpose_A (bool) – Whether to transpose A. Defaults to False.

  • transpose_B (bool) – Whether to transpose B. Defaults to False.

  • policy (tilelang.language.gemm_op.GemmWarpPolicy) – Warp partition policy for the warp-level path.

  • clear_accum (bool) – Whether to zero the accumulator before accumulating.

  • k_start (int | tvm.tirx.PrimExpr) – Logical K-axis start offset of this tile.

  • sf_a_granularity_k (int) – K elements covered by one A scale factor.

  • sf_b_granularity_k (int) – K elements covered by one B scale factor.

  • mbar (tilelang.language.gemm_op.BarrierType | None) – Completion barrier, required by TCGEN05; waited on implicitly.

  • use_2cta (bool) – Request the 2CTA TCGEN05 variant.

  • sf_layout (str | None) – SM120 scale layout, "rowmajor" or "blockscaled_chunk_kmajor".

  • annotations (dict | None) – Additional annotations.

Return type:

tvm.tirx.PrimExpr

tilelang.cuda.language.gemm_op.gemm_sp(A_sparse, E, B, C, transpose_A=False, transpose_E=False, transpose_B=False, policy=GemmWarpPolicy.Square, clear_accum=False, wg_wait=0, annotations=None)

Sparse GEMM (2:4 structured sparsity) for CUDA.

Same semantics as the common tilelang.language.experimental.gemm_sp_op.gemm_sp(). wg_wait is the Hopper warpgroup wait count consumed when the WGMMA SP lowering is selected (-1 defers the wait to an explicit T.wait_wgmma); it rides in the tile-op annotations.

Parameters:

  • A_sparse (tilelang._typing.BufferLikeType | tvm.tirx.Var) – Compressed sparse matrix containing only non-zero elements.

  • E (tilelang._typing.BufferLikeType | tvm.tirx.Var) – Metadata tensor encoding the sparsity pattern of A.

  • B (tilelang._typing.BufferLikeType | tvm.tirx.Var) – Dense input matrix.

  • C (tilelang._typing.BufferLikeType | tvm.tirx.Var) – Output accumulator matrix.

  • transpose_A (bool) – Whether to transpose A. Defaults to False.

  • transpose_E (bool) – Whether to transpose E. Defaults to False.

  • transpose_B (bool) – Whether to transpose B. Defaults to False.

  • policy (tilelang.language.gemm_op.GemmWarpPolicy) – Warp partition policy. Defaults to GemmWarpPolicy.Square.

  • clear_accum (bool) – Whether to zero the accumulator before computation. Defaults to False.

  • wg_wait (int) – Warp group wait count. Defaults to 0.

  • annotations (dict | None) – Additional annotations; values in it take precedence.

Returns:

A handle to the sparse GEMM operation.

Return type:

tirx.Call

tilelang.cuda.language.gemm_op.wgmma_gemm(A, B, C, transpose_A=False, transpose_B=False, policy=GemmWarpPolicy.Square, clear_accum=False, annotations=None)

Explicit Hopper WGMMA GEMM without an implicit wait.

This is the explicit asynchronous Hopper WGMMA counterpart to the default synchronous T.gemm(…) interface, with two stricter guarantees: - it always requests the WGMMA lowering path - it never auto-emits an inlined warpgroup_wait

If the current target or operand pattern cannot use Hopper WGMMA, compilation fails instead of silently falling back to MMA.

Parameters:

  • A (tilelang._typing.BufferLikeType)

  • B (tilelang._typing.BufferLikeType)

  • C (tilelang._typing.BufferLikeType)

  • transpose_A (bool)

  • transpose_B (bool)

  • policy (tilelang.language.gemm_op.GemmWarpPolicy)

  • clear_accum (bool)

  • annotations (dict | None)

Return type:

tvm.tirx.PrimExpr

tilelang.cuda.language.gemm_op.tcgen05_gemm(A, B, C, transpose_A=False, transpose_B=False, policy=GemmWarpPolicy.Square, clear_accum=False, *, mbar, use_2cta=False, annotations=None)

Explicit Blackwell TCGEN05 GEMM without an implicit wait.

This is the explicit asynchronous Blackwell TCGEN5MMA counterpart to the default synchronous T.gemm(…) interface, with two stricter guarantees: - it always requests the TCGEN5MMA lowering path - it never auto-emits an inlined mbarrier_wait_parity

mbar=None omits the completion arrival for an intermediate issue. A later TCGEN05 operation remains ordered in the same issue stream and may publish the completion event for the whole sequence.

When use_2cta=True, the instruction is lowered to the 2CTA variant which requires cluster_dims to be (2,1,1) or (1,2,1).

If the current target or operand pattern cannot use Blackwell TCGEN5MMA, compilation fails instead of silently falling back to another GEMM path.

Parameters:

  • A (tilelang._typing.BufferLikeType)

  • B (tilelang._typing.BufferLikeType)

  • C (tilelang._typing.BufferLikeType)

  • transpose_A (bool)

  • transpose_B (bool)

  • policy (tilelang.language.gemm_op.GemmWarpPolicy)

  • clear_accum (bool)

  • mbar (tilelang.language.gemm_op.BarrierType | None)

  • use_2cta (bool)

  • annotations (dict | None)

Return type:

tvm.tirx.PrimExpr

tilelang.cuda.language.gemm_op.tcgen05_gemm_blockscaled(A, B, C, SFA_tmem, SFB_tmem, transpose_A=False, transpose_B=False, clear_accum=False, wg_wait=0, mbar=None, *, k_start, sf_a_granularity_k, sf_b_granularity_k, use_2cta=False)

Explicit Blackwell TCGEN05 block-scaled GEMM without an implicit wait.

This is the explicit counterpart of T.gemm_blockscaled(…) for Blackwell TCGEN5MMA, with the same guarantees as T.tcgen05_gemm(…): it always requests the TCGEN5MMA lowering path and compilation fails instead of silently falling back if that path is unavailable. It never auto-emits an inlined mbarrier_wait_parity.

mbar=None omits the completion arrival for this issue. The caller must publish completion with a later TCGEN05 operation or an explicit T.tcgen05_mma_arrive before waiting and consuming the result.

With use_2cta=True, this lowers to the true 2CTA block-scaled TCGEN05 path only; there is no fallback or emulation. That mode requires cluster_dims to be (2,1,1) or (1,2,1).

A and B are FP8/FP6/FP4 mxf8f6f4 operands in shared memory, C is the accumulator in tensor memory, and SFA/SFB are E8M0 scale factors already resident in tensor memory. The API is explicit-async: it issues the MMA and leaves synchronization to the user schedule.

k_start is the logical K-axis start offset for this MMA tile. sf_a_granularity_k and sf_b_granularity_k describe how many K elements one packed scale factor covers. The compiler derives the PTX scale-factor A/B IDs for each internal K32 MMA atom from these values.

Parameters:

  • A (tilelang._typing.BufferLikeType) – FP8/FP6/FP4 input buffer A in shared memory.

  • B (tilelang._typing.BufferLikeType) – FP8/FP6/FP4 input buffer B in shared memory.

  • C (tilelang._typing.BufferLikeType) – Accumulator in tensor memory.

  • SFA_tmem (tilelang._typing.BufferLikeType) – Scale factors for A in tensor memory.

  • SFB_tmem (tilelang._typing.BufferLikeType) – Scale factors for B in tensor memory.

  • transpose_A (bool) – Whether A is MN-major. Default: False (K-major).

  • transpose_B (bool) – Whether B is K-major. Default: False (MN-major).

  • clear_accum (bool) – Whether to zero the accumulator.

  • wg_wait (int) – Warp group wait identifier.

  • mbar (tilelang.language.gemm_op.BarrierType | None) – Completion barrier, or None to defer the completion arrival.

  • k_start (int | tvm.tirx.PrimExpr) – Logical K-axis start offset for this MMA tile.

  • sf_a_granularity_k (int) – K elements covered by one A scale factor.

  • sf_b_granularity_k (int) – K elements covered by one B scale factor.

  • use_2cta (bool) – Whether to request true cta_group::2 lowering.

Return type:

tvm.tirx.PrimExpr

tilelang.cuda.language.gemm_op.mma_gemm_blockscaled(A, B, C, SFA, SFB, transpose_A=False, transpose_B=False, policy=GemmWarpPolicy.Square, clear_accum=False, *, k_start, sf_a_granularity_k, sf_b_granularity_k, sf_layout=None)

Explicit SM120 warp-level block-scaled MMA GEMM.

This is the explicit counterpart of T.gemm_blockscaled(…) for the SM120 warp-level path and follows the same scale-factor model: users pass the scale tensors, logical k_start, and K granularity, while the lowering derives the low-level scale addressing. Unlike TCGEN05, this path is synchronous warp-level mma.sync and does not use tensor memory or mbarriers, so C must be a fragment.

The current supported instruction is SM120 NVF4: m16n8k64.kind::mxf4nvf4.block_scale.scale_vec::4X with E2M1 operands, FP32 accumulation, and UE4M3 scale factors.

Parameters:

  • A (tilelang._typing.BufferLikeType)

  • B (tilelang._typing.BufferLikeType)

  • C (tilelang._typing.BufferLikeType)

  • SFA (tilelang._typing.BufferLikeType)

  • SFB (tilelang._typing.BufferLikeType)

  • transpose_A (bool)

  • transpose_B (bool)

  • policy (tilelang.language.gemm_op.GemmWarpPolicy)

  • clear_accum (bool)

  • k_start (int | tvm.tirx.PrimExpr)

  • sf_a_granularity_k (int)

  • sf_b_granularity_k (int)

  • sf_layout (str | None)

Return type:

tvm.tirx.PrimExpr

tilelang.cuda.language.gemm_op.make_blockscaled_gemm_layout(C, A, transpose_A=False)

Build the TMEM store layout for the C accumulator of a block-scaled GEMM.

Users must call T.annotate_layout({C_tmem: layout}) with the returned layout so that subsequent T.copy(C_tmem, ...) can be lowered correctly.

Parameters:

  • C (tilelang._typing.BufferLikeType) – The TMEM accumulator buffer (block_M, block_N).

  • A (tilelang._typing.BufferLikeType) – The FP8 operand A buffer (used to infer K and dtype).

  • transpose_A (bool) – Whether A is MN-major.

Returns:

A Layout object for C’s TMEM storage.

Return type:

tilelang.layout.Layout