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_and_bmm_fp8_helper

def quantize_and_bmm_fp8_helper[dtype: DType, fp8_dtype: DType, fp8_scale_dtype: DType, m_scale_granularity: Int, n_scale_granularity: Int, k_scale_granularity: Int, target: StringSlice[ImmStaticOrigin] = StringSlice("cpu")](c: TileTensor[dtype, Storage=c.Storage, linear_idx_type=c.linear_idx_type], a: TileTensor[dtype, Storage=a.Storage, linear_idx_type=a.linear_idx_type], b: TileTensor[fp8_dtype, Storage=b.Storage, linear_idx_type=b.linear_idx_type], b_scales: TileTensor[fp8_scale_dtype, Storage=b_scales.Storage, linear_idx_type=b_scales.linear_idx_type], ctx: DeviceContext)

Helper function to quantize and perform a batched matrix multiplication. This function uses the transposed view of the input tensor a.

Parameters:

dtype (DType): Data type of the input tensor a and output tensor c.
fp8_dtype (DType): Data type of the FP8 quantized tensors.
fp8_scale_dtype (DType): Data type of the FP8 scale tensors.
m_scale_granularity (Int): Granularity of the scale for the M dimension of the batched matrix multiplication.
n_scale_granularity (Int): Granularity of the scale for the N dimension of the batched matrix multiplication.
k_scale_granularity (Int): Granularity of the scale for the K dimension of the batched matrix multiplication, also the quantization group size for a.
target (StringSlice[ImmStaticOrigin]): Target device for the batched matrix multiplication (defaults to "cpu").

Args:

c (TileTensor[dtype, Storage=c.Storage, linear_idx_type=c.linear_idx_type]): Output tensor for the batched matrix multiplication result. Shape: [batch, m, n].
a (TileTensor[dtype, Storage=a.Storage, linear_idx_type=a.linear_idx_type]): Input tensor to quantize to FP8 and use as the left operand, loaded via a transposed view. Shape: [m, batch, k].
b (TileTensor[fp8_dtype, Storage=b.Storage, linear_idx_type=b.linear_idx_type]): FP8 weight tensor used as the right operand. Shape: [batch, n, k].
b_scales (TileTensor[fp8_scale_dtype, Storage=b_scales.Storage, linear_idx_type=b_scales.linear_idx_type]): Scale tensor for b.
ctx (DeviceContext): Device context for buffer allocation and kernel execution.