SpikeSkip

Bypass the Von Neumann Memory Wall for Spiking Neural Networks & Sparse Activations.

SpikeSkip is a standalone, production-grade PyTorch CUDA plugin that skips weight DRAM memory fetches for silent (zero-valued) neurons using fused Compressed Sparse Row (CSR) kernels on standard NVIDIA GPUs.

Scope & Philosophy

SpikeSkip is a standalone, zero-dependency model plugin, not a monolithic SNN framework. It seamlessly integrates into any existing PyTorch codebase as a drop-in replacement for standard dense layers (torch.nn.Linear, torch.nn.Conv2d).

• Target Sparsity: >98-99% activation sparsity (typical for SNNs & deep ReLU networks).
• DRAM Reduction: Up to 546x less memory bandwidth traffic (120 KB vs 64 MB per pass on 4096×4096 layers).
• Zero Overhead: Fused GPU-only prefix-sum (~1μs) and cached weight transposes eliminate Python-CUDA runtime latency.

Installation

git clone https://github.com/Griffith-7/spikeskip.git
cd spikeskip
pip install -e .

Requirements

• PyTorch $\ge$ 2.0
• NVIDIA GPU with Compute Capability $\ge$ 7.0
• CUDA Toolkit 11.0+
• C++ Compiler (GCC 7+ on Linux, Visual Studio 2019+ on Windows)

Quick Start

1. High-Level Drop-In Module

import torch
from spikeskip import SparseLinear, SparseConv2d

# Drop-in replacement for nn.Linear
layer = SparseLinear(4096, 4096, bias=True, device="cuda")

# Input tensor with high sparsity (>99% zero elements)
x = torch.randn(256, 4096, device="cuda")
x[x < 2.0] = 0.0

# Forward pass
output = layer(x)

2. Multi-Timestep SNN Temporal Loop

# Input shape: (T=10 timesteps, Batch=256, In_Features=4096)
spikes_t = torch.randn(10, 256, 4096, device="cuda")
spikes_t[spikes_t < 2.0] = 0.0

# Executes CSR conversion and GEMM in a tight C++ loop (zero Python overhead between timesteps)
output = layer.forward_multistep(spikes_t, T=10)

3. Low-Level Functional API

from spikeskip import alloc_csr_buffers, sparse_linear_forward

# Pre-allocate reusable scratch buffers
bufs = alloc_csr_buffers(batch_size=256, in_features=4096, device="cuda")
weightT = layer.linear.weight.t().contiguous()

output = sparse_linear_forward(x, weightT, *bufs)

Benchmark Summary

1. Spiking Language Model Benchmark (pretraining_code.jsonl)

Trained and evaluated on real code pretraining text (pretraining_code.jsonl):

2. Multi-Step Layer Benchmarks ($4096 \to 4096$, $B=256$)

Testing

Run the automated Pytest suite:

pip install -e ".[dev]"
pytest tests/ -v

Tests verify exact float32 numerical equivalence against dense PyTorch layers, automatic scratch buffer expansion for large batch sizes ($B > 1024$), and parameter version invalidation during optimizer steps.

Documentation

• Architecture & Kernel Design
• Empirical Benchmarks
• Changelog

License

MIT License. See LICENSE for details.