feat(ops): fused FFT-conv2d backend, torch.compile lowering, and CUDA 13 migration by farhadrgh · Pull Request #139 · NVIDIA-BioNeMo/nvSubquadratic

farhadrgh · 2026-09-02T18:50:57Z

Extracted from farhadr/2d_bench (PR #137), taking the library feature and the build migration but not the ND benchmark harness. Supersedes PR #138.

What's here

Fused FFT-conv2d backend. fft_backend="subq_ops_fused" on CKConvND, backed by subquadratic_ops_torch.fused_fft_conv2d — one cuFFTDx launch for the whole rfft2 → multiply → irfft2 pipeline, running natively in fp32/fp16/bf16 instead of upcasting. Restricted to data_dim=2, non-causal, fft_padding="zero", spatial extents ≤64/axis.

torch.compile lowering. nvsubquadratic.ops.fftconv_lowering — a pre-grad pass that rewrites fftconv.py's 2D chain onto the fused kernel, so a model already on torch_fft picks it up without a config change.

CUDA 12.9 → 13.0 migration. cu130 torch wheels, torch floor → >=2.14.0,<2.15.0, [cuda] → subquadratic-ops-torch-cu13, [dali] → nvidia-dali-cuda130, Dockerfile/enroot updates, and a clean-install CI matrix.

The pin change, and why it fixes CI

farhadr/2d_bench pinned subquadratic-ops-torch-cu13>=0.2.2, but 0.2.2 was published only to the internal NVIDIA GitLab registry — public PyPI topped out at 0.2.1. That is precisely why both of PR #137's checks fail today:

gpu-tests dies building the Docker .[all] layer: SUBQ_INDEX="$(cat /run/secrets/subq_index ...)" is empty on CI, so >=0.2.2 is unresolvable from PyPI.
lint fails on an mdformat table reflow in docs/ops/README.md.

0.3.0 reached public PyPI on 2026-08-28 with the same fused_fft_conv2d entry point and signature. This PR pins >=0.3.0, which lets the Dockerfile and both enroot scripts drop their GITLAB_TOKEN / BuildKit-secret plumbing entirely and resolve from PyPI. The mdformat reflow is applied. Both failures should therefore be resolved here.

Version consistency

main's install docs and pyproject.toml had drifted apart during the migration — the documented conda/venv paths still installed torch==2.10.0, which the new >=2.14.0,<2.15.0 floor rejects, so pip install -e . would silently re-resolve torch from PyPI and replace the CUDA-matched wheel after apex/mamba were compiled against it. Now aligned on torch 2.14.0 / torchvision 0.29.0 / cu130 across pyproject.toml, Dockerfile, setup_conda_env.sh, scripts/slurm/setup_env.sh, build_sqsh_slurm.sh, README.md, docs/getting_started.md, and the lean-path CI check. Also corrected the README's apex arch list to 7.5;... (Turing through Blackwell) — CUDA 13 dropped offline Volta compilation, so the previous 7.0 entry hard-fails the apex build.

Relationship to other PRs

Supersedes [draft]: integrate fused_fft_conv2d backend and torch.compile lowering #138 (cuda13.2). Its commit ee5146e is already an ancestor here; all 14 of its files are present, 6 byte-identical and 8 with later fixes on top. Note this branch corrects [draft]: integrate fused_fft_conv2d backend and torch.compile lowering #138's CUDA 13.2/cu132 base to 13.0/cu130 — apex requires the base nvcc CUDA to match torch's exactly, and a 13.2 base fails with Cuda extensions ... compiled with Cuda 13.0 vs nvcc 13.2. [draft]: integrate fused_fft_conv2d backend and torch.compile lowering #138 can be closed.
Does not include Mixer Benchmarks #137's benchmark harness: benchmark_forward_time_nd_resolution.py, benchmark_patch_size_2d.py, visualize_forward_time_nd.py, the submit_forward_time_*.sh scripts, and the FA4/FlexAttention additions to attention.py all stay on farhadr/2d_bench.

Testing

gpu-tests on the Colossus A40 runner (SM86): 1136 passed, 38 skipped, 2 xpassed. Full pre-commit run --all-files clean.

The fused path does execute and pass in CI. An earlier revision of this description claimed it was entirely unvalidated — that was written before the image resolved subquadratic-ops-torch-cu13>=0.3.0, and it is no longer true. The run covers 88 test_fused_fftconv2d, 36 test_fftconv_lowering, and 24 test_ckconv_nd_subq_fused cases against the fp32 torch_fft reference, across fp32/fp16/bf16, both BHL and BLH layouts, chunked and unchunked, per-sample (FiLM) kernels, and the lowering pass's fire/decline behaviour.

⚠️ What remains unvalidated is narrower: the 128 FFT tile. resolve_fused_fft_size escalates any extent above 32 per axis to the 128 tile, which needs more shared memory than SM80/SM86 provide — so the documented 64-per-axis cap is only reachable on SM90+. Those 4 cases skip via the requires_sm90 marker and have never run on Hopper/Blackwell. Below SM90 the arch guard raises a clear error rather than failing inside the kernel, so the path is safe by construction, but an H100/B200 run should confirm it before anyone relies on 64x64 in production.

Not covered anywhere: non-square (H != W) extents have no numeric test, and the lowering tests all run with the inductor FX cache disabled, so the cache-hit path that motivates eager op registration is untested.

🤖 Generated with Claude Code

… 13 migration Adds fft_backend="subq_ops_fused" on CKConvND, backed by subquadratic_ops_torch.fused_fft_conv2d: one cuFFTDx launch for the whole rfft2 → multiply → irfft2 pipeline, running natively in fp32/fp16/bf16 instead of upcasting. Ships with a torch.compile pre-grad pass (nvsubquadratic.ops.fftconv_lowering) that rewrites the existing torch_fft chain onto the same kernel, so models pick it up without a config change. Also migrates the CUDA runtime 12.9 → 13.0: cu130 torch wheels, torch floor raised to >=2.12.0, [cuda] extra retargeted to subquadratic-ops-torch-cu13, [dali] to nvidia-dali-cuda130, plus the Dockerfile/enroot build updates and a clean-install CI matrix. The [cuda] extra pins >=0.3.0 rather than 0.2.2: 0.2.2 was published only to the internal GitLab registry, so pinning it would have broken `pip install nvsubquadratic[cuda]` for public users. 0.3.0 reached public PyPI on 2026-08-28 with the same fused_fft_conv2d entry point, which lets the Dockerfile and enroot scripts drop their GITLAB_TOKEN/BuildKit-secret plumbing and resolve everything from PyPI. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

pyproject.toml declares the supported range, but the Docker image, the conda bootstrap, the SLURM builder and the install docs each repeat a concrete version and wheel index. Nothing links them — build_sqsh_slurm.sh replays the Dockerfile rather than parsing it, and the docs are prose — so they drift when only one is edited, which is exactly how setup_conda_env.sh ended up pointing at the cu130 index while still pinning torch 2.10.0. That drift is not cosmetic: installing a torch outside the pyproject range means the later `pip install -e .` re-resolves torch from PyPI and replaces the CUDA-matched wheel after apex/mamba/causal-conv1d were already compiled against it, leaving extensions built for a torch that is no longer installed. scripts/check_version_pins.py parses the range from pyproject.toml and asserts across all 14 pin sites that every exact pin satisfies it, that the pins agree with each other, that every PyTorch wheel index names the same cuXYZ build, that the [cuda] and [dali] extras target that same CUDA major, and that files mirroring the specifier quote it verbatim. It also fails loudly rather than silently passing when a pin is reworded out from under its pattern. Wired in as a pre-commit local hook, so CI's `pre-commit run --all-files` picks it up with no workflow change. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

farhadrgh · 2026-09-02T19:31:29Z

Correction: the `>=0.3.0` pin does not fix `gpu-tests` — it introduces a hard resolution conflict

The PR description claims bumping to subquadratic-ops-torch-cu13>=0.3.0 resolves both of #137's failing checks. lint is fixed (green here). gpu-tests is not, and I was wrong to claim it would be. Root cause, from the failed run:

torch 2.12.0                     depends on nvidia-cudnn-cu13==9.20.0.48   (exact pin)
subquadratic-ops-torch-cu13 0.3.0 depends on nvidia-cudnn-cu13>=9.24.0.43

These are mutually exclusive, so pip install nvsubquadratic[cuda] — the exact command in the README — cannot resolve on torch 2.12.x. In the Docker build this surfaces indirectly: pip backtracks through ~26 versions of datasets/jinja2/MarkupSafe looking for an escape and dies on MarkupSafe 1.0's setup.py (ImportError: cannot import name 'Feature' from 'setuptools'), which is misleading — the real conflict is cudnn.

Verified locally with pip install --dry-run on py3.12 against the cu130 index:

torch	pins `nvidia-cudnn-cu13`	vs 0.3.0's `>=9.24.0.43`
2.12.1	`==9.20.0.48`	conflict
2.13.0	`==9.20.0.48`	conflict
2.14.0	`==9.24.0.43`	satisfies

So subquadratic-ops-torch-cu13 0.3.0 was built against torch 2.14, not 2.12. Note also that nvidia-cudnn-cu13 on PyPI no longer publishes 9.20.0.48 at all (its oldest is now 9.25.1.1), so torch 2.12/2.13's exact pin is only satisfiable from the PyTorch index.

Falling back is not an option: 0.2.1 is the newest cu13 release without the cudnn floor, and I confirmed its wheel contains no fused_fft_conv2d module — it cannot back this PR's feature.

Validated fix

Raising the floor to torch>=2.14.0,<2.15.0 / torchvision>=0.29.0,<0.30.0 makes the whole closure resolve:

torch-2.14.0+cu130  torchvision-0.29.0+cu130  nvidia-cudnn-cu13-9.24.0.43
subquadratic-ops-torch-cu13-0.3.0  nvidia-dali-cuda130-1.53.0
megatron-core-0.19.0  timm-1.0.24  quack-kernels-0.6.4

torch-2.14.0+cu130 is present on the cu130 index, and it is still CUDA 13.0, so the Dockerfile's 13.0.3-devel base still satisfies the "base nvcc must match torch's CUDA exactly" constraint that apex/mamba need.

Not yet done — this widens the PR from a 2.10→2.12 bump to 2.10→2.14, and apex / mamba / DALI have not been built against 2.14. Holding for a call on scope before making that change.

Also note main's [all] set still resolves cleanly today (torch 2.10.0+cu130, subq-ops-cu12 0.2.1), so this is a regression introduced here, not pre-existing bit-rot.

…cument SM90+ `pip install nvsubquadratic[cuda]` could not resolve at all on the previous 2.12 floor. torch pins nvidia-cudnn-cu13 exactly, and subquadratic-ops-torch-cu13>=0.3.0 requires nvidia-cudnn-cu13>=9.24.0.43: torch 2.12/2.13 -> nvidia-cudnn-cu13==9.20.0.48 subq-ops 0.3.0 -> nvidia-cudnn-cu13>=9.24.0.43 torch 2.14.0 pins ==9.24.0.43 and satisfies both, so 0.3.0 was in fact built against 2.14. In the Docker build the conflict surfaced misleadingly: pip backtracked through ~26 versions of datasets/jinja2/MarkupSafe hunting for an escape and died on MarkupSafe 1.0's setup.py under --no-build-isolation. Falling back is not an option — 0.2.1 is the newest cu13 release without the cudnn floor and its wheel carries no fused_fft_conv2d module. Verified with `pip install --dry-run` on py3.12 against the cu130 index that the whole [all] closure now resolves: torch 2.14.0+cu130, torchvision 0.29.0+cu130, nvidia-cudnn-cu13 9.24.0.43, subquadratic-ops-torch-cu13 0.3.0, DALI 1.53.0, megatron-core, timm, quack-kernels. torch 2.14.0+cu130 is published for both x86_64 and aarch64, and is still CUDA 13.0, so the 13.0.3-devel base still satisfies apex/mamba's exact-CUDA-match requirement. Also drops the torchaudio pin from setup_env.sh: nothing in the repo imports it, and the cu130 index has no build past 2.11. SM90+ gating: The 128 FFT tile needs more shared memory than SM80/SM86 provide, but that gate only existed in the torch.compile lowering pass. The eager fft_backend="subq_ops_fused" path had no check, while every user-facing doc described the limit as purely spatial — so a 64x64 input on an A100/A6000 (which resolves to the 128 tile) failed inside the kernel with no explanation. Adds fused_fftconv2d_arch_supported() plus a raise at the single call site, with the constant now shared with the lowering pass rather than duplicated, and states the requirement in README, docs/ops/README.md, ops.rst and the CHANGELOG. Docs: - docs/index.rst named subquadratic-ops-torch-cu12 and CUDA Toolkit 12.0+, which this branch made wrong; also syncs the backend paragraph with README. - docs-tracker.md gains the fftconv_lowering.py row required by CONVENTIONS.md and the PR template. - scripts/slurm/enroot/README.md documents build_sqsh_slurm.sh, which the branch added without mentioning it anywhere. - FusedFFTConv2dLowering gets an API page, so the three docstrings pointing at it are no longer dead text; adds the two missing chunked variants to ops.rst. `sphinx -W --keep-going` builds clean; 276 tests pass. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

The first gpu-tests run that got far enough to execute the suite reported 1126 passed / 5 failed, and all five were the `[64]` parameter: test_ckconv_nd_subq_fused.py::test_spatial_sizes_within_cap[64] test_fftconv_lowering.py::test_rewrites_across_supported_spatial_sizes[64] test_fused_fftconv2d.py::TestForwardMatchesReference::test_double_grid_kernel[64-dtype0/1/2] The CI runner reports compute capability 8.6, and resolve_fused_fft_size escalates any extent above 32 per axis to the 128 FFT tile, which needs more shared memory than SM80/SM86 provide. So these tests encoded an assumption — that the documented 64-per-axis cap holds on every GPU — that is only true on Hopper/Blackwell. On Ampere the effective cap is 32. Four of the five surfaced as the new arch guard raising; the fifth (test_rewrites_across_supported_spatial_sizes[64]) is the lowering pass correctly declining to rewrite via its own pre-existing SM90 check, so that one would have failed on this hardware regardless of the guard. Neither is a regression — this is the first run in which the fused suite executed at all. Adds a shared `requires_sm90` marker in tests/conftest.py and applies it to the individual 64 parameters rather than whole files, so the sub-64 coverage keeps running everywhere and the 64 cases still execute on an H100/B200 runner. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

farhadrgh · 2026-09-02T22:17:01Z

The torch 2.14 bump worked — the image builds and the suite runs

The gpu-tests run on 88832cc got all the way through: the Docker image built against torch 2.14.0+cu130 (apex and mamba both compiled), and the test suite executed for the first time on this branch — 1126 passed, 33 skipped, 5 failed in 7m53s. The dependency-resolution blocker is resolved.

What the 5 failures actually mean

All five were the [64] parameter, and the runner reports compute capability 8.6:

RuntimeError: The fused 2D FFT kernel needs fft_size=128 for input (64, 64) with
kernel (127, 127), and that tile requires compute capability 9.0+
(Hopper/Blackwell); this device reports 8.6.

Four are the new arch guard firing as designed. The fifth,
test_rewrites_across_supported_spatial_sizes[64], is the lowering pass declining to rewrite via its own pre-existing SM90 check (assert lowering_stats().get("rewritten") == 1 → None) — that one would have failed on this hardware with or without the guard.

Neither is a regression. This is simply the first run in which the fused suite executed at all; before the pin fix the build died before reaching pytest.

The substantive finding is that these tests encoded an assumption that is false on the project's own CI hardware: that the documented "64 per axis" cap holds on every GPU. It doesn't. resolve_fused_fft_size escalates anything above 32/axis to the 128 tile, so on Ampere (A40 runner, A6000 workstation) the effective cap is 32, not 64.

9fa20d3 adds a shared requires_sm90 marker in tests/conftest.py and applies it to the individual 64 parameters rather than to whole files, so sub-64 coverage keeps running everywhere and the 64 cases still execute on an H100/B200 runner.

Caveat for reviewers

With this gating, gpu-tests passing on the A40 runner does not validate the fused kernel at its headline 64×64 size, or the lowering pass's rewrite at that size. Those paths remain unexercised by CI until this runs on Hopper/Blackwell.

moradza · 2026-09-04T15:12:12Z

@farhadrgh can you replace fftconv_lowering with the files below.

Follow-up to the scoped-DCE change: tightened the matcher, expanded the tests.

In addition to the DCE fix, I found that the new "is the spectrum shared?" check only controlled deletion, not whether we rewrite at all. That and a few older gaps are fixed here. Every change was verified by running the pass on Dynamo graphs with the fused kernel stubbed by the eager reference.

Matcher fixes

The pass now skips when any spectrum or the multiply result is used outside the chain (spectrum-has-other-users). Before, extra math on the spectrum was silently dropped, or the cuFFT chain was left running next to the fused kernel while we reported success. Because of this gate, the erase is now unconditional.
An in-place mul_ ordered after the irfft2 no longer matches. That graph is not a convolution in eager.
Casts are only unwrapped and erased when they are plain same-device upcasts to fp32. An fp64 or cross-device .to() is kept and the recipe check runs against it.
The norm argument is checked on all three FFT nodes. Only the default backward norm matches.
rfft2(input=...) keyword spelling is handled via a shared _fft_input helper instead of raising IndexError.
Symbolic slice bounds under dynamic shapes now return "no match" instead of raising out of the inductor pass.
Kernel dtype and device must equal x's. The fused kernel rounds the kernel to x's dtype before the FFT, which the reference never does.
The chain's output must have x's dtype and device, with or without a trailing cast. This closes the case where a bf16 graph without the final .to() would silently return bf16 instead of fp32.

Cleanup

_erase_chain erases a known consumers-first list instead of scanning the whole graph. Its docstring now gives the real reason we avoid eliminate_dead_code.
rewrote_any, the conditional lint(), and the empty-candidates early return are gone. Inductor lints and recompiles right after the pass.
uuid() hashes this file's contents via torch's get_hash_for_files, so _PASS_VERSION and the risk of forgetting to bump it are gone.

Tests

A recording wrapper snapshots the graph after the pass. Positive tests now assert that no rfft2, irfft2, or mul_ survives. A mutation check with the erase disabled fails as expected.
New positive tests: out-of-place multiply, keyword input= spelling.
New decline tests: shared spectrum, misordered mul_, fp64 upcast, ortho norm, kernel dtype mismatch, uncast bf16 output.
Removed the duplicate assert in the in-place mutation test.
L2_TOL, L2_TOL_GRAD, l2_rel, and assert_l2_close moved to tests/conftest.py. The fused op test and the ckconv test import them; ckconv keeps its own looser gradient table.
Skip-reason changes in case anything greps for them: irfft-dim-not-2-3 is folded into irfft-args-mismatch. New reasons are spectrum-has-other-users, kernel-dtype-or-device-mismatch, and output-dtype-or-device-mismatch.

To run the real suite:

python -m pytest tests/ops/test_fftconv_lowering.py tests/ops/test_fused_fftconv2d.py tests/modules/test_ckconv_nd_subq_fused.py -v -o addopts=""

test_fftconv_lowering.py
fftconv_lowering.py

Replaces fftconv_lowering.py and its test file with the revision Alireza posted on PR #139, which fixes the two correctness bugs raised in review plus a matcher gap neither of us had caught. Whole-graph `graph.eliminate_dead_code()` ran after every successful rewrite with no is_impure_node override. torch.fx's Node.is_impure() only inspects call_function schemas, so an in-place call_method elsewhere in the same compiled graph (`state.add_(1)`) is treated as pure and silently deleted when its result is only reachable through the mutated variable — wrong numerics, no error. _erase_chain now erases a known consumers-first list of just the nodes the rewrite orphaned. `any(_try_rewrite(...) for node in candidates)` short-circuited on the first success, so a graph with two independent FFT-conv chains fused only the first and silently left the rest on the eager cuFFT path (and under-counted lowering_stats()). Every candidate is now attempted. The matcher additionally declines when a spectrum or the multiply result has users outside the chain — previously that extra math was dropped, or the cuFFT chain was left running beside the fused kernel while the pass reported success. That gate is what makes the erase unconditionally safe. Also declined rather than mis-rewritten: an in-place mul_ ordered after the irfft2, non-default norm, fp64/cross-device casts, kernel dtype/device mismatches, an output whose dtype/device does not match x, and symbolic slice bounds under dynamic shapes. uuid() hashes this file's contents instead of a hand-bumped version constant. L2_TOL/L2_TOL_GRAD/l2_rel/assert_l2_close move to tests/conftest.py so the three fused test files share one definition rather than three drifting copies; test_ckconv_nd_subq_fused keeps its own looser gradient table, since gradients through the full CKConvND stack accumulate more error than the bare op. 702 passed, 177 skipped, 77 xfailed, 3 xpassed across tests/ops and tests/modules; pre-commit clean. The fused path itself still does not execute here — this env has subquadratic-ops-torch-cu12 0.1.1 (below the 0.2.2 gate) and SM86 hardware — so the rewrite still needs an H100/B200 run. Co-Authored-By: Alireza Moradzadeh <amoradzadeh@nvidia.com> Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Bumps VERSION 0.1.1 -> 0.2.0 and closes the changelog section. release.yml fails fast when the tag does not match the VERSION file, so this has to land before the tag is created. 0.2.0 rather than 0.1.2: this release is dependency-breaking for every 0.1.1 user. The torch floor moves >=2.10,<2.11 -> >=2.14,<2.15, [cuda] retargets from subquadratic-ops-torch-cu12 to -cu13 (the cu12 line tops out at 0.2.1 and has no fused_fft_conv2d), and [dali] moves to nvidia-dali-cuda130. Under PEP 440 both `~=0.1.1` and `>=0.1,<0.2` reject 0.2.0, which is what protects downstream pins from resolving into an environment that cannot build. Also records three things the changelog was missing: - SubqOpsCausalConv1d could not run under torch.autocast in 0.1.0/0.1.1. The kernel selects its specialisation from the input dtype and then demands an exact match, so bf16/fp16 activations against fp32 parameters raised ValueError. This is a genuine Fixed entry — unlike the fftconv_lowering correctness fixes, which never shipped and so are folded into that feature's own description rather than listed as regressions. - The lowering pass's decline conditions, which are user-visible through lowering_stats() skip reasons, and the scoped erase that replaced whole-graph dead-code elimination. - That CI validates the fused path only up to 32 extents per axis (64 FFT tile) on SM86; the 64-per-axis 128-tile path is guarded but has never run on Hopper/Blackwell. Better stated in the release notes than discovered by a user with an H100. scripts/check_version_pins.py: 14 pins consistent. pre-commit clean. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

…ost requirements Three things this release got wrong that become immutable once 0.2.0 is on PyPI. The [quack] extra was missed by the cu12->cu13 sweep. pyproject.toml's `cuda` and `dali` extras were retargeted at CUDA 13, but `quack = ["quack-kernels"]` still resolved nvidia-cutlass-dsl's CUDA 12 runtime libraries — libs-cu12 is an unconditional dependency, so only the [cu13] extra adds libs-cu13. Now `quack-kernels[cu13]`, with a note to install via pip rather than uv (quack's README documents a uv ordering race on that path, NVIDIA/cutlass#3259). nvsubquadratic/modules/mlp.py imports quack.linear for fused GEMM+activation, so this path is more load-bearing than the RMSNorm-only framing suggested. Upgrading a 0.1.1 GPU environment silently corrupts it. cu12 and cu13 both install a top-level subquadratic_ops_torch/ package, and both DALI builds install nvidia/, so `pip install -U` overwrites files without removing the old distribution and a later uninstall of either breaks the other. Adds a Notes block with the uninstall-first sequence. Also records that a plain `pip install nvsubquadratic[cuda]` resolves from public PyPI unaided — torch 2.14.0 there is already a CUDA 13.0 build — so the cu130 index is only needed when compiling Apex/mamba against a specific torch. The NVIDIA driver >= 580 floor was documented only in a Dockerfile comment. A host on a 12.x-era driver fails at runtime with no documented cause. Added to the requirements lists in README.md, docs/getting_started.md and docs/index.rst. While in docs/index.rst, dropped the claim that [cuda] needs nvcc to build: subquadratic-ops-torch-cu13 0.3.0 ships prebuilt manylinux wheels for x86_64 and aarch64, so nothing is compiled at install time. Not touched: the "3-4x faster" claims in README/CHANGELOG/docs/ckconv_nd, which have no reproducing benchmark in the tree. Tracked separately. pre-commit clean; 14 version pins consistent; sphinx -W --keep-going green. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

"Roughly 3-4x faster than torch_fft" appeared in four public places — README, CHANGELOG, docs/ops/README.md and the CKConvND docstring — with nothing in the tree to back it. benchmarks/ has no fused arm: bench_fftconv2d.py covers torch_fft and subq_ops only, and on farhadr/2d_bench the 2D sweep records one backend per resolution (subq_ops_fused at R=16/32/64, subq_ops at R>=128), so no same-resolution head-to-head exists there either. The numbers are not invented — they trace to ee5146e, which measured 3.6-3.9x over torch_fft and 1.2-2.4x over subq_ops on an H200 at B=8, hidden=768, forward+backward, bf16. That configuration existed only in a commit message no user will ever read, and "3-4x" alone reads as a property of the kernel rather than of one shape on one GPU. Replaces the rounded figure with the measured range plus its conditions at all four sites, and notes that speedups are shape- and hardware-dependent. Nothing here changes what ships; it makes a public performance claim checkable. An H100/B200 re-measurement is planned. If it disagrees, these four sites are where to correct it. pre-commit clean (mdformat idempotent over two passes on docs/ops/README.md); sphinx -W --keep-going green. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

…achable The comment added in 5a3a985 described [quack] as providing "fused RMSNorm + GEMM/activation kernels". The GEMM/activation half is not reachable: _validate_quack_backend (modules/mlp.py:285) raises NotImplementedError unconditionally, so MLP(backend="quack") always fails at __init__ and the quack.linear imports below it are dead code. Nothing in the tree passes backend="quack". The live consumers are RMSNorm and RMSNormChannelFirst, both of which gate on _cuda_supports_quack (SM90+) and fall back to pure PyTorch otherwise. This does not change the dependency — quack-kernels[cu13] is still the right target for a CUDA 13 release — only the comment, which ships in the sdist. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

The README, CHANGELOG, docs/ops/README.md and the CKConvND docstring all cite 3.6-3.9x over torch_fft and 1.2-2.4x over subq_ops, but nothing in benchmarks/ produced those numbers: bench_fftconv2d.py compares torch_fft against subq_ops and never touches the fused kernel, and bench_subquadratic_fftconv.py is 1D causal. The figures traced to a commit message (ee5146e), not to runnable code. This script reproduces them. It defaults to the configuration the claim was measured at — B=8, hidden=768, bf16, forward+backward, median of 100 iters — and sweeps spatial 16/32/64 with the CKConvND double-grid kernel K=2N-1, which resolves to FFT tiles 32/64/128 respectively. The 128 tile means the spatial-64 row needs SM90+. Two traps it avoids, both worth knowing: - fftconv2d_fp32_bhl exists in BOTH ops/fftconv.py and ops/fftconv_chunked.py. Importing the wrong one silently benchmarks the chunked implementation. - subq_ops takes the kernel without the leading batch-1 dim ([H, Kx, Ky]), while torch_fft and the fused wrapper take [1|B, H, Kx, Ky]. Passing the same tensor to all three either throws or measures the wrong shape. Not in the wheel: [tool.setuptools.packages.find] includes only nvsubquadratic*, so this cannot affect the published artifact. Has not been run on SM90+ yet — that is the point of adding it. If the measured numbers disagree with the citations, the four sites to correct are README.md:70, CHANGELOG.md:17, docs/ops/README.md:127 and modules/ckconv_nd.py:621. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

The script called subquadratic_ops_torch.fft_conv2d directly, which raises NotImplementedError: fft_conv2d supports float32 only on bf16/fp16 input. Since the benchmark's whole subject is native-dtype execution, the "vs subq_ops" column could never be produced — it failed at every spatial size on the first H100 run and printed n/a throughout. Uses nvsubquadratic.ops.fftconv_custom.fftconv2d_bhl instead, the repo's documented drop-in for fftconv2d_fp32_bhl. It performs the fp32 upcast/downcast that the raw op requires, which is also the honest comparison: subq_ops really does pay an upcast that the fused kernel does not, and that cost belongs in the measurement rather than being defined away. It also takes the same [1|B, H, Kx, Ky] kernel as torch_fft and the fused wrapper, so the special-cased squeeze for subq_ops is gone along with the docstring note warning about it. Found by the first H100 run of this script. The torch_fft column was unaffected and is reported separately. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

The cited 3.6-3.9x over torch_fft does not hold. Measured on an H100 (B=8, hidden=768, bf16, fwd+bwd, median of 100, via the now-fixed benchmarks/ops/bench_fused_fftconv2d.py): spatial tile torch_fft subq_ops fused vs torch vs subq 16 32 0.620ms 0.466ms 0.329ms 1.88x 1.42x 32 64 2.124ms 1.214ms 0.482ms 4.41x 2.52x 64 128 8.266ms 2.248ms 1.689ms 4.89x 1.33x Reproduced across two runs within 1-2%, so the spread is real rather than measurement noise. vs torch_fft is 1.9-4.9x, not 3.6-3.9x. The direction that matters is spatial=16: 1.88x is less than half the previously advertised floor, so the old text overclaimed for small extents — the case where the fused kernel is weakest, since its advantage grows with the FFT tile. The 3.6-3.9x figure was likely a single shape or a narrow sweep, which is why it read as a tight range. vs subq_ops is 1.3-2.5x, close enough to the previous 1.2-2.4x, but it does NOT vary monotonically with extent (1.42x, 2.52x, 1.33x), so it cannot be summarised as "larger is better" the way torch_fft can. All four sites now quote the range with its conditions and point at the script, so the claim is checkable rather than inherited. The H200 attribution is dropped: it is not reproducible from this tree, whereas these numbers are. The benchmark records the reference run in its docstring so a future regression is visible as a diff against a committed baseline. Measured on H100 80GB, torch 2.14.0+cu130, subquadratic-ops-torch-cu13 0.3.0. pre-commit clean (mdformat stable over two passes on the new CHANGELOG table); sphinx -W --keep-going green. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Dafidofff

LGTM!

The 0.2.0 section was written when #139 was the only thing in the release and dated 2026-09-08. Since then #140 and #141 merged to main and came into this branch via e505327, so both ship inside this tag and belong in the notes. - Dates the section 2026-09-16. - Adds #141: PatchMerging / PatchMerging2D / PatchEmbedHierarchical, the ViT5HierarchicalNet + StageSpec / ViT5HierarchicalClassificationNet pair built on them, the CIFAR-10 and ImageNet recipes, and the two new CIFAR-10 datamodules (torchvision and HuggingFace). - Adds #140: SpatialRecall3DMotionDataset / SpatialRecall3DMotionDataModule and the motion_3d example configs. Every class named above was checked to exist at this commit rather than transcribed from the PR descriptions. pre-commit clean, mdformat idempotent over two passes. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

The first version was written against the PR #139 branch and committed to a branch off main, so a third of it described code that does not exist here. Verified: resolve_fused_fft_size, requires_sm90, assert_l2_close, L2_TOL, torch.fx/eliminate_dead_code and scripts/check_version_pins.py each return zero hits in the tree — their only occurrence was .coderabbit.yaml citing itself. It also instructed on *.rs files, of which the repo has none. Two of those were worse than stale. tests/conftest.py on main defines only _subq_ops_version, requires_subq_ops_v2, device and dtype_fixture, so telling a reviewer to require requires_sm90 and assert_l2_close imports was advice that breaks test collection if followed. Those belong in #139, which introduces them. Also cut the SPDX rule as pure CI redundancy: scripts/license_check.py inserts the header itself and lint.yml runs pre-commit --all-files, so a missing header can never reach a human reviewer. What replaces them are invariants a general reviewer cannot infer from a diff: - COMPILE_COMPATIBLE and the chunking flags are flipped at runtime after import, so `from ... import` snapshots them and the flip silently no-ops. - fftconv_chunked.py shadows fftconv.py's public names deliberately; the defect is a new op with no matching shim, not the shadowing. - Transform length and crop offset are one contract — change one and the output is merely shifted, which a symmetric probe kernel will not catch. - OmegaConf returns ListConfig/DictConfig, neither of which subclasses list/dict, so isinstance checks on config-supplied values silently take the wrong branch. - Buffers derived from config use persistent=False so old checkpoints survive a resolution change. - QKVSequenceMixer passes cp_group positionally; a mixer reading it from **kwargs gets None and runs without cross-rank communication. - Only torch/einops/omegaconf/numpy may be imported at module scope; extras need a lazy getter and an autodoc_mock_imports entry. Every symbol and path cited above was checked to exist at this commit. Reviewed by four independent critics (over-fitting, CI redundancy, staleness, coverage gaps) and an adversarial verifier that re-derived each claim. 8 entries, 10.9k chars, well under the 20k per-entry limit. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Farhad Ramezanghorbani <farhadr@nvidia.com>

NVIDIA runs CodeRabbit self-hosted (coderabbit.nvidia.com), approved for EC 3e991, with code disposed after each review. The app is installed through https://github-onboarding.nvidia.com/install-github-app rather than the public GitHub Marketplace; this file only configures it once installed. Both top-level settings match the in-org precedent, NVIDIA-BioNeMo/bionemo-recipes, whose entire config is the same four lines: profile: chill - balanced rather than assertive, which reads as nitpicky and is the fastest way to get an AI reviewer muted. auto_review: false - on request only, via `@coderabbitai review` or `@coderabbitai full review`. Keeps the tool opt-in while the team forms a view of its signal-to-noise. Neither costs CI time: CodeRabbit is a hosted app reacting to webhooks, not a GitHub Action, so it never triggers the ~36-minute self-hosted gpu-tests job. The path_instructions are the part worth reviewing. Each encodes a defect that actually occurred in this repo rather than generic advice: - the worker_init_fn DDP-seeding trap, which silently collapses augmentation diversity to 1/world_size (PR #140, still open) - config.net needing to be a LazyConfig, and iterations-per-epoch needing the effective batch rather than the micro-batch (PR #141) - torch.fx passes calling eliminate_dead_code without an is_impure_node override, which can delete unrelated in-place ops (PR #139) - the two colliding fftconv2d_fp32_bhl definitions in ops/fftconv.py and ops/fftconv_chunked.py - the raw subquadratic_ops_torch.fft_conv2d op being fp32-only - SM90+ gating for fused spatial extents above 32 per axis - the requires_sm90 marker convention and the shared L2 tolerance helpers - display math nested under a bullet in docs/ops/, which makes mdformat non-idempotent and loops pre-commit forever - new public classes needing both a docs-tracker.md row and a Sphinx entry Ordering note: the pin-consistency instruction names scripts/check_version_pins.py, which lands with #139 and is not on main yet. The invariant it describes holds regardless; only the enforcing script is pending. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Farhad Ramezanghorbani <farhadr@nvidia.com>

The first version was written against the PR #139 branch and committed to a branch off main, so a third of it described code that does not exist here. Verified: resolve_fused_fft_size, requires_sm90, assert_l2_close, L2_TOL, torch.fx/eliminate_dead_code and scripts/check_version_pins.py each return zero hits in the tree — their only occurrence was .coderabbit.yaml citing itself. It also instructed on *.rs files, of which the repo has none. Two of those were worse than stale. tests/conftest.py on main defines only _subq_ops_version, requires_subq_ops_v2, device and dtype_fixture, so telling a reviewer to require requires_sm90 and assert_l2_close imports was advice that breaks test collection if followed. Those belong in #139, which introduces them. Also cut the SPDX rule as pure CI redundancy: scripts/license_check.py inserts the header itself and lint.yml runs pre-commit --all-files, so a missing header can never reach a human reviewer. What replaces them are invariants a general reviewer cannot infer from a diff: - COMPILE_COMPATIBLE and the chunking flags are flipped at runtime after import, so `from ... import` snapshots them and the flip silently no-ops. - fftconv_chunked.py shadows fftconv.py's public names deliberately; the defect is a new op with no matching shim, not the shadowing. - Transform length and crop offset are one contract — change one and the output is merely shifted, which a symmetric probe kernel will not catch. - OmegaConf returns ListConfig/DictConfig, neither of which subclasses list/dict, so isinstance checks on config-supplied values silently take the wrong branch. - Buffers derived from config use persistent=False so old checkpoints survive a resolution change. - QKVSequenceMixer passes cp_group positionally; a mixer reading it from **kwargs gets None and runs without cross-rank communication. - Only torch/einops/omegaconf/numpy may be imported at module scope; extras need a lazy getter and an autodoc_mock_imports entry. Every symbol and path cited above was checked to exist at this commit. Reviewed by four independent critics (over-fitting, CI redundancy, staleness, coverage gaps) and an adversarial verifier that re-derived each claim. 8 entries, 10.9k chars, well under the 20k per-entry limit. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Farhad Ramezanghorbani <farhadr@nvidia.com>

main now contains the curated extraction of this branch's library work, shipped as v0.2.0 via PR #139, plus #140 and #141. Every file both sides touched was resolved to main's version, because main is authoritative for all of it: - fftconv_lowering.py: this branch carried the original 533-line pass. main has the 631-line revision with _erase_chain, which fixes two correctness bugs — an unguarded whole-graph eliminate_dead_code that could delete unrelated in-place ops, and an any() short-circuit that fused only the first chain per graph. Keeping this branch's copy would reintroduce both. - pyproject.toml: torch >=2.12 here vs >=2.14 on main. The 2.14 floor is not a preference — torch pins nvidia-cudnn-cu13 exactly, and subquadratic-ops-torch-cu13 >=0.3.0 needs >=9.24.0.43, which only 2.14 satisfies. 2.12 makes `pip install nvsubquadratic[cuda]` unresolvable. - ckconv_nd.py, fftconv_custom.py, tests/conftest.py, the Dockerfile and enroot scripts, and the docs: same work, later revision on main. 21 files resolved this way. What this branch uniquely owns is untouched — the ND benchmark harness, the sweep JSONLs and results digest, the visualization and submit scripts, docs/mamba2_limits.md, and the FA4/FlexAttention additions to attention.py and sequence_mixer.py. Merged rather than rebased: 37 commits against 21 overlapping files would replay the same conflicts repeatedly, and Alireza has commits here, so rewriting the branch's history has a cost a merge does not. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Farhad Ramezanghorbani <farhadr@nvidia.com>

…ess (#142) * ci: add CodeRabbit configuration NVIDIA runs CodeRabbit self-hosted (coderabbit.nvidia.com), approved for EC 3e991, with code disposed after each review. The app is installed through https://github-onboarding.nvidia.com/install-github-app rather than the public GitHub Marketplace; this file only configures it once installed. Both top-level settings match the in-org precedent, NVIDIA-BioNeMo/bionemo-recipes, whose entire config is the same four lines: profile: chill - balanced rather than assertive, which reads as nitpicky and is the fastest way to get an AI reviewer muted. auto_review: false - on request only, via `@coderabbitai review` or `@coderabbitai full review`. Keeps the tool opt-in while the team forms a view of its signal-to-noise. Neither costs CI time: CodeRabbit is a hosted app reacting to webhooks, not a GitHub Action, so it never triggers the ~36-minute self-hosted gpu-tests job. The path_instructions are the part worth reviewing. Each encodes a defect that actually occurred in this repo rather than generic advice: - the worker_init_fn DDP-seeding trap, which silently collapses augmentation diversity to 1/world_size (PR #140, still open) - config.net needing to be a LazyConfig, and iterations-per-epoch needing the effective batch rather than the micro-batch (PR #141) - torch.fx passes calling eliminate_dead_code without an is_impure_node override, which can delete unrelated in-place ops (PR #139) - the two colliding fftconv2d_fp32_bhl definitions in ops/fftconv.py and ops/fftconv_chunked.py - the raw subquadratic_ops_torch.fft_conv2d op being fp32-only - SM90+ gating for fused spatial extents above 32 per axis - the requires_sm90 marker convention and the shared L2 tolerance helpers - display math nested under a bullet in docs/ops/, which makes mdformat non-idempotent and loops pre-commit forever - new public classes needing both a docs-tracker.md row and a Sphinx entry Ordering note: the pin-consistency instruction names scripts/check_version_pins.py, which lands with #139 and is not on main yet. The invariant it describes holds regardless; only the enforcing script is pending. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Farhad Ramezanghorbani <farhadr@nvidia.com> * docs: document the CodeRabbit review request process in the PR template auto_review is disabled, so CodeRabbit does nothing unless asked. Without this the config lands and nobody uses it — the trigger comments are invisible to anyone who has not read .coderabbit.yaml. Adds a Review section with two checkboxes: request the AI first pass, then request a codeowner. That order is the point — a codeowner's time is better spent on design than on what a bot would have caught. Records two ruleset facts that are easy to learn the hard way: - A CodeRabbit review does not count toward the required approving review. main requires a codeowner approval and bypass_actors is empty, so there is no override. - required_review_thread_resolution is true, so unresolved CodeRabbit threads block merge. They have to be resolved or replied to, not ignored. Also lists the other trigger commands, and points at .coderabbit.yaml as the place to correct systematic bad advice rather than repeating the correction on every PR. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Farhad Ramezanghorbani <farhadr@nvidia.com> * ci: rewrite the CodeRabbit path instructions against main The first version was written against the PR #139 branch and committed to a branch off main, so a third of it described code that does not exist here. Verified: resolve_fused_fft_size, requires_sm90, assert_l2_close, L2_TOL, torch.fx/eliminate_dead_code and scripts/check_version_pins.py each return zero hits in the tree — their only occurrence was .coderabbit.yaml citing itself. It also instructed on *.rs files, of which the repo has none. Two of those were worse than stale. tests/conftest.py on main defines only _subq_ops_version, requires_subq_ops_v2, device and dtype_fixture, so telling a reviewer to require requires_sm90 and assert_l2_close imports was advice that breaks test collection if followed. Those belong in #139, which introduces them. Also cut the SPDX rule as pure CI redundancy: scripts/license_check.py inserts the header itself and lint.yml runs pre-commit --all-files, so a missing header can never reach a human reviewer. What replaces them are invariants a general reviewer cannot infer from a diff: - COMPILE_COMPATIBLE and the chunking flags are flipped at runtime after import, so `from ... import` snapshots them and the flip silently no-ops. - fftconv_chunked.py shadows fftconv.py's public names deliberately; the defect is a new op with no matching shim, not the shadowing. - Transform length and crop offset are one contract — change one and the output is merely shifted, which a symmetric probe kernel will not catch. - OmegaConf returns ListConfig/DictConfig, neither of which subclasses list/dict, so isinstance checks on config-supplied values silently take the wrong branch. - Buffers derived from config use persistent=False so old checkpoints survive a resolution change. - QKVSequenceMixer passes cp_group positionally; a mixer reading it from **kwargs gets None and runs without cross-rank communication. - Only torch/einops/omegaconf/numpy may be imported at module scope; extras need a lazy getter and an autodoc_mock_imports entry. Every symbol and path cited above was checked to exist at this commit. Reviewed by four independent critics (over-fitting, CI redundancy, staleness, coverage gaps) and an adversarial verifier that re-derived each claim. 8 entries, 10.9k chars, well under the 20k per-entry limit. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Farhad Ramezanghorbani <farhadr@nvidia.com> --------- Signed-off-by: Farhad Ramezanghorbani <farhadr@nvidia.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

farhadrgh requested a review from saeepaliwal as a code owner September 2, 2026 18:50

farhadrgh and others added 2 commits September 2, 2026 12:43

farhadrgh requested review from Dafidofff, dwromero and moradza September 3, 2026 16:28

farhadrgh commented Sep 3, 2026

View reviewed changes

Comment thread nvsubquadratic/ops/fftconv_lowering.py Outdated

farhadrgh commented Sep 3, 2026

View reviewed changes

Comment thread nvsubquadratic/ops/fftconv_lowering.py Outdated

farhadrgh and others added 8 commits September 8, 2026 11:36

Dafidofff approved these changes Sep 10, 2026

View reviewed changes

saeepaliwal approved these changes Sep 16, 2026

View reviewed changes

moradza approved these changes Sep 16, 2026

View reviewed changes

farhadrgh and others added 2 commits September 16, 2026 11:34

Merge branch 'main' into farhadr/cuda13-fused-fftconv

e505327

farhadrgh mentioned this pull request Sep 16, 2026

ci: add CodeRabbit configuration and document the review request process #142

Merged

4 tasks

farhadrgh merged commit e83b821 into main Sep 16, 2026
11 checks passed

farhadrgh deleted the farhadr/cuda13-fused-fftconv branch September 16, 2026 18:35

Navigation Menu

feat(ops): fused FFT-conv2d backend, torch.compile lowering, and CUDA 13 migration - #139

feat(ops): fused FFT-conv2d backend, torch.compile lowering, and CUDA 13 migration#139
farhadrgh merged 14 commits into
mainfrom
farhadr/cuda13-fused-fftconv

farhadrgh commented Sep 2, 2026 •

edited

Loading

Uh oh!

farhadrgh commented Sep 2, 2026

Uh oh!

farhadrgh commented Sep 2, 2026

Uh oh!

Uh oh!

Uh oh!

moradza commented Sep 4, 2026

Uh oh!

Dafidofff left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

farhadrgh commented Sep 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What's here

The pin change, and why it fixes CI

Version consistency

Relationship to other PRs

Testing

Uh oh!

farhadrgh commented Sep 2, 2026

Correction: the >=0.3.0 pin does not fix gpu-tests — it introduces a hard resolution conflict

Validated fix

Uh oh!

farhadrgh commented Sep 2, 2026

The torch 2.14 bump worked — the image builds and the suite runs

What the 5 failures actually mean

Caveat for reviewers

Uh oh!

Uh oh!

Uh oh!

moradza commented Sep 4, 2026

Uh oh!

Dafidofff left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

farhadrgh commented Sep 2, 2026 •

edited

Loading

Correction: the `>=0.3.0` pin does not fix `gpu-tests` — it introduces a hard resolution conflict