b24c4f0d61
The FP4 implementation is specifically NVIDIA's block-scaled FP4 (e2m1 mantissa, fp32 alpha, ``layout_128x4`` scale layout, group size 16) backed by FlashInfer's ``nvfp4_quantize`` and ``mm_fp4`` kernels. Naming the public surface ``FP4`` would collide with other FP4 variants we may want to support later (OCP-FP4 / MX-FP4 / e3m0). Mechanical rename, no behavior change: * ``fp4_config.py`` → ``nvfp4_config.py`` * ``FP4Config`` → ``NVFP4Config``; ``get_name()`` returns ``"nvfp4"`` * ``FP4QuantizeMethod`` → ``NVFP4QuantizeMethod`` * ``convert_model_to_fp4`` → ``convert_model_to_nvfp4`` * ``QuantizationMethods`` literal: ``"FP4"`` → ``"NVFP4"`` * registered buffer names: ``_fp4_weight``/``_fp4_alpha`` → ``_nvfp4_weight``/``_nvfp4_alpha`` * loader helper ``_maybe_convert_model_to_fp4`` → ``_maybe_convert_model_to_nvfp4`` * test file rename + symbol updates * Dreamverse worker updates its single import + call site Internal-scope torch op namespace ``fastvideo_fp4::*`` and ``_get_ltx2_fp4_stage_profile`` left as-is (purely internal naming that mirrors FastVideo-internal). Co-authored-by: Junda (David) Su <90978028+Davids048@users.noreply.github.com> Co-authored-by: Matthew Noto <99706358+RandNMR73@users.noreply.github.com> Co-authored-by: XOR-op <17672363+XOR-op@users.noreply.github.com> Co-authored-by: Zhang Peiyuan <42993249+jzhang38@users.noreply.github.com>
72 lines
2.5 KiB
Python
72 lines
2.5 KiB
Python
from typing import Literal, get_args
|
|
|
|
from fastvideo.layers.quantization.base_config import QuantizationConfig
|
|
|
|
QuantizationMethods = Literal[None, "AbsMaxFP8", "NVFP4"]
|
|
|
|
QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods))
|
|
|
|
# The customized quantization methods which will be added to this dict.
|
|
_CUSTOMIZED_METHOD_TO_QUANT_CONFIG = {}
|
|
|
|
|
|
def register_quantization_config(quantization: str):
|
|
"""Register a customized vllm quantization config.
|
|
|
|
When a quantization method is not supported by vllm, you can register a customized
|
|
quantization config to support it.
|
|
|
|
Args:
|
|
quantization (str): The quantization method name.
|
|
|
|
Examples:
|
|
>>> from fastvideo.layers.quantization import register_quantization_config
|
|
>>> from fastvideo.layers.quantization import get_quantization_config
|
|
>>> from fastvideo.layers.quantization.base_config import QuantizationConfig
|
|
>>>
|
|
>>> @register_quantization_config("my_quant")
|
|
... class MyQuantConfig(QuantizationConfig):
|
|
... pass
|
|
>>>
|
|
>>> get_quantization_config("my_quant")
|
|
<class 'MyQuantConfig'>
|
|
""" # noqa: E501
|
|
|
|
def _wrapper(quant_config_cls):
|
|
if quantization in QUANTIZATION_METHODS:
|
|
raise ValueError(f"The quantization method `{quantization}` is already exists.")
|
|
if not issubclass(quant_config_cls, QuantizationConfig):
|
|
raise ValueError("The quantization config must be a subclass of "
|
|
"`QuantizationConfig`.")
|
|
_CUSTOMIZED_METHOD_TO_QUANT_CONFIG[quantization] = quant_config_cls
|
|
QUANTIZATION_METHODS.append(quantization)
|
|
return quant_config_cls
|
|
|
|
return _wrapper
|
|
|
|
|
|
def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
|
|
if quantization not in QUANTIZATION_METHODS:
|
|
raise ValueError(f"Invalid quantization method: {quantization}")
|
|
|
|
# lazy import to avoid triggering `torch.compile` too early
|
|
from .absmax_fp8 import AbsMaxFP8Config
|
|
from .nvfp4_config import NVFP4Config
|
|
|
|
method_to_config: dict[str, type[QuantizationConfig]] = {
|
|
"AbsMaxFP8": AbsMaxFP8Config,
|
|
"NVFP4": NVFP4Config,
|
|
}
|
|
# Update the `method_to_config` with customized quantization methods.
|
|
method_to_config.update(_CUSTOMIZED_METHOD_TO_QUANT_CONFIG)
|
|
|
|
return method_to_config[quantization]
|
|
|
|
|
|
all = [
|
|
"QuantizationMethods",
|
|
"QuantizationConfig",
|
|
"get_quantization_config",
|
|
"QUANTIZATION_METHODS",
|
|
]
|