[draft]: integrate fused_fft_conv2d backend and torch.compile lowering

[moradza]

Adds subquadratic_ops_torch.fused_fft_conv2d as a first-class 2D FFT-conv path. Unlike every other FFT backend it runs the whole rfft2/multiply/irfft2 pipeline in a single cuFFTDx launch and natively in fp32/fp16/bf16 rather than upcasting to fp32. Measured on H200 (B=8, H=768, fwd+bwd, bf16): 3.6-3.9x over torch_fft and 1.2-2.4x over the existing subq_ops path.

Two entry points:

• fft_backend="subq_ops_fused" on CKConvND — explicit and predictable. Restricted to data_dim=2, non-causal, zero padding, and spatial extents of at most 64 per axis (the kernel's largest FFT tile is 128 and it requires max(X, Y) <= fft_size // 2). The spatial cap is enforced on the first forward pass, not at construction, since input size is unknown there.

• nvsubquadratic.ops.fftconv_lowering — an inductor pre-grad pass that rewrites fftconv.py's chain onto the fused kernel, so a model already on fft_backend="torch_fft" picks it up without a config change. This matters because inductor cannot codegen complex operators and otherwise falls back to eager cuFFT for the entire chain. Enable per-callable with torch.compile(model, options=fused_fftconv2d_options()), or for a scope with the fused_fftconv2d_lowering() context manager.

  Pre-grad rather than post-grad so autograd is derived from the custom op's registered backward instead of requiring a consistent forward+backward rewrite. It fires only on an exact match of the reference recipe (padding rule, crop offset, shape limits, CUDA device, and a compute capability that supports the required tile — the 128 tile needs SM90+, as SM80/SM86 lack the shared memory). lowering_stats() reports rewrite and per-reason skip counts, since a silent pass is otherwise indistinguishable from one that never ran.

Crop-offset reconciliation: the upstream kernel crops the 'same' window at fft_size // 2 whereas fftconv.py crops at K // 2. The wrapper pre-pads the filter's top/left by the difference, making results interchangeable with the other backends (~3e-7 normwise in fp32, ~3e-3 in bf16). Without it the output is shifted by fft_size // 2 - K // 2 pixels — which reads as a ~1.41 (sqrt 2, fully decorrelated) error, not as reduced accuracy. A regression test pins this so the pre-pad cannot be optimised away.

Also registers the fused operators eagerly when the pass is constructed: inductor's on-disk FX cache lets a compiled artifact call torch.ops.subquadratic_ops_torch.* on a cache hit without any Python call having triggered the wrappers' lazy import, which failed with an opaque op-namespace AttributeError.

Fixes a pre-existing test-gate bug: tests/conftest.py resolved the kernel version from the subquadratic-ops-torch-cu12 distribution only. On the -cu13 install that pyproject.toml pins it returned (0, 0, 0), silently turning requires_subq_ops_v2 into a blanket xfail and hiding every subq_ops test result. It now checks both distributions.

Includes the in-flight CUDA 13.2 build changes already present in the working tree (Dockerfile base image, cu132 torch/DALI pins, cu13 kernel distribution).

Tests: 138 new (op-level equivalence across dtypes/shapes/layouts/FiLM, forward and backward, CKConvND integration, and lowering fire/decline behaviour).

Summary

Environment setup

Create the conda environment (required to run tests):

bash setup_conda_env.sh
conda activate nvsubquadratic

Test plan

• pre-commit run --all-files passes (pre-commit install if not yet set up).
• Existing tests pass (pytest tests/).
• New tests added, or explain why not needed:

Documentation checklist

For every new or modified public symbol in nvsubquadratic/ or experiments/:

• Every new module has a module-level docstring explaining what it contains and why.
• Every new public class has a class docstring covering purpose, math/motivation, and key attributes.
• Every new public method / function has Args: and Returns: blocks with tensor shapes where applicable.
• Math notation is consistent with the paper (or a comment explains any deviation).
• Docstrings containing backslashes use r"""...""" (required by ruff D301).
• If a new file was added, a row has been added to docs-tracker.md with status [x].

See CONVENTIONS.md for the full style guide.