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

grouped_matmul_rowwise_scaled_fp8_kernel

def grouped_matmul_rowwise_scaled_fp8_kernel[c_type: DType, a_type: DType, b_type: DType, a_scales_type: DType, b_scales_type: DType, accum_type: DType, CLayout: TensorLayout, ALayout: TensorLayout, BLayout: TensorLayout, AScalesLayout: TensorLayout, BScalesLayout: TensorLayout, AOffsetsLayout: TensorLayout, ExpertIdsLayout: TensorLayout, c_storage: TensorStorage, a_storage: TensorStorage, b_storage: TensorStorage, a_scales_storage: TensorStorage, b_scales_storage: TensorStorage, a_offsets_storage: TensorStorage, expert_ids_storage: TensorStorage, *, transpose_b: Bool = True, elementwise_lambda_fn: Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None] = None](c: TileTensor[c_type, CLayout, MutAnyOrigin, Storage=c_storage], a: TileTensor[a_type, ALayout, ImmUnsafeAnyOrigin, Storage=a_storage], b: TileTensor[b_type, BLayout, ImmUnsafeAnyOrigin, Storage=b_storage], a_scales: TileTensor[a_scales_type, AScalesLayout, ImmUnsafeAnyOrigin, Storage=a_scales_storage], b_scales: TileTensor[b_scales_type, BScalesLayout, ImmUnsafeAnyOrigin, Storage=b_scales_storage], a_offsets: TileTensor[DType.uint32, AOffsetsLayout, ImmUnsafeAnyOrigin, Storage=a_offsets_storage], expert_ids: TileTensor[DType.int32, ExpertIdsLayout, ImmUnsafeAnyOrigin, Storage=expert_ids_storage])

Computes the naive grouped FP8 matmul with rowwise weight scales and per-token activation scales, accumulating in fp32 and applying a single post-reduction scale.

For each token t in group g's row range and output channel n, computes out[t, n] = (sum_k a[t, k] * b[expert, n, k]) * a_scale[t] * b_scale[expert, n]. The a_scale is indexed by the global ragged row and b_scale by the real expert id, per the correctness invariants documented in the file header. Skips the matmul for expert == -1 (inactive LoRA blocks).