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 struct

TensorCoreAsync

struct TensorCoreAsync[c_type: DType, a_type: DType, b_type: DType, mma_shape: IndexList[Int(3)], /, a_swizzle: TensorMapSwizzle = TensorMapSwizzle.SWIZZLE_NONE, b_swizzle: TensorMapSwizzle = TensorMapSwizzle.SWIZZLE_NONE, transpose_b: Bool = False]

High-performance asynchronous tensor core operations for matrix multiplication.

This struct provides methods for utilizing NVIDIA's Tensor Cores for asynchronous matrix multiplication operations, with support for various data types and swizzling configurations.

Parameters​

  • ​c_type (DType): Data type of the output matrix C.
  • ​a_type (DType): Data type of the input matrix A.
  • ​b_type (DType): Data type of the input matrix B.
  • ​mma_shape (IndexList[Int(3)]): Dimensions for the matrix multiply-accumulate (MMA) operation as [M, N, K].
  • ​a_swizzle (TensorMapSwizzle): Swizzling mode for matrix A (default: SWIZZLE_NONE).
  • ​b_swizzle (TensorMapSwizzle): Swizzling mode for matrix B (default: SWIZZLE_NONE).
  • ​transpose_b (Bool): Whether to transpose matrix B (default: False).

Implemented traits​

AnyType, Defaultable, ImplicitlyDeletable

Methods​

__init__​

def __init__(out self)

Initialize the TensorCoreAsync instance.

Ensures that the provided MMA shape is supported.

Note: Fails to compile if mma_shape is not supported.

wgmma​

static def wgmma[num_warp_groups: Int = Int(1), scale_c: Int = Int(1), scale_a: Int = Int(1), scale_b: Int = Int(1), num_k_iters: Optional[Int] = None](a_smem_tile: LayoutTensor[a_type, address_space=AddressSpace.SHARED, element_layout=a_smem_tile.element_layout, layout_int_type=a_smem_tile.layout_int_type, linear_idx_type=a_smem_tile.linear_idx_type, masked=a_smem_tile.masked, alignment=a_smem_tile.alignment], b_smem_tile: LayoutTensor[b_type, address_space=AddressSpace.SHARED, element_layout=b_smem_tile.element_layout, layout_int_type=b_smem_tile.layout_int_type, linear_idx_type=b_smem_tile.linear_idx_type, masked=b_smem_tile.masked, alignment=b_smem_tile.alignment], c_reg_tile: LayoutTensor[c_type, address_space=AddressSpace.LOCAL, element_layout=c_reg_tile.element_layout, layout_int_type=c_reg_tile.layout_int_type, linear_idx_type=c_reg_tile.linear_idx_type, masked=c_reg_tile.masked, alignment=c_reg_tile.alignment], wg_idx: Int = Int(0))

Perform asynchronous matrix multiplication using warp group matrix multiply-accumulate (WGMMA).

This method handles the case where both A and B matrices are in shared memory.

Parameters:

  • ​num_warp_groups (Int): Number of warp groups to distribute work across (default: 1).
  • ​scale_c (Int): Scale factor for matrix C. Valid values are 1 or 0 (default: 1).
  • ​scale_a (Int): Scale factor for matrix A. Valid values are 1 or -1 (default: 1).
  • ​scale_b (Int): Scale factor for matrix B. Valid values are 1 or -1 (default: 1).
  • ​num_k_iters (Optional[Int]): Number of iterations for the K dimension. This is useful to save computation when we pad shared memory. (default: None which is just a_smem_layout[1].size() // mma_shape[2]).

Args:

static def wgmma[num_warp_groups: Int = Int(1), scale_c: Int = Int(1), scale_a: Int = Int(1), scale_b: Int = Int(1), num_k_iters: Optional[Int] = None](a_smem_tile: TileTensor[a_type, Storage=a_smem_tile.Storage, address_space=AddressSpace.SHARED, linear_idx_type=a_smem_tile.linear_idx_type], b_smem_tile: TileTensor[b_type, Storage=b_smem_tile.Storage, address_space=AddressSpace.SHARED, linear_idx_type=b_smem_tile.linear_idx_type], c_reg_tile: LayoutTensor[c_type, address_space=AddressSpace.LOCAL, element_layout=c_reg_tile.element_layout, layout_int_type=c_reg_tile.layout_int_type, linear_idx_type=c_reg_tile.linear_idx_type, masked=c_reg_tile.masked, alignment=c_reg_tile.alignment], wg_idx: Int = Int(0))

Perform asynchronous matrix multiplication with TileTensor inputs.

This overload handles the case where both A and B matrices are in shared memory as TileTensor values. It converts at this boundary so SM90 kernels can stay TileTensor-native while the WGMMA implementation continues to reuse the legacy descriptor path internally.

Parameters:

  • ​num_warp_groups (Int): Number of warp groups to distribute work across.
  • ​scale_c (Int): Scale factor for matrix C. Valid values are 1 or 0.
  • ​scale_a (Int): Scale factor for matrix A. Valid values are 1 or -1.
  • ​scale_b (Int): Scale factor for matrix B. Valid values are 1 or -1.
  • ​num_k_iters (Optional[Int]): Number of K-dimension iterations to execute.

Args:

static def wgmma(a_frag_tile: LayoutTensor[a_type, address_space=AddressSpace.LOCAL, element_layout=a_frag_tile.element_layout, layout_int_type=a_frag_tile.layout_int_type, linear_idx_type=a_frag_tile.linear_idx_type, masked=a_frag_tile.masked, alignment=a_frag_tile.alignment], b_smem_tile: LayoutTensor[b_type, address_space=AddressSpace.SHARED, element_layout=b_smem_tile.element_layout, layout_int_type=b_smem_tile.layout_int_type, linear_idx_type=b_smem_tile.linear_idx_type, masked=b_smem_tile.masked, alignment=b_smem_tile.alignment], c_reg_tile: LayoutTensor[c_type, address_space=AddressSpace.LOCAL, element_layout=c_reg_tile.element_layout, layout_int_type=c_reg_tile.layout_int_type, linear_idx_type=c_reg_tile.linear_idx_type, masked=c_reg_tile.masked, alignment=c_reg_tile.alignment])

Perform asynchronous matrix multiplication using warp group matrix multiply-accumulate (WGMMA).

This overloaded method handles the case where matrix A is in register memory and matrix B is in shared memory.

Args:

arrive​

static def arrive()

Ensures memory consistency by creating a fence for WGMMA operations.

This method should be called before committing a group to ensure all shared memory accesses are properly aligned and visible.

commit_group​

static def commit_group()

Commits the current warp group for execution.

This synchronizes the warp group and commits all pending WGMMA operations that have been previously issued.

wait_group​

static def wait_group[group: Int = Int(0)]()

Waits for the completion of a specific warp group's operations.

This method blocks until all WGMMA operations from the specified group are complete.

Parameters:

  • ​group (Int): The group ID to wait for (default: 0).

View source

View source

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community

Contact