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_block_scaled_mxfp4_kernel

def grouped_matmul_block_scaled_mxfp4_kernel[out_dtype: DType, BLOCK_N: Int](c_ptr: Pointer[Scalar[out_dtype], MutAnyOrigin, _safe=False], a_ptr: Pointer[UInt8, ImmutAnyOrigin, _safe=False], b_ptr: Pointer[UInt8, ImmutAnyOrigin, _safe=False], a_scales_ptr: Pointer[Float8_e8m0fnu, ImmutAnyOrigin, _safe=False], b_scales_ptr: Pointer[Float8_e8m0fnu, ImmutAnyOrigin, _safe=False], row_offsets_ptr: Pointer[UInt32, ImmutAnyOrigin, _safe=False], expert_ids_ptr: Pointer[Int32, ImmutAnyOrigin, _safe=False], num_active_experts: Int, M: Int, N: Int, K: Int)

GPU kernel that computes a per-expert grouped MXFP4 block-scaled matmul on AMD CDNA4.

Each thread identifies its expert via linear search over row_offsets_ptr and expert_ids_ptr, then dispatches to _mxfp4_dotprod against the expert's weight slice.

Parameters:

out_dtype (DType): Element type of the output matrix c_ptr.
BLOCK_N (Int): Number of output columns computed per thread; must divide N evenly or be 1.