Views
No views yet
torch.compile compliance via pt2_compliant_tag +
needs_fixed_stride_order.pip install kernels torch1import torch
2from kernels import get_kernel
3
4hotdog = get_kernel("chrisvoncsefalvay/hotdog")
5
6B, N, L = 4, 1024, 1024
7v = torch.randn(B, N, device="cuda", dtype=torch.float32)
8omega = torch.linspace(0.5, 1.5, L, device="cuda", dtype=torch.float32)
9lambd = torch.linspace(-1.5, -0.5, N, device="cuda", dtype=torch.float32)
10
11out = torch.empty(B, L, device="cuda", dtype=torch.float32)
12hotdog.cauchy_kernel_forward(out, v, omega, lambd)kernels client picks the pre-built binary that matches your Python,
PyTorch, and CUDA versions. For day-to-day PyTorch use prefer the top-level
Python API (see below) rather than calling the raw ops.1hotdog.cauchy_kernel(v, omega, lambd) # Cauchy kernel evaluation
2hotdog.toeplitz_matvec(c, r, x) # y = T(c, r) @ x
3hotdog.toeplitz_matvec_transpose(c, r, x) # y = T(c, r)^T @ x
4hotdog.toeplitz_solve(c, r, b) # T(c, r) x = b
5hotdog.toeplitz_solve_symmetric(c, b) # T(c, c) x = b (fast path)
6hotdog.toeplitz_inverse_apply(gs_u, gs_v, x) # T^{-1} x via Gohberg-Semencul
7hotdog.gs_compute_generators(c, r) # Gohberg-Semencul generators
8hotdog.toeplitz_to_cauchy(c, r) # C = F T F^{-1}
9hotdog.displacement_matvec(G, H, x, disp_type=0) # y = A x, generator-aware
10hotdog.cauchy_pin_l2(omega, lambd) # L2 persistence pinstream= keyword argument; forward-only
entries additionally accept an out= buffer. Calls broadcast leading
batch dimensions and stay correct under torch.compile and
torch.cuda.graph capture. Only Sylvester-type displacement
(disp_type=0) is supported today. nn.Module wrappers (CauchyKernel,
ToeplitzMV, ToeplitzSolve, ToeplitzInverse, DisplacementMV) are
provided for parameterised layers that store structured matrices as
buffers.__nv_bfloat162 pair loads on the BF16 path, $+1$ bank-padded
shared tiles, and a dedicated complex (float2) kernel for S4 poles.
Custom CUDA backward over $v$, $\omega$, $\lambda$.cudaFuncSetAttribute up to the device's
sharedMemPerBlockOptin ceiling.[U(Jv)|U(v)] then [L(u)|L(Ju)]) with shared $x$ FFT, record_stream
discipline on every intermediate, and a small-$N$ shared-memory FFT route.cauchy_pin_l2: pins $\omega$ and $\lambda$ into the L2 persistence
window on the current stream for repeated-call workloads (Ampere+).| Operation | Shape ($B\times n$) | HOTDOG | vs reference | Achieved |
|---|---|---|---|---|
cauchy_kernel | 16x4096 | 0.55 ms | 25x naive | 1.87 TF/s |
cauchy_kernel | 16x1024 | 0.08 ms | 8.3x naive | 855 GF/s |
toeplitz_matvec | 4x256 | 0.031 ms | 1.6x FFT | 9 GF/s |
toeplitz_solve_symmetric | 4x64 | 0.084 ms | 54.7x SciPy | — |
toeplitz_solve_symmetric | 4x256 | 0.278 ms | 16.5x SciPy | — |
toeplitz_inverse_apply | 4x256 | 0.344 ms | 28.4x over 4x FFT | — |
displacement_matvec ($r=2$) | 4x256 | 0.101 ms | 1.9x dense | 44 GF/s |
torch.cuda.graph, op fusion,
cauchy_pin_l2), not a kernel choice.record_stream discipline.__ldca, __ldcs, __stwt) plus
bank-conflict padding on shared tiles.pt2_compliant_tag, needs_fixed_stride_order, and
Meta (FakeTensor) implementations on every op.| Op | float32 | bfloat16 | float16 | complex64 |
|---|---|---|---|---|
cauchy_kernel (forward) | ✓ | ✓ | — | ✓ |
cauchy_kernel (backward) | ✓ | — | — | — |
toeplitz_matvec | ✓ | ✓ | ✓ | — |
toeplitz_solve_* | ✓ | — | — | — |
toeplitz_inverse_apply | ✓ | — | — | — |
displacement_matvec | ✓ | — | — | — |
toeplitz_to_cauchy | ✓ | — | — | — |
kernels-community publishes against. For local
development and additional arches, build with kernel-builder:nix run -L --max-jobs 1 --cores 4 .#build-and-copy