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>
The public LTX-2 path was previously running full bf16 even when
callers set ``pipeline_config.dit_config.quant_config = FP4Config()``,
because:
* the DiT used plain ``torch.nn.Linear`` for FP4-eligible projections
(``attn1/attn2.to_q/k/v/out``, ``audio_to_video_attn``,
``video_to_audio_attn``, FFN ``fc_in``/``fc_out``), so
``LinearBase.quant_method.apply`` was never reached;
* the loader did not call ``convert_model_to_fp4`` after weights
loaded, so even matched layers had no ``_fp4_weight*`` buffers;
* ``QuantizationMethods`` did not register ``"FP4"``, so the typed
``engine.quantization.transformer_quant`` surface had no way to
select it.
This commit completes the inference wire-up through the existing
``fastvideo.layers.quantization`` registry — no parallel pathway:
* swap ``nn.Linear`` → ``ReplicatedLinear`` for the LTX-2 attention
``to_q``/``to_k``/``to_v``/``to_out``/``to_gate_compress`` and the
FFN ``fc_in``/``fc_out`` (``GELUApprox.proj`` and ``FeedForward``
``project_out``). Other linears (timestep MLP, caption proj,
patchify, AdaLN scale/shift) stay ``nn.Linear`` to match internal.
* port ``_supports_prequantized_input`` and
``_linear_project_with_optional_prequant`` so the attention forward
quantizes input once and reuses the ``(x_fp4, x_scale, x_global_sf)``
tuple across q/k/v projections — matches internal exactly.
* plumb ``quant_config`` and ``prefix`` through
``BasicAVTransformerBlock`` → ``_init_transformer_blocks`` →
``LTXModel`` → ``LTX2Transformer3DModel`` so each
``ReplicatedLinear`` gets the correct
``ltx2.blocks.<i>.<sub>.<proj>`` prefix that
``FP4Config.get_quant_method`` matches against.
* register ``"FP4"`` in ``QuantizationMethods`` literal +
``get_quantization_config`` registry so typed
``engine.quantization.transformer_quant: "FP4"`` resolves to a
concrete ``FP4Config()`` instance.
* add ``transformer_quant`` carrier on ``FastVideoArgs`` and
``__post_init__._apply_transformer_quant`` to pin the resolved
quant config onto ``dit_config.quant_config`` (without
overwriting an explicit setter on the dit_config).
* add ``_maybe_convert_model_to_fp4`` in ``fsdp_load.py`` that
walks the loaded model once and registers
``_fp4_weight*``/``_fp4_alpha``/``_weight_global_sf`` buffers on
layers whose ``quant_method`` is ``FP4QuantizeMethod``. flashinfer
is imported lazily inside ``convert_model_to_fp4`` so this is a
no-op on hosts without the FP4 kernels.
* restore ``LinearBase`` fallback to ``UnquantizedLinearMethod`` when
``quant_config.get_quant_method`` returns ``None`` — ``FP4Config``
only tags a curated subset of layers, and the previous
``assert quant_method is not None`` would crash any non-tagged
layer that received a quant_config.
Non-FP4 callers are unaffected: ``UnquantizedLinearMethod.apply``
runs the standard ``F.linear`` path, and parameter names
(``to_q.weight``, ``to_out.0.weight`` from ``nn.ModuleList``) are
identical to the previous ``nn.Sequential`` layout.
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>
Ports the FP4 quantization config + ops from FastVideo-internal
(fastvideo.layers.quantization.fp4_config). flashinfer is imported
lazily inside _require_flashinfer() rather than at module top so
``import fastvideo`` stays cheap on hosts that don't ship flashinfer
— the FP4 kernels themselves still raise a clear ImportError pointing
at ``pip install flashinfer-python`` when invoked.
Behavior is identical to the internal version:
* FP4Config: LTX-2 layer set (48 transformer blocks across attn1/attn2
/audio_to_video_attn/video_to_audio_attn/ffn plus adaln_single.linear),
layer_profile selects between "base" and "refine" stage subsets.
* FP4QuantizeMethod: per-layer quantize_input + apply with stage-aware
routing for refine-only layers (audio_to_video_attn.to_q,
video_to_audio_attn.to_k, video_to_audio_attn.to_v fall through to
dense F.linear during the base stage).
* convert_model_to_fp4(): pre-quantize weight buffers + alpha onto
module attrs at load time so the forward path doesn't pay the
per-step weight-side quantize cost.
Required for Dreamverse integration: dreamverse-server's GPU worker
runs ``pipeline_config.dit_config.quant_config = FP4Config()`` during
LTX-2 model load, which is the dominant inference path.
test_fp4_config covers the lazy-flashinfer contract: import succeeds
without flashinfer; layer_profile round-trips via from_config; calling
into _require_flashinfer raises ImportError with an actionable
message when flashinfer is missing.
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>
