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

quantize_dynamic_scaled_async_fp4_kernel

def quantize_dynamic_scaled_async_fp4_kernel[input_dtype: DType, input_tile_rank: Int, input_tile_shape: IndexList[input_tile_rank], input_desc_shape: IndexList[input_tile_rank], output_dtype: DType, output_tile_rank: Int, output_tile_shape: IndexList[output_tile_rank], output_desc_shape: IndexList[output_tile_rank], scales_dtype: DType, scales_tile_rank: Int, scales_tile_shape: IndexList[scales_tile_rank], scales_desc_shape: IndexList[scales_tile_rank], input_swizzle_mode: TensorMapSwizzle, output_swizzle_mode: TensorMapSwizzle, scales_swizzle_mode: TensorMapSwizzle, SF_VECTOR_SIZE: Int, NUM_PIPELINES_STAGES: Int](input_tma_op: TMATensorTile[input_dtype, input_tile_rank, input_tile_shape, input_desc_shape], output_tma_op: TMATensorTile[output_dtype, output_tile_rank, output_tile_shape, output_desc_shape], scales_tma_op: TMATensorTile[scales_dtype, scales_tile_rank, scales_tile_shape, scales_desc_shape], tensor_sf: Float32)

GPU kernel that quantizes BF16 tiles to NVFP4 using TMA async copies and warp-specialized PDL.

One warpgroup issues TMA loads while another computes per-group scale factors and packs FP4 elements, then stores results back via TMA async stores.

Parameters:

input_dtype (DType): Element type of the input activation tensor.
input_tile_rank (Int): Rank of the input TMA tile descriptor.
input_tile_shape (IndexList[input_tile_rank]): Per-dimension shape of the input TMA tile.
input_desc_shape (IndexList[input_tile_rank]): Per-dimension descriptor shape of the input TMA tile.
output_dtype (DType): Element type of the quantized output tensor.
output_tile_rank (Int): Rank of the output TMA tile descriptor.
output_tile_shape (IndexList[output_tile_rank]): Per-dimension shape of the output TMA tile.
output_desc_shape (IndexList[output_tile_rank]): Per-dimension descriptor shape of the output TMA tile.
scales_dtype (DType): Element type of the block scale-factor tensor.
scales_tile_rank (Int): Rank of the scales TMA tile descriptor.
scales_tile_shape (IndexList[scales_tile_rank]): Per-dimension shape of the scales TMA tile.
scales_desc_shape (IndexList[scales_tile_rank]): Per-dimension descriptor shape of the scales TMA tile.
input_swizzle_mode (TensorMapSwizzle): TensorMapSwizzle mode for input TMA loads.
output_swizzle_mode (TensorMapSwizzle): TensorMapSwizzle mode for output TMA stores.
scales_swizzle_mode (TensorMapSwizzle): TensorMapSwizzle mode for scales TMA stores.
SF_VECTOR_SIZE (Int): Number of elements covered by each block scale factor.
NUM_PIPELINES_STAGES (Int): Number of pipeline stages for TMA double-buffered async copies.