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 module
matmul2d_fp8
Apple M5 weight-only FP8 (W8A16) tiled matmul: bf16 A x fp8 B -> fp32.
Target: Apple M5 (compute_capability == 5) only.
The DENSE W8A16 path (enqueue_matmul2d_fp8) now runs on the shared
AppleM5MatMul simdgroup-tiled GEMM body (matmul_kernel.mojo) specialized on
b_type=float8_e4m3fn -- the weight-loader seam swap -- so it inherits the dense
body's Morton schedule and epilogue with only the B operand widened. What remains
here is the GROUPED (MoE) Matmul2dFp8.run_grouped + _gemm_body, pending its
own migration onto the shared body's group-indexed wrapper.
out = activation @ weight^T where the FP8 weight stays float8_e4m3fn [N, K]
in DRAM (one byte per element -- NOT packed, NO block scales). This is the
tiled-MMA decode kernel that beats the warp-per-column GEMV's structural ceiling
at large-N / small-K (Mamba in_proj [17504, 3136], MLP up [12544, 3136]),
where the GEMV is capped ~250 GB/s while a tiled matmul amortizes the per-output
cost across a threadgroup tile.
Why this structure (NOT the FP4 matmul2d transpose_right / SMEM path)
FP4 (matmul2d_fp4.mojo) decodes through SMEM because its dequant (nibble unpack
- E2M1 LUT + per-16 block scale) is expensive and its per-lane B gather is
N-scattered -- coalescing it through SMEM took it 10 -> 38 TF/s (KB
apple-m5-gpu-perf-model). FP8 has NO such decode: the widen is native. Crucially, the Apple simdgroup MMA (_mma_apple_transposable,mma_apple.mojo:213) accepts an FP8 B operand DIRECTLY (float8_e4m3fnis in its valid-float input set; KGEN lowers the fp8 fragment to AIR's<8 x i8>), so a bf16 A x fp8 B -> fp32 MMA is a single native instruction -- no manual widen, no SMEM staging.
So this kernel mirrors the WINNING dense bf16 structure (AppleM5MatMul), which
sustains ~400 GB/s at exactly this M=1 large-N/small-K decode regime: a
simdgroup-tiled GEMM with MmaOpApple, the B operand loaded DIRECTLY from DRAM
each K-strip (Apple has no async copy and SMEM staging DEGRADES matmul -- KB
apple-m5-gpu-performance-considerations), coalesced via the col-major (K, N)
weight view (K-contiguous fast axis). The ONLY difference from the dense bf16
path is MmaOpApple's b_type = float8_e4m3fn (1 byte/weight vs 2) -- the
"weight-load seam" swap. MmaOpApple already carries a separate b_type and
loads A / B fragments independently, so this needs NO change to the shared MMA op
or the bf16 matmul.
The per-tensor scalar weight_scale folds in OUTSIDE the kernel (a post-matmul
multiply by the graph lowering), identically to the GEMV: a scalar factors out of
the fp32 sum, so the fold is EXACT.
Ragged M/N/K are handled by MmaOpApple's bounded load (zero-fill OOB) + the
< M/< N store guard -- no tile-alignment requirement (unlike the FP4
interior-only matmul2d). The 2x2 (SG=32) simdgroup subtile is the M5
register-cliff optimum (KB apple-m5-gpu-perf-model); do NOT exceed 4
accumulators.
Structs
-
Matmul2dFp8: W8A16 simdgroup-tiled GEMM: bf16 A x fp8 B -> fp32, direct DRAM B feed.
Functions
-
enqueue_grouped_matmul2d_fp8: Enqueue the tiled grouped (MoE) W8A16 GEMM (transpose_b), raw/unscaled. -
enqueue_matmul2d_fp8: Enqueue the tiled W8A16 GEMM:out = a @ W_fp8^T(raw, unscaled).
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!