For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Python class

Fp8Strategy

Fp8Strategy

class max.nn.moe.Fp8Strategy(config, dtype)

source

Bases: object

FP8 quantization for MoE.

Parameters:

• config (QuantConfig)
• dtype (DType)

fused_silu_quantize()

fused_silu_quantize(gate_up_projs, input_scales=None, expert_inputs=(), max_padded_M=0, activation=None)

source

Applies the gated activation and returns quantized activations.

The activation is plain graph ops feeding the row-bounded quantize, which the graph compiler fuses into a single kernel.

max_padded_M is accepted for QuantStrategy conformance; it only applies to the MX EP paths and is ignored here.

Parameters:

• gate_up_projs (TensorValue)
• input_scales (TensorValue | None)
• expert_inputs (tuple[TensorValue, ...])
• max_padded_M (int)
• activation (GatedActivation | None)

Return type:

tuple[TensorValue, TensorValue]

grouped_matmul()

grouped_matmul(weight, weight_scales, expert_scales=None, expert_inputs=(), estimated_total_m=None, a_scales_max_padded_m=0)

source

Runs grouped FP8 matmul for the routed experts.

a_scales_max_padded_m is accepted for QuantStrategy conformance; it only applies to the MX EP paths and is ignored here.

Parameters:

• weight (TensorValue)
• weight_scales (TensorValue)
• expert_scales (TensorValue | None)
• expert_inputs (tuple[TensorValue, ...])
• estimated_total_m (TensorValue | None)
• a_scales_max_padded_m (int)

Return type:

TensorValue

grouped_quantize()

grouped_quantize(tensor, group_size, input_scale, expert_start, scales_offset, expert_ids)

source

Quantizes to FP8, stopping at the last row expert_start covers.

FP8 needs no per-expert scale padding, so this differs from quantize() only in what it skips: tensor is the EP receive buffer, sized for the worst-case dispatch, and expert_start ends at the row count this step actually received. The rows below that count quantize identically either way.

Parameters:

• tensor (TensorValue)
• group_size (int)
• input_scale (TensorValue | None)
• expert_start (TensorValue)
• scales_offset (TensorValue | None)
• expert_ids (TensorValue)

Return type:

tuple[TensorValue, TensorValue]

prepare_weight_scales()

prepare_weight_scales(gate_up, down, device)

source

Passes FP8 weight scales through without reformatting.

Parameters:

• gate_up (TensorValue)
• down (TensorValue)
• device (DeviceRef)

Return type:

tuple[TensorValue, TensorValue]

quantize()

quantize(tensor, group_size)

source

Quantizes activations to FP8 and returns (quantized, scales).

Parameters:

• tensor (TensorValue)
• group_size (int)

Return type:

tuple[TensorValue, TensorValue]