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

get_batched_scale_factor

def get_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) -> Scalar[scales_dtype]

Loads a scale factor from 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 returns the stored scale factor.

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]): 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.

Returns:

Scalar[scales_dtype]: The scale factor stored at the translated 6D coordinate.

def get_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) -> Scalar[scales_dtype]

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

TileTensor overload of get_batched_scale_factor that translates the linear (batch_idx, row_idx, col_idx) coordinates into the 6D scale-factor layout and returns the stored scale factor.

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]): 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.

Returns:

Scalar[scales_dtype]: The scale factor stored at the translated 6D coordinate.