diff --git a/docs/operations/CausalConv1d.md b/docs/operations/CausalConv1d.md index f1b170ae8..7b69c2af7 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**: Blackwell (SM100+) -- **Data types**: FP32, FP16, BF16 +- **Architectures**: Turing (SM75) and newer, subject to available dynamic shared memory +- **Data types**: FP64, FP32, FP16, BF16 - **Activations**: `identity`, `silu` ## Python API @@ -54,7 +54,7 @@ In most cases you should use `cudnn.ops.causal_conv1d` instead, which handles au | Tensor | Device | Data Type | Shape | |--------|--------|-----------|-------| -| x | GPU | FP32, FP16, or BF16 | $(B, D, L)$ | +| x | GPU | FP64, FP32, FP16, or BF16 | $(B, D, L)$ | | weight | GPU | same as x | $(D, K)$ | | bias | GPU | same as x | $(D,)$ | @@ -78,6 +78,10 @@ In most cases you should use `cudnn.ops.causal_conv1d` instead, which handles au | dweight | GPU | same as x | $(D, K)$ | | dbias | GPU | same as x | $(D,)$ | +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. + ## Example ```python diff --git a/docs/operations/FFTCausalConv1d.md b/docs/operations/FFTCausalConv1d.md new file mode 100644 index 000000000..330f0d426 --- /dev/null +++ b/docs/operations/FFTCausalConv1d.md @@ -0,0 +1,71 @@ +# 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 521b2a4e1..7ae5f61ff 100644 --- a/include/cudnn_frontend_shim.h +++ b/include/cudnn_frontend_shim.h @@ -781,6 +781,159 @@ 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 5439d4a8f..2b9479c7c 100644 --- a/llms.txt +++ b/llms.txt @@ -22,6 +22,7 @@ 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 1a4d380dc..78cfbc811 100644 --- a/python/cudnn/__init__.py +++ b/python/cudnn/__init__.py @@ -53,6 +53,11 @@ 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 5518db097..04cc009ca 100644 --- a/python/cudnn/ops/__init__.py +++ b/python/cudnn/ops/__init__.py @@ -1 +1,2 @@ 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 74512044b..a200ba984 100644 --- a/python/cudnn/ops/causal_conv1d.py +++ b/python/cudnn/ops/causal_conv1d.py @@ -5,6 +5,7 @@ _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 } @@ -17,10 +18,16 @@ def _dtype_to_int(dtype: torch.dtype) -> int: if dtype not in _TORCH_DTYPE_TO_CUDNN: - raise ValueError(f"Unsupported dtype {dtype}. Supported: float32, float16, bfloat16.") + raise ValueError(f"Unsupported dtype {dtype}. Supported: float64, 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'.") @@ -133,8 +140,9 @@ def _bwd_primitive(grad_out: Tensor, x: Tensor, weight: Tensor, bias: Tensor, ac kernel_size = weight.shape[1] dx = torch.empty_like(x) - dweight = torch.zeros(weight.shape, device=x.device, dtype=torch.float32) - dbias = torch.zeros(bias.shape, device=x.device, dtype=torch.float32) + 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) import cudnn @@ -152,7 +160,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(torch.float32), + _dtype_to_int(grad_dtype), _activation_to_int(activation), ) return [dx, dweight.to(x.dtype), dbias.to(x.dtype)] @@ -214,7 +222,7 @@ def causal_conv1d( Args: x (torch.Tensor): Input tensor of shape ``(batch, dim, seq_len)``. - Must be BF16, FP16, or FP32. Must be contiguous and on CUDA. + Must be BF16, FP16, FP32, or FP64. Must be contiguous and on CUDA. weight (torch.Tensor): Filter tensor of shape ``(dim, kernel_size)``. Same dtype as *x*. bias (torch.Tensor | None): Optional bias of shape ``(dim,)``. @@ -342,8 +350,9 @@ def _nwh_bwd_primitive(grad_out: Tensor, x: Tensor, weight: Tensor, bias: Tensor kernel_size = weight.shape[0] dx = torch.empty_like(x) - dweight = torch.zeros(weight.shape, device=x.device, dtype=torch.float32) - dbias = torch.zeros(bias.shape, device=x.device, dtype=torch.float32) + 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) import cudnn @@ -361,7 +370,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(torch.float32), + _dtype_to_int(grad_dtype), _activation_to_int(activation), ) return [dx, dweight.to(x.dtype), dbias.to(x.dtype)] @@ -416,6 +425,7 @@ 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)``. bias (torch.Tensor | None): Optional bias of shape ``(dim,)``. activation (str): ``"identity"`` (default) or ``"silu"``. @@ -584,9 +594,10 @@ def _b2b_bwd_primitive( kernel_size_mixer = weights_mixer.shape[1] dx = torch.empty_like(x) - 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) + 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) import cudnn @@ -608,7 +619,7 @@ def _b2b_bwd_primitive( kernel_size_proj, kernel_size_mixer, _dtype_to_int(x.dtype), - _dtype_to_int(torch.float32), + _dtype_to_int(grad_dtype), ) return [ dx, @@ -685,6 +696,7 @@ 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)``. weights_mixer (torch.Tensor): Mixer filter ``(dim, kernel_size_mixer)``. skip_bias (torch.Tensor): Skip-connection bias ``(dim,)``. diff --git a/python/cudnn/ops/fft_causal_conv1d.py b/python/cudnn/ops/fft_causal_conv1d.py new file mode 100644 index 000000000..41701d250 --- /dev/null +++ b/python/cudnn/ops/fft_causal_conv1d.py @@ -0,0 +1,369 @@ +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 6b6492f61..829e894c0 100644 --- a/python/pycudnn.cpp +++ b/python/pycudnn.cpp @@ -289,6 +289,139 @@ PYBIND11_MODULE(_compiled_module, m) { if (status != 0) throw std::runtime_error("cudnnB2BCausalConv1dBackward failed with status " + std::to_string(status)); }); + +#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 20dd24752..029e9214b 100644 --- a/samples/cpp/CMakeLists.txt +++ b/samples/cpp/CMakeLists.txt @@ -76,6 +76,7 @@ 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 new file mode 100644 index 000000000..61f8727b6 --- /dev/null +++ b/samples/cpp/causal_conv1d/fft_causal_conv1d.cpp @@ -0,0 +1,162 @@ +/* + * 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 new file mode 100644 index 000000000..89a4a0c4a --- /dev/null +++ b/samples/python/66_fft_causal_conv1d_forward.ipynb @@ -0,0 +1,200 @@ +{ + "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 new file mode 100644 index 000000000..498f035e4 --- /dev/null +++ b/samples/python/67_fft_causal_conv1d_backward.ipynb @@ -0,0 +1,146 @@ +{ + "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 new file mode 100644 index 000000000..9b278c566 --- /dev/null +++ b/test/python/test_causal_conv1d.py @@ -0,0 +1,209 @@ +import math + +import pytest +import torch +import torch.nn.functional as F + +import cudnn + +_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(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}") + + +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 new file mode 100644 index 000000000..32acba7ae --- /dev/null +++ b/test/python/test_fft_causal_conv1d.py @@ -0,0 +1,99 @@ +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(not hasattr(cudnn, name) 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)