Metadata-Version: 2.4
Name: spikeskip
Version: 0.2.0
Summary: Standalone PyTorch CUDA Sparse Inference Engine Plugin for Spiking Neural Networks
Author: Sumith
License: MIT
Project-URL: Homepage, https://github.com/Griffith-7/spikeskip
Project-URL: Repository, https://github.com/Griffith-7/spikeskip.git
Project-URL: Issues, https://github.com/Griffith-7/spikeskip/issues
Keywords: spiking-neural-networks,cuda,sparse-inference,csr,pytorch-plugin,snns,memory-bandwidth
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: License :: OSI Approved :: MIT License
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.0.0
Requires-Dist: numpy>=1.22.0
Provides-Extra: dev
Requires-Dist: pytest>=7.0.0; extra == "dev"
Requires-Dist: ruff>=0.1.0; extra == "dev"
Dynamic: license-file
Dynamic: requires-python

# SpikeSkip

[![PyPI Version](https://img.shields.io/badge/version-0.2.0-blue.svg)](https://github.com/Griffith-7/spikeskip)
[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
[![CUDA](https://img.shields.io/badge/CUDA-11.0%2B-green.svg)](https://developer.nvidia.com/cuda-toolkit)
[![PyTorch](https://img.shields.io/badge/PyTorch-%E2%89%A52.0-ee4c2c.svg)](https://pytorch.org/)

**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

```bash
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

```python
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

```python
# 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

```python
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`):

| Metric / Feature | Model A (Astrocyte Baseline) | Model B (SpikeSkip LM) | Key Takeaway |
|---|---|---|---|
| **Activation Sparsity (%)** | 70.0% (unconstrained) | **99.9% (target >99%)** | **SpikeSkip (+29.9% Sparsity)** |
| **Model Perplexity (PPL)** | 25.6 (Loss 3.24) | **17.6 (Loss 2.87)** | **SpikeSkip (Better Accuracy & PPL)** |
| **DRAM Memory Traffic** | 1.0x (Baseline Reads All) | **992.8x DRAM Reduction** | **SpikeSkip (Fetches <0.1% Weights from DRAM)** |
| **Inference Decoding Speed** | 236.0 tokens/sec | 235.6 tokens/sec | Matching token generation speed |
| **Training Speed** | 60.10 ms/step | 62.95 ms/step | Parity (~2.8ms difference) |

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

| Layer Config | Sparsity | cuBLAS Dense | SpikeSkip CSR | Speedup |
|---|---|---|---|---|
| 4096→4096 (Single Pass) | 99.9% | 1.97 ms | 1.30 ms | **1.52x** |
| 2-layer SNN (T=10 Timesteps) | 99.9% | 20.19 ms | 12.79 ms | **3.24x** |
| Raw Coalesced Kernel | 99.0% | 0.54 ms | 0.25 ms | **2.13x** |

---

## Testing

Run the automated Pytest suite:

```bash
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](docs/ARCHITECTURE.md)
- [Empirical Benchmarks](docs/BENCHMARKS.md)
- [Changelog](CHANGELOG.md)

---

## License

MIT License. See [LICENSE](LICENSE) for details.
