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

fused_silu_fp8_kernel

def fused_silu_fp8_kernel[fp8_dtype: DType, scales_dtype: DType, input_dtype: DType, output_layout: TensorLayout, scales_layout: TensorLayout, input_layout: TensorLayout, offsets_layout: TensorLayout, num_threads: Int, num_sms: Int, group_size: Int = Int(128)](output_tensor: TileTensor[fp8_dtype, output_layout, MutUntrackedOrigin], scales_tensor: TileTensor[scales_dtype, scales_layout, MutUntrackedOrigin], input_tensor: TileTensor[input_dtype, input_layout, ImmUntrackedOrigin], row_offsets: TileTensor[DType.uint32, offsets_layout, ImmUntrackedOrigin])

This kernel performs the SILU operation for all the MLPs in the EP MoE module. We need to manually implement the kernel here is because after the EP dispatch phase, the actual number of received tokens is not known to the host. This kernel will read the row offsets to determine the actual number of received tokens in the input tensor.

Once the SILU operation is performed, the output tensor will be quantized to the FP8 format. The scales tensor will be stored in a transposed way.

Arguments: output_tensor: The output tensor to store the result. scales_tensor: The tensor to store the scales. input_tensor: The input tensor to perform the SILU operation. row_offsets: The row offsets to determine the actual number of received tokens.

Parameters:

fp8_dtype (DType): FP8 element type of the quantized output_tensor (e.g. DType.float8_e4m3fn).
scales_dtype (DType): Element type of the block-wise scale factors stored in scales_tensor.
input_dtype (DType): Element type of the input_tensor; its accumulation type must be floating-point.
output_layout (TensorLayout): Layout of the FP8 output_tensor TileTensor.
scales_layout (TensorLayout): Layout of the scales_tensor TileTensor; scales are stored transposed (group index in dim 0, token index in dim 1).
input_layout (TensorLayout): Layout of the input_tensor TileTensor.
offsets_layout (TensorLayout): Layout of the 1D row_offsets TileTensor.
num_threads (Int): Number of threads per block; sets the MAX_THREADS_PER_BLOCK launch metadata.
num_sms (Int): Number of streaming multiprocessors (SMs) used to scatter processing across thread blocks.
group_size (Int): Number of elements per quantization group; the output dimension must be divisible by this (defaults to 128).