fusedtok 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- fusedtok-0.1.0/.github/workflows/ci.yml +37 -0
- fusedtok-0.1.0/.github/workflows/publish.yml +76 -0
- fusedtok-0.1.0/.gitignore +35 -0
- fusedtok-0.1.0/CHANGELOG.md +26 -0
- fusedtok-0.1.0/CMakeLists.txt +56 -0
- fusedtok-0.1.0/CODE_OF_CONDUCT.md +66 -0
- fusedtok-0.1.0/CONTRIBUTING.md +51 -0
- fusedtok-0.1.0/LICENSE +21 -0
- fusedtok-0.1.0/NOTICES.md +21 -0
- fusedtok-0.1.0/PKG-INFO +235 -0
- fusedtok-0.1.0/README.md +185 -0
- fusedtok-0.1.0/README_zh.md +178 -0
- fusedtok-0.1.0/SECURITY.md +22 -0
- fusedtok-0.1.0/benchmarks/bench.py +221 -0
- fusedtok-0.1.0/docs/benchmark_results.json +228 -0
- fusedtok-0.1.0/docs/benchmark_rt3060.png +0 -0
- fusedtok-0.1.0/examples/demo.py +155 -0
- fusedtok-0.1.0/include/fusedtok/activations.hpp +78 -0
- fusedtok-0.1.0/include/fusedtok/cuda_launch.hpp +72 -0
- fusedtok-0.1.0/include/fusedtok/fusedtok.hpp +83 -0
- fusedtok-0.1.0/include/fusedtok/layernorm.hpp +22 -0
- fusedtok-0.1.0/include/fusedtok/softmax.hpp +18 -0
- fusedtok-0.1.0/pyproject.toml +46 -0
- fusedtok-0.1.0/python/fusedtok/__init__.py +517 -0
- fusedtok-0.1.0/src/activations.cu +314 -0
- fusedtok-0.1.0/src/bindings.cpp +627 -0
- fusedtok-0.1.0/src/cuda_util.cuh +103 -0
- fusedtok-0.1.0/src/fusedtok.cu +38 -0
- fusedtok-0.1.0/src/layernorm.cu +93 -0
- fusedtok-0.1.0/src/rmsnorm.cu +94 -0
- fusedtok-0.1.0/src/rope.cu +155 -0
- fusedtok-0.1.0/src/sampling.cu +76 -0
- fusedtok-0.1.0/src/softmax.cu +79 -0
- fusedtok-0.1.0/src/topk.cu +405 -0
- fusedtok-0.1.0/tests/conftest.py +11 -0
- fusedtok-0.1.0/tests/test_activations.py +144 -0
- fusedtok-0.1.0/tests/test_basic.py +62 -0
- fusedtok-0.1.0/tests/test_layernorm.py +84 -0
- fusedtok-0.1.0/tests/test_rmsnorm.py +127 -0
- fusedtok-0.1.0/tests/test_rope.py +126 -0
- fusedtok-0.1.0/tests/test_sampling.py +197 -0
- fusedtok-0.1.0/tests/test_softmax.py +72 -0
- fusedtok-0.1.0/tests/test_swiglu.py +78 -0
|
@@ -0,0 +1,37 @@
|
|
|
1
|
+
name: CI
|
|
2
|
+
|
|
3
|
+
on:
|
|
4
|
+
push:
|
|
5
|
+
branches: ["**"]
|
|
6
|
+
pull_request:
|
|
7
|
+
|
|
8
|
+
jobs:
|
|
9
|
+
build-and-test:
|
|
10
|
+
# GPU runners are not available on free plans; this job verifies that the
|
|
11
|
+
# project compiles with nvcc and that CPU-side tests pass. CUDA test cases
|
|
12
|
+
# detect the absence of a GPU and skip themselves (pytest skipif marker).
|
|
13
|
+
runs-on: ubuntu-latest
|
|
14
|
+
container:
|
|
15
|
+
image: nvidia/cuda:12.4.1-devel-ubuntu22.04
|
|
16
|
+
|
|
17
|
+
steps:
|
|
18
|
+
- uses: actions/checkout@v4
|
|
19
|
+
|
|
20
|
+
- name: Install build dependencies
|
|
21
|
+
run: |
|
|
22
|
+
apt-get update
|
|
23
|
+
apt-get install -y --no-install-recommends python3-dev python3-pip
|
|
24
|
+
pip3 install --no-cache-dir pybind11 pytest numpy cmake ninja
|
|
25
|
+
|
|
26
|
+
- name: Configure
|
|
27
|
+
run: >
|
|
28
|
+
cmake -S . -B build -G Ninja -DCMAKE_BUILD_TYPE=Release
|
|
29
|
+
-DCMAKE_PREFIX_PATH="$(python3 -m pybind11 --cmakedir)"
|
|
30
|
+
|
|
31
|
+
- name: Build
|
|
32
|
+
run: cmake --build build
|
|
33
|
+
|
|
34
|
+
- name: Test (CPU; CUDA cases auto-skip without a GPU)
|
|
35
|
+
run: |
|
|
36
|
+
export PYTHONPATH=$PWD/build
|
|
37
|
+
python3 -m pytest tests -q
|
|
@@ -0,0 +1,76 @@
|
|
|
1
|
+
name: Publish to PyPI
|
|
2
|
+
|
|
3
|
+
# Trusted Publishing (OIDC): no API tokens or passwords anywhere.
|
|
4
|
+
# The pending publisher configured on PyPI side:
|
|
5
|
+
# project=fusedtok owner=Hai-Wenxiang repo=fusedtok
|
|
6
|
+
# workflow=publish.yml environment=(none)
|
|
7
|
+
#
|
|
8
|
+
# Two jobs: the CUDA build must run inside a container, but
|
|
9
|
+
# pypa/gh-action-pypi-publish is a Docker action and Docker actions cannot
|
|
10
|
+
# execute within container jobs - so artifacts are passed through the
|
|
11
|
+
# workflow artifact store to a plain ubuntu-latest publishing job (the
|
|
12
|
+
# pattern documented by PyPA for containerized builds).
|
|
13
|
+
|
|
14
|
+
on:
|
|
15
|
+
workflow_dispatch:
|
|
16
|
+
release:
|
|
17
|
+
types: [published]
|
|
18
|
+
|
|
19
|
+
jobs:
|
|
20
|
+
build:
|
|
21
|
+
runs-on: ubuntu-latest
|
|
22
|
+
container:
|
|
23
|
+
image: nvidia/cuda:12.4.1-devel-ubuntu22.04
|
|
24
|
+
steps:
|
|
25
|
+
- uses: actions/checkout@v4
|
|
26
|
+
|
|
27
|
+
- name: Install build dependencies
|
|
28
|
+
run: |
|
|
29
|
+
apt-get update
|
|
30
|
+
apt-get install -y --no-install-recommends python3-dev python3-pip python3.10-venv
|
|
31
|
+
# patchelf via pip: ubuntu22.04's apt copy (0.14.3) is older than
|
|
32
|
+
# auditwheel's required >= 0.14.5
|
|
33
|
+
pip3 install --no-cache-dir build auditwheel patchelf
|
|
34
|
+
|
|
35
|
+
- name: Build sdist and wheel
|
|
36
|
+
run: python3 -m build --outdir dist
|
|
37
|
+
|
|
38
|
+
- name: Repair wheel to manylinux (PyPI rejects bare linux tags)
|
|
39
|
+
run: |
|
|
40
|
+
auditwheel show dist/*.whl
|
|
41
|
+
auditwheel repair dist/*.whl -w wheelhouse/
|
|
42
|
+
rm -f dist/*.whl
|
|
43
|
+
mv wheelhouse/*.whl dist/
|
|
44
|
+
|
|
45
|
+
- name: List artifacts
|
|
46
|
+
run: ls -lh dist/
|
|
47
|
+
|
|
48
|
+
- name: Store artifacts
|
|
49
|
+
uses: actions/upload-artifact@v4
|
|
50
|
+
with:
|
|
51
|
+
name: dist
|
|
52
|
+
path: dist/
|
|
53
|
+
if-no-files-found: error
|
|
54
|
+
retention-days: 5
|
|
55
|
+
|
|
56
|
+
publish:
|
|
57
|
+
needs: build
|
|
58
|
+
runs-on: ubuntu-latest
|
|
59
|
+
permissions:
|
|
60
|
+
# Required for PyPI Trusted Publishing (OIDC).
|
|
61
|
+
id-token: write
|
|
62
|
+
contents: read
|
|
63
|
+
steps:
|
|
64
|
+
- name: Download artifacts
|
|
65
|
+
uses: actions/download-artifact@v4
|
|
66
|
+
with:
|
|
67
|
+
name: dist
|
|
68
|
+
path: dist/
|
|
69
|
+
|
|
70
|
+
- name: Upload to PyPI (Trusted Publishing)
|
|
71
|
+
uses: pypa/gh-action-pypi-publish@release/v1
|
|
72
|
+
|
|
73
|
+
- name: Verify project is live
|
|
74
|
+
run: |
|
|
75
|
+
sleep 10
|
|
76
|
+
curl -fsS "https://pypi.org/pypi/fusedtok/json" | head -c 200 || true
|
|
@@ -0,0 +1,35 @@
|
|
|
1
|
+
# Build
|
|
2
|
+
build/
|
|
3
|
+
out/
|
|
4
|
+
*.obj
|
|
5
|
+
*.exe
|
|
6
|
+
*.dll
|
|
7
|
+
*.lib
|
|
8
|
+
*.exp
|
|
9
|
+
|
|
10
|
+
# Python
|
|
11
|
+
__pycache__/
|
|
12
|
+
*.pyc
|
|
13
|
+
*.egg-info/
|
|
14
|
+
dist/
|
|
15
|
+
.venv/
|
|
16
|
+
|
|
17
|
+
# IDE
|
|
18
|
+
.vs/
|
|
19
|
+
.vscode/
|
|
20
|
+
.idea/
|
|
21
|
+
|
|
22
|
+
# CUDA
|
|
23
|
+
*.cubin
|
|
24
|
+
*.fatbin
|
|
25
|
+
|
|
26
|
+
# CMake
|
|
27
|
+
CMakeFiles/
|
|
28
|
+
CMakeCache.txt
|
|
29
|
+
cmake_install.cmake
|
|
30
|
+
|
|
31
|
+
# OS
|
|
32
|
+
Thumbs.db
|
|
33
|
+
Desktop.ini
|
|
34
|
+
|
|
35
|
+
build2/
|
|
@@ -0,0 +1,26 @@
|
|
|
1
|
+
# Changelog
|
|
2
|
+
|
|
3
|
+
All notable changes to this project are documented here. The format follows
|
|
4
|
+
[Keep a Changelog](https://keepachangelog.com/en/1.1.0/) and the project
|
|
5
|
+
adheres to [Semantic Versioning](https://semver.org/).
|
|
6
|
+
|
|
7
|
+
## [0.1.0] - 2026-08-23
|
|
8
|
+
|
|
9
|
+
First public release.
|
|
10
|
+
|
|
11
|
+
### Added
|
|
12
|
+
- Operators: RMSNorm (+fused residual), LayerNorm (affine), RoPE
|
|
13
|
+
(interleaved and NeoX layouts, kv-cache `pos_offset`), SwiGLU, row-wise
|
|
14
|
+
softmax, SiLU / GeLU (erf + tanh) / ReLU / Tanh / Sigmoid, elementwise
|
|
15
|
+
add / mul, temperature scaling, repetition penalty, top-k, top-p
|
|
16
|
+
(nucleus), argmax.
|
|
17
|
+
- Three execution paths per op: CPU reference (ground truth, no GPU
|
|
18
|
+
needed), staged CUDA (numpy in / numpy out), and **zero-copy CUDA**
|
|
19
|
+
(kernels run directly in torch device buffers via `data_ptr()`).
|
|
20
|
+
- `pip install .` via scikit-build-core; sm_80/sm_86 cubins plus
|
|
21
|
+
compute_86 PTX for JIT on newer architectures.
|
|
22
|
+
- Optimized kernels: block-reduced norms/softmax, float4-vectorized
|
|
23
|
+
elementwise, exp2f/sincosf RoPE frequencies (6.2x vs eager), parallel
|
|
24
|
+
packed-key selection for top-k/top-p/argmax with deterministic ties.
|
|
25
|
+
- Benchmark suite (`benchmarks/bench.py`) with CUDA-event timing and
|
|
26
|
+
chart; bilingual README; CI on GitHub Actions (build + CPU tests).
|
|
@@ -0,0 +1,56 @@
|
|
|
1
|
+
cmake_minimum_required(VERSION 3.24)
|
|
2
|
+
project(fusedtok VERSION 0.1.0 LANGUAGES CXX CUDA)
|
|
3
|
+
|
|
4
|
+
# C++17 for host code; CUDA follows via nvcc default
|
|
5
|
+
set(CMAKE_CXX_STANDARD 17)
|
|
6
|
+
set(CMAKE_CXX_STANDARD_REQUIRED ON)
|
|
7
|
+
|
|
8
|
+
# Target GPU architectures: native sm_80 (A100) and sm_86 (RTX 30) cubins
|
|
9
|
+
# plus a compute_86 PTX embed so newer architectures (RTX 40/50, sm_89+) can
|
|
10
|
+
# JIT the kernels with their drivers.
|
|
11
|
+
# Set unconditionally (not only when undefined): build frontends may inject
|
|
12
|
+
# an empty CMAKE_CUDA_ARCHITECTURES that would otherwise fall back to the
|
|
13
|
+
# compiler default. Override with -DFUSEDTOK_CUDA_ARCHITECTURES="..." if
|
|
14
|
+
# you need something else.
|
|
15
|
+
set(FUSEDTOK_CUDA_ARCHITECTURES "80-real;86-real;86-virtual" CACHE STRING
|
|
16
|
+
"CUDA architectures fusedtok compiles for")
|
|
17
|
+
set(CMAKE_CUDA_ARCHITECTURES ${FUSEDTOK_CUDA_ARCHITECTURES})
|
|
18
|
+
|
|
19
|
+
find_package(CUDAToolkit REQUIRED)
|
|
20
|
+
# pybind11 provided by pip; pass its cmake dir via CMAKE_PREFIX_PATH
|
|
21
|
+
# (or it is injected by scikit-build-core during packaged builds)
|
|
22
|
+
find_package(pybind11 CONFIG REQUIRED)
|
|
23
|
+
|
|
24
|
+
# The compiled extension module loaded as `fusedtok._fusedtok`.
|
|
25
|
+
# Under scikit-build-core it installs into the python/fusedtok package dir.
|
|
26
|
+
pybind11_add_module(_fusedtok
|
|
27
|
+
src/fusedtok.cu
|
|
28
|
+
src/rmsnorm.cu
|
|
29
|
+
src/rope.cu
|
|
30
|
+
src/activations.cu
|
|
31
|
+
src/softmax.cu
|
|
32
|
+
src/layernorm.cu
|
|
33
|
+
src/topk.cu
|
|
34
|
+
src/sampling.cu
|
|
35
|
+
src/bindings.cpp)
|
|
36
|
+
target_include_directories(_fusedtok PRIVATE include src)
|
|
37
|
+
# Static CUDA runtime: the wheel then carries no libcudart.so dependency,
|
|
38
|
+
# letting auditwheel tag it manylinux (PyPI rejects bare linux_x86_64 tags).
|
|
39
|
+
target_link_libraries(_fusedtok PRIVATE CUDA::cudart_static)
|
|
40
|
+
|
|
41
|
+
if(SKBUILD)
|
|
42
|
+
install(TARGETS _fusedtok LIBRARY DESTINATION fusedtok)
|
|
43
|
+
endif()
|
|
44
|
+
|
|
45
|
+
# Suppress C4819 (codepage warning) on MSVC hosts, and force the conforming
|
|
46
|
+
# preprocessor (CCCL / cooperative_groups headers refuse the traditional one).
|
|
47
|
+
# Options must be routed per-language: nvcc rejects a bare "/Zc:..." or
|
|
48
|
+
# "/wd4819" on a CUDA source as a stray input file, so CUDA sources get them
|
|
49
|
+
# wrapped in -Xcompiler, CXX sources get them directly.
|
|
50
|
+
if(MSVC)
|
|
51
|
+
target_compile_options(_fusedtok PRIVATE
|
|
52
|
+
"$<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=/Zc:preprocessor>"
|
|
53
|
+
"$<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=/wd4819>"
|
|
54
|
+
"$<$<COMPILE_LANGUAGE:CXX>:/Zc:preprocessor>"
|
|
55
|
+
"$<$<COMPILE_LANGUAGE:CXX>:/wd4819>")
|
|
56
|
+
endif()
|
|
@@ -0,0 +1,66 @@
|
|
|
1
|
+
# Contributor Covenant Code of Conduct
|
|
2
|
+
|
|
3
|
+
## Our Pledge
|
|
4
|
+
|
|
5
|
+
We as members, contributors, and leaders pledge to make participation in our
|
|
6
|
+
community a harassment-free experience for everyone, regardless of age, body
|
|
7
|
+
size, visible or invisible disability, ethnicity, sex characteristics, gender
|
|
8
|
+
identity and expression, level of experience, education, socio-economic status,
|
|
9
|
+
nationality, personal appearance, race, religion, or sexual identity and
|
|
10
|
+
orientation.
|
|
11
|
+
|
|
12
|
+
We pledge to act and interact in ways that contribute to an open, welcoming,
|
|
13
|
+
diverse, inclusive, and healthy community.
|
|
14
|
+
|
|
15
|
+
## Our Standards
|
|
16
|
+
|
|
17
|
+
Examples of behavior that contributes to a positive environment include:
|
|
18
|
+
|
|
19
|
+
* Demonstrating empathy and kindness toward other people
|
|
20
|
+
* Being respectful of differing opinions, viewpoints, and experiences
|
|
21
|
+
* Giving and gracefully accepting constructive feedback
|
|
22
|
+
* Accepting responsibility and apologizing to those affected by our mistakes,
|
|
23
|
+
and learning from the experience
|
|
24
|
+
* Focusing on what is best not just for us as individuals, but for the overall
|
|
25
|
+
community
|
|
26
|
+
|
|
27
|
+
Examples of unacceptable behavior include:
|
|
28
|
+
|
|
29
|
+
* The use of sexualized language or imagery, and sexual attention or advances
|
|
30
|
+
of any kind
|
|
31
|
+
* Trolling, insulting or derogatory comments, and personal or political attacks
|
|
32
|
+
* Public or private harassment
|
|
33
|
+
* Publishing others' private information, such as a physical or email address,
|
|
34
|
+
without their explicit permission
|
|
35
|
+
* Other conduct which could reasonably be considered inappropriate in a
|
|
36
|
+
professional setting
|
|
37
|
+
|
|
38
|
+
## Enforcement Responsibilities
|
|
39
|
+
|
|
40
|
+
Community leaders are responsible for clarifying and enforcing our standards
|
|
41
|
+
of acceptable behavior and will take appropriate and fair corrective action in
|
|
42
|
+
response to any behavior that they deem inappropriate, threatening, offensive,
|
|
43
|
+
or harmful.
|
|
44
|
+
|
|
45
|
+
## Scope
|
|
46
|
+
|
|
47
|
+
This Code of Conduct applies within all community spaces (issues, PRs,
|
|
48
|
+
discussions) and also applies when an individual is officially representing
|
|
49
|
+
the community in public spaces.
|
|
50
|
+
|
|
51
|
+
## Enforcement
|
|
52
|
+
|
|
53
|
+
Instances of abusive, harassing, or otherwise unacceptable behavior may be
|
|
54
|
+
reported to the community leaders responsible for enforcement via the
|
|
55
|
+
repository's maintainer contact. All complaints will be reviewed and
|
|
56
|
+
investigated promptly and fairly.
|
|
57
|
+
|
|
58
|
+
All community leaders are obligated to respect the privacy and security of
|
|
59
|
+
the reporter of any incident.
|
|
60
|
+
|
|
61
|
+
## Attribution
|
|
62
|
+
|
|
63
|
+
This Code of Conduct is adapted from the
|
|
64
|
+
[Contributor Covenant](https://www.contributor-covenant.org), version 2.1,
|
|
65
|
+
available at
|
|
66
|
+
https://www.contributor-covenant.org/version/2/1/code_of_conduct.html.
|
|
@@ -0,0 +1,51 @@
|
|
|
1
|
+
# Contributing to fusedtok
|
|
2
|
+
|
|
3
|
+
Thanks for your interest in improving fusedtok!
|
|
4
|
+
|
|
5
|
+
## Ways to help
|
|
6
|
+
|
|
7
|
+
- **Bug reports**: open an issue with the op name, shapes, dtypes, GPU model,
|
|
8
|
+
driver/CUDA version, and a minimal reproducer. Include the CPU vs GPU
|
|
9
|
+
mismatch if it is a numerical issue.
|
|
10
|
+
- **Performance ideas**: profiles, Nsight Compute reports, or concrete kernel
|
|
11
|
+
suggestions are extremely welcome. The bar for merging a faster kernel is
|
|
12
|
+
that the full test suite stays green.
|
|
13
|
+
- **New operators**: check the roadmap in the README first; avoid operators
|
|
14
|
+
that duplicate PyTorch one-to-one without a fusion or sampling benefit.
|
|
15
|
+
|
|
16
|
+
## Development setup
|
|
17
|
+
|
|
18
|
+
```bash
|
|
19
|
+
git clone https://github.com/Hai-Wenxiang/fusedtok.git
|
|
20
|
+
cd fusedtok
|
|
21
|
+
pip install -e . # or: cmake -S . -B build -G Ninja && cmake --build build
|
|
22
|
+
pip install pytest torch # torch optional but recommended for the CUDA tests
|
|
23
|
+
pytest tests -q
|
|
24
|
+
```
|
|
25
|
+
|
|
26
|
+
Requirements: CUDA Toolkit >= 12.0, a CUDA GPU of compute capability 8.0+
|
|
27
|
+
(Ampere or newer). CPU-only machines can still run the CPU reference tests —
|
|
28
|
+
CUDA cases skip automatically.
|
|
29
|
+
|
|
30
|
+
## Rules of the road
|
|
31
|
+
|
|
32
|
+
1. **Every kernel ships with a CPU reference implementation** and parity
|
|
33
|
+
tests (multiple shapes, edge cases, error paths, GPU-vs-CPU comparison).
|
|
34
|
+
2. **Determinism where promised**: selection ops (top-k / top-p / argmax)
|
|
35
|
+
resolve ties toward the earliest index; keep that invariant.
|
|
36
|
+
3. **Error contract**: shape problems raise `ValueError`, CUDA problems
|
|
37
|
+
raise `RuntimeError` (mapped from `std::invalid_argument` /
|
|
38
|
+
`std::runtime_error` in C++).
|
|
39
|
+
4. **Comments in English**, explaining *why* (design constraints, GPU
|
|
40
|
+
micro-arch reasons), not *what*.
|
|
41
|
+
5. Benchmarks use CUDA events, never wall clock (WDDM makes host timing
|
|
42
|
+
on Windows meaningless).
|
|
43
|
+
6. Keep the CI green: `ubuntu-latest` + CUDA container build and CPU tests
|
|
44
|
+
run on every push.
|
|
45
|
+
|
|
46
|
+
## Pull requests
|
|
47
|
+
|
|
48
|
+
- One logical change per PR, with tests and README updates in the same PR.
|
|
49
|
+
- PRs are merged only with passing CI and maintainer review.
|
|
50
|
+
- By contributing you agree your contributions are licensed under the MIT
|
|
51
|
+
license (see [LICENSE](LICENSE)).
|
fusedtok-0.1.0/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Hai-Wenxiang
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
# Third-party notices
|
|
2
|
+
|
|
3
|
+
fusedtok is MIT licensed (see LICENSE). This product builds on:
|
|
4
|
+
|
|
5
|
+
## pybind11
|
|
6
|
+
|
|
7
|
+
Licensed under the BSD-style license found in
|
|
8
|
+
https://github.com/pybind/pybind11/blob/master/LICENSE.
|
|
9
|
+
Used at build time to generate the Python bindings.
|
|
10
|
+
|
|
11
|
+
## NVIDIA CUDA Toolkit
|
|
12
|
+
|
|
13
|
+
The compiled extension links the CUDA Runtime (cudart) and is built with the
|
|
14
|
+
NVIDIA CUDA Compiler. Use of the CUDA Toolkit is governed by NVIDIA's
|
|
15
|
+
software license: https://developer.nvidia.com/cuda-toolkit-software-license-agreement
|
|
16
|
+
|
|
17
|
+
## Optional runtime dependencies
|
|
18
|
+
|
|
19
|
+
- NumPy (BSD-3-Clause) - array interface for the staged path.
|
|
20
|
+
- PyTorch (BSD-3-Clause) - optional; used for the zero-copy CUDA path and in
|
|
21
|
+
the benchmark suite.
|
fusedtok-0.1.0/PKG-INFO
ADDED
|
@@ -0,0 +1,235 @@
|
|
|
1
|
+
Metadata-Version: 2.2
|
|
2
|
+
Name: fusedtok
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Fused CUDA kernels for LLM inference: RMSNorm, RoPE, SwiGLU, sampling ops, with zero-copy torch support
|
|
5
|
+
Keywords: cuda,llm,inference,kernels,deep-learning,pytorch
|
|
6
|
+
Author: Hai-Wenxiang
|
|
7
|
+
License: MIT License
|
|
8
|
+
|
|
9
|
+
Copyright (c) 2026 Hai-Wenxiang
|
|
10
|
+
|
|
11
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
12
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
13
|
+
in the Software without restriction, including without limitation the rights
|
|
14
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
15
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
16
|
+
furnished to do so, subject to the following conditions:
|
|
17
|
+
|
|
18
|
+
The above copyright notice and this permission notice shall be included in all
|
|
19
|
+
copies or substantial portions of the Software.
|
|
20
|
+
|
|
21
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
22
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
23
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
24
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
25
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
26
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
27
|
+
SOFTWARE.
|
|
28
|
+
|
|
29
|
+
Classifier: Development Status :: 4 - Beta
|
|
30
|
+
Classifier: Intended Audience :: Developers
|
|
31
|
+
Classifier: Intended Audience :: Science/Research
|
|
32
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
33
|
+
Classifier: Programming Language :: C++
|
|
34
|
+
Classifier: Programming Language :: Python :: 3
|
|
35
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
36
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
37
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
38
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
39
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
40
|
+
Project-URL: Homepage, https://github.com/Hai-Wenxiang/fusedtok
|
|
41
|
+
Project-URL: Repository, https://github.com/Hai-Wenxiang/fusedtok
|
|
42
|
+
Project-URL: Issues, https://github.com/Hai-Wenxiang/fusedtok/issues
|
|
43
|
+
Project-URL: Changelog, https://github.com/Hai-Wenxiang/fusedtok/releases
|
|
44
|
+
Requires-Python: >=3.10
|
|
45
|
+
Requires-Dist: numpy>=1.24
|
|
46
|
+
Provides-Extra: torch
|
|
47
|
+
Provides-Extra: test
|
|
48
|
+
Requires-Dist: pytest>=7; extra == "test"
|
|
49
|
+
Description-Content-Type: text/markdown
|
|
50
|
+
|
|
51
|
+
# fusedtok
|
|
52
|
+
|
|
53
|
+
[](https://github.com/Hai-Wenxiang/fusedtok/actions/workflows/ci.yml)
|
|
54
|
+
[](https://pypi.org/project/fusedtok/)
|
|
55
|
+
[](LICENSE)
|
|
56
|
+
[](pyproject.toml)
|
|
57
|
+
|
|
58
|
+
**Fused CUDA kernels for LLM inference** — RMSNorm / RoPE / SwiGLU and friends,
|
|
59
|
+
with **zero-copy torch tensor support**: up to **6.2x faster than PyTorch eager**
|
|
60
|
+
(RoPE, RTX 3060, see [Benchmarks](#benchmarks)).
|
|
61
|
+
|
|
62
|
+
**中文文档请看 [README_zh.md](README_zh.md)** | English below.
|
|
63
|
+
|
|
64
|
+
## Why
|
|
65
|
+
|
|
66
|
+
LLM inference frameworks launch many small, memory-bound operators per token. Each launch
|
|
67
|
+
round-trips through global memory. `fusedtok` fuses them into single kernels to cut memory
|
|
68
|
+
traffic and launch overhead.
|
|
69
|
+
|
|
70
|
+
## Operators
|
|
71
|
+
|
|
72
|
+
| Status | Kernel | Notes |
|
|
73
|
+
|---|---|---|
|
|
74
|
+
| ✅ | RMSNorm (+residual) | LLaMA/Qwen style, fused residual add |
|
|
75
|
+
| ✅ | LayerNorm | with affine |
|
|
76
|
+
| ✅ | RoPE | interleaved **and** NeoX layouts, kv-cache `pos_offset` |
|
|
77
|
+
| ✅ | SwiGLU | fused MLP activation |
|
|
78
|
+
| ✅ | Softmax (row-wise) | numerically stable |
|
|
79
|
+
| ✅ | SiLU / GeLU / GeLU-tanh / ReLU / Tanh / Sigmoid | elementwise |
|
|
80
|
+
| ✅ | add / mul | elementwise binary (fused add+residual pattern) |
|
|
81
|
+
| ✅ | top-k / top-p (nucleus) | deterministic ties |
|
|
82
|
+
| ✅ | argmax / temperature | greedy decoding helpers |
|
|
83
|
+
| ✅ | repetition penalty | CTRL-style, applied to sampled token ids |
|
|
84
|
+
| ⏳ | INT8/FP8 quantized path | planned v0.3 |
|
|
85
|
+
|
|
86
|
+
## Install
|
|
87
|
+
|
|
88
|
+
```bash
|
|
89
|
+
pip install fusedtok
|
|
90
|
+
```
|
|
91
|
+
|
|
92
|
+
Prebuilt Linux x86_64 wheels (manylinux, built with CUDA 12.4) are on PyPI.
|
|
93
|
+
On Windows (or any platform without a matching wheel) pip builds from
|
|
94
|
+
source automatically:
|
|
95
|
+
|
|
96
|
+
```bash
|
|
97
|
+
git clone https://github.com/Hai-Wenxiang/fusedtok.git
|
|
98
|
+
cd fusedtok
|
|
99
|
+
pip install .
|
|
100
|
+
```
|
|
101
|
+
|
|
102
|
+
**Requirements:**
|
|
103
|
+
|
|
104
|
+
- NVIDIA GPU of **RTX 30 series (Ampere) or newer** — e.g. RTX 3060/3090, RTX 4080, RTX 5090, A100, H100
|
|
105
|
+
- CUDA Toolkit >= 12.0
|
|
106
|
+
- A C++17 compiler (MSVC on Windows, GCC/Clang on Linux); Python 3.10+
|
|
107
|
+
|
|
108
|
+
<details>
|
|
109
|
+
<summary>What is "compute capability"? (click to expand)</summary>
|
|
110
|
+
|
|
111
|
+
Compute capability is NVIDIA's version number for a GPU architecture generation — not a
|
|
112
|
+
performance score. CUDA code must be compiled for a specific architecture to run on it.
|
|
113
|
+
The wheel builds native cubins for compute capability 8.0 (A100) and 8.6 (RTX 30) plus a
|
|
114
|
+
compute_86 PTX fallback, so Ampere runs natively and newer architectures (RTX 40/50, ...)
|
|
115
|
+
JIT the PTX with their driver.
|
|
116
|
+
|
|
117
|
+
| Compute capability | Architecture | Example GPUs |
|
|
118
|
+
|---|---|---|
|
|
119
|
+
| 7.5 | Turing | GTX 16xx, RTX 20xx (not supported) |
|
|
120
|
+
| 8.0 / 8.6 | Ampere | A100, RTX 30xx |
|
|
121
|
+
| 8.9 | Ada | RTX 40xx (via PTX) |
|
|
122
|
+
| 9.0 | Hopper | H100 (via PTX) |
|
|
123
|
+
| 12.0 | Blackwell | RTX 50xx (via PTX) |
|
|
124
|
+
|
|
125
|
+
Check yours: run `nvidia-smi` to see your GPU model, then look it up at
|
|
126
|
+
https://developer.nvidia.com/cuda-gpus
|
|
127
|
+
|
|
128
|
+
</details>
|
|
129
|
+
|
|
130
|
+
## Usage
|
|
131
|
+
|
|
132
|
+
numpy in / numpy out, or torch in / torch out — including **zero-copy CUDA**:
|
|
133
|
+
kernels read and write torch device buffers directly via `data_ptr()`, with
|
|
134
|
+
no staging copies and no host synchronization.
|
|
135
|
+
|
|
136
|
+
```python
|
|
137
|
+
import numpy as np
|
|
138
|
+
import torch
|
|
139
|
+
import fusedtok
|
|
140
|
+
|
|
141
|
+
x = np.random.randn(4, 1024).astype(np.float32)
|
|
142
|
+
w = np.random.rand(1024).astype(np.float32)
|
|
143
|
+
|
|
144
|
+
# CPU reference implementation (ground truth, runs anywhere)
|
|
145
|
+
y = fusedtok.rmsnorm(x, w, eps=1e-6)
|
|
146
|
+
|
|
147
|
+
# staged CUDA: copies to GPU, runs kernel, copies back
|
|
148
|
+
y = fusedtok.rmsnorm(x, w, cuda=True)
|
|
149
|
+
|
|
150
|
+
# zero-copy CUDA with torch tensors: kernels run in torch's own buffers,
|
|
151
|
+
# stream-ordered with other torch operations
|
|
152
|
+
xt, wt = torch.from_numpy(x).cuda(), torch.from_numpy(w).cuda()
|
|
153
|
+
yt = fusedtok.rmsnorm(xt, wt) # -> CUDA torch tensor
|
|
154
|
+
|
|
155
|
+
# RoPE with kv-cache position offset, NeoX (LLaMA-HF) layout
|
|
156
|
+
q = torch.randn(1, 4096, device="cuda") # new token only
|
|
157
|
+
q_rot, k_rot = fusedtok.rope(q, k=None, pos_offset=1023, neox=True)
|
|
158
|
+
|
|
159
|
+
# sampling side
|
|
160
|
+
logits = fusedtok.repetition_penalty(logits, sampled_ids, penalty=1.1)
|
|
161
|
+
values, indices = fusedtok.topk(logits, k=50)
|
|
162
|
+
```
|
|
163
|
+
|
|
164
|
+
Every function accepts float32 numpy arrays or torch tensors (other dtypes
|
|
165
|
+
are converted with a copy) and returns float32 outputs of the same family.
|
|
166
|
+
CUDA torch tensors select the zero-copy path automatically.
|
|
167
|
+
|
|
168
|
+
See `examples/demo.py` for a runnable tour of every operator.
|
|
169
|
+
|
|
170
|
+
## Correctness
|
|
171
|
+
|
|
172
|
+
Every kernel ships with a CPU reference implementation and element-wise parity tests
|
|
173
|
+
(pytest). Tests run on machines without a GPU (CUDA cases skip automatically).
|
|
174
|
+
|
|
175
|
+
## Benchmarks
|
|
176
|
+
|
|
177
|
+
RTX 3060 (sm_86), float32, zero-copy torch tensors, CUDA-event timing, vs the
|
|
178
|
+
equivalent PyTorch eager expressions (full data: `docs/benchmark_results.json`,
|
|
179
|
+
reproduce with `python benchmarks/bench.py`):
|
|
180
|
+
|
|
181
|
+
| Op | Shape | fusedtok | PyTorch eager | Speedup |
|
|
182
|
+
|---|---|---:|---:|---:|
|
|
183
|
+
| RoPE NeoX (q+k) | [2048×4096] | 416 µs | 2570 µs | **6.2x** |
|
|
184
|
+
| RMSNorm (+residual) | [1024×4096] | 260 µs | 538 µs | **2.1x** |
|
|
185
|
+
| SwiGLU | [1024×4096] | 153 µs | 257 µs | **1.7x** |
|
|
186
|
+
| LayerNorm | [1024×4096] | 168 µs | 162 µs | ~1.0x |
|
|
187
|
+
| SiLU | [1024×4096] | 105 µs | 112 µs | ~1.0x |
|
|
188
|
+
| Softmax | [1024×4096] | 159 µs | 115 µs | 0.7x |
|
|
189
|
+
| argmax | [131072] | 36 µs | 46 µs | **1.3x** |
|
|
190
|
+
| top-k (k=50) | [131072] | 168 µs | 129 µs | 0.8x |
|
|
191
|
+
|
|
192
|
+

|
|
193
|
+
|
|
194
|
+
Fusions win big (RoPE / RMSNorm / SwiGLU) because eager mode round-trips
|
|
195
|
+
intermediate tensors through global memory. Pure memory-bound elementwise ops
|
|
196
|
+
run at the same ~330-500 GB/s as PyTorch's tuned kernels (silu, gelu, add ≈
|
|
197
|
+
parity). Softmax and top-k remain behind PyTorch's CUB-based kernels —
|
|
198
|
+
honest numbers, on the v0.2 roadmap.
|
|
199
|
+
|
|
200
|
+
## Development
|
|
201
|
+
|
|
202
|
+
See [CONTRIBUTING.md](CONTRIBUTING.md) for the full guide (test rules,
|
|
203
|
+
error contract, determinism invariants). Quick start:
|
|
204
|
+
|
|
205
|
+
```bash
|
|
206
|
+
# Windows: run inside a VS developer prompt (vcvars64)
|
|
207
|
+
cmake -S . -B build -G Ninja -DCMAKE_BUILD_TYPE=Release
|
|
208
|
+
cmake --build build
|
|
209
|
+
# from repo root: PYTHONPATH picks up the built module, conftest.py adds python/
|
|
210
|
+
$env:PYTHONPATH = "$PWD/build" # Windows
|
|
211
|
+
PYTHONPATH=$PWD/build # Linux
|
|
212
|
+
python -m pytest tests -q
|
|
213
|
+
python benchmarks/bench.py # GPU benchmark + chart
|
|
214
|
+
```
|
|
215
|
+
|
|
216
|
+
Windows / Linux. Windows uses MSVC via nvcc; CI builds and runs the CPU test
|
|
217
|
+
suite on every push.
|
|
218
|
+
|
|
219
|
+
## Roadmap
|
|
220
|
+
|
|
221
|
+
- v0.2: bf16 support, radix-select top-k/top-p (CUB-class speed), fused
|
|
222
|
+
sampling (softmax+top-p+draw in one pass), CUDA graph-friendly batching
|
|
223
|
+
- v0.3: INT8/FP8 quantized paths, block-size autotuning
|
|
224
|
+
- v0.4+: lightweight fused attention; prebuilt wheels on PyPI
|
|
225
|
+
|
|
226
|
+
## Community
|
|
227
|
+
|
|
228
|
+
- [Contributing guide](CONTRIBUTING.md) — setup, rules of the road, PR process
|
|
229
|
+
- [Code of conduct](CODE_OF_CONDUCT.md)
|
|
230
|
+
- [Security policy](SECURITY.md)
|
|
231
|
+
- [Changelog](CHANGELOG.md)
|
|
232
|
+
|
|
233
|
+
## License
|
|
234
|
+
|
|
235
|
+
MIT — see [LICENSE](LICENSE). Third-party notices: [NOTICES.md](NOTICES.md).
|