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

DenseALoader

struct DenseALoader[dtype: DType, a_layout: TensorLayout, BK: Int = Int(16), SG_M: Int = Int(32), use_x2: Bool = False, b_dtype: DType = dtype]

Plain-GEMM A loader: holds the pre-tiled (SG_M, K) slab.

The run wrapper bakes the simdgroup's slab base once (see its docstring for the hoist rationale); the hot-loop .tile[SG_M, BK](0, k_strip) here only adds this strip's k_strip * BK K offset. The strip sub-tiles are type-identical to the per-strip form, so MmaOpApple.mma is unchanged.

BK/SG_M are explicit params (not re-derived from a fresh default AppleM5MatMul[dtype]) so a non-default block_k/sg_m -- e.g. use_x2's block_k=32 -- isn't silently ignored; run passes Self.BK/Self.SG_M.

b_dtype sets the B/weight operand fed to the MMA (defaults to dtype, the dense case). The W8A16 FP8 path constructs this loader with b_dtype=fp8: the A slab stays bf16 and only b_sub / mma_op widen to the native-fp8 MMA (KGEN lowers the fp8 B fragment to AIR <8 x i8>); the K-loop and accumulation order are byte-for-byte the dense path.

Lifetime: the a_slab view is held with UntrackedOrigin (struct fields cannot expose an AnyOrigin); the kernel arg it derives from outlives the K-loop, so this is the explicit-lifetime case the field-origin rule allows.

Parameters​

  • ​dtype (DType): Element type of the A operand (fp16, bf16, fp32).
  • ​a_layout (TensorLayout): TensorLayout of the pre-tiled A slab held by the loader.
  • ​BK (Int): K-strip depth per accumulate step; must match the body's BK tiling (defaults to 16).
  • ​SG_M (Int): Simdgroup subtile rows SG_M; the .tile[SG_M, BK] extent used per strip (defaults to 32).
  • ​use_x2 (Bool): Use the NT "double-strip" dense MMA (mma_dense_x2, block_k=32) instead of the single-strip mma (defaults to False).
  • ​b_dtype (DType): Element type of the B/weight operand fed to the MMA (defaults to dtype, the dense case).

Fields​

  • ​a_slab (TileTensor[dtype, a_layout, ImmUntrackedOrigin]):

Implemented traits​

AOperandLoader, AnyType, Copyable, ImplicitlyCopyable, ImplicitlyDeletable, Movable

comptime members​

b_type​

comptime b_type = b_dtype

in_type​

comptime in_type = dtype

num_m_mmas​

comptime num_m_mmas = 2

num_n_mmas​

comptime num_n_mmas = 2

Methods​

accumulate_strip​

def accumulate_strip[bounded: Bool, b_layout: TensorLayout](mut self, mut mma_op: MmaOpApple[DType.float32, Self.in_type, Int(2), Int(2), b_type=Self.b_type], mut accum: Array[SIMD[DType.float32, SIMDLength(8)], MmaOpApple[DType.float32, Self.in_type, Int(2), Int(2), b_type=Self.b_type].num_accum], b_sub: TileTensor[Self.b_type, b_layout, ImmutAnyOrigin], conv: ConvIm2colParams, k_strip: Int32, *, a_valid_rows: Int32, b_valid_cols: Int, k_valid: Int)

Edit this page

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community

Contact