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

batched_matmul_dynamic_scaled_fp8_naive

def batched_matmul_dynamic_scaled_fp8_naive[c_type: DType, a_type: DType, b_type: DType, a_scales_type: DType, b_scales_type: DType, //, *, scales_granularity_mnk: IndexList[Int(3)], transpose_b: Bool = False](c_: TileTensor[c_type, Storage=c_.Storage, address_space=c_.address_space, linear_idx_type=c_.linear_idx_type], a_: TileTensor[a_type, Storage=a_.Storage, address_space=a_.address_space, linear_idx_type=a_.linear_idx_type], b_: TileTensor[b_type, Storage=b_.Storage, address_space=b_.address_space, linear_idx_type=b_.linear_idx_type], a_scales_: TileTensor[a_scales_type, Storage=a_scales_.Storage, address_space=a_scales_.address_space, linear_idx_type=a_scales_.linear_idx_type], b_scales_: TileTensor[b_scales_type, Storage=b_scales_.Storage, address_space=b_scales_.address_space, linear_idx_type=b_scales_.linear_idx_type], ctx: DeviceContext)

Computes a batched blockwise scaled FP8 matrix multiplication using a naive per-batch loop that calls the 2D blockwise scaled FP8 kernel for each batch slice.

Parameters:

c_type (DType): Output tensor element dtype.
a_type (DType): LHS input tensor element dtype.
b_type (DType): RHS input tensor element dtype.
a_scales_type (DType): LHS scales tensor element dtype.
b_scales_type (DType): RHS scales tensor element dtype.
scales_granularity_mnk (IndexList[Int(3)]): Scale granularity (m, n, k); only (1, 128, 128) is currently supported.
transpose_b (Bool): Whether the RHS input is transposed.

Args:

c_ (TileTensor[c_type, Storage=c_.Storage, address_space=c_.address_space, linear_idx_type=c_.linear_idx_type]): Rank-3 output tensor of shape (batch, m, n).
a_ (TileTensor[a_type, Storage=a_.Storage, address_space=a_.address_space, linear_idx_type=a_.linear_idx_type]): Rank-3 LHS input tensor of shape (batch, m, k).
b_ (TileTensor[b_type, Storage=b_.Storage, address_space=b_.address_space, linear_idx_type=b_.linear_idx_type]): Rank-3 RHS input tensor of shape (batch, k, n).
a_scales_ (TileTensor[a_scales_type, Storage=a_scales_.Storage, address_space=a_scales_.address_space, linear_idx_type=a_scales_.linear_idx_type]): Rank-3 LHS scales tensor.
b_scales_ (TileTensor[b_scales_type, Storage=b_scales_.Storage, address_space=b_scales_.address_space, linear_idx_type=b_scales_.linear_idx_type]): Rank-3 RHS scales tensor.
ctx (DeviceContext): Device context used to dispatch the per-batch kernels.