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

CHANGELOG.md

@@ -4,6 +4,55 @@ All notable changes to nvSubquadratic are documented here. The format is based
 on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/) and the project
 follows [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## \[Unreleased\]
+
+### Added
+
+- **`fft_backend="subq_ops_fused"` on `CKConvND`**, backed by
+  `subquadratic_ops_torch.fused_fft_conv2d`. It runs the whole
+  rfft2 → multiply → irfft2 pipeline in a single cuFFTDx launch and, unlike
+  every other FFT path, **natively in fp32/fp16/bf16** instead of upcasting to
+  fp32. Roughly 3–4× faster than `torch_fft` and 1.2–2.4× faster than
+  `subq_ops` on forward+backward in bf16.
+
+  Restricted to `data_dim=2`, `is_causal=False`, `fft_padding="zero"`, 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 rather than at construction, because the input size
+  is not known when the module is built. Per-sample (FiLM) kernels are
+  supported. Requires `subquadratic-ops-torch >= 0.2.2`.
+
+  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 so results are interchangeable with the other backends
+  (verified to ~3e-7 normwise in fp32). Without that pre-pad the output is
+  shifted by `fft_size // 2 - K // 2` pixels.
+
+- **`nvsubquadratic.ops.fftconv_lowering`** — a `torch.compile` pre-grad pass
+  that detects `fftconv.py`'s 2D FFT-conv chain and rewrites it 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 generate code for complex
+  operators and otherwise falls back to eager cuFFT for the entire chain.
+
+  Enable it per-callable with
+  `torch.compile(model, options=fused_fftconv2d_options())`, or globally for a
+  scope with the `fused_fftconv2d_lowering()` context manager when a framework
+  owns the `torch.compile` call.
+
+  The pass only fires on an exact match of the reference recipe (padding rule,
+  crop offset, shape limits, CUDA device, and a compute capability that
+  supports the required FFT tile — the 128 tile needs SM90+). `lowering_stats()`
+  reports rewrite and per-reason skip counts, since a silent pass is otherwise
+  hard to tell apart from one that never fired.
+
+### Fixed
+
+- `tests/conftest.py` only queried the `subquadratic-ops-torch-cu12`
+  distribution when resolving the installed kernel version. On a
+  `-cu13` install (what `pyproject.toml` pins) it resolved to `(0, 0, 0)`,
+  which silently turned `requires_subq_ops_v2` into a blanket `xfail` and hid
+  every `subq_ops` test result. It now checks both distributions.
+
 ## \[0.1.1\]
 
 ### Changed

Dockerfile

@@ -1,19 +1,31 @@
-# Development Dockerfile for nvSubquadratic
+# Development Dockerfile for nvSubquadratic (CUDA 13.2 / PyTorch cu132)
 #
 # Build instructions:
-#   docker build -t nvsubquadratic:dev .
+#   DOCKER_BUILDKIT=1 docker build -t nvsubquadratic:dev .
+#
+# Requires an NVIDIA driver >= 580 on the host (CUDA 13.x minimum), or the
+# forward-compatibility package. Volta (sm_70) GPUs are no longer supported.
 #
 # Layer order is intentional for CI cache efficiency:
 #   1. Base image + conda + torch + DALI  (never changes)
 #   2. Apex build                         (changes only if apex version bumped)
 #   3. requirements-dev.txt               (changes when dev deps change)
 #   4. COPY . . + pip install             (changes on every code push — fast)
 
-FROM nvcr.io/nvidia/cuda:12.9.0-devel-ubuntu22.04
+FROM nvcr.io/nvidia/cuda:13.2.1-devel-ubuntu22.04
 
 ARG MINIFORGE_NAME=Miniforge3
 ARG MINIFORGE_VERSION=25.3.0-3
 
+# ── CUDA 13.2 toolchain pins ─────────────────────────────────────────────────
+# CUDA 13 wheels live under the cu132 index. The first PyTorch release built
+# against CUDA 13.2 is 2.12.0, so the previous 2.10.0 pin cannot be kept here.
+ARG TORCH_VERSION=2.12.1
+ARG TORCHVISION_VERSION=0.27.1
+ARG TORCH_INDEX_URL=https://download.pytorch.org/whl/cu132
+# DALI ships one build per CUDA major version; cuda130 is the CUDA 13.x build.
+ARG DALI_PACKAGE=nvidia-dali-cuda130
+
 ENV CONDA_DIR=/opt/conda
 ENV LANG=C.UTF-8 LC_ALL=C.UTF-8
 ENV PATH=${CONDA_DIR}/bin:${PATH}
@@ -40,11 +52,14 @@ RUN conda install --yes \
     python=3.12 \
     && conda clean --all --yes
 
+# The torch install is repeated after DALI on purpose: DALI pulls its own
+# nvidia-* CUDA runtime wheels and can shuffle the ones torch depends on, so we
+# re-assert the pinned cu132 build afterwards.
 RUN pip install --no-cache-dir \
-    torch==2.10.0 torchvision==0.25.0 --index-url https://download.pytorch.org/whl/cu129 \
-    && pip install --no-cache-dir nvidia-dali-cuda120 \
+    torch==${TORCH_VERSION} torchvision==${TORCHVISION_VERSION} --index-url ${TORCH_INDEX_URL} \
+    && pip install --no-cache-dir ${DALI_PACKAGE} \
     && pip install --no-cache-dir \
-       torch==2.10.0 torchvision==0.25.0 --index-url https://download.pytorch.org/whl/cu129 \
+       torch==${TORCH_VERSION} torchvision==${TORCHVISION_VERSION} --index-url ${TORCH_INDEX_URL} \
     && conda clean --all --yes
 
 # Create ubuntu user with sudo privileges
@@ -66,7 +81,11 @@ WORKDIR /workspaces/nvSubquadratic
 # ── Heavy build: Apex from source (cached until apex commit changes) ──────────
 # This layer is intentionally placed before COPY so code changes do not
 # trigger a rebuild. Apex does not depend on the project source.
-ARG TORCH_CUDA_ARCH_LIST="7.0;7.5;8.0;8.6;8.9;9.0;10.0;12.0"
+# CUDA 13 dropped offline compilation for Maxwell/Pascal/Volta — `nvcc
+# --list-gpu-arch` in 13.2 starts at compute_75, so 7.0 must be removed or the
+# apex build fails with "Unsupported gpu architecture 'compute_70'".
+# Newly available in 13.2 if you need them: 8.7, 8.8, 10.3 (B300), 11.0, 12.1.
+ARG TORCH_CUDA_ARCH_LIST="7.5;8.0;8.6;8.9;9.0;10.0;12.0"
 ENV TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST}
 ARG MAX_JOBS=""
 RUN MAX_JOBS="${MAX_JOBS}" pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation \
@@ -90,7 +109,7 @@ RUN git config --global --add safe.directory /workspaces/nvSubquadratic
 # pulls them — [all] restores the complete pre-restructure dependency set.
 RUN pip install --no-cache-dir wheel-stub \
     && pip install --no-cache-dir --no-build-isolation ".[all]" \
-       --extra-index-url https://download.pytorch.org/whl/cu129
+       --extra-index-url ${TORCH_INDEX_URL}
 
 # Set up ubuntu user's home directory and permissions
 RUN chown -R ubuntu:ubuntu /workspaces && \

README.md

@@ -13,6 +13,7 @@ nvSubquadratic consolidates efforts from across NVIDIA Research teams (nvResearc
 - **B2B CausalConv1d**: Back-to-back causal convolutions for striped Hyena architectures
 - **CausalConv1d**: Standard causal convolutions with various kernel sizes (2-256)
 - **FFT CausalConv1d**: FFT-based causal convolutions for large kernel sizes (up to 8K-16M)
+- **Fused FFT Conv2d**: single-launch 2D FFT convolution running natively in fp32/fp16/bf16 (spatial dims up to 64 per axis); requires `subquadratic-ops-torch >= 0.2.2`
 
 **Requirements**:
 
@@ -60,8 +61,15 @@ pip install "nvsubquadratic[all]"          # all of the above
 The accelerated CUDA kernels (`[cuda]`) are a source build that requires `nvcc`,
 so they are kept out of the core install — this is what lets `pip install nvsubquadratic` succeed in environments without the CUDA toolkit (e.g. a
 downstream project's CPU CI). The operators default to the portable `torch.fft`
-backend; selecting `fft_backend="subq_ops"` without `[cuda]` installed raises a
-clear `ImportError` pointing you to the extra.
+backend; selecting `fft_backend="subq_ops"` (or `"subq_ops_fused"`) without
+`[cuda]` installed raises a clear `ImportError` pointing you to the extra.
+
+On 2D problems with spatial dims of at most 64 per axis, `fft_backend="subq_ops_fused"`
+is the fastest option: it fuses the whole FFT-conv pipeline into one launch and
+runs it natively in bf16/fp16 instead of upcasting to fp32, for roughly a 3-4x
+speedup over `torch_fft`. Models already written against `torch_fft` can pick up
+the same kernel under `torch.compile` without a config change — see the
+[torch.compile lowering](docs/ops/README.md#torchcompile-lowering).
 
 ### Package Manager
 

docs/api_reference/ops.rst

@@ -49,6 +49,47 @@ fp32 reference ops above.
    ~ops.fftconv_custom.causal_fftconv1d_bhl
    ~ops.fftconv_custom.causal_fftconv1d_bhl_w_reshape
 
+.. _ops-fftconv-fused:
+
+Fused 2D FFT convolutions (native dtype)
+-----------------------------------------
+
+Wrappers around ``subquadratic_ops_torch.fused_fft_conv2d``, which fuses the
+whole rfft2/multiply/irfft2 pipeline into one launch and runs it in the input
+dtype (fp32/fp16/bf16) rather than upcasting to fp32.  Spatial dims are capped
+at 64 per axis.  Selected on ``CKConvND`` via ``fft_backend="subq_ops_fused"``.
+
+.. autosummary::
+   :toctree: generated/
+   :template: function_template.rst
+
+   ~ops.fftconv_custom.fused_fftconv2d_blh
+   ~ops.fftconv_custom.fused_fftconv2d_bhl
+   ~ops.fftconv_custom.fused_fftconv2d_bhl_w_reshape
+   ~ops.fftconv_custom.fused_fftconv2d_bhl_chunked
+   ~ops.fftconv_custom.resolve_fused_fft_size
+   ~ops.fftconv_custom.fused_fftconv2d_supported
+   ~ops.fftconv_custom.fused_fftconv2d_max_spatial
+   ~ops.fftconv_custom.load_fused_fft_conv2d
+
+.. _ops-fftconv-lowering:
+
+torch.compile lowering
+-----------------------
+
+Inductor pre-grad pass that rewrites the reference 2D FFT-conv chain onto the
+fused CUDA kernel, so an existing ``fft_backend="torch_fft"`` model picks it up
+without a config change.
+
+.. autosummary::
+   :toctree: generated/
+   :template: function_template.rst
+
+   ~ops.fftconv_lowering.fused_fftconv2d_options
+   ~ops.fftconv_lowering.fused_fftconv2d_lowering
+   ~ops.fftconv_lowering.lowering_stats
+   ~ops.fftconv_lowering.reset_lowering_stats
+
 .. _ops-causal-conv1d:
 
 Direct 1D causal convolutions (CUDA-accelerated)

docs/getting_started.md

@@ -100,7 +100,10 @@ hyena_cfg = LazyConfig(Hyena)(
         mask_cfg=LazyConfig(torch.nn.Identity)(),
         grid_type="double",  # linear (non-circular) convolution
         fft_padding="zero",
-        fft_backend="torch_fft",  # portable; "subq_ops" uses the fused 2D CUDA kernel
+        # Portable default. "subq_ops" uses the fused 2D CUDA kernel;
+        # "subq_ops_fused" is faster still and runs natively in bf16/fp16,
+        # but caps spatial dims at 64 per axis.
+        fft_backend="torch_fft",
         is_causal=False,
     ),
     # Depthwise short conv on the concatenated [Q; K; V] (3 * H channels).

docs/ops/README.md

@@ -41,7 +41,8 @@ ______________________________________________________________________
 | {ref}`circular_fftconv.py <ops-circular-fftconv>`  | fp32      | circular      | depthwise      | Periodic boundaries (e.g. PDEs, ARC grids), or when `K = N` so padding is wasteful.                                                                                                                                                              |
 | {ref}`mixed_fftconv.py <ops-mixed-fftconv>`        | fp32      | per-axis BC   | depthwise      | **Mixed boundaries**: periodic on some spatial axes, zero-padded on others (e.g. Well's `rayleigh_benard`, `viscoelastic_instability`, `turbulent_radiative_layer`). Routes to the existing linear/circular ops in the all-False/all-True cases. |
 | {ref}`fftconv_chunked.py <ops-chunking>`           | fp32      | linear        | depthwise      | Memory-constrained training; processes channels in chunks. Has a global flag so models can opt in transparently.                                                                                                                                 |
-| {ref}`fftconv_custom.py <ops-fftconv-custom>`      | fp32      | linear        | depthwise      | Wraps optional fused CUDA kernels (`subquadratic_ops_torch.fft_conv2d` for 2D non-causal, `fft_causal_conv1d` for 1D causal) behind the same API as `fftconv.py`.                                                                                |
+| {ref}`fftconv_custom.py <ops-fftconv-custom>`      | fp32 / native | linear    | depthwise      | Wraps optional fused CUDA kernels (`subquadratic_ops_torch.fft_conv2d` for 2D non-causal, `fft_causal_conv1d` for 1D causal) behind the same API as `fftconv.py`. Also hosts the `fused_fftconv2d_*` wrappers around `fused_fft_conv2d`, which run natively in fp32/fp16/bf16 but cap spatial dims at 64 per axis. |
+| {ref}`fftconv_lowering.py <ops-fftconv-lowering>`  | native    | linear        | depthwise      | `torch.compile` pre-grad pass that rewrites `fftconv.py`'s 2D chain onto `fused_fft_conv2d`. Use when you want the fused kernel without changing `fft_backend` on an existing model.                                                             |
 | {ref}`causal_conv1d_custom.py <ops-causal-conv1d>` | fp32      | direct causal | depthwise      | Non-FFT 1D causal kernels (`causal_conv1d` short conv, `b2b_causal_conv1d` fused proj-gate-mixer-gate). Use for kernels short enough that FFT overhead dominates, or as a fused-Hyena building block.                                            |
 
 The [FP16 circular FFT convolution report](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/blob/main/reports/fp16_fft_convolution/REPORT.md) contains the full derivation of the numerically stable fp16 circular conv (dual mean-centering + inclusion-exclusion geometric correction). Read it if you are touching the fp16 path or want to understand the math behind those `T1, T2, T3, T4` terms in the code.
@@ -123,14 +124,74 @@ ______________________________________________________________________
 1. **Is there a fused CUDA kernel for my shape?**
 
    - 2D non-causal or 1D causal long-conv → `fftconv_custom.py` exposes the upstream fused FFT kernels (`fft_conv2d`, `fft_causal_conv1d`) through the same API. Wire in via the `fft_backend="subq_ops"` flag on `CKConvND`. The 1D path requires `data_dim=1, is_causal=True`; the 2D path requires `data_dim=2, is_causal=False`.
+   - 2D non-causal with spatial dims **≤ 64 per axis** → prefer `fft_backend="subq_ops_fused"`, which uses the newer `fused_fft_conv2d` kernel. It runs the whole rfft2/multiply/irfft2 pipeline in one launch **natively in fp32/fp16/bf16** instead of upcasting to fp32 — roughly 3–4× faster than `torch_fft` end-to-end in bf16 and ~1.2–2.4× faster than `subq_ops`. Requires `subquadratic-ops-torch >= 0.2.2`. The 64×64 cap comes from the kernel's largest FFT tile (128) combined with its `max(X, Y) ≤ fft_size // 2` requirement, and is enforced on the first forward pass rather than at construction. Beyond that, fall back to `subq_ops` or `torch_fft`.
+   - Already have a model on `fft_backend="torch_fft"` and want the fused kernel without a config change → `fftconv_lowering.py` provides a `torch.compile` pre-grad pass that detects the FFT-conv chain and rewrites it. See [the lowering section below](#torchcompile-lowering).
    - 1D causal short conv (typical short_conv slot in a Hyena block) → `causal_conv1d_custom.py` exposes `causal_conv1d` directly, and `nvsubquadratic.modules.subq_ops_causal_conv1d.SubqOpsCausalConv1d` wraps it as a depthwise `nn.Conv1d`-compatible module.
    - 1D causal fused proj+gate+mixer+gate block → `b2b_causal_conv1d` in `causal_conv1d_custom.py`. Not yet wired into a `Hyena` variant; exposed as a building block.
 
 ______________________________________________________________________
 
+## torch.compile lowering
+
+Inductor cannot generate code for complex operators — it warns
+`Torchinductor does not support code generation for complex operators` and falls
+back to eager cuFFT for the whole `rfft2 → multiply → irfft2` chain. So
+`torch.compile` on its own buys essentially nothing for `fftconv.py`.
+
+`fftconv_lowering.py` closes that gap. It registers an inductor **pre-grad**
+pass (before AOTAutograd, so autograd comes from the custom op's registered
+backward) that detects the chain and swaps in the fused kernel:
+
+```python
+import torch
+from nvsubquadratic.ops.fftconv_lowering import fused_fftconv2d_options, lowering_stats
+
+compiled = torch.compile(model, options=fused_fftconv2d_options())
+out = compiled(x)
+
+print(lowering_stats())  # {'rewritten': 1}
+```
+
+`fused_fftconv2d_options()` is scoped to that one compiled callable and leaves
+global inductor config untouched. When a trainer or framework owns the
+`torch.compile` call and you cannot pass `options`, the
+`fused_fftconv2d_lowering()` context manager does the same thing by patching
+global config for its duration:
+
+```python
+from nvsubquadratic.ops.fftconv_lowering import fused_fftconv2d_lowering
+
+with fused_fftconv2d_lowering():
+    out = trainer.fit(model)  # whatever compiles internally
+```
+
+The pass only fires when the graph is *exactly* `fftconv.py`'s recipe — the
+`min(N + (K+1)//2, 2N)` padding, the `K // 2` crop, shapes within the fused
+kernel's limits, CUDA device, and a compute capability that supports the
+required FFT tile (the 128 tile needs SM90+; SM80/SM86 lack the shared memory).
+Anything else is left on the eager path.
+
+Two things to know:
+
+- **Setting `fft_backend="subq_ops_fused"` is the more predictable route** and
+  needs no pass at all. Reach for the lowering when you want the kernel under an
+  existing `torch_fft` model without touching its config.
+- **`lowering_stats()` only counts passes that actually ran.** Inductor caches
+  compiled artifacts on disk, and a cache hit skips every pre-grad pass, so
+  empty counters can mean "served from cache" rather than "did not fire". Set
+  `torch._inductor.config.force_disable_caches = True` when verifying.
+
+By default the pass also rewrites fp16/bf16 graphs, which moves the convolution
+from fp32-internal to native-dtype — that is where most of the speedup comes
+from, and it is a real numerics change (~2e-3 normwise in bf16). Pass
+`fused_fftconv2d_lowering(allow_reduced_precision=False)` to restrict it to
+fp32, where the rewrite is numerically neutral.
+
+______________________________________________________________________
+
 ## Numerical notes
 
-- All operators **accept any input dtype** but cast to the internal compute precision (fp32 or fp16) before the FFT. The output is returned in the **original dtype of `x`**, so no manual cast is needed on the caller side.
+- All operators **accept any input dtype** but cast to the internal compute precision (fp32 or fp16) before the FFT. The output is returned in the **original dtype of `x`**, so no manual cast is needed on the caller side. The one exception is the `fused_fftconv2d_*` family, which runs in `x.dtype` end to end — that is the point of it. Measured against the fp32 reference, its normwise relative error is ~3e-7 in fp32, ~3e-4 in fp16, and ~3e-3 in bf16.
 - The fp32 ops are correct for any input range. The fp16 ops impose two constraints: spatial dims must be powers of two (cuFFT), and the dynamic range is handled by mean-centering both `x` and `k` (see derivation doc).
 - The non-causal linear ops match a standard `torch.nn.ConvNd(padding='same')` up to floating-point rounding. The circular ops match `torch.nn.functional.conv*d` after a circular pad. Both are exercised in `tests/`.