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>
