IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
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

QuantConfig

QuantConfig​

class max.nn.QuantConfig(input_scale, weight_scale, mlp_quantized_layers, attn_quantized_layers, format, embedding_output_dtype=None, shared_experts_weight_dtype=None, bias_dtype=None, can_use_fused_mlp=False, can_use_fused_swiglu=False, scales_pre_interleaved=False, _mxfp6_element_format='e2m3', block_scaled_preshuffled_b=False)

source

Bases: object

Configures scaled quantization settings for a layer or model section.

For example, to configure NVFP4 block-scaled quantization for all layers in a 19-layer model:

from max.dtype import DType
from max.nn import QuantConfig, QuantFormat
from max.nn.quant_config import (
    InputScaleSpec,
    ScaleGranularity,
    ScaleOrigin,
    WeightScaleSpec,
)

all_layers = set(range(19))

input_spec = InputScaleSpec(
    granularity=ScaleGranularity.BLOCK,
    origin=ScaleOrigin.STATIC,
    dtype=DType.float32,
    block_size=(1, 16),
)
weight_spec = WeightScaleSpec(
    granularity=ScaleGranularity.BLOCK,
    dtype=DType.float8_e4m3fn,
    block_size=(1, 8),
)
config = QuantConfig(
    input_scale=input_spec,
    weight_scale=weight_spec,
    mlp_quantized_layers=all_layers,
    attn_quantized_layers=all_layers,
    format=QuantFormat.NVFP4,
)

Parameters:

attn_quantized_layers​

attn_quantized_layers: set[int]

source

Set of layer indices with quantized attention projections.

Attention projections are quantized on an all-or-nothing basis per layer: either all of q_proj, k_proj, v_proj, and o_proj are quantized, or all four remain in bfloat16.

bias_dtype​

bias_dtype: DType | None = None

source

The DType of bias weights.

block_scaled_preshuffled_b​

block_scaled_preshuffled_b: bool = False

source

Whether MXFP4 weight B is preshuffled into the 5D layout that the AMD preb kernel reads (produced by Shuffler.preshuffle_b_5d). When True, MoEQuantized dispatches the grouped matmul to the block_scaled_grouped_matmul_amd_preb kernel variant; when False (default) it dispatches to the dense row-major block_scaled_grouped_matmul_amd kernel. Must be set in lockstep with the weight loader actually applying the preshuffle (e.g. Kimi K2.5’s weight_adapters.py:_shuffle_group).

can_use_fused_mlp​

can_use_fused_mlp: bool = False

source

Whether the quantization scales can be used with fused MLP operations.

can_use_fused_swiglu​

can_use_fused_swiglu: bool = False

source

Whether to use the fused grouped matmul + SwiGLU + NVFP4/MXFP4/MXFP8 quant SM100 kernel for the MoE gate/up projection. When True, the MoE layer pre-permutes gate_up_proj and its scales on the N axis (sigma(2i)=i, sigma(2i+1)=D+i) and dispatches the internal grouped_matmul_blocked_swiglu kernel wrapper. Defaults to False so the chained (matmul -> BF16 -> SwiGLU+quant) path is unchanged.

embedding_output_dtype​

embedding_output_dtype: DType | None = None

source

The DType of the output from the embedding layer.

format​

format: QuantFormat

source

The QuantFormat identifying the quantization format.

input_scale​

input_scale: InputScaleSpec

source

InputScaleSpec for input activation scaling.

is_dynamic​

property is_dynamic: bool

source

True if this input scale is dynamic.

is_fp4​

property is_fp4: bool

source

True if this config represents any FP4 variant (NVFP4 or MXFP4).

is_int8_w8a8​

property is_int8_w8a8: bool

source

True if this config represents symmetric int8 W8A8.

is_mx​

property is_mx: bool

source

True for any OCP microscaling format (MXFP4, MXFP6 or MXFP8).

These share the E8M0 per-32-element block scale; NVFP4 is excluded, as it scales per 16 elements with an FP8 scale.

is_mxfp4​

property is_mxfp4: bool

source

Returns True if this config represents MXFP4 quantization.

is_mxfp6​

property is_mxfp6: bool

source

Returns True if this config represents MXFP6 quantization.

is_mxfp8​

property is_mxfp8: bool

source

Returns True if this config represents MXFP8 quantization.

is_nvfp4​

property is_nvfp4: bool

source

True if this config represents modelopt NVFP4.

is_static​

property is_static: bool

source

True if this input scale is static.

mlp_quantized_layers​

mlp_quantized_layers: set[int]

source

Set of layer indices with quantized MLPs.

MLPs are quantized on an all-or-nothing basis per layer: either all of gate_proj, down_proj, and up_proj are quantized, or all three remain in bfloat16.

mx_element_dtype​

property mx_element_dtype: DType

source

The dtype of one quantized element, for MX formats.

This is the logical element type, not the storage type: sub-byte MX weights are stored as packed uint8 bytes, and that spelling cannot say which format the bytes hold – MXFP4 and MXFP6 share it. Code that needs to branch on the format should branch on this instead of pairing dtype == DType.uint8 with an is_mxfp6 flag.

Raises:

ValueError – If this config is not an MX format.

mxfp6_format​

property mxfp6_format: str

source

The FP6 element encoding, "e2m3" or "e3m2".

E2M3 is the default: with 3 mantissa bits against E3M2’s 2, it is the better choice for weights, whose dynamic range a per-32 block scale already absorbs.

quantized_scales_type()​

quantized_scales_type(quantized_shape, device_ref)

source

The TensorType of the scales tensor after dynamic quantization.

Parameters:

Return type:

TensorType

scales_granularity_mnk​

property scales_granularity_mnk: tuple[int, int, int]

source

The weight and input scale granularities on the M, N, and K axes.

scales_pre_interleaved​

scales_pre_interleaved: bool = False

source

Whether weight scales in the checkpoint are already stored in the 5D TCGEN-interleaved layout expected by the FP4 matmul kernel (NVFP4 only). Note that scales in the 5D TCGEN-interleaved layout are typically flattened to 2D [M, K//16] in the checkpoint.

shared_experts_dtype()​

shared_experts_dtype(routed_weight_dtype)

source

Resolve weight dtype for MoE shared-expert MLPs.

Parameters:

routed_weight_dtype (DType)

Return type:

DType

shared_experts_use_quant()​

shared_experts_use_quant(routed_weight_dtype)

source

Whether shared experts use the same quantized weights as routed experts.

Parameters:

routed_weight_dtype (DType)

Return type:

bool

shared_experts_weight_dtype​

shared_experts_weight_dtype: DType | None = None

source

Weight storage dtype for MoE shared-expert MLPs when they differ from routed experts.

When None, shared experts use the same dtype and quantization as routed experts. When set (e.g. bfloat16 for mixed Kimi K2.6 NVFP4 checkpoints), shared-expert linears omit quant_config while routed experts remain quantized.

weight_scale​

weight_scale: WeightScaleSpec

source

WeightScaleSpec for weight scaling.

Models

Cookbook

Puzzles

Blog

Community

Contact