[feat]: FA4-FP4 ATTN_QAT_INFER on sm_100/sm_103 + NVFP4 weight purge

[SolitaryThinker]

Problem

Two follow-ups to the LTX-2 NVFP4 QAD deploy path (author task, one PR by request):

(A) ATTN_QAT_INFER only resolves on consumer Blackwell (sm_12x, fastvideo-kernel CUTLASS extension). On GB200 (sm_100) / GB300 (sm_103) the selector silently demotes to FlashAttention — measured on the GB200 bench: the NVFP4 arm requested ATTN_QAT_INFER, got flash, so FP4 bought nothing at attention. Meanwhile the FP4 FA4 kernel for exactly these archs is already merged (#1221, credit @Edenzzzz): flash_attn.cute plumbing with per-16 block-scaled NVFP4 Q/K + BF16 P/V, from hao-ai-lab/flash-attention-fp4 @ fp4 (kernel README).

(B) NVFP4Config's linear method retains the original bf16 layer.weight alongside _nvfp4_weight+scales on every quantized layer (+6.8 GiB measured on the GB200 LTX-2 bench) — deliberate for the stage-profile dense fallback, but only 3 of 13 curated suffixes (the cross-modal AV projections) can ever take that dense path.

Solution

(A) Arch-aware ATTN_QAT_INFER resolution (fastvideo/attention/backends/attn_qat_infer.py):

• Named capability sets as the single routing truth: _SUPPORTED_DEVICE_CAPABILITIES (CUTLASS ext, value untouched) and _FA4_FP4_CAPABILITIES = {(10,0), (10,3)}.
• (10,0)/(10,3) route AttnQatInferImpl.forward through feat: FP4 Flash Attention 4 for Blackwell GPUs #1221's plumbing verbatim (mirrors FlashAttentionImpl._forward_nvfp4: flashinfer nvfp4_quantize → flash_attn_fp4_func). Tuned defaults passed through, not re-tuned: qk_mode=nvfp4 (per-16 E4M3 SFs), pv_mode=bf16 — recorded in the resolution receipt.
• Receipt line on every resolution (ATTN_QAT_INFER resolved: arch=... kernel=... qk_mode=... pv_mode=...); the fallback log keeps its stable trailing sentence and now names arch + install pointer.
• Dependency story: sm_100/103 availability is probe-based via feat: FP4 Flash Attention 4 for Blackwell GPUs #1221's guarded flash_attn.cute import (cutlass-dsl skew maps to ImportError with a loud warning) — never a hard crash; install pointer = the fork @ fp4 branch.
• Validation callback error message is now arch-aware (was 'requires the sm120 fastvideo-kernel extension').

Train-sim mismatch (stated, measured, not assumed): ATTN_QAT_TRAIN simulates the sm_12x CUTLASS scheme. The FA4 scheme (block-scaled QK-only) differs, so sm_100/sm_103 deployment carries a train-sim mismatch. Documented in the QAD recipe header and the LTX-2 deploy example; MS-SSIM against the recorded QAT quality reference (0.846) is being measured on GB200 on this branch — in-tolerance ships with the receipt, out-of-tolerance gets surfaced before merge.

(B) Fallback-aware weight purge (fastvideo/layers/quantization/nvfp4_config.py):

• NVFP4Config(retain_original_weights=None|True|False): auto (default) purges always-FP4 layers' originals and retains the 3 refine-only suffixes; False purges everything (single-stage deploys, e.g. the QAD distilled 8-step bench); True = old behavior.
• Purge happens in convert_model_to_nvfp4 (the single conversion site); apply() takes out_dim from the packed FP4 weight when the original is gone, and a purged layer hitting the dense path raises naming the flag. FSDP-sharded (DTensor) originals are skipped (retained) until a sharded deploy needs per-shard purge.
• Receipt log: NVFP4 weight purge receipt: purged N (... GiB freed); retained M. Target on the GB200 bench config: NVFP4 peak alloc below bf16's 63.9 GiB (re-measured on GB200 for the receipt).

torch.compile interaction

The FA4-FP4 forward previously resolved its quantization/kernel entry points on every call (function-level imports + flashinfer's JIT-module lookup, whose first resolution probes the CUDA toolchain via a subprocess). Under torch.compile that pattern graph-breaks every step — fullgraph=True was impossible for the NVFP4 path, and at fullgraph=False the per-step break kept eager dispatch overhead on the hot path (measurable at short-step resolutions; documented as Finding 1 of the LTX-2.3 bf16-vs-NVFP4 compile benchmark). Two-stage fix. Round 1 memoized the resolution (lazy, module-level) — necessary for eager overhead but insufficient for compile: Dynamo traces function bodies symbolically, so a runtime cache hit is invisible at trace time and tracing still descended into the toolchain probe. Round 2 registers the whole quantize step as a torch.library custom op (fastvideo::nvfp4_quantize_fa4, forward + fake — the same boundary pattern as the masked flash-attention entry points), making it one opaque graph node. Regressions: the resolve-once-across-N-forwards test (eager) plus a torch.compile(fullgraph=True) trace-and-run test over the op-backed path with a CPU kernel registered in-test. The construct-anywhere/fail-cleanly-at-forward contract is unchanged. Round 3 (this head): the previous note called the fake's contiguous-vs-permuted-view stride mismatch "known minor — at worst an inductor copy". That was wrong: torch.compile bakes the fake's output strides into the compiled artifact and asserts at runtime (expected size 32==32, stride 16==8192 at dim=0) — tracing was fixed but all compiled NVFP4 executions crashed. The fake now constructs its scale-factor output with the impl's own base shape and permutation (new_empty((batch, nheads, rest_m, rest_k, 32, 4, 4)).permute(4, 5, 2, 6, 3, 1, 0)), matching the real strided view for general batch. Regression upgraded accordingly: torch.library.opcheck now cross-checks the registered fake against a real kernel run — shapes, dtypes, AND strides — with a batch>1, non-multiple-of-128 seqlen input so no stride is degenerate; the fullgraph trace test's CPU kernel also produces the production stride contract. Lesson encoded: shape-only fakes are insufficient for ops returning non-contiguous views.

Final performance evidence (at this head)

Matched same-node comparisons (LTX-2.3-Distilled, single GPU, identical prompts/seeds/settings, compiled inference, within-batch ratios), replicated across three nodes:

resolution	NVFP4 vs BF16 e2e	peak memory
1080p	1.47–1.56x faster (stable across all three nodes)	−15 GiB (byte-stable)
720p	node-dependent — no single-node claim holds	−15 GiB (byte-stable)

At 720p the ordering itself flips with the node: two nodes measured NVFP4 1.34x and 1.21x faster, a third measured it 1.43x slower — with identical graph-break counts in both arms, so it is not a compile artifact. The 1080p advantage and the memory reduction are the replicated claims; 720p performance should be measured on the target node before deployment decisions.

• torch.compile(fullgraph=True) now executes for the NVFP4 path (previously impossible): zero graph breaks, routing receipt confirmed in-log.
• Graph-break parity: 7/7 identical break sites in both arms at default compile mode — the NVFP4 path adds none.
• Quality: FA4-routed QAD MS-SSIM delta +1.04e-4 vs the recorded reference (zero-width self-noise band; the delta favors the model — surfaced for the maintainer/author verdict rather than self-blessed). The LTX-2.3 speed table above is an untrained-quant configuration: it carries no quality gate yet (the 2.3 QAD training campaign is separate); treat it as speed/memory evidence only.

Blast radius (two-stack)

• Inference: ATTN_QAT_INFER resolution (new archs gain the kernel; all previously-working archs unchanged — the CUTLASS set's value is untouched); NVFP4 deploys now purge by default (safe by construction: only refine-only layers can reach the dense path, and those are retained under auto).
• Training: behavior untouched. ATTN_QAT_TRAIN is not modified; the only training-side edits are the validation-swap error message and recipe comments. The receipt makes the validation-time kernel choice explicit per arch.
• Overlap note: is_attn_qat_infer_available is also touched by my open [bugfix]: complete #1601: atomic export, attn pass-through, arch gating #1646 (which narrows the CUTLASS capability set to the built arch) — same owner, whichever lands second rebases trivially; the capability-set VALUE stays [bugfix]: complete #1601: atomic export, attn pass-through, arch gating #1646/[kernel] Build + allow attn_qat_infer FP4 attention on sm_121a (DGX Spark) #1598's concern, the routing structure is this PR's.

Test evidence

• 19 passed locally (CPU): fastvideo/tests/attention/test_attn_qat_infer_arch_gate.py (resolution matrix for (12,0)/(12,1)/(10,0)/(10,3)/(9,0)/(8,9)/no-cuda + receipt-content tests) and fastvideo/tests/ops/quantization/test_nvfp4_purge.py (construction-level: auto/retain/purge matrix, out_dim after purge, dense-path-after-purge raises, receipt log content).
• tests/local_tests/test_nvfp4_fa4.py extended with the ATTN_QAT_INFER-on-sm_100/103 route: resolution + receipt, forward parity vs SDPA (cos ≥ 0.97, same bound as the sm_120 kernel test), cross-attention lengths. Runs on GB200; sm_103 (GB300) is hardware-validated: route/parity/cross-attention 3/3 PASS and an 8-shape fp32-SDPA sweep (cos 0.99044–0.99077) on a single-GPU GB300 node (NVIDIA GB300, compute_cap 10.3) — receipts in the PR comments.
• pre-commit green (yapf, ruff, codespell, mypy) on all touched files.

Context

• feat: FP4 Flash Attention 4 for Blackwell GPUs #1221 (merged FA4-FP4 plumbing, @Edenzzzz), flash-attention-fp4 @ fp4 kernel README (per-arch tuned configs; B200 2018 TFLOPS / GB300 2677 TFLOPS NVFP4+FP8 — the fp8-PV mode is faster still but not plumbed by feat: FP4 Flash Attention 4 for Blackwell GPUs #1221; follow-up candidate, not this PR), GB200 bf16-vs-NVFP4 benchmark measurements (2026-07-26), [kernel] Build + allow attn_qat_infer FP4 attention on sm_121a (DGX Spark) #1598 (sm_121 CUTLASS build — composes via the untouched capability set), [bugfix]: complete #1601: atomic export, attn pass-through, arch gating #1646 (same-file capability-set change — same owner).

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[mergify]

Merge Protections

🔴 1 of 1 protections blocking · waiting on 👀 reviews and 🤖 CI

	Protection	Waiting on
🔴	PR merge requirements	👀 reviews and 🤖 CI

🔴 PR merge requirements

Waiting for

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

This rule is failing.

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

[SolitaryThinker]

CI fix pushed in b74a1858 — microscope-unit-tests exit 1 attributed as REAL (my diff), reproduced locally before touching CI:

fastvideo/tests/train/callbacks/test_validation.py::TestAttnQatInferValidation::test_context_swaps_and_restores_qat_attention — the #1626 swap test constructs AttnQatInferImpl on a kernel-less host under a monkeypatched is_attn_qat_infer_available. My first cut resolved the kernel in __init__ and raised, breaking the pre-existing construct-anywhere/fail-at-forward contract that test pins (and that the validation callback relies on: it gates on availability before constructing).

Fix: kernel resolution moved back to per-forward; the constructor no longer raises. The unavailable-at-forward ImportError now carries the arch receipt. Local evidence: the failing test + full train/callbacks + api + arch-gate + purge suites = 297 passed; pre-commit green. The 5 pre-existing test_nvfp4_ltx2_wiring.py failures on my macOS host reproduce identically on pristine origin/main (CUDA-built-torch requirement) — not this PR.

Standing by for the GB200 parity traceback on the FA4 route — if it's a code bug in _forward_fa4_fp4 it lands on this branch same-PR.

[SolitaryThinker]

Second microscope-unit-tests failure fixed — pulled the build-4657 log (attribution: REAL, second failure mode, not the same one incompletely fixed):

fastvideo/tests/api/test_attn_qat_infer_capability_gate.py::test_sm100_host_with_bundled_extension_falls_back (807 passed / 1 failed in the lane). That #1601-era test pins the OLD contract — 'sm_100 always falls back' — which is exactly the behavior this PR intentionally changes, and it doesn't fake the FP4-FA4 probe. On the CI runner flash_attn.cute is importable, so _fa4_fp4_available() is genuinely True and sm_100 now resolves ATTN_QAT_INFER (the feature). Locally the probe is False (no flash-attn on macOS), which is why my earlier run missed it.

Fix: _fake_gpu now fakes the FA4 probe like the other two physical facts (fa4_imports, default False); test_sm100_host_with_bundled_extension_falls_back keeps pinning the originally-reviewed failure class (bundled sm_12x ext + NO FA4 → fallback), and a new parametrized test_datacenter_blackwell_with_fa4_selects_backend pins the new route for (10,0)/(10,3). Full lane selection locally: 803 passed; only the 5 pre-existing test_nvfp4_ltx2_wiring.py macOS host-artifact failures remain (identical on pristine main).

[SolitaryThinker]

Two GB200-validation follow-ups pushed in 8f371c3b:

1. Purge breadth fix (real bug caught by the GB200 purge validation). retain_original_weights=False crashed the distilled single-stage deploy: the single-stage pipeline runs under the base stage profile, whose deployment contract routes the refine-only layers (audio_to_video_attn / video_to_audio_attn projections) through the DENSE path every step — so purging their bf16 originals removes weights that are genuinely needed. The purge was too broad, not the guard: convert_model_to_nvfp4 now never purges refine-only layers regardless of the flag (retain_original_weights only widens retention: True keeps everything; default purges every always-FP4 layer). The dense-after-purge guard stays as defense against hand-purged modules. Docs/example/recipe comments updated so the documented recommendation (the default) is the usable one; unit tests updated to pin the corrected semantics. Positive receipt from the same validation run stands: the purge mechanism works (24.00 GiB of retained bf16 originals freed on the LTX-2 deploy).

2. Dependency pins declared (parity-failure attribution: NOT a code bug). The reported FA4-route parity failure on GB200 was a dependency-lattice issue, reproduced and root-caused with a full traceback: the fork's fp4 branch targets the cutlass-dsl 4.4 API surface, and 4.6-era installs fail at CuTe JIT trace (AttributeError: module 'cutlass.cute' has no attribute 'make_fragment' — removed at module level in 4.6, not shim-able). With the era-matched set, all three route tests pass 3/3 on GB200 and the receipt line confirms kernel=flash-attention-fp4 qk_mode=nvfp4(per-16-e4m3-sf) pv_mode=bf16. This PR now declares that set instead of letting every deploy rediscover it: _FA4_INSTALL_HINT names nvidia-cutlass-dsl==4.4.2, quack-kernels==0.4.1, flashinfer-python==0.6.8, FASTVIDEO_FA4=1 (the fork ships no compiled FA2, so dense paths need the FA4 opt-in); docs/inference/optimizations.md gains a branch-to-dsl compatibility table; the QAD recipe header carries the pins.

Quality measurement on this branch (GB200, distilled LTX-2 QAD): MS-SSIM 0.846894 vs reference 0.846790 (measured self-noise 0.000000) — the +1.04e-4 delta is formally out-of-band only because the reference band is zero-width; flagged for a verdict rather than self-blessed.

[SolitaryThinker]

sm_103 upgraded: code-verified → hardware-validated. The FA4-FP4 ATTN_QAT_INFER route has now run on real sm_103 silicon (a single-GPU GB300 node, nvidia-smi receipt: NVIDIA GB300, compute_cap 10.3):

• tests/local_tests/test_nvfp4_fa4.py -k AttnQatInfer — 3/3 PASS on sm_103 (backend resolution + receipt line, forward parity vs fp32 SDPA at cos ≥ 0.97, cross-attention lengths 384/512), same suite previously validated on sm_100.
• 8-shape kernel-level sweep vs an fp32 SDPA reference (CUDA-event timing, 5 warmup / 20 timed): cosine 0.99044–0.99077 across all shapes — matching the kernel repository's published precision band — including the Wan-representative 32760×12×128 (cos 0.990553, max-diff 0.011381) and a 384/512 cross-attention shape. Full per-shape table (cos / max-diff / µs / input-artifact hashes) recorded in the validation logs.

The capability gate's (10, 3) entry and the PR body's sm_103 caveat are therefore backed by hardware evidence now; the kernel install pins documented in this PR are exactly what produced these runs.

[SolitaryThinker]

Dispatch guard repaired in ad4d5e2e (review finding): AttnQatInferImpl.forward() now branches on the single per-arch resolution — fa4_fp4 → FA4 route, cutlass_sm12x → bundled extension, otherwise a clean ImportError carrying the receipt. Previously an unsupported GPU carrying an importable sm_12x wheel could reach the wrong binary through direct construction (the selector gates the common path, but the construct-anywhere/fail-cleanly-at-forward contract must hold everywhere). New CPU regression pins it: an sm_90 host with an importable extension raises at forward and the extension is never called. Also refreshed the stale sm_103 wording (PR body + GPU-test docstring) to hardware-validated per the receipts above.

[mergify]

Pre-commit checks failed

Hi @SolitaryThinker, the pre-commit checks have failed. To fix them locally:

# Install pre-commit if you haven't already
uv pip install pre-commit
pre-commit install

# Run all checks and auto-fix what's possible
pre-commit run --all-files

Common fixes:

• yapf: yapf -i <file> (formatting)
• ruff: ruff check --fix <file> (linting)
• codespell: codespell --write-changes <file> (spelling)

After fixing, commit and push the changes. The checks will re-run automatically.

For future commits, pre-commit will run automatically on changed files before each commit.

[SolitaryThinker]

Compile-interaction fix pushed in cc358f17 (plus a one-line mypy narrowing in 0984304f for the dispatch-guard commit): the FA4-FP4 forward resolved its quantization/kernel entry points on every call — function-level imports plus flashinfer's JIT-module lookup, whose first resolution probes the CUDA toolchain via a subprocess. Under torch.compile this graph-broke every step: fullgraph=True was impossible for the NVFP4 path, and even without fullgraph the per-step break kept eager dispatch overhead that the BF16 path sheds — inverting the NVFP4 speed story at short-step resolutions (documented as Finding 1 of the LTX-2.3 bf16-vs-NVFP4 compile benchmark).

Fix at both call sites (_nvfp4_quantize_for_fa4, shared with the existing nvfp4_fa4 path, and the ATTN_QAT_INFER route): resolution stays lazy — the construct-anywhere/fail-cleanly-at-forward contract from the dispatch-guard repair is unchanged — but is memoized at module level after first success, with the slow path factored into its own function so tests can pin it. New CPU regression: across 3 forward calls the resolution executes exactly once (call-recorder, same pattern as the dispatch regression).

[SolitaryThinker]

Compile fix round 2 in afb29e5b — the memoization alone was insufficient, and here is the principle it missed: Dynamo traces function bodies symbolically, so a runtime cache hit is invisible at trace time; fullgraph=True kept failing byte-identically at the toolchain probe (subprocess.Popen reached through flashinfer's JIT module lookup), and the fullgraph=False numbers were statistically unchanged. The resolution call has to leave the compiled region entirely.

Fix: the FP4 quantize step is now a torch.library custom op (fastvideo::nvfp4_quantize_fa4, forward + registered fake) — one opaque graph node, same boundary pattern this repo already uses for the masked flash-attention entry points. The round-1 memo stays (it still removes eager-path re-resolution; the subprocess runs once per process inside the op's first eager call). Public helper name/signature unchanged, so both consumers (the nvfp4_fa4 path and the ATTN_QAT_INFER route) are untouched.

New regression alongside the resolve-once test: torch.compile(fullgraph=True, backend="eager") traces AND runs the op-backed path with a CPU kernel registered in-test — any graph break fails the test by construction.

[SolitaryThinker]

Compile fix round 3 in e05849c6 — execution-gate fix. The GPU rerun proved the two gates separately: tracing was fixed at the previous head (zero graph breaks, first successful fullgraph trace of the NVFP4 path), but every compiled execution crashed at runtime — torch.compile bakes the registered fake's output STRIDES into the compiled artifact, and the fake returned a contiguous scale-factor tensor where the real kernel returns a permuted view (AssertionError: expected size 32==32, stride 16==8192 at dim=0). The earlier "known minor — at worst an inductor copy" note in this PR was wrong.

Fix: the fake now builds the scale-factor output with the impl's own base shape and permutation — new_empty((batch, nheads, rest_m, rest_k, 32, 4, 4)).permute(4, 5, 2, 6, 3, 1, 0) — matching the real strided view for general batch (a 6-dim variant matched the observed asserts only because that repro ran batch=1, where the trailing size-1 dim's stride is layout-irrelevant; the impl is 7-dim and its docstring was correct).

Regression upgraded so this class is mechanically caught: torch.library.opcheck cross-checks the registered fake against a real kernel run — shapes, dtypes, and strides — using a batch>1, seqlen-not-multiple-of-128 input so no stride is degenerate (test_utils restricted to schema/faketensor/aot-dispatch; the op is forward-only). The fullgraph trace test's CPU kernel now produces the production stride contract too. Lesson encoded in-test: shape-only fakes are insufficient for ops returning non-contiguous views.