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

NvMxf4f8Strategy

NvMxf4f8Strategy

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

source

Bases: object

NVIDIA NVFP4/MXFP4/MXFP8 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 then quantizes the result.

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 NVIDIA block-scaled matmul with per-expert scales.

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_matmul_swiglu()

grouped_matmul_swiglu(weight, weight_scales, *, expert_scales=None, input_scales=None, expert_inputs, estimated_total_m=None, activation=None, a_row_scales=None)

source

Runs the fused quantized grouped matmul + SwiGLU + quant kernel.

Equivalent to grouped_matmul followed by fused_silu_quantize, but folds both into a single SM100 kernel. The caller must pass a sigma-permuted weight and weight_scales (see max.nn.kernels.grouped_matmul_blocked_swiglu); the layout is produced by MoE.gate_up_proj() and MoEQuantized.gate_up_proj_scales() when quant_config.can_use_fused_swiglu is set, or ships natively (gate, up)-interleaved in the checkpoint (the TP path gated by MoEQuantized._can_fuse_swiglu_interleaved()).

Parameters:

• weight (TensorValue) – Sigma-permuted gate/up projection weights.
• weight_scales (TensorValue) – Sigma-permuted gate/up projection scales (6D).
• expert_scales (TensorValue | None) – Per-expert matmul-epilogue scaling factors. Typically nvfp4.gate_up_expert.
• input_scales (TensorValue | None) – Raw per-expert SiLU-output scale (nvfp4.down_input). The fused kernel consumes the inverted scale internally; this method inverts here to match the chained fused_silu_quantize() convention.
• expert_inputs (tuple[TensorValue, ...]) – Same tuple shape as grouped_matmul(): (hidden, hidden_scales, expert_start, scales_offsets, expert_ids, usage_stats).
• estimated_total_m (TensorValue | None) – Estimated total non-padded token count.
• activation (GatedActivation | None) – The gated activation the kernel fuses, forwarded to grouped_matmul_blocked_swiglu.
• a_row_scales (TensorValue | None) – Optional BF16 scale per hidden row, applied with the expert scale before the SwiGLU. NVFP4 only.

Returns:

Tuple (c_packed, c_swiglu_scales) matching the chained reference path byte-for-byte under the kernel’s default match_bf16=True setting.

Return type:

tuple[TensorValue, TensorValue]

grouped_quantize()

grouped_quantize(tensor, group_size, input_scale, expert_start, scales_offset, expert_ids, indices=None, indices_divisor=1)

source

Quantizes activations per-expert with padded scale alignment.

Parameters:

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

Return type:

tuple[TensorValue, TensorValue]

is_nvfp4

property is_nvfp4: bool

source

Whether this strategy is handling NVFP4 rather than MXFP4 or MXFP8.

prepare_swiglu_operands()

prepare_swiglu_operands(expert_inputs, *, expert_scales=None, input_scales=None)

source

Normalizes the operands every fused-SwiGLU kernel call needs.

Both the chained path and the EP-combine-send path hand the same operands to the same kernel, so the two derivations that are easy to get wrong live here instead of at each call site: the per-expert SiLU-output scale must be INVERTED, and the usage stats must be on the host. A call site that re-derived either one silently disagreed with this one.

Parameters:

• expert_inputs (tuple[TensorValue, ...]) – (hidden, hidden_scales, expert_start, scales_offsets, expert_ids, usage_stats).
• expert_scales (TensorValue | None) – Per-expert matmul-epilogue scaling factors.
• input_scales (TensorValue | None) – Raw per-expert SiLU-output scale.

Returns:

The normalized operands.

Raises:

ValueError – If NVFP4 is missing a required scale tensor.

Return type:

_PreparedSwigluOperands

prepare_weight_scales()

prepare_weight_scales(gate_up, down, device)

source

Interleaves NVIDIA block scales for kernel layout.

Parameters:

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

Return type:

tuple[TensorValue, TensorValue]

quantize()

quantize(tensor, group_size)

source

Parameters:

• tensor (TensorValue)
• group_size (int)

Return type:

tuple[TensorValue, TensorValue]