fusedtok

Fused CUDA kernels for LLM inference 鈥?RMSNorm / RoPE / SwiGLU and friends, with zero-copy torch tensor support: up to 6.2x faster than PyTorch eager (RoPE, RTX 3060, see Benchmarks).

涓枃鏂囨。璇风湅 README_zh.md | English below.

Why

LLM inference frameworks launch many small, memory-bound operators per token. Each launch round-trips through global memory. fusedtok fuses them into single kernels to cut memory traffic and launch overhead.

Operators

Install

pip install fusedtok

Prebuilt Linux x86_64 wheels (manylinux, built with CUDA 12.4) are on PyPI. On Windows (or any platform without a matching wheel) pip builds from source automatically:

git clone https://github.com/Hai-Wenxiang/fusedtok.git
cd fusedtok
pip install .

Requirements:

• NVIDIA GPU of RTX 30 series (Ampere) or newer 鈥?e.g. RTX 3060/3090, RTX 4080, RTX 5090, A100, H100
• CUDA Toolkit >= 12.0
• A C++17 compiler (MSVC on Windows, GCC/Clang on Linux); Python 3.10+

Usage

numpy in / numpy out, or torch in / torch out 鈥?including zero-copy CUDA: kernels read and write torch device buffers directly via data_ptr(), with no staging copies and no host synchronization.

import numpy as np
import torch
import fusedtok

x = np.random.randn(4, 1024).astype(np.float32)
w = np.random.rand(1024).astype(np.float32)

# CPU reference implementation (ground truth, runs anywhere)
y = fusedtok.rmsnorm(x, w, eps=1e-6)

# staged CUDA: copies to GPU, runs kernel, copies back
y = fusedtok.rmsnorm(x, w, cuda=True)

# zero-copy CUDA with torch tensors: kernels run in torch's own buffers,
# stream-ordered with other torch operations
xt, wt = torch.from_numpy(x).cuda(), torch.from_numpy(w).cuda()
yt = fusedtok.rmsnorm(xt, wt)          # -> CUDA torch tensor

# RoPE with kv-cache position offset, NeoX (LLaMA-HF) layout
q = torch.randn(1, 4096, device="cuda")          # new token only
q_rot, k_rot = fusedtok.rope(q, k=None, pos_offset=1023, neox=True)

# sampling side
logits = fusedtok.repetition_penalty(logits, sampled_ids, penalty=1.1)
values, indices = fusedtok.topk(logits, k=50)

Every function accepts float32 numpy arrays or torch tensors (other dtypes are converted with a copy) and returns float32 outputs of the same family. CUDA torch tensors select the zero-copy path automatically.

See examples/demo.py for a runnable tour of every operator.

Correctness

Every kernel ships with a CPU reference implementation and element-wise parity tests (pytest). Tests run on machines without a GPU (CUDA cases skip automatically).

Benchmarks

RTX 3060 (sm_86), float32, zero-copy torch tensors, CUDA-event timing, vs the equivalent PyTorch eager expressions (full data: docs/benchmark_results.json, reproduce with python benchmarks/bench.py):

Fusions win big (RoPE / RMSNorm / SwiGLU) because eager mode round-trips intermediate tensors through global memory. Pure memory-bound elementwise ops run at the same ~330-500 GB/s as PyTorch's tuned kernels (silu, gelu, add 鈮? parity). Softmax and top-k remain behind PyTorch's CUB-based kernels 鈥? honest numbers, on the v0.2 roadmap.

Development

See CONTRIBUTING.md for the full guide (test rules, error contract, determinism invariants). Quick start:

# Windows: run inside a VS developer prompt (vcvars64)
cmake -S . -B build -G Ninja -DCMAKE_BUILD_TYPE=Release
cmake --build build
# from repo root: PYTHONPATH picks up the built module, conftest.py adds python/
$env:PYTHONPATH = "$PWD/build"        # Windows
PYTHONPATH=$PWD/build                 # Linux
python -m pytest tests -q
python benchmarks/bench.py            # GPU benchmark + chart

Windows / Linux. Windows uses MSVC via nvcc; CI builds and runs the CPU test suite on every push.

Roadmap

• v0.2: bf16 support, radix-select top-k/top-p (CUB-class speed), fused sampling (softmax+top-p+draw in one pass), CUDA graph-friendly batching
• v0.3: INT8/FP8 quantized paths, block-size autotuning
• v0.4+: lightweight fused attention; prebuilt wheels on PyPI

Community

• Contributing guide 鈥?setup, rules of the road, PR process
• Code of conduct
• Security policy
• Changelog

License

MIT 鈥?see LICENSE. Third-party notices: NOTICES.md.