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_quantize_dynamic_scaled_fp4_async_kernel

def grouped_quantize_dynamic_scaled_fp4_async_kernel[output_dtype: DType, scales_dtype: DType, input_dtype: DType, scales_tile_rank: Int, scales_tile_shape: IndexList[scales_tile_rank], scales_desc_shape: IndexList[scales_tile_rank], scales_swizzle_mode: TensorMapSwizzle, output_layout: TensorLayout, input_layout: TensorLayout, row_offsets_layout: TensorLayout, scales_offsets_layout: TensorLayout, expert_ids_layout: TensorLayout, sf_layout: TensorLayout, num_threads: Int = Int(128)](output_tensor: TileTensor[output_dtype, output_layout, MutAnyOrigin], scales_tma_op: TMATensorTile[scales_dtype, scales_tile_rank, scales_tile_shape, scales_desc_shape], input_tensor: TileTensor[input_dtype, input_layout, ImmutAnyOrigin], row_offsets: TileTensor[DType.uint32, row_offsets_layout, ImmutAnyOrigin], scales_offsets: TileTensor[DType.uint32, scales_offsets_layout, ImmutAnyOrigin], expert_ids: TileTensor[DType.int32, expert_ids_layout, ImmutAnyOrigin], sf_tensor: TileTensor[DType.float32, sf_layout, ImmutAnyOrigin])

GPU kernel that quantizes per-expert BF16 activation tiles to NVFP4/MXFP4/MXFP8 with TMA-based scale-factor stores.

Each block locates its assigned expert via binary search on row_offsets and scales_offsets, then quantizes the expert's activation tile and writes scale factors back through TMA async stores.

Parameters:

output_dtype (DType): Element type of the quantized output tensor (inferred).
scales_dtype (DType): Element type of the block scale-factor tensor (inferred).
input_dtype (DType): Element type of the input activation tensor (inferred).
scales_tile_rank (Int): Rank of the scales TMA tile descriptor (inferred).
scales_tile_shape (IndexList[scales_tile_rank]): Per-dimension tile shape of the scales TMA descriptor (inferred).
scales_desc_shape (IndexList[scales_tile_rank]): Per-dimension descriptor shape of the scales TMA descriptor (inferred).
scales_swizzle_mode (TensorMapSwizzle): Swizzle mode applied to the scales TMA descriptor (inferred).
output_layout (TensorLayout): TileTensor layout of the quantized output tensor (inferred).
input_layout (TensorLayout): TileTensor layout of the input activation tensor (inferred).
row_offsets_layout (TensorLayout): TileTensor layout of the per-expert row offsets tensor (inferred).
scales_offsets_layout (TensorLayout): TileTensor layout of the per-expert scales offsets tensor (inferred).
expert_ids_layout (TensorLayout): TileTensor layout of the expert IDs tensor (inferred).
sf_layout (TensorLayout): TileTensor layout of the per-expert tensor-wise scale factor tensor (inferred).
num_threads (Int): Number of threads per block in the launch grid (defaults to 128).