tilelang.language.gemm_op

GEMM (General Matrix Multiplication) operators exposed on the TileLang language surface.

Functions

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

TileLang GEMM operator.

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

Target-neutral block-scaled GEMM: C (+)= (A * SFA) @ (B * SFB).

Module Contents

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

TileLang GEMM operator.

This is the default synchronous GEMM interface. 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.

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

Return type:

tvm.tirx.PrimExpr

Backend dialects extend this signature with their hardware’s knobs: tilelang.cuda.language.gemm adds mbar (Blackwell TCGEN5MMA barrier), tilelang.rocm.language.gemm adds k_pack (packed MFMA).

Returns:

A handle to the GEMM operation.

Return type:

tirx.Call

Parameters:

  • A (tilelang._typing.BufferLikeType)

  • B (tilelang._typing.BufferLikeType)

  • C (tilelang._typing.BufferLikeType)

  • transpose_A (bool)

  • transpose_B (bool)

  • policy (tilelang.tileop.base.GemmWarpPolicy)

  • clear_accum (bool)

  • annotations (dict | None)

tilelang.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, annotations=None)

Target-neutral block-scaled GEMM: C (+)= (A * SFA) @ (B * SFB).

Scale factors apply to blocks along the reduction axis: k_start is the logical K-axis start of this tile and sf_*_granularity_k gives the number of K elements covered by one scale factor. The backend owns the supported dtypes, operand scopes, scale representation and lowering. Compilation fails when the backend has no block-scaled implementation; lowering to an unscaled GEMM would change the result.

Like T.gemm(…), this is the synchronous interface: the result is complete when the call returns. On Blackwell TCGEN5MMA, TileLang inserts the corresponding mbarrier_wait_parity(…) implicitly after issue, so that path needs a completion barrier, which the CUDA dialect accepts as mbar (it also adds use_2cta and sf_layout). For manual asynchronous scheduling use T.tcgen05_gemm_blockscaled(…).

Parameters:

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

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

  • C (tilelang._typing.BufferLikeType) – Accumulator tile.

  • 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.tileop.base.GemmWarpPolicy) – GEMM warp partition policy.

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

  • k_start (int | tvm.tirx.PrimExpr) – Logical K-axis start offset for 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.

  • annotations (dict | None) – Additional annotations.

Return type:

tvm.tirx.PrimExpr