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_scale_factor

def set_scale_factor[scales_dtype: DType, scales_layout: Layout, //, SF_VECTOR_SIZE: Int, width: SIMDLength](scales_tensor: LayoutTensor[scales_dtype, scales_layout, address_space=scales_tensor.address_space, element_layout=scales_tensor.element_layout, layout_int_type=scales_tensor.layout_int_type, linear_idx_type=scales_tensor.linear_idx_type, masked=scales_tensor.masked, alignment=scales_tensor.alignment], row_idx: Int, col_idx: Int, scale_value: SIMD[scales_dtype, width])

Stores a scale factor into a 5D non-batched LayoutTensor at the given row and column.

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

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.
width (SIMDLength): SIMD width of the value to store; must not exceed SF_ATOM_K.

Args:

scales_tensor (LayoutTensor[scales_dtype, scales_layout, address_space=scales_tensor.address_space, element_layout=scales_tensor.element_layout, layout_int_type=scales_tensor.layout_int_type, linear_idx_type=scales_tensor.linear_idx_type, masked=scales_tensor.masked, alignment=scales_tensor.alignment]): Mutable 5D LayoutTensor holding the scale factors.
row_idx (Int): Row index in the original (unscaled) tensor coordinates.
col_idx (Int): Column index in the original (unscaled) tensor coordinates.
scale_value (SIMD[scales_dtype, width]): Scale factor value(s) to store.

def set_scale_factor[scales_dtype: DType, width: SIMDLength, //, 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], row_idx: Int, col_idx: Int, scale_value: SIMD[scales_dtype, width])

Stores a scale factor into a 5D non-batched TileTensor at the given row and column.

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

Parameters:

scales_dtype (DType): Element type of the scales tensor.
width (SIMDLength): SIMD width of the value to store; must not exceed SF_ATOM_K.
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 >= 5) holding the scale factors.
row_idx (Int): Row index in the original (unscaled) tensor coordinates.
col_idx (Int): Column index in the original (unscaled) tensor coordinates.
scale_value (SIMD[scales_dtype, width]): Scale factor value(s) to store.