FlashKAN

Fast B-spline Kolmogorov-Arnold Network layers for PyTorch.

6-15x faster than standard Cox-de Boor implementations (PyKAN, efficient-kan), faster than Gaussian RBF alternatives (FastKAN) — while producing exact B-spline basis values with compact support, C2 continuity, and partition of unity.

How it works

Standard KAN implementations compute B-spline basis functions using the Cox-de Boor recursion — 3 sequential passes for cubic splines, each creating intermediate tensors. FlashKAN replaces this with the truncated power closed form:

N(u) = (1/6) [relu(u)³ - 4·relu(u-1)³ + 6·relu(u-2)³ - 4·relu(u-3)³ + relu(u-4)³]

This single expression computes exact B-spline values with no recursion, no span lookups, and no gather operations. torch.compile fuses all elementwise ops into one GPU kernel.

Installation

pip install flashkan

Requirements: Python >= 3.9, PyTorch >= 2.0

Supported devices: CPU, CUDA (NVIDIA), MPS (Apple Silicon)

From source

git clone https://github.com/NAVEENMN/flashkan.git
cd flashkan
pip install -e .

Quick start

import torch
from flashkan import KANLayer, KANNetwork

# Drop-in replacement for nn.Linear
layer = KANLayer(784, 64)
x = torch.randn(32, 784)
y = layer(x)  # [32, 64]

# Multi-layer network
net = KANNetwork([784, 64, 10])
y = net(torch.randn(32, 784))  # [32, 10]

MNIST example

import torch
import torch.nn as nn
from flashkan import KANNetwork

model = KANNetwork([784, 64, 10])
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()

# Standard PyTorch training loop
for images, labels in train_loader:
    output = model(images.view(-1, 784))
    loss = criterion(output, labels)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

See examples/ for complete runnable scripts.

API

KANLayer(in_features, out_features, grid_size=5, spline_order=3)

A single KAN layer. Drop-in replacement for nn.Linear.

KANNetwork(layer_dims, grid_size=5, spline_order=3)

Stack of KAN layers.

# 3-layer KAN: 784 -> 128 -> 64 -> 10
net = KANNetwork([784, 128, 64, 10])

Benchmarks

Forward pass time (ms) on Apple M-series GPU (MPS), batch=256:

FlashKAN is 6.3x faster than vanilla Cox-de Boor on MNIST and 15.3x faster on CIFAR-10.

Why it's fast

91% of a standard KAN forward pass is spent computing B-spline basis functions. FlashKAN eliminates this bottleneck:

B-spline properties preserved

Unlike Gaussian RBF approximations, FlashKAN computes exact B-spline basis values:

• Compact support — each basis function is exactly zero outside its knot span window
• C2 continuity — second derivatives are continuous at every knot
• Partition of unity — basis values sum to 1 at every point in the interior
• Non-negativity — all basis values are >= 0

Verified: max difference vs Cox-de Boor reference is < 5e-5 in float32.

Project structure

src/flashkan/
├── __init__.py     # Public API
├── basis.py        # Truncated power B-spline + torch.compile (core math)
├── layer.py        # KANLayer
└── network.py      # KANNetwork

4 source files. The core innovation is in basis.py — 30 lines of math.

Citation

If you use FlashKAN in your research, please cite:

@software{flashkan2026,
  title={FlashKAN: Fast B-spline KAN Layers via Truncated Power Basis},
  author={Mysore, Naveen},
  year={2026},
  url={https://github.com/NAVEENMN/flashkan}
}

License

MIT