perf(FourierUnit): single fp32 region — +8% fps, −20% VRAM, weight-compatible

[NevermindNilas]

What

FourierUnit.forward runs the FFT block as a single fp32 region instead of casting the spectrum down to fp16 for the intervening norm/conv ops and back up for the inverse FFT.

The FFT mandates fp32. The old path did 4 casts per FourierUnit (fp32→fp16 around the convs, fp16→fp32 around the FFTs). Keeping the norm/convs in the same fp32 region needs only 2 casts (entry/exit), cutting memory traffic and peak VRAM.

Why it's safe (weight-compatible, no retrain)

• Existing checkpoints load unchanged — no params added/removed/reshaped.
• FP32 inference path is bit-identical to before.
• FP16 path differs only by running the FFT-block convs in fp32, which is strictly more accurate.

Performance (RTX 3090, x4, autocast FP16, drift-canceled interleaved bench)

input	fps	peak VRAM
360p (640×360)	6.68 → 7.49 (+8.7%)	551 → 436 MB (−115)
480p (854×480)	3.84 → 4.13 (+7.7%)	892 → 686 MB (−206)

FP32: speed/VRAM neutral, output bit-identical. VRAM saving grows with resolution (larger spectrum tensors).

Quality (Urban100, 5 frames, x4, edited vs original, identical weights)

Verified on real periodic structure (building facades), where FFT ringing would surface — not synthetic noise.

	PSNR vs fp32 GT	SSIM
original arch fp16 (noise floor)	74.60 dB	0.99999
edited arch fp16	74.70 dB	0.99999

Δ +0.10 dB — within fp16 noise (in fact fractionally better, since the convs now run fp32). FP32 output bit-exact (PSNR ∞ / SSIM 1.000000). Edited meets-or-beats the noise floor on all 5 frames.

Notes

No torch.compile / CUDA graphs / channels_last — those were explicitly out of scope, and channels_last in particular measured −26% here because the FFT path forces NCHW and fights it. This is a pure arithmetic-region change.

🤖 Generated with Claude Code

[Copilot]

Pull request overview

Note

Copilot was unable to run its full agentic suite in this review.

This PR optimizes the FFT-based forward pass by keeping FFT-adjacent operations in fp32 to reduce dtype cast round-trips and memory traffic under mixed precision.

Changes:

• Disable AMP autocast around the FFT block and run intermediate ops in fp32.
• Replace contiguous().view(...) with reshape(...) after a permute(...).
• Move final dtype restoration to a single cast at the end of forward.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[NevermindNilas]

Fixed in cb01d85 — autocast now uses x.device.type. (Note: torch.autocast("cuda", enabled=False) is actually a no-op on CPU and doesn't raise — verified — but device-agnostic is cleaner for the export path.)

[NevermindNilas]

Made explicit in cb01d85 (contiguous().view(...)). For reference these are equivalent here — after a permute, reshape performs the same single copy as contiguous().view, so no perf difference — but explicit is fine.

[NevermindNilas]

Intentional and measured. post_norm is a custom RMSNorm (norm + addcmul), not a fused fp16 kernel — and AMP's autocast policy already runs normalization in fp32, so keeping it in this fp32 region matches prior numerics rather than changing them.

The suggested alternative (post_norm outside the disabled block) re-introduces a cast: irfft returns fp32, so post_norm in fp16 needs an extra fp32→fp16 round-trip — exactly the round-trip this PR removes (4 casts → 2). The +8% fps / −20% VRAM benchmark already has post_norm inside the region, so its effect is included and the net is positive. Moving it out measured slower.