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 function

enqueue_apple_fp4_matmul

def enqueue_apple_fp4_matmul[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, *, use_matmul2d: Bool = False](c: TileTensor[c_type, Storage=c.Storage, address_space=c.address_space, linear_idx_type=c.linear_idx_type], a: TileTensor[DType.bfloat16, Storage=a.Storage, address_space=a.address_space, linear_idx_type=a.linear_idx_type], packed: TileTensor[DType.uint8, Storage=packed.Storage, address_space=packed.address_space, linear_idx_type=packed.linear_idx_type], scales: TileTensor[DType.float8_e4m3fn, Storage=scales.Storage, address_space=scales.address_space, linear_idx_type=scales.linear_idx_type], ctx: DeviceContext)

Enqueue the W4A16 matmul: out = a @ dequant(packed, scales)^T.

a is the bf16 activation (M, K), packed the FP4 weight (N, K//2) (lo-nibble first), scales the FP8-E4M3 block scales (N, ceil(K/16)). C is (M, N).

The default strategy (use_matmul2d=False) is chosen by (K, M) -- all paths produce bit-identical dequant (the E2M1 * |scale| arithmetic is the same on every path; the GEMV differs only by its fp32 reduction order vs the MMA):

  • Batch-1 decode (M == 1): the register-resident W4A16 GEMV (enqueue_apple_fp4_gemv), no MMA to feed. See fp4_gemv.mojo for the rationale (bandwidth win, ~1.53x at the Llama-8B down-proj shape).
  • Deep-K niche (K >= 18432 AND M >= 1024, aligned interior): the matmul2d W4A16 kernel (enqueue_matmul2d_fp4_smem). At deep K the MAT->DENSE path below must read a large (>=226 MB) materialized bf16 weight and hits the M5 DRAM wall, while matmul2d stays 4-bit-in-DRAM and holds ~42 TF/s -- measured 1.02-1.31x over the walled MAT at M in {1024,2048,4096} (thermal-fair A/B). This is the real FLUX.2-dev FFN-down (N=6144, K=18432) at prefill. Only when there is no fused epilogue (the matmul2d interior store does not yet apply the lambda).
  • Large M (M >= 1536): MATERIALIZE the weight to a transient (N, K) bf16 buffer (enqueue_fp4_materialize), then run the stock dense bf16 MMA (enqueue_apple_matmul). The FP4 decode is off the MMA critical path, so this beats the fused path ~1.25-1.34x on M5 (which serializes the scalar-ALU decode against the matrix MMA). The transient buffer is execution-time only -- it is NOT a persistent graph constant, so the packed weight stays 4-bit in DRAM at rest. The 1536 crossover (vs a naive ~512) accounts for the per-call buffer alloc/free the launcher pays.
  • Mid M (256 <= M < 1536): the FUSED cooperative-SMEM path with BM=128 (deep dequant amortization across 8 co-resident simdgroups). Faster than materialize here once the per-call alloc is paid.
  • Small M (M < 256): the FUSED path with BM=64 (keeps occupancy up when a 128-tall tile would waste rows).

In the fused bands each (BN, BK) weight sub-tile is cooperatively dequantized to bf16 in threadgroup memory once per K-strip, then the dense bf16 MMA reads B from SMEM; weights stay packed in DRAM the whole time.

On pre-M5 Apple silicon GPUs the M5-only fused / matmul2d paths are skipped: the call routes unconditionally to materialize->dense (FP4 decode to a transient bf16 buffer, then the pre-M5 dense bf16 gemm_kernel_apple_8x8).

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), threaded to whichever path runs.
  • ​use_matmul2d (Bool): OPT-IN OVERRIDE. When True, route EVERY aligned interior (N % _M2D_TG_N=32 == 0, K % _M2D_SMEM_BK=256 == 0) to the native matmul2d-tile W4A16 kernel (enqueue_matmul2d_fp4_smem), falling back to the incumbent only off that aligned interior. DEFAULT False -- off the deep-K niche the P1 crossover measured matmul2d dominated at every production M by the incumbent bands (FUSED at M<=512, MAT->DENSE at mid-K M>=1024), so it is not forced in general; the flag keeps it reachable for A/B and future geometry. NOTE this override is ORTHOGONAL to the deep-K niche below, which routes to matmul2d BY DEFAULT (use_matmul2d=False) where it actually wins. See the _M2D_* threshold comments.

Args:

Edit this page

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community

Contact