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_grouped_matmul2d_fp8

def enqueue_grouped_matmul2d_fp8[c_type: DType = DType.float32, 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)](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], b: TileTensor[DType.float8_e4m3fn, Storage=b.Storage, address_space=b.address_space, linear_idx_type=b.linear_idx_type], a_offsets: TileTensor[DType.uint32, Storage=a_offsets.Storage, address_space=a_offsets.address_space, linear_idx_type=a_offsets.linear_idx_type], expert_ids: TileTensor[DType.int32, Storage=expert_ids.Storage, address_space=expert_ids.address_space, linear_idx_type=expert_ids.linear_idx_type], max_num_tokens_per_expert: Int, num_active_experts: Int, ctx: DeviceContext)

Enqueue the tiled grouped (MoE) W8A16 GEMM (transpose_b), raw/unscaled.

The FP8 analog of naive_grouped_matmul on the tiled simdgroup MMA: ONE kernel dispatch handles all expert groups (grid block_idx.z selects the group), so there is no host-side per-expert launch loop. a is the ragged bf16 activation (total_M, K), b the FP8-E4M3 weight stack (num_experts, N, K), c the ragged output (total_M, N). Each group z reads its expert weight slab b[expert_ids[z]] and its token rows [a_offsets[z], a_offsets[z+1]). bf16 A x fp8 B -> fp32 on the native Apple MMA, direct DRAM B feed (no SMEM). The per-expert weight_scale folds post-matmul (graph lowering), matching the dense Apple FP8 Linear.

Raises:

If the attached GPU is not Apple M5 (compute_capability == 5).