Skip to content
FastVideo
Esc
↑↓navigate↵open⌘Jpreview
On this page

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&#91;:-1&#93;, out_dim)</code></pre></details>


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() &gt; 2:
        x_fp8 = x_fp8.reshape(-1, x_fp8.shape&#91;-1&#93;)
    if x_scale.dim() &gt; 2:
        x_scale = x_scale.reshape(-1, x_scale.shape&#91;-1&#93;)
elif self.granularity == &quot;channel&quot;:
    x_fp8, x_scale = _quantize_rowwise(x.reshape(-1, x.shape&#91;-1&#93;))
else:
    x_fp8, x_scale = _quantize_tensorwise(x.reshape(-1, x.shape&#91;-1&#93;))

w_fp8 = layer._fp8_weight
w_scale = layer._fp8_weight_scale
if self.granularity == &quot;channel&quot; 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&#91;:-1&#93;, out_dim)
scale_b = w_scale.view(1, -1) if self.granularity == &quot;channel&quot; 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&#91;0&#93;
if bias is not None:
    out = out + bias
return out.view(*original_shape&#91;:-1&#93;, out_dim)</code></pre></details>


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 = &quot;tensor&quot;):
    super().__init__()
    if granularity not in (&quot;tensor&quot;, &quot;channel&quot;):
        raise ValueError(f&quot;granularity must be &#x27;tensor&#x27; or &#x27;channel&#x27;, got &#123;granularity!r&#125;&quot;)
    self.granularity = granularity

def get_name(self) -&gt; str:
    return &quot;FP8&quot;

def get_supported_act_dtypes(self) -&gt; list&#91;torch.dtype&#93;:
    return &#91;torch.bfloat16, torch.float16&#93;

@classmethod
def get_min_capability(cls) -&gt; int:
    return 89

@staticmethod
def get_config_filenames() -&gt; list&#91;str&#93;:
    return &#91;&#93;

@classmethod
def from_config(cls, config: dict&#91;str, Any&#93;) -&gt; FP8Config:
    return cls(granularity=config.get(&quot;granularity&quot;, &quot;tensor&quot;))

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</code></pre></details>


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</code></pre></details>


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, &quot;quant_method&quot;, None)
        if not isinstance(qm, FP8QuantizeMethod):
            continue
        weight = getattr(mod, &quot;weight&quot;, None)
        if weight is None:
            continue
        weight_local = weight.to_local() if isinstance(weight, DTensor) else weight  # type: ignore&#91;arg-type&#93;
        if getattr(qm, &quot;granularity&quot;, &quot;tensor&quot;) == &quot;channel&quot;:
            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(&quot;_fp8_weight&quot;, w_fp8.contiguous(), persistent=False)
        mod.register_buffer(&quot;_fp8_weight_scale&quot;, w_scale.to(torch.float32), persistent=False)
        removed_weight = mod._parameters.pop(&quot;weight&quot;, 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()</code></pre></details>

Was this page helpful?