Ops#

Low-level convolution primitives. Pure-PyTorch reference implementations double as the spec the CUDA kernels must match; the subquadratic_ops_torch wrappers are the production path on GPU.

FFT convolutions (reference fp32)#

Use these for correctness and as the spec for the CUDA kernels below.

fftconv1d_fp32_blh(x, kernel[, shortcut])

Non-causal 1D FFT convolution (BLH layout, channels-last) with optional shortcut.

fftconv2d_fp32_blh(x, kernel[, shortcut])

2D FFT convolution with optional shortcut.

fftconv3d_fp32_blh(x, kernel[, shortcut])

3D FFT convolution with optional shortcut.

causal_fftconv1d_fp32_blh(x, kernel[, shortcut])

Causal 1D FFT convolution (BLH layout, channels-last) with optional shortcut.

fftconv1d_fp32_bhl(x, kernel[, shortcut])

1D FFT convolution with optional shortcut, for inputs with layout (batch, hidden, length).

fftconv2d_fp32_bhl(x, kernel[, shortcut])

2D FFT convolution with optional shortcut, for inputs with layout (batch, hidden, height, width).

fftconv3d_fp32_bhl(x, kernel[, shortcut])

3D FFT convolution with optional shortcut, for inputs with layout (batch, hidden, depth, height, width).

causal_fftconv1d_fp32_bhl(x, kernel[, shortcut])

1D FFT convolution with optional shortcut, for inputs with layout (batch, hidden, length).

FFT convolutions (CUDA-accelerated)#

Drop-in wrappers around the subquadratic_ops_torch fused CUDA kernels. 2D non-causal and 1D causal long-conv variants share the same API as the fp32 reference ops above.

fftconv2d_blh(x, kernel[, shortcut])

Alias for fftconv2d_bhl_w_reshape().

fftconv2d_bhl(x, kernel[, shortcut])

2D FFT convolution via subq_ops CUDA kernel, BHL layout [B, H, X, Y].

fftconv2d_bhl_w_reshape(x, kernel[, shortcut])

2D FFT convolution via subq_ops for BLH inputs [B, X, Y, H].

causal_fftconv1d_blh(x, kernel[, shortcut])

Alias for causal_fftconv1d_bhl_w_reshape().

causal_fftconv1d_bhl(x, kernel[, shortcut])

1D causal FFT convolution via subq_ops CUDA kernel, BHL layout [B, H, L].

causal_fftconv1d_bhl_w_reshape(x, kernel[, ...])

1D causal FFT convolution via subq_ops for BLH inputs [B, L, H].

Fused 2D FFT convolutions (native dtype)#

Wrappers around subquadratic_ops_torch.fused_fft_conv2d, which fuses the whole rfft2/multiply/irfft2 pipeline into one launch and runs it in the input dtype (fp32/fp16/bf16) rather than upcasting to fp32. Spatial dims are capped at 64 per axis. Selected on CKConvND via fft_backend="subq_ops_fused".

Note

Spatial extents above 32 per axis select the 128 FFT tile, which requires compute capability 9.0+ (Hopper/Blackwell); SM80/SM86 raise a clear error.

fused_fftconv2d_blh(x, kernel[, shortcut, ...])

Alias for fused_fftconv2d_bhl_w_reshape().

fused_fftconv2d_bhl(x, kernel[, shortcut, ...])

Fused 2D FFT convolution, BHL layout [B, H, X, Y], native dtype.

fused_fftconv2d_bhl_w_reshape(x, kernel[, ...])

Fused 2D FFT convolution for BLH inputs [B, X, Y, H].

fused_fftconv2d_bhl_chunked(x, kernel[, ...])

Channel-chunked fused 2D FFT convolution, BHL layout.

fused_fftconv2d_blh_chunked(x, kernel[, ...])

Alias for fused_fftconv2d_bhl_w_reshape_chunked().

fused_fftconv2d_bhl_w_reshape_chunked(x, kernel)

Channel-chunked fused 2D FFT convolution for BLH inputs [B, X, Y, H].

resolve_fused_fft_size(x_dim, y_dim, k_x, k_y)

Pick the FFT tile size for a fused 2D conv, or validate an explicit one.

fused_fftconv2d_supported(x_dim, y_dim, k_x, k_y)

Whether the fused 2D kernel can handle this input/kernel shape combination.

fused_fftconv2d_max_spatial()

Largest supported spatial extent per axis for the fused 2D kernel.

fused_fftconv2d_arch_supported([device, ...])

Whether device can run the fused 2D kernel at this FFT tile size.

load_fused_fft_conv2d()

Import the fused CUDA kernel now instead of on first call.

torch.compile lowering#

Inductor pre-grad pass that rewrites the reference 2D FFT-conv chain onto the fused CUDA kernel, so an existing fft_backend="torch_fft" model picks it up without a config change.

fused_fftconv2d_options([...])

Return a torch.compile(options=...) dict that enables the lowering.

fused_fftconv2d_lowering([...])

Context manager that installs FusedFFTConv2dLowering globally.

lowering_stats()

Return a snapshot of rewrite counters since the last reset.

reset_lowering_stats()

Clear the rewrite counters returned by lowering_stats().

FusedFFTConv2dLowering([allow_reduced_precision])

Inductor pre-grad pass replacing the 2D FFT-conv chain with the fused kernel.

Direct 1D causal convolutions (CUDA-accelerated)#

Non-FFT CUDA kernels for short and fused 1D causal convolutions. Useful for small kernel sizes (where FFT overhead dominates) and as building blocks for fused Hyena variants.

causal_conv1d(x, weight[, bias, activation])

Depthwise causal 1D conv via the subq_ops CUDA kernel.

b2b_causal_conv1d(x, weight_proj, ...)

Back-to-back fused causal 1D conv via the subq_ops CUDA kernel.

Circular FFT convolutions#

Periodic-boundary FFT convolutions for global mixing without zero padding.

circular_fftconv1d_fp32_bhl(x, kernel[, ...])

1D circular FFT convolution with optional shortcut (BHL layout).

circular_fftconv2d_fp32_bhl(x, kernel[, ...])

2D circular FFT convolution with optional shortcut (BHL layout).

circular_fftconv3d_fp32_bhl(x, kernel[, ...])

3D circular FFT convolution with optional shortcut (BHL layout).

Chunking utilities#

Helpers to bound the FFT working-set memory by processing along the sequence axis in chunks.

enable_chunking([module_or_flag, chunk_size])

Enable chunked FFT conv globally, as decorator, or as context manager.

chunking_enabled([enabled, chunk_size])

Context manager to temporarily enable/disable chunked FFT conv.

set_default_chunk_size(chunk_size)

Set the default chunk size for chunked FFT convolutions.

get_default_chunk_size()

Return the default chunk size for chunked FFT convolutions.

Mixed boundary-condition FFT convolutions#

FFT convolutions with per-axis boundary conditions: periodic on some spatial axes, zero-padded on others. See Mixed Boundary-Condition FFT Convolution for the per-axis algorithm and the fft_padding API.

mixed_fftconv1d_fp32_bhl(x, kernel, periodic)

1D mixed-BC FFT convolution (BHL layout).

mixed_fftconv2d_fp32_bhl(x, kernel, periodic)

2D mixed-BC FFT convolution (BHL layout).

mixed_fftconv3d_fp32_bhl(x, kernel, periodic)

3D mixed-BC FFT convolution (BHL layout).

mixed_fftconv1d_fp32_bhl_w_reshape(x, ...[, ...])

1D mixed-BC FFT conv wrapper for BLH layout (batch, length, hidden).

mixed_fftconv2d_fp32_bhl_w_reshape(x, ...[, ...])

2D mixed-BC FFT conv wrapper for BLH layout (batch, X, Y, hidden).

mixed_fftconv3d_fp32_bhl_w_reshape(x, ...[, ...])

3D mixed-BC FFT conv wrapper for BLH layout (batch, X, Y, Z, hidden).

mixed_fftconv1d_fp32_bhl_chunked(x, kernel, ...)

Memory-efficient 1D mixed-BC FFT conv (BHL) via channel chunking.

mixed_fftconv2d_fp32_bhl_chunked(x, kernel, ...)

Memory-efficient 2D mixed-BC FFT conv (BHL) via channel chunking.

mixed_fftconv3d_fp32_bhl_chunked(x, kernel, ...)

Memory-efficient 3D mixed-BC FFT conv (BHL) via channel chunking.