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

set_batched_scale_factor

def set_batched_scale_factor[scales_dtype: DType, scales_layout: Layout, //, SF_VECTOR_SIZE: Int](scales_tensor: LayoutTensor[scales_dtype, scales_layout, MutAnyOrigin], batch_idx: Int, row_idx: Int, col_idx: Int, scale_value: Scalar[scales_dtype])

Stores a scale factor into a 6D batched LayoutTensor at the given batch, row, and column.

Translates the linear (batch_idx, row_idx, col_idx) coordinates into the 6D scale-factor layout used by batched MXFP4/NVFP4 block-scaled tensors and stores scale_value there.

Parameters:

scales_dtype (DType): Element type of the scales tensor.
scales_layout (Layout): Layout of the scales LayoutTensor.
SF_VECTOR_SIZE (Int): Number of elements each scale factor covers.

Args:

scales_tensor (LayoutTensor[scales_dtype, scales_layout, MutAnyOrigin]): Mutable 6D LayoutTensor holding the batched scale factors.
batch_idx (Int): Batch index in the original (unscaled) tensor coordinates.
row_idx (Int): Row index in the original (unscaled) tensor coordinates.
col_idx (Int): Column index in the original (unscaled) tensor coordinates.
scale_value (Scalar[scales_dtype]): Scale factor value to store.

def set_batched_scale_factor[scales_dtype: DType, //, SF_VECTOR_SIZE: Int](scales_tensor: TileTensor[scales_dtype, Storage=scales_tensor.Storage, address_space=scales_tensor.address_space, linear_idx_type=scales_tensor.linear_idx_type], batch_idx: Int, row_idx: Int, col_idx: Int, scale_value: Scalar[scales_dtype])

Stores a scale factor into a 6D batched TileTensor at the given batch, row, and column.

TileTensor overload of set_batched_scale_factor that translates the linear (batch_idx, row_idx, col_idx) coordinates into the 6D scale-factor layout and stores scale_value there.

Parameters:

scales_dtype (DType): Element type of the scales tensor.
SF_VECTOR_SIZE (Int): Number of elements each scale factor covers.

Args:

scales_tensor (TileTensor[scales_dtype, Storage=scales_tensor.Storage, address_space=scales_tensor.address_space, linear_idx_type=scales_tensor.linear_idx_type]): Mutable TileTensor (flat rank == 6) holding the batched scale factors.
batch_idx (Int): Batch index in the original (unscaled) tensor coordinates.
row_idx (Int): Row index in the original (unscaled) tensor coordinates.
col_idx (Int): Column index in the original (unscaled) tensor coordinates.
scale_value (Scalar[scales_dtype]): Scale factor value to store.