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

Matmul2dFp8

struct Matmul2dFp8[c_type: DType = DType.float32, elementwise_lambda_fn: Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None] = None, block_m: Int = Int(64), block_n: Int = Int(64), block_k: Int = Int(16), sg_m: Int = Int(32), sg_n: Int = Int(32)]

W8A16 simdgroup-tiled GEMM: bf16 A x fp8 B -> fp32, direct DRAM B feed.

Mirrors AppleM5MatMul (the ~400 GB/s dense bf16 decode path) with the sole change that the B (weight) operand is float8_e4m3fn via MmaOpApple's b_type. A is bf16, C is c_type, accumulation fp32. No SMEM, no barrier: each K-strip's B fragments load straight from DRAM (coalesced K-contiguous via the col-major (K, N) view) and feed the native fp8-capable MMA.

Parameters​

  • ​c_type (DType): Output element type (fp16, bf16, fp32). Accumulation is fp32.
  • ​elementwise_lambda_fn (Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None]): Optional fused epilogue (AMD's (row, col) contract). Currently unused on the interior store; the launcher routes fused-epilogue shapes elsewhere.
  • ​block_m (Int): Threadgroup block rows BM (multiple of SG_M).
  • ​block_n (Int): Threadgroup block cols BN (multiple of SG_N).
  • ​block_k (Int): K-strip depth BK per accumulate step (multiple of 16).
  • ​sg_m (Int): Simdgroup subtile rows SG_M (multiple of 16).
  • ​sg_n (Int): Simdgroup subtile cols SG_N (multiple of 16). sg_m/16 * sg_n/16 is the accumulator count; 2x2 == 4 is the M5 optimum (do NOT exceed).

Implemented traits​

AnyType, ImplicitlyDeletable

comptime members​

b_type​

comptime b_type = DType.float8_e4m3fn

BK​

comptime BK = block_k

BM​

comptime BM = block_m

BN​

comptime BN = block_n

in_type​

comptime in_type = DType.bfloat16

Mma​

comptime Mma = MmaOpApple[DType.float32, DType.bfloat16, Matmul2dFp8[c_type, elementwise_lambda_fn, block_m, block_n, block_k, sg_m, sg_n].NUM_MMA_M, Matmul2dFp8[c_type, elementwise_lambda_fn, block_m, block_n, block_k, sg_m, sg_n].NUM_MMA_N, b_type=DType.float8_e4m3fn]

NUM_MMA_M​

comptime NUM_MMA_M = (Matmul2dFp8[c_type, elementwise_lambda_fn, block_m, block_n, block_k, sg_m, sg_n].SG_M // Int(16))

NUM_MMA_N​

comptime NUM_MMA_N = (Matmul2dFp8[c_type, elementwise_lambda_fn, block_m, block_n, block_k, sg_m, sg_n].SG_N // Int(16))

NUM_SG​

comptime NUM_SG = (Matmul2dFp8[c_type, elementwise_lambda_fn, block_m, block_n, block_k, sg_m, sg_n].NUM_SG_M * Matmul2dFp8[c_type, elementwise_lambda_fn, block_m, block_n, block_k, sg_m, sg_n].NUM_SG_N)

NUM_SG_M​

comptime NUM_SG_M = (Matmul2dFp8[c_type, elementwise_lambda_fn, block_m, block_n, block_k, sg_m, sg_n].BM // Matmul2dFp8[c_type, elementwise_lambda_fn, block_m, block_n, block_k, sg_m, sg_n].SG_M)

NUM_SG_N​

comptime NUM_SG_N = (Matmul2dFp8[c_type, elementwise_lambda_fn, block_m, block_n, block_k, sg_m, sg_n].BN // Matmul2dFp8[c_type, elementwise_lambda_fn, block_m, block_n, block_k, sg_m, sg_n].SG_N)

SG_M​

comptime SG_M = sg_m

SG_N​

comptime SG_N = sg_n

THREADS_PER_BLOCK​

comptime THREADS_PER_BLOCK = (Matmul2dFp8[c_type, elementwise_lambda_fn, block_m, block_n, block_k, sg_m, sg_n].NUM_SG * _resolve_warp_size())

Methods​

run_grouped​

static def run_grouped[c_layout: TensorLayout, a_layout: TensorLayout, b_layout: TensorLayout, ao_layout: TensorLayout, ei_layout: TensorLayout](c: TileTensor[c_type, c_layout, MutAnyOrigin], a: TileTensor[DType.bfloat16, a_layout, ImmutAnyOrigin], b: TileTensor[DType.float8_e4m3fn, b_layout, ImmutAnyOrigin], a_offsets: TileTensor[DType.uint32, ao_layout, ImmUnsafeAnyOrigin], expert_ids: TileTensor[DType.int32, ei_layout, ImmUnsafeAnyOrigin], N: Int, K: Int)

Grouped (MoE) W8A16 kernel entry: one expert group per block_idx.z.

Computes, for group z = block_idx.z, C[a_offsets[z]:a_offsets[z+1], :] = A[a_offsets[z]:a_offsets[z+1], :] @ b[expert_ids[z], :, :]^T i.e. the FP8 analog of naive_grouped_matmul_kernel but on the tiled simdgroup MMA. c is the ragged output (total_M, N), a the ragged bf16 activation (total_M, K), b the FP8 weight stack (num_experts, N, K). Produces the RAW (unscaled) matmul; the per-expert weight_scale is folded post-matmul by the graph lowering, exactly as on the dense Apple FP8 Linear (a per-tensor scalar factors out of the sum).

Grid (ceil(N/BN), ceil(max_tokens_per_expert/BM), num_active_experts). Per-group ragged M is a_offsets[z+1] - a_offsets[z]; empty groups (M == 0) and fully-OOB M-tiles early-return in _gemm_body. An inactive LoRA group (expert_ids[z] == -1) writes zeros for its rows.

The per-group base views use a base-pointer offset (a.ptr + a_start*K, b.ptr + expert*N*K) -- the ragged-group addressing idiom shared by every grouped matmul kernel in the tree (naive_grouped_matmul_kernel, grouped_matmul_amd_kernel_launcher); TileTensor has no runtime-offset ragged-slice primitive that keeps the inner (., K) dims static. All element access inside _gemm_body is TileTensor-indexed (MMA fragment loaders + guarded c.store).

Edit this page

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community

Contact