fusedtok

Fused CUDA kernels for LLM inference — RMSNorm / RoPE / SwiGLU / attention decode and friends, with zero-copy torch tensor support: up to 8.9x 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-cp313) and Windows x86_64 (cp311-cp313). 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"))
# ...or over a paged (vLLM-style block-pool) cache: pools [Nb, Hkv, P, D]
# + a per-sequence block table; append each new token, then decode
fusedtok.kv_append_paged(k_pool, v_pool, block_table, k_new, v_new, lens)
out = fusedtok.attention_decode_paged(q_heads, k_pool, v_pool,
                                      block_table, lens + 1)
# 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)
# top-k sampling variant (renormalizes within the k survivors)
token = fusedtok.sample_topk(logits, k=50, 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 be bfloat16 on every operator that moves tensor data (elementwise / norms / RoPE / attention), and the attention operators additionally take float16 — 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, and the usage guide for the topic-structured manual (data flow, attention, INT8 workflow, sampling contract, CUDA graphs, error contract) — also available in 中文.

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).

API stability

1.0 freezes the public surface: the names in fusedtok.__all__ (32 operators + helpers) keep their signatures across the 1.x series. Type stubs (__init__.pyi, PEP 561 py.typed) ship with the package. New operators arrive in minor releases; breaking changes require a new major version and a deprecation window. Determinism promises: selection ties resolve to the earliest index; sampling is deterministic per seed.

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/benchmarks/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/benchmarks/benchmark_rtx5060ti.json):

On smaller shapes the Blackwell card shows bigger wins (softmax 2.5x, RMSNorm 3.2x at 256 rows, attention decode 3.78x at T=4096 running 187 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; the v1.0 retune (in-block-sort threshold and sort chunk both dropped 2048 -> 1024 - a single block bitonic-sorting 2048 keys was the whole mid-k regression) brings the mid-k window to parity-or-winning as well (k=4096 @131k: 1.36x / 1.09x). The fused samplers win against the eager composites when the logits look like real decode output (sample_topp peaked: 2.69x / 2.48x; sample_topk: 1.91x / 2.04x); on a FLAT distribution sample_topp is honestly 0.10-0.17x — the nucleus then spans ~90% of the vocabulary and the pipeline must effectively order the whole thing. v1.2 cut that worst case ~8.5x (18.2ms -> 2.2ms at n=131072 on a 3060) with three token-preserving changes - an adaptive window jump driven by a p*total mass bound, a full-vocabulary fast path that skips the selection stages, and a batched-load serial sampling walk (the strictly sequential float adds ARE the CPU-parity determinism contract; only the loads got pipelined) - but torch's fully parallel sort still owns that regime. attention_decode wins big at decode (one launch streams the GQA cache once at up to ~156 GB/s effective while SDPA pays head expansion or small-query inefficiency); attention_decode_paged (v1.2) pays only 1.06-1.07x for the block-table indirection of the fragmentation-free vLLM-style cache layout, with bit-identical output on matching slice schedules; 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 pipelined IMMA GEMM (v1.0 rework: cp.async double-buffered slabs, runtime-tuned 64x64 / 128x128 tiles) reaches ~39 TOPS on a 3060 and ~67 TOPS on a 5060 Ti — 2x-4x the v0.4 kernel — but cuBLASLt (torch._int_mm) still holds a ~2.2-2.6x lead: its tiles pipeline deeper and its epilogue is tuned per-arch. For now qgemm is the exact / graph-capturable / zero-copy INT8 path, not the fastest one; honest numbers, a CUTLASS-class schedule stays future work. The per-channel variant (qgemm_perchannel, the W8A8 layout INT8 inference actually uses) fuses the per-output-channel scale multiply into the same epilogue at zero kernel cost — the composite torch reference pays for that broadcast separately, which is where its 0.55-0.58x comes from.

All sampling rows above measure FIXED logits since 1.1.1 (the bench seeds torch's RNG); since v1.2 the peaked row spikes +20 and the flat row uses near-uniform logits — both regimes are now deterministically what their labels say (the v1.1 peaked row sat on the coverage boundary at n=131072 and flipped regimes per seed). The argmax rows are event-timed over a host-synchronized call and swing on WDDM (0.73-1.24x across runs); wall-clock probes with the sync excluded from the timed loop measure 1.12x (3060) / 0.96x (5060 Ti) — v1.2 removed one CUDA submission and one allocation per call.

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
• 1.0 (released): pipelined tensor-core INT8 GEMM (cp.async double-buffering, runtime tile tuning; 17 -> 39 TOPS on a 3060) with per-channel weight scales (W8A8), fused top-k sampling (2.1x vs the topk+multinomial composite), top-k mid-range-k parity, text hygiene gate, wheel matrix expansion (Linux cp310-313, Windows cp311-313), API freeze
• 1.1 (released): half-precision attention — attention_decode / attention_prefill accept bfloat16 and float16 caches (float32 compute, half the decode bytes); parallel exp precompute halves the flat-distribution sampling worst case with bit-identical tokens
• 1.2 (released): paged kv-cache attention — attention_decode_paged over a vLLM-style block pool [Nb, Hkv, P, D] + per-sequence block tables (1.06-1.07x the contiguous op, any valid table honored) and kv_append_paged (the in-place cache-write side); flat-distribution sampling worst case cut ~8.5x (adaptive widening jump + full-vocab fast path + batched-load serial walk, tokens bit-identical); argmax launch diet (one submission and one allocation less per call)

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.