diff --git a/docs/operations/CausalConv1d.md b/docs/operations/CausalConv1d.md index 8af400c0f..5f3bd782b 100644 --- a/docs/operations/CausalConv1d.md +++ b/docs/operations/CausalConv1d.md @@ -10,8 +10,8 @@ Supports forward and backward passes with `torch.autograd` and `torch.compile`. ## Support -- **Architectures**: Turing (SM75) and newer, subject to available dynamic shared memory -- **Data types**: FP64, FP32, FP16, BF16 +- **Architectures**: Turing (SM75) or later +- **Data types**: FP32, FP16, BF16 - **Activations**: `identity` and `silu` for NHW and NWH; B2B uses fixed gating ### Kernel sizes @@ -171,7 +171,3 @@ B2B, requiring cuDNN 9.24.0 or later, with projection kernel size 2–32 and mix The low-level B2B forward binding writes both the mixer-plus-skip intermediate `y` and the final post-gated `y_gated`; the high-level API returns only `y_gated`. In most cases, use the corresponding `cudnn.ops` API, which handles autograd, `torch.compile`, and tensor management automatically. - -The backend accumulation buffers for `dweight` and `dbias` are FP32 for -FP16, BF16, and FP32 inputs, and FP64 for FP64 inputs. The Python autograd -results are returned in the corresponding input tensor dtype. diff --git a/docs/operations/FFTCausalConv1d.md b/docs/operations/FFTCausalConv1d.md deleted file mode 100644 index 330f0d426..000000000 --- a/docs/operations/FFTCausalConv1d.md +++ /dev/null @@ -1,71 +0,0 @@ -# FFT Causal Conv1d - -The cuDNN Frontend Python API exposes FFT-based depthwise causal convolution -through: - -```python -y = cudnn.ops.fft_causal_conv1d(x, weight) -``` - -This API requires cuDNN 9.26.0 or newer. It supports FP16, BF16, FP32, and FP64 -CUDA tensors, `torch.autograd`, and `torch.compile`. - -## Operation - -For `x` shaped `(batch, dim, seq_len)` and `weight` shaped -`(dim, kernel_size)`, the operation computes: - -```text -y[b, c, t] = sum(x[b, c, t - j] * weight[c, j], j=0..kernel_size-1) -``` - -Here, `x[b, c, t - j]` is zero when `t - j` is outside the input sequence. -This is FIR-order weight storage: -`weight[0]` multiplies the current sample. It is reversed relative to -`cudnn.ops.causal_conv1d`, whose first stored weight multiplies the oldest -sample in the causal window. - -## Path Selection - -The convenience wrapper follows cuhyena's Python API: - -- It rounds `kernel_size` up to a power of two, with a minimum of 128. -- It selects the medium FFT path when the filter fits that path's dtype and - device limit. -- It right-pads the input to a multiple of the medium filter length. -- Otherwise it right-pads input and filter to one common power-of-two length - and selects the long FFT path. -- It trims the output and autograd gradients back to the caller's shapes. - -The medium path supports FP64 filters through 4096. Other dtypes support -filters through 8192 before SM90 and through 16384 on SM90 or newer. - -The raw long backend path requires power-of-two `seq_len == kernel_size` in -`[4096, 16777216]`. FP64 supports lengths through 8388608; length 16777216 for -other dtypes requires SM90 or newer. - -## Long FFT Buffers - -Long FFT forward queries and allocates two opaque byte buffers: - -- Workspace is temporary scratch for the current forward or backward call. -- Reserve space stores transformed signal and filter state from forward and is - retained by the autograd context for the matching backward call. - -Applications calling the C APIs directly must preserve this same reserve-space -lifetime. See -`samples/cpp/causal_conv1d/fft_causal_conv1d.cpp` for medium and long C API -examples. - -## Example - -```python -import cudnn -import torch - -x = torch.randn(2, 16, 4096, device="cuda", requires_grad=True) -weight = torch.randn(16, 256, device="cuda", requires_grad=True) - -y = cudnn.ops.fft_causal_conv1d(x, weight) -y.sum().backward() -``` diff --git a/include/cudnn_frontend_shim.h b/include/cudnn_frontend_shim.h index 7525690b4..269ef0104 100644 --- a/include/cudnn_frontend_shim.h +++ b/include/cudnn_frontend_shim.h @@ -764,159 +764,6 @@ b2b_causal_conv1d_backward(cudaStream_t stream, } #endif -#if CUDNN_VERSION >= 92600 -inline cudnnStatus_t -fft_causal_conv1d_forward(cudaStream_t stream, - const void *x, - const void *weight, - void *y, - int batch, - int dim, - int seq_len, - int kernel_size, - cudnnDataType_t data_type) { - auto effective_cudnn_ver = std::min(detail::get_compiled_version(), detail::get_backend_version()); - if (effective_cudnn_ver < 92600) { - return CUDNN_STATUS_NOT_SUPPORTED; - } - NV_FE_CALL_TO_BACKEND(fft_causal_conv1d_forward, - cudnnFFTCausalConv1dForward, - stream, - x, - weight, - y, - batch, - dim, - seq_len, - kernel_size, - data_type); -} - -inline cudnnStatus_t -fft_causal_conv1d_backward(cudaStream_t stream, - const void *x, - const void *weight, - const void *dy, - void *dx, - void *dweight, - int batch, - int dim, - int seq_len, - int kernel_size, - cudnnDataType_t data_type) { - auto effective_cudnn_ver = std::min(detail::get_compiled_version(), detail::get_backend_version()); - if (effective_cudnn_ver < 92600) { - return CUDNN_STATUS_NOT_SUPPORTED; - } - NV_FE_CALL_TO_BACKEND(fft_causal_conv1d_backward, - cudnnFFTCausalConv1dBackward, - stream, - x, - weight, - dy, - dx, - dweight, - batch, - dim, - seq_len, - kernel_size, - data_type); -} - -inline cudnnStatus_t -long_fft_causal_conv1d_get_buffer_sizes(int batch, - int dim, - int seq_len, - int kernel_size, - cudnnDataType_t data_type, - size_t *workspace_size_in_bytes, - size_t *reserve_space_size_in_bytes) { - auto effective_cudnn_ver = std::min(detail::get_compiled_version(), detail::get_backend_version()); - if (effective_cudnn_ver < 92600) { - return CUDNN_STATUS_NOT_SUPPORTED; - } - NV_FE_CALL_TO_BACKEND(long_fft_causal_conv1d_get_buffer_sizes, - cudnnLongFFTCausalConv1dGetBufferSizes, - batch, - dim, - seq_len, - kernel_size, - data_type, - workspace_size_in_bytes, - reserve_space_size_in_bytes); -} - -inline cudnnStatus_t -long_fft_causal_conv1d_forward(cudaStream_t stream, - const void *x, - const void *weight, - void *y, - int batch, - int dim, - int seq_len, - int kernel_size, - cudnnDataType_t data_type, - void *workspace, - size_t workspace_size_in_bytes, - void *reserve_space, - size_t reserve_space_size_in_bytes) { - auto effective_cudnn_ver = std::min(detail::get_compiled_version(), detail::get_backend_version()); - if (effective_cudnn_ver < 92600) { - return CUDNN_STATUS_NOT_SUPPORTED; - } - NV_FE_CALL_TO_BACKEND(long_fft_causal_conv1d_forward, - cudnnLongFFTCausalConv1dForward, - stream, - x, - weight, - y, - batch, - dim, - seq_len, - kernel_size, - data_type, - workspace, - workspace_size_in_bytes, - reserve_space, - reserve_space_size_in_bytes); -} - -inline cudnnStatus_t -long_fft_causal_conv1d_backward(cudaStream_t stream, - const void *dy, - void *dx, - void *dweight, - int batch, - int dim, - int seq_len, - int kernel_size, - cudnnDataType_t data_type, - void *workspace, - size_t workspace_size_in_bytes, - void *reserve_space, - size_t reserve_space_size_in_bytes) { - auto effective_cudnn_ver = std::min(detail::get_compiled_version(), detail::get_backend_version()); - if (effective_cudnn_ver < 92600) { - return CUDNN_STATUS_NOT_SUPPORTED; - } - NV_FE_CALL_TO_BACKEND(long_fft_causal_conv1d_backward, - cudnnLongFFTCausalConv1dBackward, - stream, - dy, - dx, - dweight, - batch, - dim, - seq_len, - kernel_size, - data_type, - workspace, - workspace_size_in_bytes, - reserve_space, - reserve_space_size_in_bytes); -} -#endif - inline std::string convert_version_to_str(size_t const version) { // The multiplier for major version pre-v9 and post-v9 are different. diff --git a/llms.txt b/llms.txt index 2b9479c7c..5439d4a8f 100644 --- a/llms.txt +++ b/llms.txt @@ -22,7 +22,6 @@ Published documentation: https://docs.nvidia.com/deeplearning/cudnn/latest/devel - [Block Scaling (MXFP8/NVFP4 quantization)](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/BlockScaling.md) - [RoPE](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/RoPE.md) - [Causal Conv1d](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/CausalConv1d.md) -- [FFT Causal Conv1d](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/FFTCausalConv1d.md) - [Concatenate](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/Concatenate.md), [Reshape](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/Reshape.md), [Slice](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/Slice.md), [Transpose](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/Transpose.md), [Resampling](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/Resampling.md) ## Open-source (frontend-only) kernel APIs diff --git a/python/cudnn/__init__.py b/python/cudnn/__init__.py index 2411fdfeb..3d23d3471 100644 --- a/python/cudnn/__init__.py +++ b/python/cudnn/__init__.py @@ -56,11 +56,6 @@ def is_windows(): "causal_conv1d_nwh_backward", "b2b_causal_conv1d_forward", "b2b_causal_conv1d_backward", - "fft_causal_conv1d_forward", - "fft_causal_conv1d_backward", - "long_fft_causal_conv1d_get_buffer_sizes", - "long_fft_causal_conv1d_forward", - "long_fft_causal_conv1d_backward", ]: if hasattr(_pybind_module, _optional_symbol): globals()[_optional_symbol] = getattr(_pybind_module, _optional_symbol) diff --git a/python/cudnn/ops/__init__.py b/python/cudnn/ops/__init__.py index a6c54a3d3..857081820 100644 --- a/python/cudnn/ops/__init__.py +++ b/python/cudnn/ops/__init__.py @@ -2,4 +2,3 @@ # SPDX-License-Identifier: Apache-2.0 from .causal_conv1d import causal_conv1d, causal_conv1d_nwh, b2b_causal_conv1d -from .fft_causal_conv1d import fft_causal_conv1d diff --git a/python/cudnn/ops/causal_conv1d.py b/python/cudnn/ops/causal_conv1d.py index 91cfbce66..9f3731e8e 100644 --- a/python/cudnn/ops/causal_conv1d.py +++ b/python/cudnn/ops/causal_conv1d.py @@ -8,7 +8,6 @@ _TORCH_DTYPE_TO_CUDNN = { torch.float32: 0, # CUDNN_DATA_FLOAT - torch.float64: 1, # CUDNN_DATA_DOUBLE torch.float16: 2, # CUDNN_DATA_HALF torch.bfloat16: 9, # CUDNN_DATA_BFLOAT16 } @@ -21,16 +20,10 @@ def _dtype_to_int(dtype: torch.dtype) -> int: if dtype not in _TORCH_DTYPE_TO_CUDNN: - raise ValueError(f"Unsupported dtype {dtype}. Supported: float64, float32, float16, bfloat16.") + raise ValueError(f"Unsupported dtype {dtype}. Supported: float32, float16, bfloat16.") return _TORCH_DTYPE_TO_CUDNN[dtype] -def _gradient_dtype(dtype: torch.dtype) -> torch.dtype: - # Match cuhyena: FP16/BF16 parameter gradients accumulate in FP32, - # while FP32 and FP64 parameter gradients accumulate in their input type. - return torch.float32 if dtype in (torch.float16, torch.bfloat16) else dtype - - def _activation_to_int(activation: str) -> int: if activation not in _ACTIVATION_TO_INT: raise ValueError(f"Unsupported activation '{activation}'. Supported: 'identity', 'silu'.") @@ -143,9 +136,8 @@ def _bwd_primitive(grad_out: Tensor, x: Tensor, weight: Tensor, bias: Tensor, ac kernel_size = weight.shape[1] dx = torch.empty_like(x) - grad_dtype = _gradient_dtype(x.dtype) - dweight = torch.zeros(weight.shape, device=x.device, dtype=grad_dtype) - dbias = torch.zeros(bias.shape, device=x.device, dtype=grad_dtype) + dweight = torch.zeros(weight.shape, device=x.device, dtype=torch.float32) + dbias = torch.zeros(bias.shape, device=x.device, dtype=torch.float32) import cudnn @@ -163,7 +155,7 @@ def _bwd_primitive(grad_out: Tensor, x: Tensor, weight: Tensor, bias: Tensor, ac seq_len, kernel_size, _dtype_to_int(x.dtype), - _dtype_to_int(grad_dtype), + _dtype_to_int(torch.float32), _activation_to_int(activation), ) return [dx, dweight.to(x.dtype), dbias.to(x.dtype)] @@ -225,7 +217,7 @@ def causal_conv1d( Args: x (torch.Tensor): Input tensor of shape ``(batch, dim, seq_len)``. - Must be BF16, FP16, FP32, or FP64. Must be contiguous and on CUDA. + Must be BF16, FP16, or FP32 and on CUDA. weight (torch.Tensor): Filter tensor of shape ``(dim, kernel_size)``. Same dtype as *x*. ``kernel_size`` must be between 2 and 256, inclusive. @@ -354,9 +346,8 @@ def _nwh_bwd_primitive(grad_out: Tensor, x: Tensor, weight: Tensor, bias: Tensor kernel_size = weight.shape[0] dx = torch.empty_like(x) - grad_dtype = _gradient_dtype(x.dtype) - dweight = torch.zeros(weight.shape, device=x.device, dtype=grad_dtype) - dbias = torch.zeros(bias.shape, device=x.device, dtype=grad_dtype) + dweight = torch.zeros(weight.shape, device=x.device, dtype=torch.float32) + dbias = torch.zeros(bias.shape, device=x.device, dtype=torch.float32) import cudnn @@ -374,7 +365,7 @@ def _nwh_bwd_primitive(grad_out: Tensor, x: Tensor, weight: Tensor, bias: Tensor seq_len, kernel_size, _dtype_to_int(x.dtype), - _dtype_to_int(grad_dtype), + _dtype_to_int(torch.float32), _activation_to_int(activation), ) return [dx, dweight.to(x.dtype), dbias.to(x.dtype)] @@ -432,7 +423,6 @@ def causal_conv1d_nwh( Args: x (torch.Tensor): Input tensor of shape ``(batch, seq_len, dim)``. - Must be BF16, FP16, FP32, or FP64. weight (torch.Tensor): Filter tensor of shape ``(kernel_size, dim)``. ``kernel_size`` must be between 2 and 128, inclusive. bias (torch.Tensor | None): Optional bias of shape ``(dim,)``. @@ -602,10 +592,9 @@ def _b2b_bwd_primitive( kernel_size_mixer = weights_mixer.shape[1] dx = torch.empty_like(x) - grad_dtype = _gradient_dtype(x.dtype) - dweights_proj = torch.zeros(weights_proj.shape, device=x.device, dtype=grad_dtype) - dweights_mixer = torch.zeros(weights_mixer.shape, device=x.device, dtype=grad_dtype) - dskip_bias = torch.zeros(skip_bias.shape, device=x.device, dtype=grad_dtype) + dweights_proj = torch.zeros(weights_proj.shape, device=x.device, dtype=torch.float32) + dweights_mixer = torch.zeros(weights_mixer.shape, device=x.device, dtype=torch.float32) + dskip_bias = torch.zeros(skip_bias.shape, device=x.device, dtype=torch.float32) import cudnn @@ -627,7 +616,7 @@ def _b2b_bwd_primitive( kernel_size_proj, kernel_size_mixer, _dtype_to_int(x.dtype), - _dtype_to_int(grad_dtype), + _dtype_to_int(torch.float32), ) return [ dx, @@ -707,7 +696,6 @@ def b2b_causal_conv1d( Args: x (torch.Tensor): Input tensor of shape ``(batch, 3*dim, seq_len)``. - Must be BF16, FP16, FP32, or FP64. weights_proj (torch.Tensor): Projection filter ``(3*dim, kernel_size_proj)``. ``kernel_size_proj`` must be between 2 and 32, inclusive. weights_mixer (torch.Tensor): Mixer filter ``(dim, kernel_size_mixer)``. diff --git a/python/cudnn/ops/fft_causal_conv1d.py b/python/cudnn/ops/fft_causal_conv1d.py deleted file mode 100644 index 41701d250..000000000 --- a/python/cudnn/ops/fft_causal_conv1d.py +++ /dev/null @@ -1,369 +0,0 @@ -from typing import List, Tuple - -import torch -import torch.nn.functional as F -from torch import Tensor - -_TORCH_DTYPE_TO_CUDNN = { - torch.float32: 0, # CUDNN_DATA_FLOAT - torch.float64: 1, # CUDNN_DATA_DOUBLE - torch.float16: 2, # CUDNN_DATA_HALF - torch.bfloat16: 9, # CUDNN_DATA_BFLOAT16 -} - -_REQUIRED_CUDNN_VERSION = 92600 -_LONG_FFT_MIN_LENGTH = 4096 -_LONG_FFT_MAX_LENGTH = 1 << 24 -_BUFFER_ALIGNMENT = 256 - - -def _dtype_to_int(dtype: torch.dtype) -> int: - if dtype not in _TORCH_DTYPE_TO_CUDNN: - supported = ", ".join(str(value) for value in _TORCH_DTYPE_TO_CUDNN) - raise ValueError(f"Unsupported dtype {dtype}. Supported: {supported}.") - return _TORCH_DTYPE_TO_CUDNN[dtype] - - -def _require_fft_causal_conv1d_api() -> None: - import cudnn - - required_symbols = ( - "fft_causal_conv1d_forward", - "fft_causal_conv1d_backward", - "long_fft_causal_conv1d_get_buffer_sizes", - "long_fft_causal_conv1d_forward", - "long_fft_causal_conv1d_backward", - ) - missing = [name for name in required_symbols if not hasattr(cudnn, name)] - if missing or cudnn.backend_version() < _REQUIRED_CUDNN_VERSION: - raise RuntimeError("FFT causal conv1d requires cuDNN 9.26.0 or newer and Frontend bindings built against cuDNN 9.26.0 or newer.") - - -def _validate_tensors(x: Tensor, weight: Tensor) -> Tuple[int, int, int, int]: - if x.dim() != 3 or weight.dim() != 2: - raise ValueError(f"Expected x(3D) and weight(2D); got x.shape={x.shape}, weight.shape={weight.shape}.") - if not (x.is_cuda and weight.is_cuda): - raise ValueError(f"x and weight must be CUDA tensors; got x.device={x.device}, weight.device={weight.device}.") - if x.device != weight.device: - raise ValueError(f"x and weight must be on the same device; got x.device={x.device}, weight.device={weight.device}.") - if x.dtype != weight.dtype: - raise TypeError(f"x and weight must have the same dtype; got x.dtype={x.dtype}, weight.dtype={weight.dtype}.") - _dtype_to_int(x.dtype) - - batch, dim, seq_len = x.shape - weight_dim, kernel_size = weight.shape - if batch <= 0 or dim <= 0 or seq_len <= 0 or kernel_size <= 0: - raise ValueError(f"All tensor dimensions must be positive; got x.shape={x.shape}, weight.shape={weight.shape}.") - if weight_dim != dim: - raise ValueError(f"Channel mismatch: x has dim={dim}, but weight.shape[0]={weight_dim}.") - return batch, dim, seq_len, kernel_size - - -def _is_power_of_two(value: int) -> bool: - return value > 0 and (value & (value - 1)) == 0 - - -def _next_power_of_two(value: int) -> int: - return 1 << (value - 1).bit_length() - - -def _compatible_filter_length(filter_length: int) -> int: - return max(128, _next_power_of_two(filter_length)) - - -def _cuda_arch(device: torch.device) -> int: - properties = torch.cuda.get_device_properties(device) - return properties.major * 100 + properties.minor * 10 - - -def _medium_fft_is_available(filter_length: int, dtype: torch.dtype, device: torch.device) -> bool: - # Match cuhyena's public Python selector: FP64 ends at K=4096; other - # dtypes end at K=8192 before SM90 and K=16384 on SM90 or newer. - if dtype == torch.float64: - return filter_length <= 4096 - return filter_length <= (16384 if _cuda_arch(device) >= 900 else 8192) - - -def _validate_medium_shape(seq_len: int, kernel_size: int) -> None: - if not _is_power_of_two(kernel_size) or not 128 <= kernel_size <= 16384: - raise ValueError(f"Medium FFT kernel_size must be a power of two in [128, 16384]; got {kernel_size}.") - if seq_len < kernel_size or seq_len % kernel_size != 0: - raise ValueError(f"Medium FFT requires seq_len >= kernel_size and seq_len % kernel_size == 0; got {seq_len} and {kernel_size}.") - - -def _validate_long_shape(seq_len: int, kernel_size: int, dtype: torch.dtype, device: torch.device) -> None: - if seq_len != kernel_size or not _is_power_of_two(kernel_size): - raise ValueError(f"Long FFT requires power-of-two seq_len == kernel_size; got {seq_len} and {kernel_size}.") - if not _LONG_FFT_MIN_LENGTH <= kernel_size <= _LONG_FFT_MAX_LENGTH: - raise ValueError(f"Long FFT kernel_size must be in [{_LONG_FFT_MIN_LENGTH}, {_LONG_FFT_MAX_LENGTH}]; got {kernel_size}.") - if dtype == torch.float64 and kernel_size == _LONG_FFT_MAX_LENGTH: - raise ValueError("Long FFT FP64 supports kernel_size through 8388608.") - 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.") - - -def _align_buffer_size(size: int) -> int: - return ((size + _BUFFER_ALIGNMENT - 1) // _BUFFER_ALIGNMENT) * _BUFFER_ALIGNMENT - - -def _long_buffer_size_bytes(batch: int, dim: int, kernel_size: int, dtype: torch.dtype) -> int: - # This is the public backend buffer-size formula used only by the fake - # implementation. Runtime allocations always use the backend size query. - fft_size = 2 * kernel_size - bits = fft_size.bit_length() - 1 - m = 1 << ((bits + 1) // 2) - n = 1 << (bits - ((bits + 1) // 2)) - intermediate_elements = 2 * (m // 2 + 1) * n - scratch_element_size = 8 if dtype == torch.float64 else 4 - signal_bytes = _align_buffer_size(batch * dim * intermediate_elements * scratch_element_size) - filter_bytes = _align_buffer_size(dim * intermediate_elements * scratch_element_size) - return signal_bytes + filter_bytes - - -@torch.library.custom_op("cudnn::fft_causal_conv1d_fwd_primitive", mutates_args=(), device_types="cuda") -def _medium_fwd_primitive(x: Tensor, weight: Tensor) -> Tensor: - _require_fft_causal_conv1d_api() - batch, dim, seq_len, kernel_size = _validate_tensors(x, weight) - _validate_medium_shape(seq_len, kernel_size) - - x = x.contiguous() - weight = weight.contiguous() - y = torch.empty_like(x) - - import cudnn - - cudnn.fft_causal_conv1d_forward( - torch.cuda.current_stream(x.device).cuda_stream, - x.data_ptr(), - weight.data_ptr(), - y.data_ptr(), - batch, - dim, - seq_len, - kernel_size, - _dtype_to_int(x.dtype), - ) - return y - - -@torch.library.register_fake("cudnn::fft_causal_conv1d_fwd_primitive") -def _medium_fwd_fake(x: Tensor, weight: Tensor) -> Tensor: - return torch.empty_like(x) - - -@torch.library.custom_op("cudnn::fft_causal_conv1d_bwd_primitive", mutates_args=(), device_types="cuda") -def _medium_bwd_primitive(grad_out: Tensor, x: Tensor, weight: Tensor) -> List[Tensor]: - _require_fft_causal_conv1d_api() - batch, dim, seq_len, kernel_size = _validate_tensors(x, weight) - _validate_medium_shape(seq_len, kernel_size) - if grad_out.shape != x.shape or grad_out.device != x.device or grad_out.dtype != x.dtype: - raise ValueError("grad_out must match x in shape, device, and dtype.") - - x = x.contiguous() - weight = weight.contiguous() - grad_out = grad_out.contiguous() - grad_x = torch.empty_like(x) - # The medium cuhyena kernel accumulates channel/batch contributions into dw. - grad_weight = torch.zeros_like(weight) - - import cudnn - - cudnn.fft_causal_conv1d_backward( - torch.cuda.current_stream(x.device).cuda_stream, - x.data_ptr(), - weight.data_ptr(), - grad_out.data_ptr(), - grad_x.data_ptr(), - grad_weight.data_ptr(), - batch, - dim, - seq_len, - kernel_size, - _dtype_to_int(x.dtype), - ) - return [grad_x, grad_weight] - - -@torch.library.register_fake("cudnn::fft_causal_conv1d_bwd_primitive") -def _medium_bwd_fake(grad_out: Tensor, x: Tensor, weight: Tensor) -> List[Tensor]: - return [torch.empty_like(x), torch.empty_like(weight)] - - -def _medium_setup_context(ctx, inputs, output) -> None: - del output - x, weight = inputs - ctx.save_for_backward(x, weight) - - -@torch.compiler.allow_in_graph -def _medium_autograd_bwd(ctx, grad_out: Tensor): - x, weight = ctx.saved_tensors - grad_x, grad_weight = torch.ops.cudnn.fft_causal_conv1d_bwd_primitive(grad_out, x, weight) - return grad_x, grad_weight - - -torch.library.register_autograd( - "cudnn::fft_causal_conv1d_fwd_primitive", - _medium_autograd_bwd, - setup_context=_medium_setup_context, -) - - -@torch.library.custom_op("cudnn::long_fft_causal_conv1d_fwd_primitive", mutates_args=(), device_types="cuda") -def _long_fwd_primitive(x: Tensor, weight: Tensor) -> List[Tensor]: - _require_fft_causal_conv1d_api() - batch, dim, seq_len, kernel_size = _validate_tensors(x, weight) - _validate_long_shape(seq_len, kernel_size, x.dtype, x.device) - - x = x.contiguous() - weight = weight.contiguous() - y = torch.empty_like(x) - - import cudnn - - workspace_size, reserve_size = cudnn.long_fft_causal_conv1d_get_buffer_sizes(batch, dim, seq_len, kernel_size, _dtype_to_int(x.dtype)) - workspace = torch.empty(workspace_size, device=x.device, dtype=torch.uint8) - reserve_space = torch.empty(reserve_size, device=x.device, dtype=torch.uint8) - cudnn.long_fft_causal_conv1d_forward( - torch.cuda.current_stream(x.device).cuda_stream, - x.data_ptr(), - weight.data_ptr(), - y.data_ptr(), - batch, - dim, - seq_len, - kernel_size, - _dtype_to_int(x.dtype), - workspace.data_ptr(), - workspace_size, - reserve_space.data_ptr(), - reserve_size, - ) - # Reserve space contains transformed x/weight state and must survive until - # autograd invokes the matching long backward primitive. - return [y, reserve_space] - - -@torch.library.register_fake("cudnn::long_fft_causal_conv1d_fwd_primitive") -def _long_fwd_fake(x: Tensor, weight: Tensor) -> List[Tensor]: - batch, dim, _ = x.shape - kernel_size = weight.shape[1] - reserve_size = _long_buffer_size_bytes(batch, dim, kernel_size, x.dtype) - return [torch.empty_like(x), torch.empty(reserve_size, device=x.device, dtype=torch.uint8)] - - -@torch.library.custom_op("cudnn::long_fft_causal_conv1d_bwd_primitive", mutates_args=(), device_types="cuda") -def _long_bwd_primitive(grad_out: Tensor, x: Tensor, weight: Tensor, reserve_space: Tensor) -> List[Tensor]: - _require_fft_causal_conv1d_api() - batch, dim, seq_len, kernel_size = _validate_tensors(x, weight) - _validate_long_shape(seq_len, kernel_size, x.dtype, x.device) - if grad_out.shape != x.shape or grad_out.device != x.device or grad_out.dtype != x.dtype: - raise ValueError("grad_out must match x in shape, device, and dtype.") - if reserve_space.device != x.device or reserve_space.dtype != torch.uint8: - raise ValueError("reserve_space must be a CUDA uint8 tensor on the same device as x.") - - grad_out = grad_out.contiguous() - reserve_space = reserve_space.contiguous() - grad_x = torch.empty_like(x) - grad_weight = torch.empty_like(weight) - - import cudnn - - workspace_size, reserve_size = cudnn.long_fft_causal_conv1d_get_buffer_sizes(batch, dim, seq_len, kernel_size, _dtype_to_int(x.dtype)) - if reserve_space.numel() < reserve_size: - raise ValueError(f"reserve_space has {reserve_space.numel()} bytes, but the backend requires {reserve_size}.") - workspace = torch.empty(workspace_size, device=x.device, dtype=torch.uint8) - cudnn.long_fft_causal_conv1d_backward( - torch.cuda.current_stream(x.device).cuda_stream, - grad_out.data_ptr(), - grad_x.data_ptr(), - grad_weight.data_ptr(), - batch, - dim, - seq_len, - kernel_size, - _dtype_to_int(x.dtype), - workspace.data_ptr(), - workspace_size, - reserve_space.data_ptr(), - reserve_size, - ) - return [grad_x, grad_weight] - - -@torch.library.register_fake("cudnn::long_fft_causal_conv1d_bwd_primitive") -def _long_bwd_fake(grad_out: Tensor, x: Tensor, weight: Tensor, reserve_space: Tensor) -> List[Tensor]: - return [torch.empty_like(x), torch.empty_like(weight)] - - -def _long_setup_context(ctx, inputs, output) -> None: - x, weight = inputs - _, reserve_space = output - ctx.save_for_backward(x, weight, reserve_space) - - -@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.library.register_autograd( - "cudnn::long_fft_causal_conv1d_fwd_primitive", - _long_autograd_bwd, - setup_context=_long_setup_context, -) - - -def fft_causal_conv1d(x: Tensor, weight: Tensor) -> Tensor: - r"""Compute depthwise causal 1D convolution with FFT kernels. - - The public signature, path selection, and right-padding behavior match - cuhyena's ``fft_causal_conv1d(x, weight)`` wrapper. Medium FFT kernels are - used when the filter fits their dtype/device limit; otherwise the full - sequence long FFT path is used. Padding introduced for either raw backend - API is removed from the returned tensor and from autograd gradients. - - The FFT weight convention is FIR order:: - - y[t] = sum(weight[j] * x[t - j], j=0..kernel_size-1) - - This is reversed relative to :func:`causal_conv1d`, whose first stored - weight multiplies the oldest sample in the causal window. - - Requires cuDNN 9.26.0 or newer. Supports FP16, BF16, FP32, and FP64 CUDA - tensors and integrates with ``torch.autograd`` and ``torch.compile``. - - Args: - x (torch.Tensor): Input tensor shaped ``(batch, dim, seq_len)``. - weight (torch.Tensor): Per-channel filters shaped - ``(dim, kernel_size)``. - - Returns: - torch.Tensor: Output shaped ``(batch, dim, seq_len)``. - """ - _, _, seq_len, weight_length = _validate_tensors(x, weight) - - if _medium_fft_is_available(weight_length, x.dtype, x.device): - filter_length = _compatible_filter_length(weight_length) - if weight_length != filter_length: - weight = F.pad(weight, (0, filter_length - weight_length)) - - padded_seq_len = max(filter_length, ((seq_len + filter_length - 1) // filter_length) * filter_length) - if seq_len != padded_seq_len: - x = F.pad(x, (0, padded_seq_len - seq_len)) - - y = torch.ops.cudnn.fft_causal_conv1d_fwd_primitive(x, weight) - return y[..., :seq_len] - - filter_length = max(_compatible_filter_length(weight_length), _compatible_filter_length(seq_len)) - _validate_long_shape(filter_length, filter_length, x.dtype, x.device) - if weight_length != filter_length: - weight = F.pad(weight, (0, filter_length - weight_length)) - if seq_len != filter_length: - x = F.pad(x, (0, filter_length - seq_len)) - - y, _ = torch.ops.cudnn.long_fft_causal_conv1d_fwd_primitive(x, weight) - return y[..., :seq_len] diff --git a/python/pycudnn.cpp b/python/pycudnn.cpp index e3fab8558..b334dc7d2 100644 --- a/python/pycudnn.cpp +++ b/python/pycudnn.cpp @@ -391,138 +391,6 @@ PYBIND11_MODULE(_compiled_module, m) { throw_if_b2b_causal_conv1d_failed( status, "cudnnB2BCausalConv1dBackward", kernel_size_proj, kernel_size_mixer); }); -#if CUDNN_VERSION >= 92600 - m.def("fft_causal_conv1d_forward", - [](std::intptr_t stream, - std::intptr_t x_ptr, - std::intptr_t weight_ptr, - std::intptr_t out_ptr, - int batch, - int dim, - int seq_len, - int kernel_size, - int data_type) { - auto status = detail::fft_causal_conv1d_forward(reinterpret_cast(stream), - reinterpret_cast(x_ptr), - reinterpret_cast(weight_ptr), - reinterpret_cast(out_ptr), - batch, - dim, - seq_len, - kernel_size, - static_cast(data_type)); - if (status != 0) - throw std::runtime_error("cudnnFFTCausalConv1dForward failed with status " + std::to_string(status)); - }); - - m.def("fft_causal_conv1d_backward", - [](std::intptr_t stream, - std::intptr_t x_ptr, - std::intptr_t weight_ptr, - std::intptr_t dy_ptr, - std::intptr_t dx_ptr, - std::intptr_t dweight_ptr, - int batch, - int dim, - int seq_len, - int kernel_size, - int data_type) { - auto status = detail::fft_causal_conv1d_backward(reinterpret_cast(stream), - reinterpret_cast(x_ptr), - reinterpret_cast(weight_ptr), - reinterpret_cast(dy_ptr), - reinterpret_cast(dx_ptr), - reinterpret_cast(dweight_ptr), - batch, - dim, - seq_len, - kernel_size, - static_cast(data_type)); - if (status != 0) - throw std::runtime_error("cudnnFFTCausalConv1dBackward failed with status " + std::to_string(status)); - }); - - m.def("long_fft_causal_conv1d_get_buffer_sizes", - [](int batch, int dim, int seq_len, int kernel_size, int data_type) { - size_t workspace_size_in_bytes = 0; - size_t reserve_space_size_in_bytes = 0; - auto status = detail::long_fft_causal_conv1d_get_buffer_sizes(batch, - dim, - seq_len, - kernel_size, - static_cast(data_type), - &workspace_size_in_bytes, - &reserve_space_size_in_bytes); - if (status != 0) - throw std::runtime_error("cudnnLongFFTCausalConv1dGetBufferSizes failed with status " + - std::to_string(status)); - return std::make_pair(workspace_size_in_bytes, reserve_space_size_in_bytes); - }); - - m.def("long_fft_causal_conv1d_forward", - [](std::intptr_t stream, - std::intptr_t x_ptr, - std::intptr_t weight_ptr, - std::intptr_t out_ptr, - int batch, - int dim, - int seq_len, - int kernel_size, - int data_type, - std::intptr_t workspace_ptr, - size_t workspace_size_in_bytes, - std::intptr_t reserve_space_ptr, - size_t reserve_space_size_in_bytes) { - auto status = detail::long_fft_causal_conv1d_forward(reinterpret_cast(stream), - reinterpret_cast(x_ptr), - reinterpret_cast(weight_ptr), - reinterpret_cast(out_ptr), - batch, - dim, - seq_len, - kernel_size, - static_cast(data_type), - reinterpret_cast(workspace_ptr), - workspace_size_in_bytes, - reinterpret_cast(reserve_space_ptr), - reserve_space_size_in_bytes); - if (status != 0) - throw std::runtime_error("cudnnLongFFTCausalConv1dForward failed with status " + - std::to_string(status)); - }); - - m.def("long_fft_causal_conv1d_backward", - [](std::intptr_t stream, - std::intptr_t dy_ptr, - std::intptr_t dx_ptr, - std::intptr_t dweight_ptr, - int batch, - int dim, - int seq_len, - int kernel_size, - int data_type, - std::intptr_t workspace_ptr, - size_t workspace_size_in_bytes, - std::intptr_t reserve_space_ptr, - size_t reserve_space_size_in_bytes) { - auto status = detail::long_fft_causal_conv1d_backward(reinterpret_cast(stream), - reinterpret_cast(dy_ptr), - reinterpret_cast(dx_ptr), - reinterpret_cast(dweight_ptr), - batch, - dim, - seq_len, - kernel_size, - static_cast(data_type), - reinterpret_cast(workspace_ptr), - workspace_size_in_bytes, - reinterpret_cast(reserve_space_ptr), - reserve_space_size_in_bytes); - if (status != 0) - throw std::runtime_error("cudnnLongFFTCausalConv1dBackward failed with status " + - std::to_string(status)); - }); -#endif #endif } diff --git a/samples/cpp/CMakeLists.txt b/samples/cpp/CMakeLists.txt index 87a316124..81db82352 100644 --- a/samples/cpp/CMakeLists.txt +++ b/samples/cpp/CMakeLists.txt @@ -79,7 +79,6 @@ add_executable( causal_conv1d/causal_conv1d.cpp causal_conv1d/causal_conv1d_nwh.cpp causal_conv1d/b2b_causal_conv1d.cpp - causal_conv1d/fft_causal_conv1d.cpp ) # target flags diff --git a/samples/cpp/causal_conv1d/fft_causal_conv1d.cpp b/samples/cpp/causal_conv1d/fft_causal_conv1d.cpp deleted file mode 100644 index 61f8727b6..000000000 --- a/samples/cpp/causal_conv1d/fft_causal_conv1d.cpp +++ /dev/null @@ -1,162 +0,0 @@ -/* - * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. - * - * Permission is hereby granted, free of charge, to any person obtaining a - * copy of this software and associated documentation files (the "Software"), - * to deal in the Software without restriction, including without limitation - * the rights to use, copy, modify, merge, publish, distribute, sublicense, - * and/or sell copies of the Software, and to permit persons to whom the - * Software is furnished to do so, subject to the following conditions: - * - * The above copyright notice and this permission notice shall be included in - * all copies or substantial portions of the Software. - * - * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR - * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, - * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL - * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER - * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING - * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER - * DEALINGS IN THE SOFTWARE. - */ - -#include -#include "../utils/helpers.h" - -#include - -#if defined(CUDNN_SUBQUADRATIC_OPS_H_) || __has_include() -#if !defined(CUDNN_SUBQUADRATIC_OPS_H_) -#include -#endif -#define HAS_SUBQUADRATIC_OPS 1 -#else -#define HAS_SUBQUADRATIC_OPS 0 -#endif - -#if HAS_SUBQUADRATIC_OPS && CUDNN_VERSION >= 92600 -namespace { - -void -skip_if_fft_causal_conv1d_is_unavailable() { - if (cudnnGetVersion() < 92600) { - SKIP("FFT causal conv1d requires cuDNN 9.26.0 or newer"); - } - if (!is_arch_supported_by_cudnn()) { - SKIP("Architecture is not supported by the current cuDNN version"); - } -} - -} // namespace -#endif - -TEST_CASE("FFT causal conv1d medium forward and backward", "[fft_causal_conv1d][medium]") { -#if !HAS_SUBQUADRATIC_OPS || CUDNN_VERSION < 92600 - SKIP("FFT causal conv1d APIs require cuDNN 9.26.0 headers"); -#elif defined(_MSC_VER) - SKIP("FFT causal conv1d kernels are not supported on Windows (MSVC)"); -#else - skip_if_fft_causal_conv1d_is_unavailable(); - - constexpr int batch = 2; - constexpr int dim = 4; - constexpr int seq_len = 256; - constexpr int kernel_size = 128; - - cudaStream_t stream = nullptr; - - Surface x_tensor(batch * dim * seq_len); - Surface weight_tensor(dim * kernel_size); - Surface y_tensor(batch * dim * seq_len); - Surface dy_tensor(batch * dim * seq_len); - Surface dx_tensor(batch * dim * seq_len); - Surface dweight_tensor(dim * kernel_size, 0.0f); - - 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()); -#endif -} - -TEST_CASE("FFT causal conv1d long forward and backward", "[fft_causal_conv1d][long]") { -#if !HAS_SUBQUADRATIC_OPS || CUDNN_VERSION < 92600 - SKIP("FFT causal conv1d APIs require cuDNN 9.26.0 headers"); -#elif defined(_MSC_VER) - SKIP("FFT causal conv1d kernels are not supported on Windows (MSVC)"); -#else - skip_if_fft_causal_conv1d_is_unavailable(); - - constexpr int batch = 1; - constexpr int dim = 1; - constexpr int seq_len = 4096; - constexpr int kernel_size = 4096; - - size_t workspace_size = 0; - size_t reserve_size = 0; - CUDNN_CHECK(cudnnLongFFTCausalConv1dGetBufferSizes( - batch, dim, seq_len, kernel_size, CUDNN_DATA_FLOAT, &workspace_size, &reserve_size)); - - cudaStream_t stream = nullptr; - - Surface x_tensor(batch * dim * seq_len); - Surface weight_tensor(dim * kernel_size); - Surface y_tensor(batch * dim * seq_len); - Surface dy_tensor(batch * dim * seq_len); - Surface dx_tensor(batch * dim * seq_len); - Surface dweight_tensor(dim * kernel_size, 0.0f); - Surface workspace_tensor(workspace_size); - Surface reserve_tensor(reserve_size); - - CUDNN_CHECK(cudnnLongFFTCausalConv1dForward(stream, - x_tensor.devPtr, - weight_tensor.devPtr, - y_tensor.devPtr, - batch, - dim, - seq_len, - kernel_size, - CUDNN_DATA_FLOAT, - workspace_tensor.devPtr, - workspace_size, - reserve_tensor.devPtr, - reserve_size)); - - // Backward consumes the frequency-domain state written to reserve_tensor - // by this matching forward call. workspace_tensor is temporary scratch. - CUDNN_CHECK(cudnnLongFFTCausalConv1dBackward(stream, - dy_tensor.devPtr, - dx_tensor.devPtr, - dweight_tensor.devPtr, - batch, - dim, - seq_len, - kernel_size, - CUDNN_DATA_FLOAT, - workspace_tensor.devPtr, - workspace_size, - reserve_tensor.devPtr, - reserve_size)); - - CUDA_CHECK(cudaDeviceSynchronize()); -#endif -} diff --git a/samples/python/66_fft_causal_conv1d_forward.ipynb b/samples/python/66_fft_causal_conv1d_forward.ipynb deleted file mode 100644 index 89a4a0c4a..000000000 --- a/samples/python/66_fft_causal_conv1d_forward.ipynb +++ /dev/null @@ -1,200 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# FFT Causal Conv1D Forward\n", - "\n", - "This notebook exercises the cuDNN Frontend `fft_causal_conv1d(x, weight)` wrapper through both its medium and long FFT paths." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Prerequisites\n", - "\n", - "An NVIDIA GPU, cuDNN 9.26.0 or newer, and a cuDNN Frontend Python binding built against cuDNN 9.26.0 or newer are required." - ] - }, - { - "cell_type": "code", - "execution_count": 1, - "metadata": { - "execution": { - "iopub.execute_input": "2026-07-24T17:42:07.068707Z", - "iopub.status.busy": "2026-07-24T17:42:07.068455Z", - "iopub.status.idle": "2026-07-24T17:42:10.359705Z", - "shell.execute_reply": "2026-07-24T17:42:10.358876Z" - } - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "cuDNN backend version: 92600\n" - ] - } - ], - "source": [ - "import math\n", - "\n", - "import cudnn\n", - "import torch\n", - "import torch.nn.functional as F\n", - "\n", - "assert torch.cuda.is_available(), \"This sample requires a CUDA device\"\n", - "assert cudnn.backend_version() >= 92600, \"FFT causal conv1d requires cuDNN 9.26.0 or newer\"\n", - "print(\"cuDNN backend version:\", cudnn.backend_version())" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Reference\n", - "\n", - "FFT causal conv1d uses FIR-order weights: `weight[0]` multiplies the current sample. PyTorch `conv1d` is cross-correlation, so the reference reverses the filter before applying left-only padding." - ] - }, - { - "cell_type": "code", - "execution_count": 2, - "metadata": { - "execution": { - "iopub.execute_input": "2026-07-24T17:42:10.385555Z", - "iopub.status.busy": "2026-07-24T17:42:10.385307Z", - "iopub.status.idle": "2026-07-24T17:42:10.388310Z", - "shell.execute_reply": "2026-07-24T17:42:10.387722Z" - } - }, - "outputs": [], - "source": [ - "def fft_causal_conv1d_reference(x, weight):\n", - " dim, kernel_size = weight.shape\n", - " x_padded = F.pad(x, (kernel_size - 1, 0))\n", - " return F.conv1d(x_padded, weight.flip(-1).unsqueeze(1), groups=dim)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Medium and long path checks\n", - "\n", - "The medium case intentionally uses non-power-of-two `seq_len` and `kernel_size`, which validates the wrapper's padding and trimming. The FP64 case uses `kernel_size=8192`, above FP64's medium limit, to force the public wrapper through the long FFT path." - ] - }, - { - "cell_type": "code", - "execution_count": 3, - "metadata": { - "execution": { - "iopub.execute_input": "2026-07-24T17:42:10.389570Z", - "iopub.status.busy": "2026-07-24T17:42:10.389383Z", - "iopub.status.idle": "2026-07-24T17:42:48.772224Z", - "shell.execute_reply": "2026-07-24T17:42:48.771509Z" - } - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "medium-padded: shape=(2, 4, 750), max_abs=1.043081e-07\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "long: shape=(1, 1, 8192), max_abs=1.776357e-15\n" - ] - } - ], - "source": [ - "torch.manual_seed(42)\n", - "configs = [\n", - " {\"name\": \"medium-padded\", \"batch\": 2, \"dim\": 4, \"seq_len\": 750, \"kernel_size\": 192, \"dtype\": torch.float32, \"atol\": 2e-6, \"rtol\": 2e-6},\n", - " {\"name\": \"long\", \"batch\": 1, \"dim\": 1, \"seq_len\": 8192, \"kernel_size\": 8192, \"dtype\": torch.float64, \"atol\": 5e-11, \"rtol\": 5e-11},\n", - "]\n", - "\n", - "for config in configs:\n", - " x = 0.1 * torch.randn(config[\"batch\"], config[\"dim\"], config[\"seq_len\"], device=\"cuda\", dtype=config[\"dtype\"])\n", - " weight = torch.randn(config[\"dim\"], config[\"kernel_size\"], device=\"cuda\", dtype=config[\"dtype\"]) / math.sqrt(config[\"kernel_size\"])\n", - "\n", - " actual = cudnn.ops.fft_causal_conv1d(x, weight)\n", - " expected = fft_causal_conv1d_reference(x.double(), weight.double()).to(config[\"dtype\"])\n", - " max_abs = (actual - expected).abs().max().item()\n", - " print(f'{config[\"name\"]}: shape={tuple(actual.shape)}, max_abs={max_abs:.6e}')\n", - " torch.testing.assert_close(actual, expected, atol=config[\"atol\"], rtol=config[\"rtol\"])" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## `torch.compile`\n", - "\n", - "The registered custom operators preserve the same result under `torch.compile`." - ] - }, - { - "cell_type": "code", - "execution_count": 4, - "metadata": { - "execution": { - "iopub.execute_input": "2026-07-24T17:42:48.774184Z", - "iopub.status.busy": "2026-07-24T17:42:48.773955Z", - "iopub.status.idle": "2026-07-24T17:43:02.389981Z", - "shell.execute_reply": "2026-07-24T17:43:02.389343Z" - } - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "torch.compile output matches eager output\n" - ] - } - ], - "source": [ - "@torch.compile\n", - "def compiled_fft_causal_conv1d(x, weight):\n", - " return cudnn.ops.fft_causal_conv1d(x, weight)\n", - "\n", - "\n", - "x = torch.randn(1, 2, 512, device=\"cuda\")\n", - "weight = torch.randn(2, 128, device=\"cuda\")\n", - "eager = cudnn.ops.fft_causal_conv1d(x, weight)\n", - "compiled = compiled_fft_causal_conv1d(x, weight)\n", - "torch.testing.assert_close(compiled, eager, atol=0.0, rtol=0.0)\n", - "print(\"torch.compile output matches eager output\")" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3 (ipykernel)", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.14.0" - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} diff --git a/samples/python/67_fft_causal_conv1d_backward.ipynb b/samples/python/67_fft_causal_conv1d_backward.ipynb deleted file mode 100644 index 498f035e4..000000000 --- a/samples/python/67_fft_causal_conv1d_backward.ipynb +++ /dev/null @@ -1,146 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# FFT Causal Conv1D Backward\n", - "\n", - "This notebook verifies input and filter gradients for the medium and long FFT paths exposed by `cudnn.ops.fft_causal_conv1d`." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Prerequisites\n", - "\n", - "An NVIDIA GPU, cuDNN 9.26.0 or newer, and a cuDNN Frontend Python binding built against cuDNN 9.26.0 or newer are required." - ] - }, - { - "cell_type": "code", - "execution_count": 1, - "metadata": { - "execution": { - "iopub.execute_input": "2026-07-24T17:43:20.257664Z", - "iopub.status.busy": "2026-07-24T17:43:20.257549Z", - "iopub.status.idle": "2026-07-24T17:43:22.491087Z", - "shell.execute_reply": "2026-07-24T17:43:22.490263Z" - } - }, - "outputs": [], - "source": [ - "import math\n", - "\n", - "import cudnn\n", - "import torch\n", - "import torch.nn.functional as F\n", - "\n", - "assert torch.cuda.is_available(), \"This sample requires a CUDA device\"\n", - "assert cudnn.backend_version() >= 92600, \"FFT causal conv1d requires cuDNN 9.26.0 or newer\"" - ] - }, - { - "cell_type": "code", - "execution_count": 2, - "metadata": { - "execution": { - "iopub.execute_input": "2026-07-24T17:43:22.492808Z", - "iopub.status.busy": "2026-07-24T17:43:22.492607Z", - "iopub.status.idle": "2026-07-24T17:43:22.495427Z", - "shell.execute_reply": "2026-07-24T17:43:22.494984Z" - } - }, - "outputs": [], - "source": [ - "def fft_causal_conv1d_reference(x, weight):\n", - " dim, kernel_size = weight.shape\n", - " x_padded = F.pad(x, (kernel_size - 1, 0))\n", - " return F.conv1d(x_padded, weight.flip(-1).unsqueeze(1), groups=dim)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The long backward call consumes opaque frequency-domain reserve space saved by its matching forward call. The Python autograd registration manages that reserve-space lifetime automatically." - ] - }, - { - "cell_type": "code", - "execution_count": 3, - "metadata": { - "execution": { - "iopub.execute_input": "2026-07-24T17:43:22.496766Z", - "iopub.status.busy": "2026-07-24T17:43:22.496648Z", - "iopub.status.idle": "2026-07-24T17:44:13.624996Z", - "shell.execute_reply": "2026-07-24T17:44:13.624307Z" - } - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "medium-padded: dx_max_abs=9.984049e-08, dw_max_abs=3.097257e-07\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "long: dx_max_abs=1.804112e-15, dw_max_abs=1.132427e-14\n" - ] - } - ], - "source": [ - "torch.manual_seed(42)\n", - "configs = [\n", - " {\"name\": \"medium-padded\", \"batch\": 2, \"dim\": 4, \"seq_len\": 750, \"kernel_size\": 192, \"dtype\": torch.float32, \"atol\": 2e-6, \"rtol\": 2e-6},\n", - " {\"name\": \"long\", \"batch\": 1, \"dim\": 1, \"seq_len\": 8192, \"kernel_size\": 8192, \"dtype\": torch.float64, \"atol\": 5e-11, \"rtol\": 5e-11},\n", - "]\n", - "\n", - "for config in configs:\n", - " x_data = 0.1 * torch.randn(config[\"batch\"], config[\"dim\"], config[\"seq_len\"], device=\"cuda\", dtype=config[\"dtype\"])\n", - " weight_data = torch.randn(config[\"dim\"], config[\"kernel_size\"], device=\"cuda\", dtype=config[\"dtype\"]) / math.sqrt(config[\"kernel_size\"])\n", - " grad_out = 0.1 * torch.randn_like(x_data)\n", - "\n", - " x = x_data.detach().requires_grad_(True)\n", - " weight = weight_data.detach().requires_grad_(True)\n", - " cudnn.ops.fft_causal_conv1d(x, weight).backward(grad_out)\n", - "\n", - " x_ref = x_data.double().detach().requires_grad_(True)\n", - " weight_ref = weight_data.double().detach().requires_grad_(True)\n", - " fft_causal_conv1d_reference(x_ref, weight_ref).backward(grad_out.double())\n", - "\n", - " dx_max_abs = (x.grad.double() - x_ref.grad).abs().max().item()\n", - " dw_max_abs = (weight.grad.double() - weight_ref.grad).abs().max().item()\n", - " print(f'{config[\"name\"]}: dx_max_abs={dx_max_abs:.6e}, dw_max_abs={dw_max_abs:.6e}')\n", - " torch.testing.assert_close(x.grad, x_ref.grad.to(config[\"dtype\"]), atol=config[\"atol\"], rtol=config[\"rtol\"])\n", - " torch.testing.assert_close(weight.grad, weight_ref.grad.to(config[\"dtype\"]), atol=config[\"atol\"], rtol=config[\"rtol\"])" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3 (ipykernel)", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.14.0" - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} diff --git a/test/python/test_causal_conv1d.py b/test/python/test_causal_conv1d.py index 1aa31eaad..8d22ca4f1 100644 --- a/test/python/test_causal_conv1d.py +++ b/test/python/test_causal_conv1d.py @@ -1,11 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -import math - import pytest -import torch -import torch.nn.functional as F import cudnn @@ -94,206 +90,3 @@ def test_causal_conv1d_rejects_unsupported_kernel_size(binding_name, args, kerne message = rf"{api_name} {argument_name} must be between 2 and {max_kernel_size}, inclusive; got {kernel_size}" with pytest.raises(ValueError, match=message): binding(*args) - - -_TOLERANCES = { - torch.float64: (5e-12, 5e-12), - torch.float32: (5e-6, 5e-6), - torch.float16: (5e-3, 5e-3), - torch.bfloat16: (1e-2, 1e-2), -} - -# Mirrors the dtype, activation, and odd/even filter-width coverage in the -# causal conv1d NHW/NWH notebooks, extended with the new FP64 support. -_CAUSAL_CONV1D_CASES = [ - pytest.param(torch.float32, 4, "silu", id="fp32-k4-silu"), - pytest.param(torch.float16, 4, "silu", id="fp16-k4-silu"), - pytest.param(torch.bfloat16, 4, "silu", id="bf16-k4-silu"), - pytest.param(torch.bfloat16, 3, "silu", id="bf16-k3-silu"), - pytest.param(torch.float16, 7, "silu", id="fp16-k7-silu"), - pytest.param(torch.bfloat16, 4, "identity", id="bf16-k4-identity"), - pytest.param(torch.float64, 4, "silu", id="fp64-k4-silu"), - pytest.param(torch.float64, 4, "identity", id="fp64-k4-identity"), -] - - -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(getattr(cudnn, name, None) is None for name in required_symbols): - pytest.skip(f"SE causal conv1d case requires cuDNN {required_version // 10000}.{required_version // 100 % 100}") - - -def _make_tensor(shape, dtype, scale=1.0): - return (scale * torch.randn(*shape, device="cuda")).to(dtype) - - -def _causal_conv1d_reference(x, weight, bias=None, activation="identity"): - dim, kernel_size = weight.shape - y = F.conv1d(F.pad(x, (kernel_size - 1, 0)), weight.unsqueeze(1), bias=bias, groups=dim) - return F.silu(y) if activation == "silu" else y - - -def _causal_conv1d_nwh_reference(x, weight, bias=None, activation="identity"): - return _causal_conv1d_reference( - x.transpose(1, 2), - weight.transpose(0, 1), - bias, - activation, - ).transpose(1, 2) - - -def _b2b_causal_conv1d_reference(x, weights_proj, weights_mixer, skip_bias): - projected = _causal_conv1d_reference(x, weights_proj) - x1, x2, value = projected[:, 0::3], projected[:, 1::3], projected[:, 2::3] - gated = x2 * value - mixed = _causal_conv1d_reference(gated, weights_mixer) + skip_bias[None, :, None] * gated - return mixed * x1 - - -def _assert_close(actual, expected, dtype): - atol, rtol = _TOLERANCES[dtype] - torch.testing.assert_close(actual, expected.to(dtype), atol=atol, rtol=rtol) - - -@pytest.mark.L0 -@pytest.mark.parametrize("dtype,kernel_size,activation", _CAUSAL_CONV1D_CASES) -def test_causal_conv1d_autograd(dtype, kernel_size, activation): - _require_variant(("causal_conv1d_forward", "causal_conv1d_backward"), dtype, 92200) - torch.manual_seed(42) - batch, dim, seq_len = 2, 16, 128 - x_data = _make_tensor((batch, dim, seq_len), dtype, scale=0.1) - weight_data = _make_tensor((dim, kernel_size), dtype, scale=1.0 / math.sqrt(kernel_size)) - bias_data = _make_tensor((dim,), dtype, scale=0.1) - grad_out = _make_tensor((batch, dim, seq_len), dtype, scale=0.1) - - x = x_data.detach().requires_grad_(True) - weight = weight_data.detach().requires_grad_(True) - bias = bias_data.detach().requires_grad_(True) - actual = cudnn.ops.causal_conv1d(x, weight, bias, activation=activation) - actual.backward(grad_out) - - x_ref = x_data.double().detach().requires_grad_(True) - weight_ref = weight_data.double().detach().requires_grad_(True) - bias_ref = bias_data.double().detach().requires_grad_(True) - expected = _causal_conv1d_reference(x_ref, weight_ref, bias_ref, activation) - expected.backward(grad_out.double()) - - _assert_close(actual, expected, dtype) - _assert_close(x.grad, x_ref.grad, dtype) - _assert_close(weight.grad, weight_ref.grad, dtype) - _assert_close(bias.grad, bias_ref.grad, dtype) - - -@pytest.mark.L0 -@pytest.mark.parametrize("dtype,kernel_size,activation", _CAUSAL_CONV1D_CASES) -def test_causal_conv1d_nwh_autograd(dtype, kernel_size, activation): - _require_variant(("causal_conv1d_nwh_forward", "causal_conv1d_nwh_backward"), dtype, 92400) - torch.manual_seed(42) - batch, dim, seq_len = 2, 16, 128 - x_data = _make_tensor((batch, seq_len, dim), dtype, scale=0.1) - weight_data = _make_tensor((kernel_size, dim), dtype, scale=1.0 / math.sqrt(kernel_size)) - bias_data = _make_tensor((dim,), dtype, scale=0.1) - grad_out = _make_tensor((batch, seq_len, dim), dtype, scale=0.1) - - x = x_data.detach().requires_grad_(True) - weight = weight_data.detach().requires_grad_(True) - bias = bias_data.detach().requires_grad_(True) - actual = cudnn.ops.causal_conv1d_nwh(x, weight, bias, activation=activation) - actual.backward(grad_out) - - x_ref = x_data.double().detach().requires_grad_(True) - weight_ref = weight_data.double().detach().requires_grad_(True) - bias_ref = bias_data.double().detach().requires_grad_(True) - expected = _causal_conv1d_nwh_reference(x_ref, weight_ref, bias_ref, activation) - expected.backward(grad_out.double()) - - _assert_close(actual, expected, dtype) - _assert_close(x.grad, x_ref.grad, dtype) - _assert_close(weight.grad, weight_ref.grad, dtype) - _assert_close(bias.grad, bias_ref.grad, dtype) - - -@pytest.mark.L0 -@pytest.mark.parametrize( - "dtype", - [ - pytest.param(torch.float32, id="fp32"), - pytest.param(torch.float16, id="fp16"), - pytest.param(torch.bfloat16, id="bf16"), - pytest.param(torch.float64, id="fp64"), - ], -) -def test_b2b_causal_conv1d_autograd(dtype): - _require_variant(("b2b_causal_conv1d_forward", "b2b_causal_conv1d_backward"), dtype, 92400) - torch.manual_seed(42) - batch, dim, seq_len = 2, 8, 128 - kernel_size_proj, kernel_size_mixer = 4, 7 - x_data = _make_tensor((batch, 3 * dim, seq_len), dtype, scale=0.1) - weights_proj_data = _make_tensor( - (3 * dim, kernel_size_proj), - dtype, - scale=1.0 / math.sqrt(kernel_size_proj), - ) - weights_mixer_data = _make_tensor( - (dim, kernel_size_mixer), - dtype, - scale=1.0 / math.sqrt(kernel_size_mixer), - ) - skip_bias_data = _make_tensor((dim,), dtype, scale=0.1) - grad_out = _make_tensor((batch, dim, seq_len), dtype, scale=0.1) - - x = x_data.detach().requires_grad_(True) - weights_proj = weights_proj_data.detach().requires_grad_(True) - weights_mixer = weights_mixer_data.detach().requires_grad_(True) - skip_bias = skip_bias_data.detach().requires_grad_(True) - actual = cudnn.ops.b2b_causal_conv1d(x, weights_proj, weights_mixer, skip_bias) - actual.backward(grad_out) - - x_ref = x_data.double().detach().requires_grad_(True) - weights_proj_ref = weights_proj_data.double().detach().requires_grad_(True) - weights_mixer_ref = weights_mixer_data.double().detach().requires_grad_(True) - skip_bias_ref = skip_bias_data.double().detach().requires_grad_(True) - expected = _b2b_causal_conv1d_reference( - x_ref, - weights_proj_ref, - weights_mixer_ref, - skip_bias_ref, - ) - expected.backward(grad_out.double()) - - _assert_close(actual, expected, dtype) - _assert_close(x.grad, x_ref.grad, dtype) - _assert_close(weights_proj.grad, weights_proj_ref.grad, dtype) - _assert_close(weights_mixer.grad, weights_mixer_ref.grad, dtype) - _assert_close(skip_bias.grad, skip_bias_ref.grad, dtype) - - -@pytest.mark.L0 -@pytest.mark.parametrize("layout", ["nhw", "nwh"]) -def test_causal_conv1d_compiled_autograd(layout): - symbols = ("causal_conv1d_forward", "causal_conv1d_backward") if layout == "nhw" else ("causal_conv1d_nwh_forward", "causal_conv1d_nwh_backward") - _require_variant(symbols, torch.bfloat16, 92200 if layout == "nhw" else 92400) - torch.manual_seed(42) - batch, dim, seq_len, kernel_size = 1, 16, 64, 4 - if layout == "nhw": - x_data = _make_tensor((batch, dim, seq_len), torch.bfloat16, scale=0.1) - weight_data = _make_tensor((dim, kernel_size), torch.bfloat16, scale=0.5) - op = cudnn.ops.causal_conv1d - else: - x_data = _make_tensor((batch, seq_len, dim), torch.bfloat16, scale=0.1) - weight_data = _make_tensor((kernel_size, dim), torch.bfloat16, scale=0.5) - op = cudnn.ops.causal_conv1d_nwh - bias_data = _make_tensor((dim,), torch.bfloat16, scale=0.1) - grad_out = _make_tensor(x_data.shape, torch.bfloat16, scale=0.1) - - def train_step(x, weight, bias, dy): - x = x.detach().requires_grad_(True) - weight = weight.detach().requires_grad_(True) - bias = bias.detach().requires_grad_(True) - op(x, weight, bias, activation="silu").backward(dy) - return x.grad, weight.grad, bias.grad - - eager = train_step(x_data, weight_data, bias_data, grad_out) - compiled = torch.compile(train_step)(x_data, weight_data, bias_data, grad_out) - for eager_grad, compiled_grad in zip(eager, compiled): - torch.testing.assert_close(compiled_grad, eager_grad, atol=0, rtol=0) diff --git a/test/python/test_fft_causal_conv1d.py b/test/python/test_fft_causal_conv1d.py deleted file mode 100644 index f71890239..000000000 --- a/test/python/test_fft_causal_conv1d.py +++ /dev/null @@ -1,99 +0,0 @@ -import math - -import pytest -import torch -import torch.nn.functional as F - -import cudnn -from cudnn.ops.fft_causal_conv1d import _long_buffer_size_bytes, fft_causal_conv1d - - -def _require_fft_causal_conv1d(): - required_symbols = ( - "fft_causal_conv1d_forward", - "fft_causal_conv1d_backward", - "long_fft_causal_conv1d_get_buffer_sizes", - "long_fft_causal_conv1d_forward", - "long_fft_causal_conv1d_backward", - ) - if cudnn.backend_version() < 92600 or any(getattr(cudnn, name, None) is None for name in required_symbols): - pytest.skip("FFT causal conv1d requires cuDNN 9.26.0 or newer bindings and backend") - - -def _reference(x, weight): - dim, kernel_size = weight.shape - return F.conv1d(F.pad(x, (kernel_size - 1, 0)), weight.flip(-1).unsqueeze(1), groups=dim) - - -def _make_inputs(batch, dim, seq_len, kernel_size, dtype): - x = (0.1 * torch.randn(batch, dim, seq_len, device="cuda")).to(dtype) - weight = (torch.randn(dim, kernel_size, device="cuda") / math.sqrt(kernel_size)).to(dtype) - return x, weight - - -@pytest.mark.L0 -@pytest.mark.parametrize( - "batch,dim,seq_len,kernel_size,dtype,atol,rtol", - [ - (2, 4, 750, 192, torch.float32, 2e-6, 2e-6), # Medium path pads both input and filter. - # FP64 filters above K=4096 select long FFT on every supported architecture. - (1, 1, 8192, 8192, torch.float64, 5e-11, 5e-11), - ], -) -def test_fft_causal_conv1d_forward_and_backward(batch, dim, seq_len, kernel_size, dtype, atol, rtol): - _require_fft_causal_conv1d() - x_data, weight_data = _make_inputs(batch, dim, seq_len, kernel_size, dtype) - grad_out = 0.1 * torch.randn_like(x_data) - - x = x_data.detach().requires_grad_(True) - weight = weight_data.detach().requires_grad_(True) - actual = fft_causal_conv1d(x, weight) - actual.backward(grad_out) - - # FP64 avoids TF32 convolution/reduction error masking the FFT error. - x_ref = x_data.double().detach().requires_grad_(True) - weight_ref = weight_data.double().detach().requires_grad_(True) - expected = _reference(x_ref, weight_ref) - expected.backward(grad_out.double()) - - torch.testing.assert_close(actual, expected.to(dtype), atol=atol, rtol=rtol) - torch.testing.assert_close(x.grad, x_ref.grad.to(dtype), atol=atol, rtol=rtol) - torch.testing.assert_close(weight.grad, weight_ref.grad.to(dtype), atol=atol, rtol=rtol) - - -@pytest.mark.L0 -@pytest.mark.parametrize( - "dtype,atol,rtol", - [ - (torch.float64, 5e-12, 5e-12), - (torch.float32, 2e-6, 2e-6), - (torch.float16, 5e-3, 5e-3), - (torch.bfloat16, 1e-2, 1e-2), - ], -) -def test_medium_fft_supported_dtypes(dtype, atol, rtol): - _require_fft_causal_conv1d() - x_data, weight_data = _make_inputs(1, 2, 256, 128, dtype) - grad_out = (0.1 * torch.randn(1, 2, 256, device="cuda")).to(dtype) - - x = x_data.detach().requires_grad_(True) - weight = weight_data.detach().requires_grad_(True) - actual = fft_causal_conv1d(x, weight) - actual.backward(grad_out) - - x_ref = x_data.double().detach().requires_grad_(True) - weight_ref = weight_data.double().detach().requires_grad_(True) - expected = _reference(x_ref, weight_ref) - expected.backward(grad_out.double()) - - torch.testing.assert_close(actual, expected.to(dtype), atol=atol, rtol=rtol) - torch.testing.assert_close(x.grad, x_ref.grad.to(dtype), atol=atol, rtol=rtol) - torch.testing.assert_close(weight.grad, weight_ref.grad.to(dtype), atol=atol, rtol=rtol) - - -@pytest.mark.L0 -def test_long_fft_fake_buffer_formula_matches_backend_query(): - _require_fft_causal_conv1d() - batch, dim, seq_len, kernel_size = 2, 3, 4096, 4096 - _, reserve_size = cudnn.long_fft_causal_conv1d_get_buffer_sizes(batch, dim, seq_len, kernel_size, 0) # CUDNN_DATA_FLOAT - assert reserve_size == _long_buffer_size_bytes(batch, dim, kernel_size, torch.float32)