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

QuantStrategy

QuantStrategy

class max.nn.moe.QuantStrategy(*args, **kwargs)

source

Bases: Protocol

Quantization strategy for MoE layers.

fused_silu_quantize()

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

source

Applies gating and quantizes activations for the down proj.

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 matmul for routed experts.

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 activations with per-expert scales and padding.

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

Prepares weight scales for kernel consumption.

Parameters:

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

Return type:

tuple[TensorValue, TensorValue]

quantize()

quantize(tensor, group_size)

source

Quantizes activations and returns (quantized, scales).

Parameters:

• tensor (TensorValue)
• group_size (int)

Return type:

tuple[TensorValue, TensorValue]