--- title: fp8_config ---

Generic FP8 quantization backed by torch._scaled_mm.

Matches linear layers by suffix (to_q/k/v/to_out, ffn.fc_in/fc_out). Supports per-tensor (default, fast) and per-channel (higher accuracy) granularity. On sm89 the per-channel path runs the per-tensor GEMM plus a Triton scale epilogue (fp8_kernels), since torch's rowwise-scaled kernel there is slower than bf16. Falls back to bf16 dequant on GPUs older than sm89.

fastvideo.layers.quantization.fp8_config.QuantizationConfig

QuantizationConfig()

Base class for quantization configs.

fastvideo.layers.quantization.fp8_config.QuantizationConfig.get_name

get_name() -> QuantizationMethods

Name of the quantization method.

fastvideo.layers.quantization.fp8_config.QuantizationConfig.get_supported_act_dtypes

get_supported_act_dtypes() -> list[torch.dtype]

List of supported activation dtypes.

fastvideo.layers.quantization.fp8_config.QuantizationConfig.get_min_capability

get_min_capability() -> int

Minimum GPU capability to support the quantization method.

E.g., 70 for Volta, 75 for Turing, 80 for Ampere. This requirement is due to the custom CUDA kernels used by the quantization method.

fastvideo.layers.quantization.fp8_config.QuantizationConfig.get_config_filenames

get_config_filenames() -> list[str]

List of filenames to search for in the model directory.

fastvideo.layers.quantization.fp8_config.QuantizationConfig.from_config

from_config(config: dict[str, Any]) -> QuantizationConfig

Create a config class from the model's quantization config.

fastvideo.layers.quantization.fp8_config.QuantizationConfig.override_quantization_method

override_quantization_method(hf_quant_cfg, user_quant) -> QuantizationMethods | None

Detects if this quantization method can support a given checkpoint format by overriding the user specified quantization method -- this method should only be overwritten by subclasses in exceptional circumstances

fastvideo.layers.quantization.fp8_config.QuantizationConfig.get_from_keys

get_from_keys(config: dict[str, Any], keys: list[str]) -> Any

Get a value from the model's quantization config.

fastvideo.layers.quantization.fp8_config.QuantizationConfig.get_from_keys_or

get_from_keys_or(config: dict[str, Any], keys: list[str], default: Any) -> Any

Get a optional value from the model's quantization config.

fastvideo.layers.quantization.fp8_config.QuantizationConfig.get_quant_method

get_quant_method(layer: torch.nn.Module, prefix: str) -> QuantizeMethodBase | None

Get the quantize method to use for the quantized layer.

Parameters

layer: torch.nn.Module

The layer for the quant method.

prefix: str

The full name of the layer in the state dict

Returns: The quantize method. None if the given layer doesn't support quant method.

fastvideo.layers.quantization.fp8_config.QuantizationConfig.validate_runtime

validate_runtime(device: torch.device) -> None

Fail fast on a device the scheme cannot execute on.

Called once by loaders that select a serialized checkpoint scheme, before any weight is read. The default accepts every device; schemes with kernel or capability requirements override it.

fastvideo.layers.quantization.fp8_config.QuantizationConfig.get_cache_scale

get_cache_scale(name: str) -> str | None

fastvideo.layers.quantization.fp8_config.QuantizeMethodBase

QuantizeMethodBase()

Base class for different quantized methods.

fastvideo.layers.quantization.fp8_config.QuantizeMethodBase.create_weights

create_weights(layer: torch.nn.Module, *weight_args, **extra_weight_attrs)

Create weights for a layer.

The weights will be set as attributes of the layer.

fastvideo.layers.quantization.fp8_config.QuantizeMethodBase.apply

apply(layer: torch.nn.Module, *args, **kwargs) -> torch.Tensor

Apply the weights in layer to the input tensor.

Expects create_weights to have been called before on the layer.

fastvideo.layers.quantization.fp8_config.QuantizeMethodBase.embedding

embedding(layer: torch.nn.Module, *args, **kwargs) -> torch.Tensor

Gather embeddings in the layer based on indices in the input tensor.

Expects create_weights to have been called before on the layer.

fastvideo.layers.quantization.fp8_config.QuantizeMethodBase.process_weights_after_loading

process_weights_after_loading(layer: nn.Module) -> None

Process the weight after loading.

This can be used for example, to transpose weights for computation.

fastvideo.layers.quantization.fp8_config.quantize_rowwise_fused

quantize_rowwise_fused(x_2d: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]

Returns (x_fp8 [M, K], x_scale [M, 1] float32); one kernel launch.

fastvideo.layers.quantization.fp8_config.quantize_tensorwise_fused

quantize_tensorwise_fused(x_2d: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]

Returns (x_fp8 [M, K], x_scale [1] float32); an aminmax reduction + one pass.

fastvideo.layers.quantization.fp8_config.triton_quant_available

triton_quant_available(x: torch.Tensor) -> bool

fastvideo.layers.quantization.fp8_config.set_weight_attrs

set_weight_attrs(weight: torch.Tensor, weight_attrs: dict[str, Any] | None)

Set attributes on a weight tensor.

This method is used to set attributes on a weight tensor. This method will not overwrite existing attributes.

Parameters

weight: torch.Tensor

The weight tensor.

weight_attrs: dict[str, Any] | None

A dictionary of attributes to set on the weight tensor.

fastvideo.layers.quantization.fp8_config.FP8QuantizeMethod

FP8QuantizeMethod(granularity: str = 'tensor')

FP8 linear method.

granularity='tensor' (default): per-tensor weight + per-tensor dynamic activation scales — the fast tensorwise _scaled_mm path. granularity='channel': per-output-channel weight + per-token activation scales (rowwise) — higher accuracy but slower _scaled_mm.

Source
class FP8QuantizeMethod(QuantizeMethodBase):
    """FP8 linear method.

    ``granularity='tensor'`` (default): per-tensor weight + per-tensor
    dynamic activation scales — the fast tensorwise ``_scaled_mm`` path.
    ``granularity='channel'``: per-output-channel weight + per-token
    activation scales (rowwise) — higher accuracy but slower ``_scaled_mm``.
    """

    def __init__(self, granularity: str = "tensor"):
        super().__init__()
        self.granularity = granularity

    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,
    ):
        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)

    def quantize_input(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, None]:
        """Pre-quantize an activation for reuse across q/k/v projections."""
        assert x.dtype in (torch.bfloat16, torch.float16), (f"only allow bf16/fp16 inputs to fp8 linear, got {x.dtype}")
        x_2d = x.view(-1, x.shape[-1])
        if self.granularity == "channel":
            x_fp8, x_scale = _quantize_rowwise(x_2d)
        else:
            x_fp8, x_scale = _quantize_tensorwise(x_2d)
        return x_fp8, x_scale, None

    def wants_prequantized_input(self) -> bool:
        return True

    def apply(
        self,
        layer: torch.nn.Module,
        x: torch.Tensor,
        bias: torch.Tensor | None = None,
        pre_quantized: tuple[torch.Tensor, torch.Tensor, Any] | None = None,
    ) -> torch.Tensor:
        out_dim = layer._fp8_weight.shape[0]
        original_shape = x.shape

        if not _supports_fp8_compute():
            return self._apply_dequant(layer, x, bias)

        if pre_quantized is not None:
            x_fp8, x_scale, _ = pre_quantized
            if x_fp8.dim() > 2:
                x_fp8 = x_fp8.reshape(-1, x_fp8.shape[-1])
            if x_scale.dim() > 2:
                x_scale = x_scale.reshape(-1, x_scale.shape[-1])
        elif self.granularity == "channel":
            x_fp8, x_scale = _quantize_rowwise(x.reshape(-1, x.shape[-1]))
        else:
            x_fp8, x_scale = _quantize_tensorwise(x.reshape(-1, x.shape[-1]))

        w_fp8 = layer._fp8_weight
        w_scale = layer._fp8_weight_scale
        if self.granularity == "channel" and _rowwise_scaled_mm_is_slow():
            from fastvideo.layers.quantization.fp8_kernels import scaled_mm_token_channel
            out = scaled_mm_token_channel(x_fp8, x_scale, w_fp8.t(), w_scale)
            if bias is not None:
                out = out + bias
            return out.view(*original_shape[:-1], out_dim)
        scale_b = w_scale.view(1, -1) if self.granularity == "channel" else w_scale

        out = torch._scaled_mm(
            x_fp8,
            w_fp8.t(),
            scale_a=x_scale,
            scale_b=scale_b,
            out_dtype=torch.bfloat16,
        )
        if isinstance(out, tuple):
            out = out[0]
        if bias is not None:
            out = out + bias
        return out.view(*original_shape[:-1], out_dim)

    def _apply_dequant(
        self,
        layer: torch.nn.Module,
        x: torch.Tensor,
        bias: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """bf16 fallback for pre-sm89 GPUs."""
        out_dim = layer._fp8_weight.shape[0]
        original_shape = x.shape
        w_fp8 = layer._fp8_weight
        w_scale = layer._fp8_weight_scale.to(x.dtype)
        weight = w_fp8.to(x.dtype) * w_scale.unsqueeze(1)
        out = F.linear(x, weight, bias)
        return out.view(*original_shape[:-1], out_dim)

fastvideo.layers.quantization.fp8_config.FP8QuantizeMethod.create_weights

create_weights(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)
Source
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,
):
    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)

fastvideo.layers.quantization.fp8_config.FP8QuantizeMethod.quantize_input

quantize_input(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, None]

Pre-quantize an activation for reuse across q/k/v projections.

Source
def quantize_input(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, None]:
    """Pre-quantize an activation for reuse across q/k/v projections."""
    assert x.dtype in (torch.bfloat16, torch.float16), (f"only allow bf16/fp16 inputs to fp8 linear, got {x.dtype}")
    x_2d = x.view(-1, x.shape[-1])
    if self.granularity == "channel":
        x_fp8, x_scale = _quantize_rowwise(x_2d)
    else:
        x_fp8, x_scale = _quantize_tensorwise(x_2d)
    return x_fp8, x_scale, None

fastvideo.layers.quantization.fp8_config.FP8QuantizeMethod.wants_prequantized_input

wants_prequantized_input() -> bool
Source
def wants_prequantized_input(self) -> bool:
    return True

fastvideo.layers.quantization.fp8_config.FP8QuantizeMethod.apply

apply(layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None = None, pre_quantized: tuple[torch.Tensor, torch.Tensor, Any] | None = None) -> torch.Tensor
Source
def apply(
    self,
    layer: torch.nn.Module,
    x: torch.Tensor,
    bias: torch.Tensor | None = None,
    pre_quantized: tuple[torch.Tensor, torch.Tensor, Any] | None = None,
) -> torch.Tensor:
    out_dim = layer._fp8_weight.shape[0]
    original_shape = x.shape

    if not _supports_fp8_compute():
        return self._apply_dequant(layer, x, bias)

    if pre_quantized is not None:
        x_fp8, x_scale, _ = pre_quantized
        if x_fp8.dim() > 2:
            x_fp8 = x_fp8.reshape(-1, x_fp8.shape[-1])
        if x_scale.dim() > 2:
            x_scale = x_scale.reshape(-1, x_scale.shape[-1])
    elif self.granularity == "channel":
        x_fp8, x_scale = _quantize_rowwise(x.reshape(-1, x.shape[-1]))
    else:
        x_fp8, x_scale = _quantize_tensorwise(x.reshape(-1, x.shape[-1]))

    w_fp8 = layer._fp8_weight
    w_scale = layer._fp8_weight_scale
    if self.granularity == "channel" and _rowwise_scaled_mm_is_slow():
        from fastvideo.layers.quantization.fp8_kernels import scaled_mm_token_channel
        out = scaled_mm_token_channel(x_fp8, x_scale, w_fp8.t(), w_scale)
        if bias is not None:
            out = out + bias
        return out.view(*original_shape[:-1], out_dim)
    scale_b = w_scale.view(1, -1) if self.granularity == "channel" else w_scale

    out = torch._scaled_mm(
        x_fp8,
        w_fp8.t(),
        scale_a=x_scale,
        scale_b=scale_b,
        out_dtype=torch.bfloat16,
    )
    if isinstance(out, tuple):
        out = out[0]
    if bias is not None:
        out = out + bias
    return out.view(*original_shape[:-1], out_dim)

fastvideo.layers.quantization.fp8_config.FP8Config

FP8Config(granularity: str = 'tensor')

FP8 (e4m3) quantization via suffix matching on standard linear layer names.

Source
class FP8Config(QuantizationConfig):
    """FP8 (e4m3) quantization via suffix matching on standard linear layer names."""

    def __init__(self, granularity: str = "tensor"):
        super().__init__()
        if granularity not in ("tensor", "channel"):
            raise ValueError(f"granularity must be 'tensor' or 'channel', got {granularity!r}")
        self.granularity = granularity

    def get_name(self) -> str:
        return "FP8"

    def get_supported_act_dtypes(self) -> list[torch.dtype]:
        return [torch.bfloat16, torch.float16]

    @classmethod
    def get_min_capability(cls) -> int:
        return 89

    @staticmethod
    def get_config_filenames() -> list[str]:
        return []

    @classmethod
    def from_config(cls, config: dict[str, Any]) -> FP8Config:
        return cls(granularity=config.get("granularity", "tensor"))

    def get_quant_method(self, layer: torch.nn.Module, prefix: str):
        from fastvideo.layers.linear import LinearBase

        if isinstance(layer, LinearBase) and any(s in prefix for s in _FP8_SUFFIXES):
            return FP8QuantizeMethod(granularity=self.granularity)
        return None

fastvideo.layers.quantization.fp8_config.FP8Config.get_name

get_name() -> str
Source
def get_name(self) -> str:
    return "FP8"

fastvideo.layers.quantization.fp8_config.FP8Config.get_supported_act_dtypes

get_supported_act_dtypes() -> list[torch.dtype]
Source
def get_supported_act_dtypes(self) -> list[torch.dtype]:
    return [torch.bfloat16, torch.float16]

fastvideo.layers.quantization.fp8_config.FP8Config.get_min_capability

get_min_capability() -> int
Source
@classmethod
def get_min_capability(cls) -> int:
    return 89

fastvideo.layers.quantization.fp8_config.FP8Config.get_config_filenames

get_config_filenames() -> list[str]
Source
@staticmethod
def get_config_filenames() -> list[str]:
    return []

fastvideo.layers.quantization.fp8_config.FP8Config.from_config

from_config(config: dict[str, Any]) -> FP8Config
Source
@classmethod
def from_config(cls, config: dict[str, Any]) -> FP8Config:
    return cls(granularity=config.get("granularity", "tensor"))

fastvideo.layers.quantization.fp8_config.FP8Config.get_quant_method

get_quant_method(layer: torch.nn.Module, prefix: str)
Source
def get_quant_method(self, layer: torch.nn.Module, prefix: str):
    from fastvideo.layers.linear import LinearBase

    if isinstance(layer, LinearBase) and any(s in prefix for s in _FP8_SUFFIXES):
        return FP8QuantizeMethod(granularity=self.granularity)
    return None

fastvideo.layers.quantization.fp8_config.convert_model_to_fp8

convert_model_to_fp8(model: torch.nn.Module) -> None

Quantize all FP8-tagged linear layers in-place after weights are loaded.

Source
def convert_model_to_fp8(model: torch.nn.Module) -> None:
    """Quantize all FP8-tagged linear layers in-place after weights are loaded."""
    import gc
    from torch.distributed.tensor import DTensor  # type: ignore

    with torch.no_grad():
        for mod in model.modules():
            qm = getattr(mod, "quant_method", None)
            if not isinstance(qm, FP8QuantizeMethod):
                continue
            weight = getattr(mod, "weight", None)
            if weight is None:
                continue
            weight_local = weight.to_local() if isinstance(weight, DTensor) else weight  # type: ignore[arg-type]
            if getattr(qm, "granularity", "tensor") == "channel":
                w_absmax = weight_local.detach().abs().amax(dim=1).nan_to_num().float()
                w_scale = (w_absmax / FP8_MAX).clamp(min=FP8_MIN_SCALE)
                w_fp8 = (weight_local / w_scale.to(weight_local.dtype).unsqueeze(1)).clamp(-FP8_MAX,
                                                                                           FP8_MAX).to(FP8_DTYPE)
            else:
                w_absmax = weight_local.detach().abs().amax().nan_to_num().to(torch.float32)
                w_scale = (w_absmax / FP8_MAX).clamp(min=FP8_MIN_SCALE).view(1)
                w_fp8 = (weight_local / w_scale.to(weight_local.dtype)).clamp(-FP8_MAX, FP8_MAX).to(FP8_DTYPE)
            mod.register_buffer("_fp8_weight", w_fp8.contiguous(), persistent=False)
            mod.register_buffer("_fp8_weight_scale", w_scale.to(torch.float32), persistent=False)
            removed_weight = mod._parameters.pop("weight", None)
            if removed_weight is not None:
                removed_weight.grad = None
            del removed_weight, weight, weight_local, w_absmax, w_scale, w_fp8

    gc.collect()
    torch.cuda.empty_cache()