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. Falls back to bf16 dequant on GPUs older than sm89.

Classes¶

fastvideo.layers.quantization.fp8_config.FP8Config ¶

FP8Config(granularity: str = 'tensor')

Bases: QuantizationConfig

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

Source code in fastvideo/layers/quantization/fp8_config.py

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

fastvideo.layers.quantization.fp8_config.FP8QuantizeMethod ¶

FP8QuantizeMethod(granularity: str = 'tensor')

Bases: 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.

Source code in fastvideo/layers/quantization/fp8_config.py

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

Methods:¶

fastvideo.layers.quantization.fp8_config.FP8QuantizeMethod.quantize_input ¶

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

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

Source code in fastvideo/layers/quantization/fp8_config.py

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

Functions:¶

fastvideo.layers.quantization.fp8_config.convert_model_to_fp8 ¶

convert_model_to_fp8(model: Module) -> None

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

Source code in fastvideo/layers/quantization/fp8_config.py

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