Skip to content

mxfp8_config

MXFP8 quantization for MiniMax-H3 transformer-block feed-forward layers.

Classes

fastvideo.layers.quantization.mxfp8_config.MXFP8Config 

MXFP8Config()

Bases: QuantizationConfig

Select MXFP8 for the main MiniMax-H3 transformer-block FFN linears.

Source code in fastvideo/layers/quantization/base_config.py 
64
65
66
67
def __init__(self) -> None:
    super().__init__()
    # mapping is updated by models as they initialize
    self.packed_modules_mapping: dict[str, list[str]] = dict()

fastvideo.layers.quantization.mxfp8_config.MXFP8QuantizeMethod

Bases: QuantizeMethodBase

Dynamically quantize activations against prequantized MXFP8 weights.

Methods:

fastvideo.layers.quantization.mxfp8_config.MXFP8QuantizeMethod.apply 
apply(layer: Module, hidden_states: Tensor, bias: Tensor | None = None) -> Tensor

Quantize one BF16 activation and apply the prequantized linear.

Source code in fastvideo/layers/quantization/mxfp8_config.py 
58
59
60
61
62
63
64
65
66
67
68
69
70
71
def apply(
    self,
    layer: torch.nn.Module,
    hidden_states: torch.Tensor,
    bias: torch.Tensor | None = None,
) -> torch.Tensor:
    """Quantize one BF16 activation and apply the prequantized linear."""
    from fastvideo.layers.mxfp8linear import quantize_mxfp8_blockwise

    original_shape = hidden_states.shape
    hidden_states_2d = hidden_states.reshape(-1, hidden_states.shape[-1])
    activation_values, activation_scales = quantize_mxfp8_blockwise(hidden_states_2d)
    output_2d = self.apply_quantized(layer, activation_values, activation_scales, bias)
    return output_2d.reshape(*original_shape[:-1], layer.output_size)
fastvideo.layers.quantization.mxfp8_config.MXFP8QuantizeMethod.apply_quantized 
apply_quantized(layer: Module, activation_values: Tensor, activation_scales: Tensor, bias: Tensor | None = None) -> Tensor

Apply one linear to an activation that is already in MXFP8.

Source code in fastvideo/layers/quantization/mxfp8_config.py 
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
def apply_quantized(
    self,
    layer: torch.nn.Module,
    activation_values: torch.Tensor,
    activation_scales: torch.Tensor,
    bias: torch.Tensor | None = None,
) -> torch.Tensor:
    """Apply one linear to an activation that is already in MXFP8."""
    from fastvideo.layers.mxfp8linear import mxfp8_scaled_mm

    if layer._mxfp8_weight is None or layer._mxfp8_weight_scale is None:
        raise RuntimeError(f"MXFP8 weight buffers are not initialized for {layer.prefix}.")
    return mxfp8_scaled_mm(
        activation_values,
        activation_scales,
        layer._mxfp8_weight,
        layer._mxfp8_weight_scale,
        bias,
    )
fastvideo.layers.quantization.mxfp8_config.MXFP8QuantizeMethod.create_weights 
create_weights(layer: Module, input_size_per_partition: int, output_partition_sizes: list[int], input_size: int, output_size: int, params_dtype: dtype, **extra_weight_attrs: Any) -> None

Create the BF16 checkpoint weight and non-persistent MXFP8 buffers.

Source code in fastvideo/layers/quantization/mxfp8_config.py 
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
def create_weights(
    self,
    layer: torch.nn.Module,
    input_size_per_partition: int,
    output_partition_sizes: list[int],
    input_size: int,
    output_size: int,
    params_dtype: torch.dtype,
    **extra_weight_attrs: Any,
) -> None:
    """Create the BF16 checkpoint weight and non-persistent MXFP8 buffers."""
    del input_size, output_size
    weight = Parameter(
        torch.empty(
            sum(output_partition_sizes),
            input_size_per_partition,
            dtype=params_dtype,
        ),
        requires_grad=False,
    )
    set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0})
    layer.register_parameter("weight", weight)
    set_weight_attrs(weight, extra_weight_attrs)
    layer.register_buffer("_mxfp8_weight", None, persistent=False)
    layer.register_buffer("_mxfp8_weight_scale", None, persistent=False)
fastvideo.layers.quantization.mxfp8_config.MXFP8QuantizeMethod.process_weights_after_loading 
process_weights_after_loading(layer: Module) -> None

Prequantize one adapter-merged BF16 linear weight.

Source code in fastvideo/layers/quantization/mxfp8_config.py 
50
51
52
53
54
55
56
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
    """Prequantize one adapter-merged BF16 linear weight."""
    from fastvideo.layers.mxfp8linear import quantize_mxfp8_weight_blockwise

    quantized_weight, blocked_scales = quantize_mxfp8_weight_blockwise(layer.weight.detach())
    layer._mxfp8_weight = quantized_weight
    layer._mxfp8_weight_scale = blocked_scales

Functions:

fastvideo.layers.quantization.mxfp8_config.convert_model_to_mxfp8 

convert_model_to_mxfp8(model: Module) -> int

Prequantize every MXFP8-tagged weight and return the converted count.

Source code in fastvideo/layers/quantization/mxfp8_config.py 
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
def convert_model_to_mxfp8(model: torch.nn.Module) -> int:
    """Prequantize every MXFP8-tagged weight and return the converted count."""
    converted_count = 0
    with torch.no_grad():
        for module in model.modules():
            quant_method = getattr(module, "quant_method", None)
            if not isinstance(quant_method, MXFP8QuantizeMethod):
                continue
            if converted_count == 0 and (not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 10):
                raise RuntimeError("MXFP8 inference requires an NVIDIA Blackwell GPU with compute capability 10.0+.")
            quant_method.process_weights_after_loading(module)
            converted_count += 1

    if converted_count:
        logger.info("Prequantized %d MiniMax-H3 feed-forward linear weights to MXFP8", converted_count)
        torch.cuda.empty_cache()
    return converted_count