Reintroduce FFT causal conv1d frontend bindings - #542
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review. 📝 WalkthroughWalkthroughThe PR adds medium- and long-FFT causal Conv1d support. It adds backend shims, Python bindings, validation, autograd integration, tests, samples, and documentation. Existing causal Conv1d variants gain FP64 and dtype-specific gradient accumulation. ChangesCausal Conv1d Features
Estimated code review effort: 4 (Complex) | ~60 minutes Merge Risk: ⚪ Minimal · up to The PR adds the FFT causal convolution frontend and documents the supported cuDNN requirements and validation; no actionable merge-blocking risk remains beyond normal checks and review. Sequence Diagram(s)sequenceDiagram
participant PythonAPI
participant FFTCausalConv1d
participant cuDNNBindings
participant cuDNNBackend
PythonAPI->>FFTCausalConv1d: Validate tensors and select medium or long path
FFTCausalConv1d->>cuDNNBindings: Invoke forward or backward operation
cuDNNBindings->>cuDNNBackend: Dispatch cuDNN 9.26+ operation
cuDNNBackend-->>cuDNNBindings: Return output, gradients, or buffer sizes
cuDNNBindings-->>FFTCausalConv1d: Return backend result
FFTCausalConv1d-->>PythonAPI: Trim padding and return result
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (9)
test/python/test_fft_causal_conv1d.py (3)
11-20: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winGate the tests on CUDA availability and on device capability.
_require_fft_causal_conv1dchecks the backend version and the binding symbols. It does not checktorch.cuda.is_available()._make_inputsthen creates tensors withdevice="cuda"at Line 29, so a CPU-only runner raises aRuntimeErrorinstead of skipping.The FFT path selection in
python/cudnn/ops/fft_causal_conv1d.pyalso depends on the architecture through_cuda_arch. Add a capability gate so unsupported architectures skip instead of failing.As per coding guidelines: "Gate tests on supported capabilities and skip unsupported architecture, dtype, or backend-version combinations using support checks,
cudnn.backend_version(), andtorch.cuda.get_device_capability()."♻️ Proposed change
def _require_fft_causal_conv1d(): + if not torch.cuda.is_available(): + pytest.skip("FFT causal conv1d requires a CUDA device") required_symbols = (🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@test/python/test_fft_causal_conv1d.py` around lines 11 - 20, Update _require_fft_causal_conv1d to skip unless CUDA is available, the existing cuDNN version and symbol checks pass, and torch.cuda.get_device_capability() is supported by the FFT path’s _cuda_arch requirements. Keep the existing skip behavior and use the capability check before _make_inputs creates CUDA tensors.Source: Coding guidelines
94-99: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExtend coverage for the long path, the buffer formula, and
torch.compile.Three gaps remain in this file:
- The long FFT path runs for FP64 only, at Line 40. FP16, BF16, and FP32 never reach
_long_fwd_primitiveor_long_bwd_primitive. This gap hides thegrad_weightinitialization question raised onpython/cudnn/ops/fft_causal_conv1d.pyLine 268.- This test compares the mirrored buffer formula for
CUDNN_DATA_FLOATonly. Parameterize it over every supported dtype, so a wrong scratch-element assumption atpython/cudnn/ops/fft_causal_conv1d.pyLine 117 fails loudly.- No test compiles the operation. The PR claims
torch.compilesupport, andregister_fakeplus@torch.compiler.allow_in_graphexist for that purpose. Add one test that wrapsfft_causal_conv1dwithtorch.compileand compares against the eager result.Place the dtype sweeps at a level above
L0to keepL0fast.As per coding guidelines: "Mark every new Python test with a level from
L0throughL4; keepL0tests fast and place large parameter sweeps at higher levels."🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@test/python/test_fft_causal_conv1d.py` around lines 94 - 99, Extend test coverage in test_long_fft_fake_buffer_formula_matches_backend_query and the existing long-path tests: parameterize supported FP16, BF16, FP32, and FP64 dtypes so both long forward/backward primitives and grad_weight initialization are exercised, and move these broad sweeps above L0. Parameterize the buffer-size assertion across every supported dtype while retaining the matching dtype-to-backend mapping. Add a higher-level marked test that wraps fft_causal_conv1d with torch.compile and compares its output with eager execution.Source: Coding guidelines
28-31: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMake the numerical checks reproducible and use gradient-specific tolerances.
test/python/conftest.pydoes not seed the RNG. Seed a CUDA generator before creatingx,weight, andgrad_out, or usetorch_fork_set_rng.Use a separate tolerance for
weight.grad. Its convolution gradient reduces over batch and sequence dimensions and can have a different error scale from the output andx.grad.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@test/python/test_fft_causal_conv1d.py` around lines 28 - 31, Update _make_inputs and the test’s grad_out setup to use a reproducible CUDA RNG, either by seeding a CUDA generator before creating x, weight, and grad_out or by wrapping generation with torch_fork_set_rng. Add a separate tolerance for weight.grad comparisons, reflecting its distinct reduction error scale, while retaining the existing tolerance for the output and x.grad checks.python/cudnn/ops/fft_causal_conv1d.py (4)
7-24: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
_dtype_to_intduplicates the helper incausal_conv1d.py.
python/cudnn/ops/causal_conv1d.pyLine 22 defines a function with the same name and purpose. Move the dtype map and the conversion into a shared private module underpython/cudnn/ops/, and let each operation declare its own supported subset. This keeps one source for the cuDNN enum values.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/ops/fft_causal_conv1d.py` around lines 7 - 24, Move the shared cuDNN dtype-to-enum map and conversion logic from _dtype_to_int in fft_causal_conv1d.py and its duplicate in causal_conv1d.py into a private module under python/cudnn/ops/. Update both operations to reuse that shared helper while retaining each operation’s own supported-dtype subset and validation behavior.
165-167: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winConfirm that the backend fully overwrites
dx.Line 167 uses
zeros_likeforgrad_weightbecause the kernel accumulates. Line 165 usesempty_likeforgrad_x, which assumes the kernel writes every element ofdx. Confirm that assumption for both the medium and long kernels. If either kernel accumulates intodx, the gradient contains uninitialized memory.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/ops/fft_causal_conv1d.py` around lines 165 - 167, Inspect the medium and long cuhyena kernel implementations invoked by the surrounding backward path and verify whether each fully overwrites every element of grad_x. If either kernel accumulates into dx or leaves any elements untouched, initialize grad_x with zeros_like(x) before dispatch; otherwise retain empty_like only when both kernels guarantee complete writes.
109-120: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winThe mirrored buffer formula is verified for one dtype only.
_long_buffer_size_bytesreimplements a cuDNN internal layout. Line 117 assumes a 4-byte scratch element for FP16 and BF16, and an 8-byte element for FP64. The only test that compares this formula against the backend query istest/python/test_fft_causal_conv1d.pyLine 98, which passes0(CUDNN_DATA_FLOAT).If the assumption is wrong for FP16, BF16, or FP64,
_long_fwd_fakereports a wrong reserve size. Undertorch.compile, downstream shape reasoning then uses that wrong size. Extend the comparison test to every supported dtype.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/ops/fft_causal_conv1d.py` around lines 109 - 120, Extend the backend-versus-formula buffer-size comparison in test_fft_causal_conv1d.py around the existing test to run for every supported dtype, including FP16, BF16, FP32, and FP64, rather than only CUDNN_DATA_FLOAT. Reuse the existing test structure and dtype-to-backend mapping while validating _long_buffer_size_bytes for each case.
99-102: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueThe
== _LONG_FFT_MAX_LENGTHbounds depend on the power-of-two precondition.Both checks compare
kernel_sizeagainst_LONG_FFT_MAX_LENGTHwith==. They are correct only because Line 95 forces a power of two and Line 97 caps the value at_LONG_FFT_MAX_LENGTH, which makes16777216the single value above8388608. If_LONG_FFT_MAX_LENGTHchanges, the FP64 check silently stops rejecting8388608 < kernel_size < _LONG_FFT_MAX_LENGTH.Also, the FP64 message states the supported bound but omits the rejected value.
♻️ Proposed change
- if dtype == torch.float64 and kernel_size == _LONG_FFT_MAX_LENGTH: - raise ValueError("Long FFT FP64 supports kernel_size through 8388608.") + if dtype == torch.float64 and kernel_size > _LONG_FFT_FP64_MAX_LENGTH: + raise ValueError(f"Long FFT FP64 supports kernel_size through {_LONG_FFT_FP64_MAX_LENGTH}; got {kernel_size}.") - if kernel_size == _LONG_FFT_MAX_LENGTH and _cuda_arch(device) < 900: - raise ValueError("Long FFT kernel_size 16777216 requires compute capability 9.0 or newer.") + if kernel_size >= _LONG_FFT_MAX_LENGTH and _cuda_arch(device) < 900: + raise ValueError(f"Long FFT kernel_size {_LONG_FFT_MAX_LENGTH} requires compute capability 9.0 or newer.")Add the new constant next to the existing bounds:
_LONG_FFT_FP64_MAX_LENGTH = 1 << 23🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/ops/fft_causal_conv1d.py` around lines 99 - 102, Replace the FP64 equality check in the FFT validation flow with an explicit upper-bound check using a new _LONG_FFT_FP64_MAX_LENGTH constant set to 1 << 23, while retaining the existing overall _LONG_FFT_MAX_LENGTH and CUDA-architecture checks. Update the FP64 error message to state the supported maximum and the rejected kernel size range/value clearly.python/cudnn/ops/__init__.py (1)
4-5: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd
__all__to declare the exported wrappers.The coding guidelines require that frontend kernel packages export their API class and wrapper through
__all__. This file re-exports four wrappers with no__all__. The gap predates this change, and the newfft_causal_conv1dexport extends it.As per coding guidelines: "Frontend kernel packages must export their API class and wrapper through
__all__."♻️ Proposed change
from .causal_conv1d import causal_conv1d, causal_conv1d_nwh, b2b_causal_conv1d from .fft_causal_conv1d import fft_causal_conv1d + +__all__ = [ + "causal_conv1d", + "causal_conv1d_nwh", + "b2b_causal_conv1d", + "fft_causal_conv1d", +]🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/ops/__init__.py` around lines 4 - 5, Define an __all__ declaration in the package initializer listing the four re-exported wrappers: causal_conv1d, causal_conv1d_nwh, b2b_causal_conv1d, and fft_causal_conv1d, matching the existing public imports.Source: Coding guidelines
python/pycudnn.cpp (1)
414-415: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueCompare the status against
CUDNN_STATUS_SUCCESS.The five new bindings use
status != 0. The existing helper at Line 62 compares againstCUDNN_STATUS_SUCCESS. Use the enum constant so the new code matches the established pattern in this file.♻️ Proposed change, shown for the forward binding
- if (status != 0) + if (status != CUDNN_STATUS_SUCCESS) throw std::runtime_error("cudnnFFTCausalConv1dForward failed with status " + std::to_string(status));Also applies to: 441-442, 456-458, 489-491, 521-523
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/pycudnn.cpp` around lines 414 - 415, Update the status checks in all five new cuDNN bindings, including the forward binding and the locations corresponding to the other reported ranges, to compare against CUDNN_STATUS_SUCCESS instead of the literal 0. Preserve the existing runtime_error messages and failure handling.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python/cudnn/__init__.py`:
- Around line 59-63: Update python/cudnn/ops/__init__.py to register both
Torch-dependent convolution modules or APIs in _LAZY_OPTIONAL_IMPORTS and define
the module’s __all__ entries for the exposed convolution symbols. Ensure
importing cudnn.ops stays Torch-free and missing Torch raises the established
install-hint ImportError when those APIs are accessed.
In `@python/cudnn/ops/fft_causal_conv1d.py`:
- Around line 265-268: Confirm the gradient-buffer contract for
cudnnLongFFTCausalConv1dBackward at
python/cudnn/ops/fft_causal_conv1d.py:265-268 and update grad_weight to
zeros_like(weight) if the kernel accumulates, adding the explanatory comment
used at the medium path; also confirm whether both backward kernels fully write
dx, and change grad_x to zeros_like(x) wherever accumulation occurs. Record the
confirmed contract beside each allocation, including the sibling site at
python/cudnn/ops/fft_causal_conv1d.py:165-167, so both paths remain consistent.
In `@samples/cpp/causal_conv1d/fft_causal_conv1d.cpp`:
- Around line 53-99: Update the FFT causal convolution test cases, including
“FFT causal conv1d medium forward and backward,” to initialize deterministic
inputs and verify y_tensor, dx_tensor, and dweight_tensor against a reference
calculation or expected values after execution. Follow the required sample
sequence by validating and building the operation graph, creating plans,
checking support with a graceful skip when unsupported, then executing and
asserting results for both forward and backward paths.
---
Nitpick comments:
In `@python/cudnn/ops/__init__.py`:
- Around line 4-5: Define an __all__ declaration in the package initializer
listing the four re-exported wrappers: causal_conv1d, causal_conv1d_nwh,
b2b_causal_conv1d, and fft_causal_conv1d, matching the existing public imports.
In `@python/cudnn/ops/fft_causal_conv1d.py`:
- Around line 7-24: Move the shared cuDNN dtype-to-enum map and conversion logic
from _dtype_to_int in fft_causal_conv1d.py and its duplicate in causal_conv1d.py
into a private module under python/cudnn/ops/. Update both operations to reuse
that shared helper while retaining each operation’s own supported-dtype subset
and validation behavior.
- Around line 165-167: Inspect the medium and long cuhyena kernel
implementations invoked by the surrounding backward path and verify whether each
fully overwrites every element of grad_x. If either kernel accumulates into dx
or leaves any elements untouched, initialize grad_x with zeros_like(x) before
dispatch; otherwise retain empty_like only when both kernels guarantee complete
writes.
- Around line 109-120: Extend the backend-versus-formula buffer-size comparison
in test_fft_causal_conv1d.py around the existing test to run for every supported
dtype, including FP16, BF16, FP32, and FP64, rather than only CUDNN_DATA_FLOAT.
Reuse the existing test structure and dtype-to-backend mapping while validating
_long_buffer_size_bytes for each case.
- Around line 99-102: Replace the FP64 equality check in the FFT validation flow
with an explicit upper-bound check using a new _LONG_FFT_FP64_MAX_LENGTH
constant set to 1 << 23, while retaining the existing overall
_LONG_FFT_MAX_LENGTH and CUDA-architecture checks. Update the FP64 error message
to state the supported maximum and the rejected kernel size range/value clearly.
In `@python/pycudnn.cpp`:
- Around line 414-415: Update the status checks in all five new cuDNN bindings,
including the forward binding and the locations corresponding to the other
reported ranges, to compare against CUDNN_STATUS_SUCCESS instead of the literal
0. Preserve the existing runtime_error messages and failure handling.
In `@test/python/test_fft_causal_conv1d.py`:
- Around line 11-20: Update _require_fft_causal_conv1d to skip unless CUDA is
available, the existing cuDNN version and symbol checks pass, and
torch.cuda.get_device_capability() is supported by the FFT path’s _cuda_arch
requirements. Keep the existing skip behavior and use the capability check
before _make_inputs creates CUDA tensors.
- Around line 94-99: Extend test coverage in
test_long_fft_fake_buffer_formula_matches_backend_query and the existing
long-path tests: parameterize supported FP16, BF16, FP32, and FP64 dtypes so
both long forward/backward primitives and grad_weight initialization are
exercised, and move these broad sweeps above L0. Parameterize the buffer-size
assertion across every supported dtype while retaining the matching
dtype-to-backend mapping. Add a higher-level marked test that wraps
fft_causal_conv1d with torch.compile and compares its output with eager
execution.
- Around line 28-31: Update _make_inputs and the test’s grad_out setup to use a
reproducible CUDA RNG, either by seeding a CUDA generator before creating x,
weight, and grad_out or by wrapping generation with torch_fork_set_rng. Add a
separate tolerance for weight.grad comparisons, reflecting its distinct
reduction error scale, while retaining the existing tolerance for the output and
x.grad checks.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 9774b3f2-db7d-4143-ada9-9b2c528c296e
📒 Files selected for processing (15)
docs/operations/CausalConv1d.mddocs/operations/FFTCausalConv1d.mdinclude/cudnn_frontend_shim.hllms.txtpython/cudnn/__init__.pypython/cudnn/ops/__init__.pypython/cudnn/ops/causal_conv1d.pypython/cudnn/ops/fft_causal_conv1d.pypython/pycudnn.cppsamples/cpp/CMakeLists.txtsamples/cpp/causal_conv1d/fft_causal_conv1d.cppsamples/python/66_fft_causal_conv1d_forward.ipynbsamples/python/67_fft_causal_conv1d_backward.ipynbtest/python/test_causal_conv1d.pytest/python/test_fft_causal_conv1d.py
15daaaa to
926794c
Compare
|
@cudnn-ci-bot run python_tests |
Anerudhan
added
mod-backend
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-542-73b6beb |
3 participants
Before submitting
pre-commit runand committed any formatting changes.Affected area
Summary
Reintroduces the FFT causal conv1d frontend support originally merged in #437 and reverted in #479:
cudnn.ops.fft_causal_conv1d(x, weight), following cuhyena's medium/long selection, padding, trimming, dtype support, and autograd behavior.Why
The cuDNN 9.26 release branch now contains the backend declarations and implementations required by this frontend API. That resolves the compatibility issue that motivated #479, so the original feature can again be built and run with its documented cuDNN 9.26 minimum version.
Related issues
Related to #437 and #479.
API and compatibility impact
Adds the public Python API:
It also exposes the low-level long-path buffer query:
torch.compile.weight[0]multiplies the current sample. This is reversed relative to the existing direct causal conv1d wrapper.CUDNN_STATUS_NOT_SUPPORTED.Testing
Validated locally on a compute-capability 10.0 GPU with CUDA 13.4.25 and cuDNN 9.26.0.27 from the merged 9.26 backend pipeline. The rebuilt Python module reported backend version
92600and loadedlibcudnn.so.9.26.0from that artifact./path/to/python -m pre_commit run --from-ref origin/develop --to-ref HEAD: passed (clang-format,black, andblack-jupyter).cmake -S . -B /tmp/cudnn-frontend-437-cpp -DCUDNN_INCLUDE_PATH=/tmp/cudnn-mr4168-9.26.0.27/cudnn/include -DCUDNN_LIBRARY_PATH=/tmp/cudnn-mr4168-9.26.0.27/cudnn/lib -DCUDAToolkit_ROOT=/path/to/cuda-13.4: passed.cmake --build /tmp/cudnn-frontend-437-cpp --parallel 32 --target samples: passed; the complete C++ samples target compiled.LD_LIBRARY_PATH=<cuDNN-9.26-and-CUDA-libs> /tmp/cudnn-frontend-437-cpp/bin/samples '[fft_causal_conv1d]': 2 passed (medium and long forward/backward).cmake -S . -B /tmp/cudnn-frontend-437-pybind -DCUDNN_FRONTEND_BUILD_PYTHON_BINDINGS=ON -DCUDNN_FRONTEND_BUILD_SAMPLES=OFF -DCUDNN_FRONTEND_BUILD_TESTS=OFF -DCUDAToolkit_ROOT=/path/to/cuda-13.4 -DCUDNN_PATH=/tmp/cudnn-mr4168-9.26.0.27/cudnn -DCUDNN_LIBRARY_PATH=/tmp/cudnn-mr4168-9.26.0.27/cudnn/lib: passed.cmake --build /tmp/cudnn-frontend-437-pybind --parallel 32: passed.PYTHONPATH=/tmp/cudnn-frontend-437-pybind python -m pytest -q -ra test/python/test_fft_causal_conv1d.py: 7 passed in 110.93s.python -m jupyter nbconvert --to notebook --execute samples/python/66_fft_causal_conv1d_forward.ipynb --ExecutePreprocessor.kernel_name=cudnn-fe-437 --ExecutePreprocessor.timeout=1200: passed with zero cell errors; medium FP32, long FP64, andtorch.compilechecks passed.python -m jupyter nbconvert --to notebook --execute samples/python/67_fft_causal_conv1d_backward.ipynb --ExecutePreprocessor.kernel_name=cudnn-fe-437 --ExecutePreprocessor.timeout=1200: passed with zero cell errors; medium and long backward/reference checks passed.Summary by CodeRabbit
New Features
Documentation
Tests