[bugfix] QAD 5090: emit sm_120a in build.sh so attn_qat_infer kernels build

build.sh auto-detected Blackwell (sm_120) and exported TORCH_CUDA_ARCH_LIST=12.0 without the arch-conditional 'a' suffix. CMake's AUTO gate for the attn_qat_infer (modified SageAttention3 FP4) kernels only matches 12.0a/120a/sm_120a, so fp4attn_cuda/fp4quant_cuda were silently skipped and the ATTN_QAT_INFER backend fell back to Flash Attention at runtime. Exporting the env var also bypassed CMake's local-GPU fallback that would otherwise have enabled them.

Mirror the existing 9.0 -> 9.0a Hopper handling for 12.0 -> 12.0a.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

examples/inference/optimizations/FastWan_QAD_TAEHV.py

@@ -0,0 +1,286 @@
+"""Fast NVFP4 linear inference for Wan2.1-T2V-1.3B with TAEHV decoding.
+
+This is the FP4-linear fast path from ``fp4_linear_wan2_1_1_3b.py`` with the
+heavy Wan VAE swapped out for TAEHV -- a tiny autoencoder that decodes Wan2.1
+latents directly (no denormalization) and is dramatically faster / lighter.
+
+How it works: the generator runs with ``output_type="latent"`` so the pipeline
+returns raw denoised latents instead of pixels (the Wan VAE is offloaded and
+never used). We then decode those latents with TAEHV in this script and save
+the frames ourselves. This mirrors the FastVideo-Quantization
+``quantization_example_taehv.py`` proof-of-concept, but kept clean: TAEHV is a
+pip package (no ``sys.path`` hacks), the latent->uint8 conversion is vectorized,
+and there is no dead profiler / sanitization code.
+
+Requirements:
+    - Blackwell GPU (B200/B300, sm100a/sm103a) for the FP4 linear path
+    - flashinfer (``pip install flashinfer-python``)
+    - TAEHV weights ``taew2_1.pth`` (https://github.com/madebyollin/taehv)
+
+Usage:
+    python fp4_linear_taehv_wan2_1_1_3b.py                    # FP4 + TAEHV + compile
+    python fp4_linear_taehv_wan2_1_1_3b.py --no-taehv         # FP4 + full Wan VAE
+    python fp4_linear_taehv_wan2_1_1_3b.py --no-compile       # eager
+    python fp4_linear_taehv_wan2_1_1_3b.py --baseline         # dense bf16 reference
+    python fp4_linear_taehv_wan2_1_1_3b.py --distilled_model ''  # base Wan2.1 weights
+"""
+
+import argparse
+import contextlib
+import logging
+import os
+import time
+
+import imageio
+import torch
+
+from fastvideo import VideoGenerator
+from fastvideo.configs.pipelines.base import PipelineConfig
+from fastvideo.layers.quantization.nvfp4_qat_config import NVFP4QATConfig
+
+OUTPUT_PATH = "video_samples"
+
+# Distilled, quantization-aware (QAD) transformer for Wan2.1-1.3B (3 steps,
+# guidance 1.0). Loaded on top of the base Wan2.1 pipeline; pass
+# ``--distilled_model ''`` to run the base weights instead.
+DEFAULT_DISTILLED_MODEL = "FastVideo/FastWan-QAD-1.3B"
+DISTILLED_WEIGHTS_FILE = (
+    "generator_inference_transformer/diffusion_pytorch_model.safetensors"
+)
+
+# TAEHV checkpoint for Wan2.1. Clone https://github.com/madebyollin/taehv to get
+# ``taew2_1.pth`` (Wan 2.1 / Wan 2.2-14B / Qwen-Image all use this VAE).
+DEFAULT_TAEHV_CHECKPOINT = "/root/taehv/taew2_1.pth"
+
+PROMPT = (
+    "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
+    "wide with interest. The playful yet serene atmosphere is complemented by soft "
+    "natural light filtering through the petals. Mid-shot, warm and cheerful tones."
+)
+
+
+class TaehvDecoder:
+    """Thin wrapper around the TAEHV tiny autoencoder for Wan2.1 latents.
+
+    TAEHV consumes the *normalized* latents the diffusion model produces (the
+    same representation FastVideo carries internally), so no denormalization is
+    needed -- unlike the full Wan VAE path.
+    """
+
+    def __init__(self, checkpoint_path: str, device: str = "cuda",
+                 dtype: torch.dtype = torch.float16) -> None:
+        from taehv import TAEHV  # pip-installed; no sys.path manipulation
+        self.device = device
+        self.dtype = dtype
+        print(f"Loading TAEHV from {checkpoint_path} ...")
+        self.model = TAEHV(checkpoint_path=checkpoint_path).to(device, dtype).eval()
+
+    @torch.no_grad()
+    def decode(self, latents: torch.Tensor):
+        """Decode FastVideo latents into uint8 RGB frames.
+
+        Args:
+            latents: ``[B, C, T, H, W]`` (NCTHW) normalized latent tensor.
+
+        Returns:
+            A ``(T, H, W, 3)`` uint8 numpy array ready for ``imageio.mimsave``.
+        """
+        # NCTHW -> NTCHW (TAEHV's expected layout), on the TAEHV device/dtype.
+        latents = latents.permute(0, 2, 1, 3, 4).to(self.device, self.dtype)
+        decoded = self.model.decode_video(
+            latents, parallel=True, show_progress_bar=False)
+        # decoded: [B, T, 3, H, W] in [0, 1]. Take batch 0, vectorize to uint8.
+        frames = (decoded[0].clamp(0, 1) * 255).to(torch.uint8)
+        return frames.permute(0, 2, 3, 1).cpu().numpy()
+
+
+def resolve_distilled_weights(hf_id: str) -> str:
+    """Return a local path to the distilled transformer safetensors."""
+    if os.path.exists(hf_id):
+        return hf_id
+    from huggingface_hub import hf_hub_download
+    return hf_hub_download(repo_id=hf_id, filename=DISTILLED_WEIGHTS_FILE)
+
+
+@contextlib.contextmanager
+def silence_request_log():
+    """Quiet ``VideoGenerator.generate``'s per-request config printout.
+
+    Each ``generate(...)`` call logs a multi-line debug block (height/width/
+    prompt/steps/...) at INFO via ``logger.info`` in
+    ``fastvideo.entrypoints.video_generator``. There is no built-in switch,
+    so this context manager raises that logger's level to WARNING while the
+    warmup calls run, then restores it for the timed run.
+    """
+    vg_logger = logging.getLogger("fastvideo.entrypoints.video_generator")
+    prev_level = vg_logger.level
+    vg_logger.setLevel(logging.WARNING)
+    try:
+        yield
+    finally:
+        vg_logger.setLevel(prev_level)
+
+
+def resolve_taehv_checkpoint(path: str) -> str:
+    """Validate the TAEHV checkpoint path, with a helpful error if missing."""
+    if os.path.exists(path):
+        return path
+    raise FileNotFoundError(
+        f"TAEHV checkpoint not found at {path!r}. Clone the weights with:\n"
+        "    git clone https://github.com/madebyollin/taehv\n"
+        "and pass --taehv_checkpoint <repo>/taew2_1.pth")
+
+
+def build_generator(args: argparse.Namespace) -> VideoGenerator:
+    model_id = args.model
+
+    # Half precision everywhere; DiT linears are additionally NVFP4-quantized
+    # via dit_config.quant_config below.
+    pipeline_config = PipelineConfig.from_pretrained(model_id)
+    pipeline_config.dit_precision = "bf16"
+    pipeline_config.vae_precision = "bf16"
+    pipeline_config.text_encoder_precisions = ("bf16",)
+
+    if not args.baseline:
+        pipeline_config.dit_config.quant_config = NVFP4QATConfig()
+
+    compile_enabled = not args.no_compile
+
+    extra_kwargs = {}
+    if args.distilled_model:
+        weights_path = resolve_distilled_weights(args.distilled_model)
+        print(f"Using distilled weights: {args.distilled_model} -> {weights_path}")
+        extra_kwargs["init_weights_from_safetensors"] = weights_path
+
+    if args.taehv:
+        # Skip the in-pipeline VAE decode entirely: the pipeline returns raw
+        # latents, the Wan VAE is offloaded to CPU (and not compiled) since we
+        # decode with TAEHV in this script instead.
+        extra_kwargs["output_type"] = "latent"
+
+    generator = VideoGenerator.from_pretrained(
+        model_id,
+        pipeline_config=pipeline_config,
+        num_gpus=args.num_gpus,
+        # Keep everything resident on the GPU -- no offloading, except the
+        # unused Wan VAE when TAEHV handles decoding.
+        use_fsdp_inference=False,
+        dit_cpu_offload=False,
+        dit_layerwise_offload=False,
+        vae_cpu_offload=args.taehv,
+        text_encoder_cpu_offload=False,
+        pin_cpu_memory=False,
+        enable_torch_compile=compile_enabled,
+        enable_torch_compile_text_encoder=compile_enabled,
+        enable_torch_compile_vae=compile_enabled and not args.taehv,
+        **extra_kwargs,
+    )
+    return generator
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="FP4 linear Wan2.1-1.3B with TAEHV decoding benchmark")
+    parser.add_argument("--model", default="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+                        help="Model path or HuggingFace ID")
+    parser.add_argument("--baseline", action="store_true",
+                        help="Run dense bf16 instead of FP4 linear")
+    parser.add_argument("--no-compile", action="store_true",
+                        help="Disable torch.compile (eager)")
+    parser.add_argument("--taehv", action=argparse.BooleanOptionalAction,
+                        default=True,
+                        help="Decode with TAEHV instead of the full Wan VAE "
+                             "(use --no-taehv for the Wan VAE path)")
+    parser.add_argument("--taehv_checkpoint", default=DEFAULT_TAEHV_CHECKPOINT,
+                        help="Path to the TAEHV taew2_1.pth checkpoint")
+    parser.add_argument("--distilled_model", default=DEFAULT_DISTILLED_MODEL,
+                        help="HuggingFace ID (or local path) of a distilled "
+                             "transformer checkpoint to load on top of --model. "
+                             "Pass '' to use the base --model weights instead.")
+    parser.add_argument("--num_gpus", type=int, default=1)
+    parser.add_argument("--infer_steps", type=int, default=3)
+    parser.add_argument("--guidance_scale", type=float, default=1.0)
+    args = parser.parse_args()
+
+    if not torch.cuda.is_available():
+        raise SystemExit("CUDA is required for FP4 inference.")
+
+    cap = torch.cuda.get_device_capability()
+    print(f"GPU: {torch.cuda.get_device_name()} (capability {cap[0]}.{cap[1]})")
+    if not args.baseline and cap[0] < 10:
+        print("Warning: NVFP4 requires Blackwell (capability 10.0+); "
+              "FP4 kernels may be unavailable on this GPU.")
+
+    mode = "bf16" if args.baseline else "fp4_linear"
+    mode += "_taehv" if args.taehv else "_wanvae"
+    if not args.no_compile:
+        mode += "_compile"
+    print(f"Mode: {mode.upper()}")
+
+    # Load TAEHV before the (slow) generator build so a bad checkpoint path
+    # fails fast.
+    taehv = TaehvDecoder(resolve_taehv_checkpoint(args.taehv_checkpoint)) \
+        if args.taehv else None
+
+    generator = build_generator(args)
+
+    os.makedirs(OUTPUT_PATH, exist_ok=True)
+
+    # Warmup: with compile enabled the first call(s) pay the DiT compilation
+    # cost. When using TAEHV we also decode the warmup latents so the timed
+    # decode below is warm -- TAEHV's decoder is all conv/upsample, so the
+    # first call otherwise pays cuDNN algo selection + allocator growth
+    # (~0.2s), which is exactly the cold-start overhead we want to exclude.
+    n_warmup = 2 if not args.no_compile else 1
+    with silence_request_log():
+        for _ in range(n_warmup):
+            warm = generator.generate(request={
+                "prompt": PROMPT,
+                "sampling": {"num_inference_steps": 2, "guidance_scale": args.guidance_scale},
+                "output": {"save_video": False, "return_frames": args.taehv},
+            })
+            if args.taehv:
+                taehv.decode(warm.samples)
+
+    output_path = os.path.join(OUTPUT_PATH, f"raccoon_{mode}.mp4")
+    torch.cuda.synchronize()
+    start = time.perf_counter()
+    result = generator.generate(request={
+        "prompt": PROMPT,
+        "sampling": {
+            "num_inference_steps": args.infer_steps,
+            "guidance_scale": args.guidance_scale,
+        },
+        # When using TAEHV we need the latents back and save manually; the Wan
+        # VAE path lets the pipeline decode and save the mp4 itself.
+        "output": {
+            "save_video": not args.taehv,
+            "return_frames": args.taehv,
+            "output_path": output_path,
+        },
+    })
+    torch.cuda.synchronize()
+    denoise_elapsed = time.perf_counter() - start
+
+    if args.taehv:
+        torch.cuda.synchronize()
+        decode_start = time.perf_counter()
+        frames = taehv.decode(result.samples)
+        torch.cuda.synchronize()
+        decode_elapsed = time.perf_counter() - decode_start
+
+        imageio.mimsave(output_path, frames, fps=16, format="mp4")
+        total = denoise_elapsed + decode_elapsed
+        print(f"[{mode.upper()}] denoise {denoise_elapsed:.2f}s + TAEHV decode "
+              f"{decode_elapsed:.2f}s = {total:.2f}s "
+              f"({frames.shape[0]} frames @ {tuple(frames.shape[1:3])})")
+        print(f"Saved video to {output_path}")
+    else:
+        print(f"[{mode.upper()}] {args.infer_steps} steps in {denoise_elapsed:.2f}s "
+              f"({args.infer_steps / denoise_elapsed:.2f} it/s)")
+
+    generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()

fastvideo-kernel/build.sh

@@ -103,6 +103,11 @@ if [ "${GPU_BACKEND}" = "CUDA" ]; then
     if [ -z "${TORCH_CUDA_ARCH_LIST:-}" ]; then
         if [ "${cc_major}" = "9" ] && [ "${cc_minor}" = "0" ]; then
             export TORCH_CUDA_ARCH_LIST="9.0a"
+        elif [ "${cc_major}" = "12" ] && [ "${cc_minor}" = "0" ]; then
+            # Blackwell sm_120 needs the arch-conditional 'a' suffix so CMake's
+            # AUTO gate (matches 12.0a/120a/sm_120a) builds the attn_qat_infer
+            # (modified SageAttention3 FP4) kernels instead of silently skipping.
+            export TORCH_CUDA_ARCH_LIST="12.0a"
         else
             export TORCH_CUDA_ARCH_LIST="${cc_major}.${cc_minor}"
         fi

fastvideo/attention/backends/sage_attn3.py

@@ -55,6 +55,20 @@ class SageAttention3Impl(AttentionImpl):
         self.softmax_scale = softmax_scale
         self.dropout = extra_impl_args.get("dropout_p", 0.0)
 
+    def preprocess_qkv(
+        self,
+        qkv: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        """Transpose stacked QKV from [3B, L, H, D] to [3B, H, L, D].
+
+        Single bulk permute+contiguous on the entire stacked tensor rather than
+        three separate transposed views for Q, K, V. The .contiguous() is
+        required: sageattn_blackwell's fake kernel returns empty_like(q), so the
+        op's output strides must match contiguous q under torch.compile.
+        """
+        return qkv.permute(0, 2, 1, 3).contiguous()
+
     def forward(
         self,
         query: torch.Tensor,
@@ -62,9 +76,15 @@ class SageAttention3Impl(AttentionImpl):
         value: torch.Tensor,
         attn_metadata: AttentionMetadata,
     ) -> torch.Tensor:
-        query = query.transpose(1, 2)
-        key = key.transpose(1, 2)
-        value = value.transpose(1, 2)
+        """Call sageattn3_blackwell directly. Input is already [B, H, L, D]
+        and contiguous from preprocess_qkv."""
         output = sageattn3_blackwell(query, key, value, is_causal=self.causal)
-        output = output.transpose(1, 2)
         return output
+
+    def postprocess_output(
+        self,
+        output: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        """Transpose output from [B, H, L, D] back to [B, L, H, D]."""
+        return output.permute(0, 2, 1, 3).contiguous()

fastvideo/attention/layer.py

@@ -1,5 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
+import os
+
 import torch
 import torch.nn as nn
 
@@ -13,6 +15,26 @@ from fastvideo.utils import get_compute_dtype
 from fastvideo.layers.rotary_embedding import _apply_rotary_emb
 
 
+def _attention_compile_disabled() -> bool:
+    """Whether to keep attention ``forward`` out of the torch.compile graph.
+
+    Defaults to ``True`` (the historical behavior: attention runs eager via
+    ``torch.compiler.disable``). Set ``FASTVIDEO_DISABLE_ATTENTION_COMPILE=0``
+    to let attention be traced/compiled into the surrounding graph.
+    """
+    val = os.environ.get("FASTVIDEO_DISABLE_ATTENTION_COMPILE")
+    if val is None:
+        return True
+    return val.strip().lower() not in ("0", "false", "no", "off", "")
+
+
+def _maybe_compiler_disable(fn):
+    """Apply ``torch.compiler.disable`` unless disabled via env var."""
+    if _attention_compile_disabled():
+        return torch.compiler.disable(fn)
+    return fn
+
+
 class DistributedAttention(nn.Module):
     """Distributed attention layer.
     """
@@ -56,7 +78,7 @@ class DistributedAttention(nn.Module):
         self.backend = backend_name_to_enum(attn_backend.get_name())
         self.dtype = dtype
 
-    @torch.compiler.disable
+    @_maybe_compiler_disable
     def forward(
         self,
         q: torch.Tensor,
@@ -146,7 +168,7 @@ class DistributedAttention_VSA(DistributedAttention):
     """Distributed attention layer with VSA support.
     """
 
-    @torch.compiler.disable
+    @_maybe_compiler_disable
     def forward(
         self,
         q: torch.Tensor,

fastvideo/layers/quantization/nvfp4_qat_config.py

@@ -1,38 +1,90 @@
 # SPDX-License-Identifier: Apache-2.0
+"""NVFP4 quantization-aware (QAD) linear method, inference path.
+
+Quantizes every targeted linear's weight to NVFP4 once at load time and
+runs each forward as a registered flashinfer-backed FP4 matmul. The
+original fp16/bf16 weight is *popped* immediately after quantization so
+the half-precision copy does not keep occupying GPU memory — that's
+what lets a Wan-2.1 pipeline stay fully resident on a single GPU
+without any CPU offloading.
+
+The quantize / matmul custom ops are owned by
+:mod:`fastvideo.layers.quantization.nvfp4_config` and registered under
+the ``fastvideo_fp4::`` namespace. We reuse them here for two reasons:
+
+1. Re-registering the same op name in a second module would raise.
+2. The registered ops have ``register_fake`` shape/dtype kernels, which
+   is what makes the inference pipeline's per-block ``torch.compile``
+   trace through without graph breaks. Calling raw flashinfer functions
+   (the old behavior of this file, plus a ``@torch.compile`` on
+   ``apply``) graph-breaks at every quantize and every matmul.
+
+For QAT *training*, see ``nvfp4_qat_train_config`` which keeps the
+weight trainable and fake-quantizes on the fly via a straight-through
+estimator.
+"""
+from __future__ import annotations
+
+import gc
 import logging
 from typing import Any
 
 import torch
 from torch.nn.parameter import Parameter
 
-from fastvideo.layers.quantization.base_config import QuantizationConfig, QuantizeMethodBase
+from fastvideo.layers.quantization.base_config import (
+    QuantizationConfig,
+    QuantizeMethodBase,
+)
+from fastvideo.layers.quantization.nvfp4_config import (
+    _mm_fp4,
+    _nvfp4_quantize,
+    _require_flashinfer,
+)
 from fastvideo.models.utils import set_weight_attrs
 
-try:
-    import flashinfer
-except ImportError:
-    flashinfer = None
-
 logger = logging.getLogger(__name__)
 
+# Wan-style attention + FFN projection layers. Matched as substrings of the
+# layer prefix (e.g. "blocks.0.attn1.to_q" contains "to_q").
+DEFAULT_FP4_LAYERS = (
+    "ffn.fc_in",
+    "ffn.fc_out",
+    "to_q",
+    "to_k",
+    "to_v",
+    "to_out",
+)
 
-def _require_flashinfer() -> Any:
-    if flashinfer is None:
-        raise ImportError("flashinfer is required for NVFP4 QAT quantization. "
-                          "Please install flashinfer to use the nvfp4_qat quantization backend.")
-    return flashinfer
+
+def _layout_128x4() -> Any:
+    SfLayout, _, _ = _require_flashinfer()
+    return SfLayout.layout_128x4
 
 
 class NVFP4QATQuantizeMethod(QuantizeMethodBase):
+    """Inference-only NVFP4 linear method with weight popping.
+
+    The dense ``weight`` parameter is materialized at load time only so
+    that :func:`convert_model_to_fp4` can read it once; the loader then
+    removes it via ``mod._parameters.pop('weight')``. From that point
+    forward, ``apply`` reads only ``_fp4_weight`` / ``_fp4_weight_scale``
+    / ``_weight_global_sf``.
+    """
 
     def __init__(self) -> None:
         super().__init__()
         self.weight_fp4 = None
         self.weight_scale = None
+        # Static input global scale factor. Matches the FastVideo-Quantization
+        # production path; recomputing it per-call via a ``.max()`` reduction
+        # (the previous behavior) adds a sync point, costs a kernel launch,
+        # and produces a data-dependent value that prevents CUDA-graph
+        # capture under ``torch.compile(mode='reduce-overhead')``.
+        self.x_global_sf = torch.tensor(1.0, device="cuda", dtype=torch.float32)
 
     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):
-        """Create weights for a linear layer. Note the corrected signature to match LinearMethodBase."""
+                       input_size: int, output_size: int, params_dtype: torch.dtype, **extra_weight_attrs) -> None:
         weight = Parameter(torch.empty(
             sum(output_partition_sizes),
             input_size_per_partition,
@@ -43,28 +95,27 @@ class NVFP4QATQuantizeMethod(QuantizeMethodBase):
         layer.register_parameter("weight", weight)
         set_weight_attrs(weight, extra_weight_attrs)
 
-    @torch.compile
     def apply(self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor:
-        """Apply NVFP4 QAT quantized computation."""
-        flashinfer_mod = _require_flashinfer()
-        out_dim = layer.weight.shape[0]
+        # ``_fp4_weight`` carries the (out, in/2) packed fp4 weight, so its
+        # row count is the output dim even after the dense weight is popped.
+        out_dim = layer._fp4_weight.shape[0]
         original_shape = x.shape
-        assert x.dtype == torch.bfloat16 or x.dtype == torch.float16, f"only allow bf16/fp16 inputs to fp4 linear, got {x.dtype}"
 
+        assert x.dtype in (torch.bfloat16, torch.float16), (f"only allow bf16/fp16 inputs to fp4 linear, got {x.dtype}")
         x = x.view(-1, x.shape[-1])
-
-        x_global_sf = (448 * 6) / x.float().abs().nan_to_num().max()
-        x_fp4, x_scale = flashinfer_mod.nvfp4_quantize(
+        x_global_sf = self.x_global_sf
+        x_fp4, x_scale = _nvfp4_quantize(
             x,
             x_global_sf,
-            sfLayout=flashinfer_mod.SfLayout.layout_128x4,
+            sfLayout=_layout_128x4(),
             do_shuffle=False,
         )
+
         weight_fp4 = layer._fp4_weight
         weight_scale = layer._fp4_weight_scale
         weight_global_sf = layer._weight_global_sf
 
-        out = flashinfer_mod.mm_fp4(
+        out = _mm_fp4(
             x_fp4,
             weight_fp4.T,
             x_scale,
@@ -76,67 +127,106 @@ class NVFP4QATQuantizeMethod(QuantizeMethodBase):
         )
 
         if bias is not None:
-            if bias.device != out.device or bias.dtype != out.dtype:
-                bias = bias.to(device=out.device, dtype=out.dtype)
             out = out + bias
-
-        if len(original_shape) == 3:
-            out = out.view(original_shape[0], original_shape[1], out_dim)
-
+        out = out.view(*original_shape[:-1], out_dim)
         return out
 
 
 class NVFP4QATConfig(QuantizationConfig):
+    """NVFP4 (Wan-style) linear quantization, inference.
 
-    def __init__(self) -> None:
+    Args:
+        target_layers: Substrings matched against each linear layer's
+            prefix. A layer is quantized if any substring is contained in
+            its prefix. Defaults to the standard Wan attention + FFN
+            projections (:data:`DEFAULT_FP4_LAYERS`).
+    """
+
+    def __init__(self, target_layers: tuple[str, ...] | None = None) -> None:
         super().__init__()
+        self.target_layers = (tuple(target_layers) if target_layers else DEFAULT_FP4_LAYERS)
 
-    def get_name(self):
+    def get_name(self) -> str:
         return "nvfp4_qat"
 
-    def get_supported_act_dtypes(self):
+    def get_supported_act_dtypes(self) -> list[torch.dtype]:
         return [torch.bfloat16, torch.float16]
 
     @classmethod
-    def get_min_capability(cls):
+    def get_min_capability(cls) -> int:
         return 100
 
     @staticmethod
-    def get_config_filenames():
+    def get_config_filenames() -> list[str]:
         return []
 
     @classmethod
-    def from_config(cls, config: dict[str, Any]) -> "NVFP4QATConfig":
-        return cls()
+    def from_config(cls, config: dict[str, Any]) -> NVFP4QATConfig:
+        target_layers = config.get("target_layers")
+        if target_layers is not None:
+            target_layers = tuple(target_layers)
+        return cls(target_layers=target_layers)
 
     def get_quant_method(self, layer: torch.nn.Module, prefix: str):
         from fastvideo.layers.linear import LinearBase
-        fp4_layers = ["ffn.fc_in", "ffn.fc_out", "to_q", "to_k", "to_v", "to_out"]
-        if isinstance(layer, LinearBase) and any(layer_name in prefix for layer_name in fp4_layers):
+        if isinstance(layer, LinearBase) and any(name in prefix for name in self.target_layers):
             return NVFP4QATQuantizeMethod()
         return None
 
 
-@torch.compile
-def convert_model_to_fp4(model: torch.nn.Module):
-    flashinfer_mod = _require_flashinfer()
+def convert_model_to_fp4(model: torch.nn.Module) -> None:
+    """Prequantize every FP4-tagged linear and drop its dense weight.
+
+    Walks the module tree, and for each layer whose ``quant_method`` is
+    an :class:`NVFP4QATQuantizeMethod`, computes the NVFP4 packed weight
+    / scale / global-scale buffers, then pops the original fp16/bf16
+    ``weight`` parameter so it no longer occupies GPU memory.
+    """
+    SfLayout, _, _ = _require_flashinfer()
     from torch.distributed.tensor import DTensor  # type: ignore
-    for mod in model.modules():
-        qm = getattr(mod, "quant_method", None)
-        if isinstance(qm, NVFP4QATQuantizeMethod):
+
+    with torch.no_grad():
+        for mod in model.modules():
+            qm = getattr(mod, "quant_method", None)
+            if not isinstance(qm, NVFP4QATQuantizeMethod):
+                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]
-            weight_global_sf = (448 * 6) / weight_local.float().abs().nan_to_num().max()
-            fp4_w, fp4_s = flashinfer_mod.nvfp4_quantize(
+
+            # Only the reduced scalar needs fp32; avoid a full fp32 copy.
+            weight_absmax = (weight_local.detach().abs().nan_to_num().amax().to(dtype=torch.float32))
+            weight_global_sf = (448 * 6) / weight_absmax
+            fp4_w, fp4_s = _nvfp4_quantize(
                 weight_local,
                 weight_global_sf,
-                sfLayout=flashinfer_mod.SfLayout.layout_128x4,
+                sfLayout=SfLayout.layout_128x4,
                 do_shuffle=False,
             )
             mod.register_buffer("_fp4_weight", fp4_w, persistent=False)
             mod.register_buffer("_fp4_weight_scale", fp4_s, persistent=False)
-            mod.register_buffer("_weight_global_sf",
-                                torch.tensor(weight_global_sf, dtype=torch.bfloat16),
-                                persistent=False)
+            mod.register_buffer(
+                "_weight_global_sf",
+                weight_global_sf.to(dtype=torch.bfloat16),
+                persistent=False,
+            )
+
+            # Drop the dense weight as soon as the fp4 buffers are installed
+            # so it cannot keep occupying GPU memory.
+            removed_weight = mod._parameters.pop("weight", None)
+            if removed_weight is not None:
+                removed_weight.grad = None
+            del removed_weight, weight, weight_local, weight_absmax
+    gc.collect()
+    torch.cuda.empty_cache()
+
+
+__all__ = [
+    "NVFP4QATConfig",
+    "NVFP4QATQuantizeMethod",
+    "convert_model_to_fp4",
+    "DEFAULT_FP4_LAYERS",
+]