Add FFT causal conv1d frontend bindings - #437
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 (4)
🚧 Files skipped from review as they are similar to previous changes (4)
📝 WalkthroughWalkthroughAdds ChangesFFT Causal Conv1d
Causal Conv1d FP64 support
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant fft_causal_conv1d
participant CustomCUDAOp
participant cuDNNBackend
Caller->>fft_causal_conv1d: Provide x and weight
fft_causal_conv1d->>CustomCUDAOp: Select medium or long path
CustomCUDAOp->>cuDNNBackend: Run FFT causal convolution
cuDNNBackend-->>CustomCUDAOp: Return output and reserve state
CustomCUDAOp-->>Caller: Return trimmed output
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: 4
🤖 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 `@docs/operations/FFTCausalConv1d.md`:
- Around line 15-25: Correct the operation equation in the FFTCausalConv1d
documentation to include batch and channel indices, indexing weights as
weight[c, j] for tensor shape (dim, kernel_size). State the elementwise contract
for y[b, c, t] and explicitly define x[b, c, t - j] as zero when the index is
out of range.
In `@python/cudnn/ops/fft_causal_conv1d.py`:
- Around line 305-310: Update _long_autograd_bwd to accept both gradient
arguments produced by _long_fwd_primitive, including the gradient for
reserve_space, while continuing to use the output gradient for the backward
primitive. Preserve the existing returned gradients for x and weight.
In `@samples/cpp/causal_conv1d/fft_causal_conv1d.cpp`:
- Around line 75-97: Add numerical verification after execution in the sample’s
forward and backward paths, covering y_tensor, dx_tensor, and dweight_tensor
against a causal-convolution reference with dtype-appropriate tolerances. Keep
the existing CUDNN_CHECK calls and synchronization, and apply the same
execute-then-verify sequence to the code around the additional indicated
section.
In `@test/python/test_fft_causal_conv1d.py`:
- Around line 11-13: Update _require_fft_causal_conv1d to also verify that
cudnn.long_fft_causal_conv1d_get_buffer_sizes exists, and skip when either
required FFT symbol is unavailable while retaining the existing backend-version
check.
🪄 Autofix (Beta)
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: 63d73f12-c0d5-4b9f-ae72-a27e92cb6362
📒 Files selected for processing (12)
docs/operations/FFTCausalConv1d.mdinclude/cudnn_frontend_shim.hllms.txtpython/cudnn/__init__.pypython/cudnn/ops/__init__.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_fft_causal_conv1d.py
| @torch.compiler.allow_in_graph | ||
| def _long_autograd_bwd(ctx, grad_out): | ||
| x, weight, reserve_space = ctx.saved_tensors | ||
| grad_y = grad_out[0] | ||
| grad_x, grad_weight = torch.ops.cudnn.long_fft_causal_conv1d_bwd_primitive(grad_y, x, weight, reserve_space) | ||
| return grad_x, grad_weight |
There was a problem hiding this comment.
🎯 Functional Correctness | 🔴 Critical | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the relevant sections and nearby tests/helpers.
sed -n '220,330p' python/cudnn/ops/fft_causal_conv1d.py
printf '\n--- tests ---\n'
sed -n '1,260p' test/python/test_fft_causal_conv1d.pyRepository: NVIDIA/cudnn-frontend
Length of output: 8202
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Check the exact autograd registration signature semantics from installed PyTorch docs if available,
# and probe the long-path backward call shape with a small standalone Python snippet.
python3 - <<'PY'
import inspect
import torch
print("torch version:", torch.__version__)
print("register_autograd:", inspect.signature(torch.library.register_autograd))
print("autograd.Function.backward docs available:", hasattr(torch.autograd.Function, 'backward'))
PYRepository: NVIDIA/cudnn-frontend
Length of output: 277
🌐 Web query:
PyTorch torch.library.register_autograd backward signature number of gradients matches number of outputs official docs
💡 Result:
Yes, the official PyTorch documentation explicitly states that the number of gradients passed to the backward function matches the number of outputs of the operator [1][2][3]. When registering a backward formula using torch.library.register_autograd, the backward function accepts the signature (ctx, *grads) [1][2][3]. In this signature, the grads parameter represents one or more gradients, and the number of these gradients corresponds exactly to the number of outputs produced by the operator [1][2][3]. The semantics of this backward_fn are consistent with those of torch.autograd.Function.backward() [1][2][3].
Citations:
- 1: https://docs.pytorch.org/docs/stable/library.md
- 2: https://docs.pytorch.org/docs/2.12/library.html
- 3: https://docs.pytorch.org/docs/2.11/library.html
_long_autograd_bwd must accept the second output grad.
_long_fwd_primitive returns y and reserve_space, so the registered backward receives two grad args. The current (ctx, grad_out) signature will fail on the long FFT path.
🐛 Proposed fix
`@torch.compiler.allow_in_graph`
-def _long_autograd_bwd(ctx, grad_out):
+def _long_autograd_bwd(ctx, grad_y, _grad_reserve_space):
x, weight, reserve_space = ctx.saved_tensors
- grad_y = grad_out[0]
grad_x, grad_weight = torch.ops.cudnn.long_fft_causal_conv1d_bwd_primitive(grad_y, x, weight, reserve_space)
return grad_x, grad_weight📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| @torch.compiler.allow_in_graph | |
| def _long_autograd_bwd(ctx, grad_out): | |
| x, weight, reserve_space = ctx.saved_tensors | |
| grad_y = grad_out[0] | |
| grad_x, grad_weight = torch.ops.cudnn.long_fft_causal_conv1d_bwd_primitive(grad_y, x, weight, reserve_space) | |
| return grad_x, grad_weight | |
| `@torch.compiler.allow_in_graph` | |
| def _long_autograd_bwd(ctx, grad_y, _grad_reserve_space): | |
| x, weight, reserve_space = ctx.saved_tensors | |
| grad_x, grad_weight = torch.ops.cudnn.long_fft_causal_conv1d_bwd_primitive(grad_y, x, weight, reserve_space) | |
| return grad_x, grad_weight |
🤖 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 305 - 310, Update
_long_autograd_bwd to accept both gradient arguments produced by
_long_fwd_primitive, including the gradient for reserve_space, while continuing
to use the output gradient for the backward primitive. Preserve the existing
returned gradients for x and weight.
| CUDNN_CHECK(cudnnFFTCausalConv1dForward(stream, | ||
| x_tensor.devPtr, | ||
| weight_tensor.devPtr, | ||
| y_tensor.devPtr, | ||
| batch, | ||
| dim, | ||
| seq_len, | ||
| kernel_size, | ||
| CUDNN_DATA_HALF)); | ||
|
|
||
| CUDNN_CHECK(cudnnFFTCausalConv1dBackward(stream, | ||
| x_tensor.devPtr, | ||
| weight_tensor.devPtr, | ||
| dy_tensor.devPtr, | ||
| dx_tensor.devPtr, | ||
| dweight_tensor.devPtr, | ||
| batch, | ||
| dim, | ||
| seq_len, | ||
| kernel_size, | ||
| CUDNN_DATA_HALF)); | ||
|
|
||
| CUDA_CHECK(cudaDeviceSynchronize()); |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
Verify numerical results for both paths.
CUDNN_CHECK and synchronization only prove the calls completed. Compare y_tensor, dx_tensor, and dweight_tensor against a causal-convolution reference with dtype-appropriate tolerances; otherwise layout, weight-order, and gradient regressions pass this sample silently.
As per coding guidelines, “C++ samples should follow the sequence: … execute, then verify results.”
Also applies to: 130-160
🤖 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 `@samples/cpp/causal_conv1d/fft_causal_conv1d.cpp` around lines 75 - 97, Add
numerical verification after execution in the sample’s forward and backward
paths, covering y_tensor, dx_tensor, and dweight_tensor against a
causal-convolution reference with dtype-appropriate tolerances. Keep the
existing CUDNN_CHECK calls and synchronization, and apply the same
execute-then-verify sequence to the code around the additional indicated
section.
Source: Coding guidelines
|
|
@yeliu-oss for ^^^ |
Anerudhan
added
mod-backend
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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 `@test/python/test_causal_conv1d.py`:
- Around line 30-33: Update _require_variant to query
torch.cuda.get_device_capability() and skip the causal-conv1d cases when the
device is below SM100, in addition to the existing cuDNN version and symbol
checks. Keep the current version requirement and skip message behavior intact
for other unsupported configurations.
- Around line 9-13: Update the torch.float64 entry in _TOLERANCES to use the
tighter FP64 tolerance established by related causal-conv tests, while leaving
the float32, float16, and bfloat16 tolerances unchanged.
🪄 Autofix (Beta)
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: b9b13106-e4fa-4481-8277-69bf7465dc28
📒 Files selected for processing (4)
docs/operations/CausalConv1d.mdpython/cudnn/ops/causal_conv1d.pytest/python/test_causal_conv1d.pytest/python/test_fft_causal_conv1d.py
🚧 Files skipped from review as they are similar to previous changes (1)
- test/python/test_fft_causal_conv1d.py
| def _require_variant(required_symbols, dtype, minimum_version): | ||
| required_version = max(minimum_version, 92600 if dtype == torch.float64 else 0) | ||
| if cudnn.backend_version() < required_version or any(not hasattr(cudnn, name) for name in required_symbols): | ||
| pytest.skip(f"SE causal conv1d case requires cuDNN {required_version // 10000}.{required_version // 100 % 100}") |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
git ls-files test/python/test_causal_conv1d.py test/AGENTS.md
printf '\n--- outline ---\n'
ast-grep outline test/python/test_causal_conv1d.py --view expanded
printf '\n--- file excerpt ---\n'
cat -n test/python/test_causal_conv1d.py | sed -n '1,220p'
printf '\n--- capability search ---\n'
rg -n "get_device_capability|SM100|sm100|cudnn.backend_version|_require_variant|causal_conv1d" test/python -SRepository: NVIDIA/cudnn-frontend
Length of output: 50377
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n--- causal conv1d references ---\n'
rg -n "causal_conv1d|b2b_causal_conv1d|fft_causal_conv1d" -S . --glob '!**/node_modules/**' --glob '!**/dist/**' --glob '!**/build/**'
printf '\n--- Blackwell/SM100 helper patterns in nearby tests ---\n'
sed -n '1,120p' test/python/test_sm100_prefill_oss_engine.py
printf '\n---\n'
sed -n '1,120p' test/python/test_sm100_rms_norm_silu_graph_api.py
printf '\n---\n'
sed -n '1,220p' test/python/test_cudnn_sdpa_op.pyRepository: NVIDIA/cudnn-frontend
Length of output: 48222
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n--- python/cudnn/ops/causal_conv1d.py excerpt ---\n'
cat -n python/cudnn/ops/causal_conv1d.py | sed -n '1,260p'
printf '\n--- docs/operations/CausalConv1d.md excerpt ---\n'
cat -n docs/operations/CausalConv1d.md | sed -n '1,220p'
printf '\n--- test/python/test_fft_causal_conv1d.py excerpt ---\n'
cat -n test/python/test_fft_causal_conv1d.py | sed -n '1,180p'Repository: NVIDIA/cudnn-frontend
Length of output: 20391
Gate causal-conv1d tests on SM100+. _require_variant() only checks backend version and symbol presence; these cases are documented as Blackwell-only, so add a torch.cuda.get_device_capability() guard before running them to avoid hitting unsupported kernels on older GPUs.
🤖 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_causal_conv1d.py` around lines 30 - 33, Update
_require_variant to query torch.cuda.get_device_capability() and skip the
causal-conv1d cases when the device is below SM100, in addition to the existing
cuDNN version and symbol checks. Keep the current version requirement and skip
message behavior intact for other unsupported configurations.
Source: Coding guidelines
FFT causal conv1d python tests took ~1.5 mins, the non-fft causal conv1d python tests took ~3.5 mins. I marked them L0 in the latest commit. |
2 participants
Before submitting
pre-commit runand committed any formatting changes.Affected area
Summary
cudnn.ops.fft_causal_conv1d(x, weight), following cuhyena's medium/long selection, padding, trimming, dtype support, and autograd behavior.Why
cuDNN 9.26 adds backend FFT causal conv1d kernels for filters beyond the efficient range of the existing direct causal conv1d kernels. This change makes those APIs usable from cuDNN Frontend while retaining the upstream cuhyena Python signature and selection policy.
Related issues
Depends on the FFT causal conv1d backend APIs targeted for cuDNN 9.26.
API and compatibility impact
Adds the public Python API:
It also exposes the low-level long-path buffer query:
torch.compile, and warmed-up CUDA Graph capture.weight[0]multiplies the current sample. This is reversed relative to the existing direct causal conv1d wrapper.CUDNN_STATUS_NOT_SUPPORTED.Testing
pre-commit run --files <all changed files>: passed (clang-format,black,black-jupyter).python -m pytest -v -m 'L1 or L2' test/python/test_fft_causal_conv1d.py: 7 passed.cmake --build /tmp/cudnn-frontend-fft-cpp-build -j16 --target samples: passed.samples '[fft_causal_conv1d]': 2 passed (medium forward/backward and long forward/backward).66_fft_causal_conv1d_forward.ipynband67_fft_causal_conv1d_backward.ipynb: passed for padded medium and genuine long paths.Summary by CodeRabbit
fft_causal_conv1d) with medium/long-kernel forward and backward support.torch.compile.fft_causal_conv1d.fft_causal_conv1d.