For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /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, OutputEngine: TensorEngine, InputEngine: TensorEngine, RowOffsetsEngine: TensorEngine, ScalesOffsetsEngine: TensorEngine, ExpertIdsEngine: TensorEngine, SfEngine: TensorEngine, row_indices_layout: TensorLayout = Layout[TypeList[Int64](), TypeList[ComptimeInt[Int(1)]]()], RowIndicesEngine: TensorEngine = DefaultEngine, num_threads: Int = Int(128), k_tiles_per_block: Int = Int(1)](output_tensor: TileTensor[output_dtype, output_layout, MutAnyOrigin, Engine=OutputEngine], scales_tma_op: TMATensorTile[scales_dtype, scales_tile_rank, scales_tile_shape, scales_desc_shape], input_tensor: TileTensor[input_dtype, input_layout, ImmutAnyOrigin, Engine=InputEngine], row_offsets: TileTensor[.uint32, row_offsets_layout, ImmutAnyOrigin, Engine=RowOffsetsEngine], scales_offsets: TileTensor[.uint32, scales_offsets_layout, ImmutAnyOrigin, Engine=ScalesOffsetsEngine], expert_ids: TileTensor[.int32, expert_ids_layout, ImmutAnyOrigin, Engine=ExpertIdsEngine], sf_tensor: TileTensor[.float32, sf_layout, ImmutAnyOrigin, Engine=SfEngine], row_indices: OptionalReg[TileTensor[.int32, row_indices_layout, ImmutAnyOrigin, Engine=RowIndicesEngine]] = None)

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 with a warp-parallel search over 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).
• ​OutputEngine (TensorEngine): Engine policy of the output_tensor.
• ​InputEngine (TensorEngine): Engine policy of the input_tensor.
• ​RowOffsetsEngine (TensorEngine): Engine policy of the row_offsets tensor.
• ​ScalesOffsetsEngine (TensorEngine): Engine policy of the scales_offsets tensor.
• ​ExpertIdsEngine (TensorEngine): Engine policy of the expert_ids tensor.
• ​SfEngine (TensorEngine): Engine policy of the sf_tensor.
• ​row_indices_layout (TensorLayout): TileTensor layout of the optional source row indices tensor.
• ​RowIndicesEngine (TensorEngine): Engine policy of the row_indices tensor.
• ​num_threads (Int): Number of threads per block in the launch grid (defaults to 128).
• ​k_tiles_per_block (Int): Column tiles handled by one block. Batching them amortizes the per-block expert search and shrinks the grid, which dominates at decode where a block's payload is one routed row.