[Bugfix] QAD 5090: Torch.compile and other optimizations (15/12)

[loaydatrain]

Purpose

Continues the QAD 5090 stack (#1225). Inference-side companion to #1463 (FP4
QAT training): wires the FP4 linear path for Wan2.1 1.3B end-to-end under
torch.compile, pops the bf16 weight after FP4 packing, removes x_global_sf, and adds TAEHV scripts

With Sage3, running in 2it/s, 1.78s e2e.

Changes

• layers/quantization/nvfp4_qat_config.py (+ loader/fsdp_load.py):
  inference path made torch.compile-friendly — precomputed x_global_sf
  (no per-call sync), single-level quant, and the dense bf16 weight is popped
  after FP4 pack (only _fp4_weight + scales survive).
• attention/backends/sage_attn3.py: removed data-dependent control flow on
  tensor metadata so the backend traces cleanly under torch.compile.
• examples/inference/optimizations/fp4_linear_wan2_1_1_3b.py and
  FastWan_QAD_TAEHV.py: end-to-end Wan2.1 1.3B inference with the FP4 linear
  path + sage_attn3 + torch.compile; the TAEHV variant swaps in the TAEHV
  decoder.

No new CLI flags; reuses --transformer-quant nvfp4_qat from #1463.

Test results (RTX 5090 / sm_120)

• ✅ Inference: --transformer-quant nvfp4_qat Wan2.1 1.3B runs end-to-end
  with torch.compile enabled; sample output
  video_samples/raccoon_fp4_linear_taehv_compile.mp4.
• ✅ Step time: DiT time: 2it/s // e2e time: 1.78s for a 5s 480p video
• ✅ Memory: dense bf16 weight popping drops DiT resident memory after the
  first forward by TODO MB.

Depends on #1463 (--transformer-quant CLI + nvfp4_qat config). Part of #1225.

[mergify]

Merge Protections

Your pull request matches the following merge protections and will not be merged until they are valid.

🔴 PR merge requirements

Waiting for

• #approved-reviews-by>=1
• check-success=fastcheck-passed
• check-success=full-suite-passed

This rule is failing.

• #approved-reviews-by>=1
• check-success=fastcheck-passed
• check-success=full-suite-passed
• check-success~=pre-commit
• title~=(?i)^\[(feat|feature|bugfix|fix|refactor|perf|ci|doc|docs|misc|chore|kernel|new.?model|skill|skills|infra)\]

[gemini-code-assist]

Code Review

This pull request introduces fast NVFP4 linear inference for Wan2.1 with TAEHV decoding, optimizes attention preprocessing/postprocessing in SageAttn3, and adds support for quantization-aware training (QAD) with a straight-through estimator. Key feedback includes resolving an environment-specific absolute path, rounding instead of truncating during float-to-uint8 conversion, avoiding backslashes for line continuation, and preventing division-by-zero errors when computing scale factors. Additionally, it is recommended to make NVFP4QATTrainConfig customizable to match its inference counterpart and to refactor the FSDP loader to handle mixed quantization methods without returning early.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

The absolute path /root/taehv/taew2_1.pth is environment-specific and will fail on other systems. It is better to use a relative path or default to a standard location, or allow it to be resolved dynamically.

Suggested change

	DEFAULT_TAEHV_CHECKPOINT = "/root/taehv/taew2_1.pth"
	DEFAULT_TAEHV_CHECKPOINT = "taew2_1.pth"

[gemini-code-assist]

Directly casting a float tensor in [0, 1] to uint8 after multiplying by 255 performs a floor operation (truncation), which can cause precision loss and slight color shifts. Rounding to the nearest integer before casting is recommended.

Suggested change

	frames = (decoded[0].clamp(0, 1) * 255).to(torch.uint8)
	frames = (decoded[0].clamp(0, 1) * 255).round().to(torch.uint8)

[gemini-code-assist]

PEP 8 discourages the use of backslashes for line continuation. Wrapping the expression in parentheses is the preferred way to break long lines.

Suggested change

	taehv = TaehvDecoder(resolve_taehv_checkpoint(args.taehv_checkpoint)) \
	if args.taehv else None
	taehv = (
	TaehvDecoder(resolve_taehv_checkpoint(args.taehv_checkpoint))
	if args.taehv
	else None
	)

References

1. PEP 8: Avoid using backslashes for line continuation. Use parentheses or standard block structures instead. ^(link)

[gemini-code-assist]

If weight_absmax is 0 (e.g., for zero-initialized or pruned weights), weight_global_sf = (448 * 6) / weight_absmax will result in a division by zero, producing inf or NaN values. Clamping weight_absmax to a small positive epsilon prevents this.

Suggested change

	weight_absmax = (weight_local.detach().abs().nan_to_num().amax().to(dtype=torch.float32))
	weight_absmax = (weight_local.detach().abs().nan_to_num().amax().to(dtype=torch.float32)).clamp(min=1e-12)

[gemini-code-assist]

NVFP4QATTrainConfig hardcodes fp4_layers and does not accept target_layers in its constructor, unlike NVFP4QATConfig. This inconsistency prevents users from customizing target layers during QAT training. Refactoring it to match NVFP4QATConfig ensures consistency and flexibility.

class NVFP4QATTrainConfig(QuantizationConfig):

    def __init__(self, target_layers: tuple[str, ...] | None = None) -> None: 
        super().__init__()
        from fastvideo.layers.quantization.nvfp4_qat_config import DEFAULT_FP4_LAYERS
        self.target_layers = (tuple(target_layers) if target_layers else DEFAULT_FP4_LAYERS)

    def get_name(self) -> str:
        return "nvfp4_qat_train"

    def get_supported_act_dtypes(self) -> list[torch.dtype]:
        return [torch.bfloat16, torch.float16]

    @classmethod
    def get_min_capability(cls) -> int:
        return 100

    @staticmethod
    def get_config_filenames() -> list[str]:
        return []

    @classmethod
    def from_config(cls, config: dict[str, Any]) -> NVFP4QATTrainConfig:
        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
        if isinstance(layer, LinearBase) and any(name in prefix for name in self.target_layers):
            return NVFP4QATTrainQuantizeMethod()
        return None

[gemini-code-assist]

The current loop returns immediately after finding the first quantization method. If a model contains mixed quantization (e.g., both NVFP4QATQuantizeMethod and NVFP4QuantizeMethod), only one of them will be converted, and the other will be silently skipped. Tracking which conversions have run instead of returning early makes the loader much more robust.

Suggested change

	for mod in model.modules():
	if isinstance(getattr(mod, "quant_method", None),
	NVFP4QuantizeMethod):
	quant_method = getattr(mod, "quant_method", None)
	# Wan-style NVFP4 linear: prequantize and pop dense weights.
	if isinstance(quant_method, NVFP4QATQuantizeMethod):
	logger.info(
	"Converting loaded model weights for NVFP4 (QAT) linear layers")
	convert_model_to_fp4(model)
	return
	# LTX-2 NVFP4 linear.
	if isinstance(quant_method, NVFP4QuantizeMethod):
	logger.info("Converting loaded model weights for NVFP4 linear layers")
	convert_model_to_nvfp4(model)
	return
	converted_qat = False
	converted_nvfp4 = False
	for mod in model.modules():
	quant_method = getattr(mod, "quant_method", None)
	# Wan-style NVFP4 linear: prequantize and pop dense weights.
	if isinstance(quant_method, NVFP4QATQuantizeMethod) and not converted_qat:
	logger.info(
	"Converting loaded model weights for NVFP4 (QAT) linear layers")
	convert_model_to_fp4(model)
	converted_qat = True
	# LTX-2 NVFP4 linear.
	elif isinstance(quant_method, NVFP4QuantizeMethod) and not converted_nvfp4:
	logger.info("Converting loaded model weights for NVFP4 linear layers")
	convert_model_to_nvfp4(model)
	converted_nvfp4 = True

[mergify]

This PR has merge conflicts with the base branch. Please rebase:

git fetch origin main
git rebase origin/main
# Resolve any conflicts, then:
git push --force-with-lease

[SolitaryThinker]

Hi @loaydatrain — automated review from Gob, one of @SolitaryThinker's AI reviewers. Findings aren't all human-verified; ping @SolitaryThinker if anything looks off.

Scope: training/quant + perf (the FP4-linear torch.compile enablement + optimizations). Open PR, mergeable=CONFLICTING / mergeStateStatus=DIRTY + needs-rebase. Verdict: COMMENT — the compile fix is correct and the core logic lands on files main hasn't touched, but rebase first and confirm two things below before merge.

Major

• Needs rebase — the branch is stacked on the pre-merge #1463 commit (09f9b56). #1463 already squash-merged to main (6da206e1), and that base is also ~13k lines behind current main (the big deletion count GitHub shows is this artifact, not intended changes). The actual textual conflict is just fastvideo/models/loader/fsdp_load.py and fastvideo/fastvideo_args.py, and both PR versions are semantically equivalent to what #1463 already put on main (variable name / branch order / comment wording in fsdp_load.py; pure formatting reflow of the existing --transformer-quant arg in fastvideo_args.py). Rebase onto main and drop those two near-dup hunks — the substantive logic (sage_attn3, attention/layer.py, nvfp4_qat_config.py) is conflict-free.

• Activation-quant numerics change with no quality evidence — nvfp4_qat_config.py:84,106. apply() switched the input scale from per-call dynamic (448*6)/x.absmax() to static x_global_sf = torch.tensor(1.0). The comment ("matches the production path") is correct — the LTX-2 NVFP4 method uses the same static 1.0 (nvfp4_config.py:270,338), and this is exactly what removes the per-call .max() sync blocking CUDA-graph capture. But it's still a real FP4 activation-quant change and the body ships no SSIM/quality number (the memory claim is a literal "TODO MB"). Per the repo's model-PR convention, please post an SSIM delta vs dense bf16. Also note it diverges from the training STE, which still quantizes activations with the dynamic scale (fp4linear.py:50) — confirm the train-vs-inference input-scale skew is intended for QAT.

• Latent layout-contract break in the sage3 hook refactor. The old impl transposed q/k/v inside forward() after the chunk; the new preprocess_qkv (sage_attn3.py:58-70) does permute(0,2,1,3) on the stacked tensor, which layer.py:134 / ltx2.py:1303 call before the replicated-QKV concat (dim=1) and the post-attention replicated split. After the permute dim-1 is heads, not sequence, so the replicated-QKV concat/split no longer mean what they did. Safe for this PR's Wan target (replicated_q is never passed there, and it's never populated anywhere in-tree today), but LTX-2's non-VSA path carries the replicated-QKV plumbing and can select SAGE_ATTN_THREE, so the contract is silently wrong for any future caller. Either preserve the per-tensor transpose for the replicated path, or assert replicated_q is None in the sage3 path so a break is loud, not silent.

Minor

• No div-by-zero guard — nvfp4_qat_config.py:201-202: weight_global_sf = (448*6) / weight_absmax divides by an unclamped reduced absmax; an all-zero tile gives inf/nan. A .clamp_min(eps) is cheap insurance.
• Hardcoded path — FastWan_QAD_TAEHV.py:53: DEFAULT_TAEHV_CHECKPOINT = "/root/taehv/taew2_1.pth" is environment-specific; default it to a repo-relative / env-driven value so the example runs out of the box.

Nits

• TAEHV docstring (FastWan_QAD_TAEHV.py:3,21-25) references a sibling fp4_linear_wan2_1_1_3b.py and a python fp4_linear_taehv_wan2_1_1_3b.py usage line that don't match the actual filename (FastWan_QAD_TAEHV.py).
• PR body "drops DiT resident memory ... by TODO MB" — fill in the number (it's the headline benefit of the weight-pop).

Verified positives: reusing the register_fake-backed fastvideo_fp4:: custom ops instead of the old @torch.compile + raw flashinfer is the correct graph-break fix (no silently-swallowed breaks); the weight-pop memory logic is sound (persistent=False buffers + _parameters.pop + grad=None + gc/empty_cache); the FASTVIDEO_DISABLE_ATTENTION_COMPILE env-gate defaults to the historical eager behavior; title matches the merge regex; pre-commit is green. CI microscope failures are uniform exit 124 timeouts consistent with the stale/conflicting base — rebase and re-run.

— Gob (@SolitaryThinker's AI reviewer).

[SolitaryThinker]

/merge