dspk

[!WARNING] Early development (0.0.x). APIs change between releases.

Documentation · Python API · Rust API · Changelog

Multi-channel signal processing written once in CubeCL and run on any GPU (wgpu: Vulkan, Metal, DirectX 12; CUDA; ROCm / HIP) or the CPU, with the device chosen at run time. Processing stages keep data on the device between them, so a recording crosses disk → RAM → device memory as few times as possible.

dspk-app (desktop workbench, cargo install dspk-app): a Neuropixels recording (385 channels, 30 kHz) drawn from min/max pyramids built by dspk-view.

Crates and folders

Every crate builds with wgpu, CUDA and the CPU runtime by default. Development and tests use one profile:

cargo build --workspace --profile validate
cargo test  --workspace --exclude dspk-app --exclude dspk-py --profile validate

Each kernel is tested on every runtime compiled in (DSPK_TEST_RUNTIMES=cuda limits a run to one). Timings against scipy and kernel benchmarks: benchmarks.

Runtimes

Runtimes are crate features; the device is picked at run time with dspk::core::ComputeTarget (or the DSPK_RUNTIME environment variable) among those compiled in.

[dependencies]
dspk = "0.0.10"                                                                  # wgpu, cuda and cpu (default)
# dspk = { version = "0.0.10", default-features = false, features = ["wgpu", "zarr"] }  # wgpu only

CUDA and HIP are loaded at run time, so building with them needs no toolkit installed. The CPU runtime downloads a prebuilt LLVM on its first build (cached in ~/.local/share/tracel); leave it out with default-features = false for a lighter build. All three are the default because the library, its tests, the CLI and the Python package are built and tested together that way.

Performance

Device time of one recording-sized block (384 channels × 60 000 samples, float32: 2 s of Neuropixels at 30 kHz), data already on the device, autotuned kernel; RTX 2070 Max-Q and an i7-9750H (6 cores, 12 threads) through CubeCL's CPU runtime:

Where a kernel has several exact implementations (direct, shared-memory tiles, sorting networks, vectors of 2–16 values per unit that the CPU runtime turns into SIMD), the autotuner times them on the device and keeps the fastest. The vector versions are bit-identical to the scalar ones; other alternatives (FIR tiles, IIR time blocks) agree to rounding; dspk_core::compute::pin_tuned_choices fixes the IIR block split and the matrix-multiply choice for runs that must reproduce bit for bit. Reproduce with cargo test -p dspk-benchmarks --profile validate --test bench_kernels -- --ignored --nocapture (and bench_vector for every choice); timings against scipy from Python are in benchmarks.

Kernel and autotune cache

Autotune decisions and compiled kernels are cached across processes: the first run on a device tunes and compiles, later runs start in about a second. The cache is one SQLite database, default.db (with its -wal / -tshm journal files), in CubeCL's folder of the user cache directory:

• DSPK_CACHE_DIR=/some/dir puts it in that directory instead (shared machines, CI caches, read-only home directories).
• DSPK_KERNEL_CACHE=0 turns the compiled-kernel cache off; autotune decisions still persist.
• Deleting the folder is safe: the next run tunes and compiles again. It grows with every runtime, device and kernel variant used (about 130 MB after the full test suite on wgpu, CUDA and CPU).
• A cubecl.toml in the working directory, or CubeCL's own environment variables, still apply.
• Code that opens a device through CubeCL directly instead of dspk gets CubeCL's default: target/environment/ of the nearest Cargo project, or the user cache directory outside one.

[!NOTE] Abort at exit (upstream, worked around). CubeCL drops its kernel-cache database (turso) while a device thread ends, which can abort a process after its work: fatal runtime error: thread local panicked on drop, aborting (exit code 134; tracel-ai/cubecl#1981). dspk avoids it by ending only wgpu devices at exit (their Vulkan drivers need it) and leaving CUDA and CPU device threads to the process exit. Calling dspk_core::compute::shutdown_devices yourself can still hit it; so can code that opens devices through CubeCL directly. DSPK_KERNEL_CACHE=0 avoids it at the cost of compiling kernels in every process.

Python

pip install dspk

import dspk
from dspk.ops import Pipeline, stage

dspk.runtime.available()                      # e.g. ['wgpu', 'cuda', 'cpu']
pipe = Pipeline([stage.car(), stage.iir("bandpass", (300, 6000), mode="forward")],
                sample_rate=30_000, stream=True)
out = pipe.process(chunk)                     # (channels, samples) in, same shape out, on the GPU

Wheels for Linux x86_64 and Windows x64 (wgpu, CUDA, CPU runtime) and macOS arm64 (wgpu, CPU); CPython ≥ 3.10. See bindings/python.

Device size limit. Kernels index with 32-bit integers (WebGPU has no 64-bit integers, and one kernel source serves every runtime): one window or batch on the device holds at most 2³² − 1 values (e.g. 384 channels × 11.1 million samples). Larger requests are refused with an error.

Spike sorting

Kilosort4, EMUsort, MountainSort 5, SpyKING CIRCUS 2 and Tridesclous 2 run on the GPU (or the CPU) in one separate program, dspk-cli-sorter, released with each version. The dspk command line runs it as its own process:

dspk sorters install                                   # this version's build, SHA-256 checked
dspk sort kilosort4 recording.bin -o sorted --format phy
dspk sort kilosort4 --help                             # every setting and its default
dspk-cli-sorter licenses                                  # each sorter's licence and source

It is installed in the user data folder (~/.local/share/dspk/sorters/<version>/ on Linux, ~/Library/Application Support/dspk/… on macOS, %LOCALAPPDATA%\dspk\… on Windows; DSPK_DATA_DIR moves it, DSPK_CLI_SORTER_BIN names another build). A source build finds the dspk-cli-sorter built next to dspk (cargo build -p dspk-cli-sorter).

License

MIT OR Apache-2.0, at your option (LICENSE-MIT, LICENSE-APACHE): the published crates, the dspk command line and the Python package.

The sorters follow their upstreams' licences:

Each sorter crate carries its upstream licence and a NOTICE.md. Nothing MIT / Apache-licensed links dspk-sorter-kilosort: dspk sort runs dspk-cli-sorter as a separate program, whose source is this repository at the version's tag.