IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /max/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. /max/get-started.md).

Mojo struct

SM100TensorAccumulator

struct SM100TensorAccumulator[operand_type: DType, accum_dtype: DType, MMA_M: Int, MMA_N: Int, BK: Int, *, a_tmem: Bool, mma_kind: UMMAKind = UMMAKind.KIND_F16, swizzle_a: TensorMapSwizzle = TensorMapSwizzle.SWIZZLE_128B, swizzle_b: TensorMapSwizzle = TensorMapSwizzle.SWIZZLE_128B, transpose_b: Bool = True, cta_group: Int = Int(1), num_stages: Int = Int(1), b_page_dense: Bool = False, allow_3_then_1_split: Bool = True]

Performs the C = A @ B tensor contraction on SM100 using tcgen05.mma instructions.

The A operand is either an SMEM tile (a_tmem=False, the "SS" contraction) or a TMEM tile (a_tmem=True, the "TS" contraction); B is always an SMEM descriptor. When cta_group == 1 and MMA_M <= 64, the warp-specialized .ws datapath is used.

Parameters​

  • ​operand_type (DType): Element dtype of the A and B input operands.
  • ​accum_dtype (DType): Element dtype of the output accumulator C.
  • ​MMA_M (Int): M dimension (rows) of the output tile in MMA units.
  • ​MMA_N (Int): N dimension (columns) of the output tile in MMA units.
  • ​BK (Int): K dimension (contraction axis) tile size in elements.
  • ​a_tmem (Bool): Whether the A operand resides in TMEM (True, the TS contraction) or SMEM (False, the SS contraction).
  • ​mma_kind (UMMAKind): UMMAKind selecting the tcgen05.mma instruction variant (defaults to UMMAKind.KIND_F16).
  • ​swizzle_a (TensorMapSwizzle): SMEM swizzle mode for the A tile (defaults to SWIZZLE_128B); meaningful only when a_tmem is False.
  • ​swizzle_b (TensorMapSwizzle): SMEM swizzle mode for the B tile (defaults to SWIZZLE_128B).
  • ​transpose_b (Bool): Whether B is stored k-major (True) or mn-major (False) (defaults to True).
  • ​cta_group (Int): Number of cooperating CTAs, 1 or 2 (defaults to 1). When 1 and MMA_M <= 64, the .ws datapath is used.
  • ​num_stages (Int): Number of K-dimension pipeline stages for latency hiding (defaults to 1).
  • ​b_page_dense (Bool): Whether B uses the row-major page-fold layout (defaults to False).
  • ​allow_3_then_1_split (Bool): Whether a num_stages == 2 A-in-TMEM contraction with num_k_blocks % 4 == 0 may use the 3-then-1 K sub-stage split (defaults to True, the historical behavior). Set False by Layout-E's P@V reduction-split accumulator so the P sub-stage chunks stay an even 2-then-2 aligned with the V reduction chunks.

Implemented traits​

AnyType, Copyable, ImplicitlyCopyable, ImplicitlyDeletable, Movable, RegisterPassable, TrivialRegisterPassable

comptime members​

a_bmn​

comptime a_bmn = align_up((MMA_M // cta_group), Int(8))

a_layout​

comptime a_layout = tile_layout_k_major[SM100TensorAccumulator[operand_type, accum_dtype, MMA_M, MMA_N, BK, a_tmem=a_tmem, mma_kind=mma_kind, swizzle_a=swizzle_a, swizzle_b=swizzle_b, transpose_b=transpose_b, cta_group=cta_group, num_stages=num_stages, b_page_dense=b_page_dense, allow_3_then_1_split=allow_3_then_1_split].operand_t, align_up((MMA_M // cta_group), Int(8)), SM100TensorAccumulator[operand_type, accum_dtype, MMA_M, MMA_N, BK, a_tmem=a_tmem, mma_kind=mma_kind, swizzle_a=swizzle_a, swizzle_b=swizzle_b, transpose_b=transpose_b, cta_group=cta_group, num_stages=num_stages, b_page_dense=b_page_dense, allow_3_then_1_split=allow_3_then_1_split].padded_BK, swizzle_a]()

accum_t​

comptime accum_t = accum_dtype

AInput​

comptime AInput = UInt32 if a_tmem else MMASmemDescriptorPair

AType​

comptime AType = TMemTile[operand_type, MMA_M, BK] if a_tmem else MMASmemDescriptorPair

b_bmn​

comptime b_bmn = MMA_N if a_tmem and not transpose_b else (MMA_N // cta_group)

b_layout​

comptime b_layout = tile_layout_k_major[SM100TensorAccumulator[operand_type, accum_dtype, MMA_M, MMA_N, BK, a_tmem=a_tmem, mma_kind=mma_kind, swizzle_a=swizzle_a, swizzle_b=swizzle_b, transpose_b=transpose_b, cta_group=cta_group, num_stages=num_stages, b_page_dense=b_page_dense, allow_3_then_1_split=allow_3_then_1_split].operand_t, MMA_N if a_tmem and (xor transpose_b, True) else (MMA_N // cta_group), SM100TensorAccumulator[operand_type, accum_dtype, MMA_M, MMA_N, BK, a_tmem=a_tmem, mma_kind=mma_kind, swizzle_a=swizzle_a, swizzle_b=swizzle_b, transpose_b=transpose_b, cta_group=cta_group, num_stages=num_stages, b_page_dense=b_page_dense, allow_3_then_1_split=allow_3_then_1_split].padded_BK, swizzle_b, b_page_dense]() if transpose_b else tile_layout_mn_major[SM100TensorAccumulator[operand_type, accum_dtype, MMA_M, MMA_N, BK, a_tmem=a_tmem, mma_kind=mma_kind, swizzle_a=swizzle_a, swizzle_b=swizzle_b, transpose_b=transpose_b, cta_group=cta_group, num_stages=num_stages, b_page_dense=b_page_dense, allow_3_then_1_split=allow_3_then_1_split].operand_t, MMA_N if a_tmem and (xor transpose_b, True) else (MMA_N // cta_group), SM100TensorAccumulator[operand_type, accum_dtype, MMA_M, MMA_N, BK, a_tmem=a_tmem, mma_kind=mma_kind, swizzle_a=swizzle_a, swizzle_b=swizzle_b, transpose_b=transpose_b, cta_group=cta_group, num_stages=num_stages, b_page_dense=b_page_dense, allow_3_then_1_split=allow_3_then_1_split].padded_BK, swizzle_b, b_page_dense]()

BType​

comptime BType = MMASmemDescriptorPair

CType​

comptime CType = TMemTile[SM100TensorAccumulator[operand_type, accum_dtype, MMA_M, MMA_N, BK, a_tmem=a_tmem, mma_kind=mma_kind, swizzle_a=swizzle_a, swizzle_b=swizzle_b, transpose_b=transpose_b, cta_group=cta_group, num_stages=num_stages, b_page_dense=b_page_dense, allow_3_then_1_split=allow_3_then_1_split].accum_t, MMA_M, MMA_N]

idesc​

comptime idesc = UMMAInsDescriptor.create[SM100TensorAccumulator[operand_type, accum_dtype, MMA_M, MMA_N, BK, a_tmem=a_tmem, mma_kind=mma_kind, swizzle_a=swizzle_a, swizzle_b=swizzle_b, transpose_b=transpose_b, cta_group=cta_group, num_stages=num_stages, b_page_dense=b_page_dense, allow_3_then_1_split=allow_3_then_1_split].accum_t, SM100TensorAccumulator[operand_type, accum_dtype, MMA_M, MMA_N, BK, a_tmem=a_tmem, mma_kind=mma_kind, swizzle_a=swizzle_a, swizzle_b=swizzle_b, transpose_b=transpose_b, cta_group=cta_group, num_stages=num_stages, b_page_dense=b_page_dense, allow_3_then_1_split=allow_3_then_1_split].operand_t, SM100TensorAccumulator[operand_type, accum_dtype, MMA_M, MMA_N, BK, a_tmem=a_tmem, mma_kind=mma_kind, swizzle_a=swizzle_a, swizzle_b=swizzle_b, transpose_b=transpose_b, cta_group=cta_group, num_stages=num_stages, b_page_dense=b_page_dense, allow_3_then_1_split=allow_3_then_1_split].operand_t, Index[Int, Int, dtype=DType.uint32](MMA_M, MMA_N), transpose_b=transpose_b]()

MMA_K​

comptime MMA_K = Int(16) if operand_type.is_half_float() else Int(32)

num_k_blocks​

comptime num_k_blocks = (SM100TensorAccumulator[operand_type, accum_dtype, MMA_M, MMA_N, BK, a_tmem=a_tmem, mma_kind=mma_kind, swizzle_a=swizzle_a, swizzle_b=swizzle_b, transpose_b=transpose_b, cta_group=cta_group, num_stages=num_stages, b_page_dense=b_page_dense, allow_3_then_1_split=allow_3_then_1_split].padded_BK // Int(16) if (eq #pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<@std::@builtin::@dtype::@DType> operand_type, "_mlir_value">>, 80) if (eq #pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<@std::@builtin::@dtype::@DType> operand_type, "_mlir_value">>, 80) else (eq #pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<@std::@builtin::@dtype::@DType> operand_type, "_mlir_value">>, 79) else Int(32))

num_k_blocks_per_stage​

comptime num_k_blocks_per_stage = ((BK if a_tmem else align_up(BK, (swizzle_b.bytes() if a_tmem else max(swizzle_a.bytes(), swizzle_b.bytes()) // size_of[operand_type]())) // Int(16) if (eq #pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<@std::@builtin::@dtype::@DType> operand_type, "_mlir_value">>, 80) if (eq #pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<@std::@builtin::@dtype::@DType> operand_type, "_mlir_value">>, 80) else (eq #pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<@std::@builtin::@dtype::@DType> operand_type, "_mlir_value">>, 79) else Int(32)) // Int(4) if SM100TensorAccumulator[operand_type, accum_dtype, MMA_M, MMA_N, BK, a_tmem=a_tmem, mma_kind=mma_kind, swizzle_a=swizzle_a, swizzle_b=swizzle_b, transpose_b=transpose_b, cta_group=cta_group, num_stages=num_stages, b_page_dense=b_page_dense, allow_3_then_1_split=allow_3_then_1_split].use_3_then_1_split else num_stages)

num_k_mmas​

comptime num_k_mmas = (BK // Int(16) if (eq #pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<@std::@builtin::@dtype::@DType> operand_type, "_mlir_value">>, 80) if (eq #pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<@std::@builtin::@dtype::@DType> operand_type, "_mlir_value">>, 80) else (eq #pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<@std::@builtin::@dtype::@DType> operand_type, "_mlir_value">>, 79) else Int(32)) if a_tmem else ceildiv(BK, Int(16) if (eq #pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<@std::@builtin::@dtype::@DType> operand_type, "_mlir_value">>, 80) if (eq #pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<@std::@builtin::@dtype::@DType> operand_type, "_mlir_value">>, 80) else (eq #pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<@std::@builtin::@dtype::@DType> operand_type, "_mlir_value">>, 79) else Int(32))

operand_size​

comptime operand_size = size_of[SM100TensorAccumulator[operand_type, accum_dtype, MMA_M, MMA_N, BK, a_tmem=a_tmem, mma_kind=mma_kind, swizzle_a=swizzle_a, swizzle_b=swizzle_b, transpose_b=transpose_b, cta_group=cta_group, num_stages=num_stages, b_page_dense=b_page_dense, allow_3_then_1_split=allow_3_then_1_split].operand_t]()

operand_t​

comptime operand_t = operand_type

padded_BK​

comptime padded_BK = BK if a_tmem else align_up(BK, (swizzle_b.bytes() if a_tmem else max(swizzle_a.bytes(), swizzle_b.bytes()) // size_of[operand_type]()))

swizzle_granularity​

comptime swizzle_granularity = (swizzle_b.bytes() if a_tmem else max(swizzle_a.bytes(), swizzle_b.bytes()) // size_of[operand_type]())

tcgen05_mma_type​

comptime tcgen05_mma_type = "tcgen05.mma.ws.cta_group::1."

use_3_then_1_split​

comptime use_3_then_1_split = allow_3_then_1_split and a_tmem and (num_stages == Int(2)) and (((BK if a_tmem else align_up(BK, (swizzle_b.bytes() if a_tmem else max(swizzle_a.bytes(), swizzle_b.bytes()) // size_of[operand_type]())) // Int(16) if (eq #pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<@std::@builtin::@dtype::@DType> operand_type, "_mlir_value">>, 80) if (eq #pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<@std::@builtin::@dtype::@DType> operand_type, "_mlir_value">>, 80) else (eq #pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<@std::@builtin::@dtype::@DType> operand_type, "_mlir_value">>, 79) else Int(32)) % Int(4)) == Int(0))

use_ws​

comptime use_ws = (MMA_M <= Int(64)) if (cta_group == Int(1)) else (cta_group == Int(1))

Methods​

mma​

static def mma[*, stage_idx: Int = Int(0)](a: Self.AInput, b: MMASmemDescriptorPair, c: UInt32, *, c_scale: UInt32, elect: Int32)

mma_maybe_partial_k​

static def mma_maybe_partial_k[*, stage_idx: Int = Int(0)](a: Self.AInput, b: MMASmemDescriptorPair, c: UInt32, *, c_scale: UInt32, elect: Int32, valid_k_mmas: UInt32)

Edit this page

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community

Contact