# SPDX-License-Identifier: Apache-2.0 """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. Falls back to bf16 dequant on GPUs older than sm89. """ from __future__ import annotations import logging from typing import Any import torch import torch.nn.functional as F from torch.nn.parameter import Parameter from fastvideo.layers.quantization.base_config import ( QuantizationConfig, QuantizeMethodBase, ) from fastvideo.models.utils import set_weight_attrs logger = logging.getLogger(__name__) FP8_DTYPE = torch.float8_e4m3fn FP8_MAX = float(torch.finfo(FP8_DTYPE).max) # 448.0 FP8_MIN_SCALE = 1.0 / (FP8_MAX * 512.0) # Wan's "to_q"/"to_k"/"to_v" also substring-match Kandinsky5's "to_query"/ # "to_key"/"to_value", so only Kandinsky5's out-projection and FFN names need # to be listed explicitly below. _FP8_SUFFIXES = ( "ffn.fc_in", "ffn.fc_out", "to_q", "to_k", "to_v", "to_out", # Kandinsky5 "self_attention.out_layer", "cross_attention.out_layer", "feed_forward.mlp.fc_in", "feed_forward.mlp.fc_out", ) def _supports_fp8_compute() -> bool: """Whether the active device supports FP8 ``_scaled_mm`` (sm89+).""" if not torch.cuda.is_available(): return False cap = torch.cuda.get_device_capability() return cap[0] > 8 or (cap[0] == 8 and cap[1] >= 9) def _quantize_tensorwise(x_2d: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """Returns ``(x_fp8 [M, K], x_scale [1] float32)``.""" x_absmax = x_2d.abs().amax().float() x_scale = (x_absmax / FP8_MAX).clamp(min=FP8_MIN_SCALE) x_fp8 = (x_2d / x_scale.to(x_2d.dtype)).clamp(-FP8_MAX, FP8_MAX).to(FP8_DTYPE) return x_fp8, x_scale.view(1) def _quantize_rowwise(x_2d: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """Returns ``(x_fp8 [M, K], x_scale [M, 1] float32)``.""" x_absmax = x_2d.abs().amax(dim=-1, keepdim=True).float() x_scale = (x_absmax / FP8_MAX).clamp(min=FP8_MIN_SCALE) x_fp8 = (x_2d / x_scale.to(x_2d.dtype)).clamp(-FP8_MAX, FP8_MAX).to(FP8_DTYPE) return x_fp8, x_scale 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 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) 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 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() __all__ = [ "FP8Config", "FP8QuantizeMethod", "convert_model_to_fp8", ]