fusedtok

Fused CUDA kernels for LLM inference — RMSNorm / RoPE / SwiGLU / attention decode and friends, with zero-copy torch tensor support: up to 9.3x faster than PyTorch SDPA (attention decode, 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 wheels on PyPI (built with CUDA 12.4): Linux x86_64 (manylinux, cp310) and Windows x86_64 (cp312). On other platforms or Python versions 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)

# attention over a GQA kv-cache: one call per decode step, no score
# materialization, variable-length batches share one cache tensor
out = fusedtok.attention_decode(
    q_heads,                                    # [B, Hq, D] new token
    k_cache, v_cache,                           # [B, Hkv, T, D]
    lens=torch.tensor([1023, 512], dtype=torch.int32, device="cuda"))
# fresh-sequence prefill (causal by default; convenience path)
ctx = fusedtok.attention_prefill(q_all, k_all, v_all, causal=True)

# sampling side: the whole decode step in one fused call
token = fusedtok.decode_step(logits, sampled_ids, penalty=1.1,
                             p=0.9, temperature=0.8, seed=step)
# or step by step:
logits = fusedtok.repetition_penalty(logits, sampled_ids, penalty=1.1)
token = fusedtok.sample_topp(logits, p=0.9, temperature=0.8, seed=step)

A minimal per-token sampling loop:

import torch, fusedtok as ft

h = torch.zeros(1, 4096, device="cuda")            # decoder state
w = torch.load("rms_weight.pt").cuda()             # float32 weights
wq, wscale = ft.quantize_int8(weight_f32.ravel())  # int8 weights
generated = []
for step in range(256):
    h = ft.rmsnorm(h, w, residual=h)               # fused add + norm
    q = ft.rope(q, k=None, pos_offset=step, neox=True)
    logits = model_output(h)                       # your model
    tok = ft.decode_step(logits, generated, penalty=1.1,
                         p=0.9, temperature=0.8, seed=step)
    generated.append(int(tok))

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 may also be bfloat16 - the kernels compute in float32 and convert at the load/store boundary (norm weights are upcast to float32 automatically; sampling/selection ops stay float32). 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 over 3 independent rounds (means below; per-round values in the JSON), vs the equivalent PyTorch reference (composite eager expressions; attention references use pre-expanded heads - repeat_interleave outside the timed region). Largest shape per op; full data: docs/benchmark_rtx3060.json, reproduce with python benchmarks/bench.py:

Row-wise kernels (norms, softmax) autotune their thread-block size per shape at first call (v0.4.1); the table reflects the tuned choices.

RTX 5060 Ti (Blackwell, sm_120) — same suite, largest shape per op (full data: docs/benchmark_rtx5060ti.json):

On smaller shapes the Blackwell card shows bigger wins (softmax 2.5x, RMSNorm 3.2x at 256 rows, attention decode 3.8x at T=4096 running 235 GB/s) - the launch-overhead share shrinks as shapes grow; full sweep in the JSON.

The PyPI wheel ships sm_80/sm_86 cubins plus a compute_86 PTX fallback — verified to JIT and run correctly on Blackwell (sm_120) drivers.

Fusions win big (RoPE / RMSNorm / SwiGLU) because eager mode round-trips intermediate tensors through global memory. The v0.4 selection pipeline (arrival-ticket radix rounds + early-exit compaction, replayed from a cached CUDA graph) beats torch's CUB radix select at small k on both GPUs; mid-range k (2048..n) stays at or below parity — honest numbers, a pipelined tensor-core sort stays future work. attention_decode wins big at decode (one launch streams the GQA cache once at up to ~157 GB/s effective while SDPA pays head expansion or small-query inefficiency); attention_prefill is the honest convenience path at ~0.45x of SDPA's flash backend — no tensor cores by design, so heavyweight prefill stays with SDPA/FlashAttention. The INT8 decode GEMV moves half the bytes of an fp16 projection and runs at full memory bandwidth (2x); the IMMA GEMM path (~17 TOPS) is correctness-first — cuBLASLt (torch._int_mm) remains faster for large prefill matmuls.

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 (done): bf16 zero-copy, radix-select top-k/top-p, fused nucleus sampling, single-read softmax, CUDA-graph verified
• v0.3 (done): chunk-merge selection sort + parallel nucleus count, bf16x4/x8 vectorized elementwise, INT8 quantize/dequantize utilities
• v0.4 (done): arrival-ticket selection pipeline (no cooperative launch, early-exit compaction, cached CUDA graphs), stream-aware launchers everywhere (real CUDA-graph capture), INT8 compute path (IMMA qgemm + decode GEMV), fused decode_step sampling
• v0.4.1 (done): runtime block-size autotuning for the row-wise kernels (norms/softmax pick 128..1024 threads per shape at first call)
• v0.5 (done): attention - GQA decode attention over a contiguous kv-cache (flash-decoding split over long caches, per-sequence lengths) and a tiled prefill path (honest ~0.45x of SDPA flash - the convenience path); single-chart-per-GPU benchmarks; Windows wheels in the PyPI publish pipeline
• future: pipelined tensor-core INT8 GEMM; top-k mid-range parity; API freeze for 1.0

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.