nextok 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.
@@ -0,0 +1,8 @@
1
+ build/
2
+ .venv/
3
+ __pycache__/
4
+ *.egg-info/
5
+ *.so
6
+ *.pyd
7
+ dist/
8
+ .pytest_cache/
@@ -0,0 +1,121 @@
1
+ cmake_minimum_required(VERSION 3.18...3.30)
2
+ # The version lives in include/nextok.h (single source of truth).
3
+ file(STRINGS include/nextok.h _nextok_version_lines
4
+ REGEX "^#define NEXTOK_VERSION_(MAJOR|MINOR|PATCH) [0-9]+$")
5
+ foreach(_part MAJOR MINOR PATCH)
6
+ string(REGEX MATCH "NEXTOK_VERSION_${_part} ([0-9]+)" _ "${_nextok_version_lines}")
7
+ set(_nextok_version_${_part} ${CMAKE_MATCH_1})
8
+ endforeach()
9
+ project(nextok
10
+ VERSION ${_nextok_version_MAJOR}.${_nextok_version_MINOR}.${_nextok_version_PATCH}
11
+ LANGUAGES C CXX)
12
+
13
+ set(CMAKE_CXX_STANDARD 17)
14
+ set(CMAKE_CXX_STANDARD_REQUIRED ON)
15
+ set(CMAKE_CXX_EXTENSIONS OFF)
16
+
17
+ if(NOT CMAKE_BUILD_TYPE AND NOT CMAKE_CONFIGURATION_TYPES)
18
+ set(CMAKE_BUILD_TYPE Release CACHE STRING "" FORCE)
19
+ endif()
20
+
21
+ if(DEFINED SKBUILD)
22
+ set(_nextok_python_default ON)
23
+ else()
24
+ set(_nextok_python_default OFF)
25
+ endif()
26
+ option(NEXTOK_BUILD_PYTHON "Build the Python extension" ${_nextok_python_default})
27
+ option(NEXTOK_BUILD_TESTS "Build native tests" OFF)
28
+ option(NEXTOK_BUILD_BENCH "Build native benchmarks" OFF)
29
+ option(NEXTOK_SIMD "Build the Highway SIMD backends (runtime dispatch)" ON)
30
+ option(NEXTOK_USE_SYSTEM_HWY "Use an installed Highway instead of fetching it" OFF)
31
+
32
+ if(NEXTOK_SIMD)
33
+ if(NEXTOK_USE_SYSTEM_HWY)
34
+ find_package(hwy 1.1 CONFIG REQUIRED)
35
+ else()
36
+ # Pinned release. Set FETCHCONTENT_SOURCE_DIR_HIGHWAY to build offline.
37
+ include(FetchContent)
38
+ set(HWY_ENABLE_TESTS OFF CACHE BOOL "" FORCE)
39
+ set(HWY_ENABLE_EXAMPLES OFF CACHE BOOL "" FORCE)
40
+ set(HWY_ENABLE_CONTRIB OFF CACHE BOOL "" FORCE)
41
+ set(HWY_ENABLE_INSTALL OFF CACHE BOOL "" FORCE)
42
+ set(HWY_FORCE_STATIC_LIBS ON CACHE BOOL "" FORCE)
43
+ FetchContent_Declare(highway
44
+ URL https://github.com/google/highway/archive/refs/tags/1.3.0.tar.gz
45
+ URL_HASH SHA256=07b3c1ba2c1096878a85a31a5b9b3757427af963b1141ca904db2f9f4afe0bc2
46
+ DOWNLOAD_EXTRACT_TIMESTAMP TRUE)
47
+ set(_nextok_pic ${CMAKE_POSITION_INDEPENDENT_CODE})
48
+ set(CMAKE_POSITION_INDEPENDENT_CODE ON)
49
+ FetchContent_MakeAvailable(highway)
50
+ set(CMAKE_POSITION_INDEPENDENT_CODE ${_nextok_pic})
51
+ if(NOT TARGET hwy::hwy)
52
+ add_library(hwy::hwy ALIAS hwy)
53
+ endif()
54
+ endif()
55
+ endif()
56
+
57
+ add_library(nextok_core STATIC
58
+ src/core/fast.cpp
59
+ src/core/sampler.cpp
60
+ src/core/thread_pool.cpp
61
+ src/kernels/dispatch.cpp
62
+ src/kernels/scalar.cpp
63
+ src/reference/reference.cpp
64
+ )
65
+ if(NEXTOK_SIMD)
66
+ target_sources(nextok_core PRIVATE src/kernels/simd.cpp)
67
+ target_compile_definitions(nextok_core PRIVATE NEXTOK_HAVE_SIMD)
68
+ target_link_libraries(nextok_core PRIVATE hwy::hwy)
69
+ endif()
70
+ target_include_directories(nextok_core
71
+ PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include
72
+ PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src
73
+ )
74
+ find_package(Threads REQUIRED)
75
+ target_link_libraries(nextok_core PUBLIC Threads::Threads)
76
+ set_target_properties(nextok_core PROPERTIES POSITION_INDEPENDENT_CODE ON)
77
+ if(MSVC)
78
+ target_compile_options(nextok_core PRIVATE /W4 /fp:precise)
79
+ else()
80
+ # No FMA contraction: keeps results bit-identical across x86 and ARM.
81
+ target_compile_options(nextok_core PRIVATE -Wall -Wextra -ffp-contract=off)
82
+ endif()
83
+
84
+ if(NEXTOK_BUILD_PYTHON)
85
+ find_package(Python 3.10 REQUIRED
86
+ COMPONENTS Interpreter Development.Module
87
+ OPTIONAL_COMPONENTS Development.SABIModule)
88
+ execute_process(
89
+ COMMAND "${Python_EXECUTABLE}" -m nanobind --cmake_dir
90
+ OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE nanobind_ROOT)
91
+ find_package(nanobind CONFIG REQUIRED)
92
+
93
+ nanobind_add_module(_core STABLE_ABI NB_STATIC python/src/bindings.cpp)
94
+ target_link_libraries(_core PRIVATE nextok_core)
95
+ target_include_directories(_core PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src)
96
+ install(TARGETS _core LIBRARY DESTINATION nextok)
97
+
98
+ nanobind_add_stub(_core_stub
99
+ MODULE _core
100
+ OUTPUT _core.pyi
101
+ PYTHON_PATH $<TARGET_FILE_DIR:_core>
102
+ MARKER_FILE py.typed
103
+ DEPENDS _core)
104
+ install(FILES ${CMAKE_CURRENT_BINARY_DIR}/_core.pyi
105
+ ${CMAKE_CURRENT_BINARY_DIR}/py.typed
106
+ DESTINATION nextok)
107
+ endif()
108
+
109
+ if(NEXTOK_BUILD_TESTS)
110
+ enable_testing()
111
+ add_executable(test_capi tests/c/test_capi.c)
112
+ target_link_libraries(test_capi PRIVATE nextok_core)
113
+ set_target_properties(test_capi PROPERTIES LINKER_LANGUAGE CXX)
114
+ add_test(NAME test_capi COMMAND test_capi)
115
+ endif()
116
+
117
+ if(NEXTOK_BUILD_BENCH)
118
+ add_executable(bench_sampler bench/bench_sampler.cpp)
119
+ target_link_libraries(bench_sampler PRIVATE nextok_core)
120
+ target_include_directories(bench_sampler PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src)
121
+ endif()
nextok-0.1.0/LICENSE ADDED
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 sjjeong94
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.
nextok-0.1.0/PKG-INFO ADDED
@@ -0,0 +1,108 @@
1
+ Metadata-Version: 2.1
2
+ Name: nextok
3
+ Version: 0.1.0
4
+ Summary: Fast cross-platform LLM token sampler for CPUs (temperature, top-k/p, min-p, penalties, logit bias)
5
+ Keywords: llm,sampling,top-p,top-k,inference,simd
6
+ Author: sjjeong94
7
+ License: MIT
8
+ Classifier: Development Status :: 3 - Alpha
9
+ Classifier: Intended Audience :: Developers
10
+ Classifier: Intended Audience :: Science/Research
11
+ Classifier: License :: OSI Approved :: MIT License
12
+ Classifier: Operating System :: MacOS
13
+ Classifier: Operating System :: Microsoft :: Windows
14
+ Classifier: Operating System :: POSIX :: Linux
15
+ Classifier: Programming Language :: C++
16
+ Classifier: Programming Language :: Python :: 3
17
+ Classifier: Programming Language :: Python :: 3 :: Only
18
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
19
+ Requires-Python: >=3.10
20
+ Requires-Dist: numpy>=1.21
21
+ Description-Content-Type: text/markdown
22
+
23
+ # nextok
24
+
25
+ Fast, cross-platform LLM token sampler for CPUs, with a NumPy-friendly Python API.
26
+
27
+ Supported options: temperature, top-k, top-p, min-p, repetition / frequency / presence
28
+ penalty, logit_bias, seed. Exact semantics: [docs/semantics.md](docs/semantics.md).
29
+ Benchmarks: [docs/performance.md](docs/performance.md). For example, at vocab 128k a
30
+ top-p=0.95 draw takes ~50 µs on one core, ~190× faster than a sort-based sampler.
31
+ Roadmap: [plan.md](plan.md).
32
+
33
+ ```bash
34
+ pip install nextok # Python >= 3.10; wheels for Linux, macOS and Windows (x86-64 / arm64)
35
+ ```
36
+
37
+ ```python
38
+ import numpy as np
39
+ import nextok
40
+
41
+ sampler = nextok.Sampler(vocab_size=32000)
42
+ seq = sampler.new_sequence(prompt_ids)
43
+ token = sampler.sample(logits, temperature=0.8, top_p=0.95, seed=42, seq=seq)
44
+ ```
45
+
46
+ ## CPU backends
47
+
48
+ SIMD kernels are built with [Highway](https://github.com/google/highway) for every
49
+ supported instruction set (SSE2 through AVX-512/AVX10 on x86, NEON/SVE on ARM), and
50
+ the best one the CPU supports is picked at runtime:
51
+
52
+ ```python
53
+ >>> nextok.cpu_info()
54
+ {'backend': 'AVX2', 'available': ['AVX2', 'SSE4', 'SSSE3', 'SSE2', 'scalar']}
55
+ ```
56
+
57
+ Set `NEXTOK_BACKEND=<name>` (for example `scalar`) to override the choice.
58
+
59
+ ## Threads
60
+
61
+ `Sampler(vocab_size, num_threads=0)` samples the rows of a batch in parallel
62
+ (`num_threads=0` means up to min(physical cores, 8) threads). A single row stays on
63
+ the calling thread unless `num_threads > 1` is given explicitly. In that case a large
64
+ row (vocab >= 64k) is also split across threads, which helps when calls come back to
65
+ back. Waking idle workers can cost more than the split saves when calls are spaced
66
+ out. Results never depend on the thread count.
67
+
68
+ ## Development
69
+
70
+ Requires [uv](https://docs.astral.sh/uv/), CMake >= 3.18, and a C++17 compiler.
71
+
72
+ ```bash
73
+ uv sync # create .venv, build the extension (editable), install dev deps
74
+ uv run pytest # run tests; edits to C++ sources trigger a rebuild automatically
75
+ uv build # build sdist + wheel into dist/
76
+
77
+ # Native C ABI test
78
+ cmake -S . -B build/native -DNEXTOK_BUILD_TESTS=ON && cmake --build build/native && ctest --test-dir build/native
79
+
80
+ # Benchmarks
81
+ cmake -S . -B build/bench -DNEXTOK_BUILD_BENCH=ON && cmake --build build/bench
82
+ ./build/bench/bench_sampler --reference # C++ latency per preset, vs. scalar reference
83
+ ./build/bench/bench_sampler --backend SSE4 # force a kernel backend
84
+ ./build/bench/bench_sampler --threads 4 --batch 16 # parallel batch
85
+ uv run python bench/bench.py # Python latency vs. a NumPy sampler
86
+ ```
87
+
88
+ Highway is downloaded at configure time (pinned release). To build offline, point
89
+ `FETCHCONTENT_SOURCE_DIR_HIGHWAY` at a local checkout, or use an installed copy with
90
+ `-DNEXTOK_USE_SYSTEM_HWY=ON`. `-DNEXTOK_SIMD=OFF` builds the scalar backend only.
91
+
92
+ ### Releasing
93
+
94
+ `.github/workflows/wheels.yml` builds the sdist and wheels with cibuildwheel. That means
95
+ abi3 wheels for CPython >= 3.12, per-version wheels for 3.10–3.11, manylinux_2_28 and
96
+ musllinux for x86_64 and aarch64, macOS arm64 and x86_64, and Windows AMD64. Every wheel
97
+ runs the test suite before it is uploaded. Pushing a `v*` tag publishes to PyPI through
98
+ Trusted Publishing. Before the first release:
99
+
100
+ 1. Create a GitHub environment named `pypi`. Adding required reviewers is recommended.
101
+ 2. Register this repository as a trusted publisher of project `nextok` on pypi.org
102
+ (workflow `wheels.yml`, environment `pypi`).
103
+ 3. Bump `NEXTOK_VERSION_MAJOR/MINOR/PATCH` in `include/nextok.h`, the single source of the
104
+ version (CMake and `pyproject.toml` read it from there). Then tag:
105
+ `git tag v0.1.0 && git push --tags`.
106
+
107
+ The Python version is pinned in `.python-version`; uv prefers its managed interpreters,
108
+ which include the headers needed to build the extension.
nextok-0.1.0/README.md ADDED
@@ -0,0 +1,86 @@
1
+ # nextok
2
+
3
+ Fast, cross-platform LLM token sampler for CPUs, with a NumPy-friendly Python API.
4
+
5
+ Supported options: temperature, top-k, top-p, min-p, repetition / frequency / presence
6
+ penalty, logit_bias, seed. Exact semantics: [docs/semantics.md](docs/semantics.md).
7
+ Benchmarks: [docs/performance.md](docs/performance.md). For example, at vocab 128k a
8
+ top-p=0.95 draw takes ~50 µs on one core, ~190× faster than a sort-based sampler.
9
+ Roadmap: [plan.md](plan.md).
10
+
11
+ ```bash
12
+ pip install nextok # Python >= 3.10; wheels for Linux, macOS and Windows (x86-64 / arm64)
13
+ ```
14
+
15
+ ```python
16
+ import numpy as np
17
+ import nextok
18
+
19
+ sampler = nextok.Sampler(vocab_size=32000)
20
+ seq = sampler.new_sequence(prompt_ids)
21
+ token = sampler.sample(logits, temperature=0.8, top_p=0.95, seed=42, seq=seq)
22
+ ```
23
+
24
+ ## CPU backends
25
+
26
+ SIMD kernels are built with [Highway](https://github.com/google/highway) for every
27
+ supported instruction set (SSE2 through AVX-512/AVX10 on x86, NEON/SVE on ARM), and
28
+ the best one the CPU supports is picked at runtime:
29
+
30
+ ```python
31
+ >>> nextok.cpu_info()
32
+ {'backend': 'AVX2', 'available': ['AVX2', 'SSE4', 'SSSE3', 'SSE2', 'scalar']}
33
+ ```
34
+
35
+ Set `NEXTOK_BACKEND=<name>` (for example `scalar`) to override the choice.
36
+
37
+ ## Threads
38
+
39
+ `Sampler(vocab_size, num_threads=0)` samples the rows of a batch in parallel
40
+ (`num_threads=0` means up to min(physical cores, 8) threads). A single row stays on
41
+ the calling thread unless `num_threads > 1` is given explicitly. In that case a large
42
+ row (vocab >= 64k) is also split across threads, which helps when calls come back to
43
+ back. Waking idle workers can cost more than the split saves when calls are spaced
44
+ out. Results never depend on the thread count.
45
+
46
+ ## Development
47
+
48
+ Requires [uv](https://docs.astral.sh/uv/), CMake >= 3.18, and a C++17 compiler.
49
+
50
+ ```bash
51
+ uv sync # create .venv, build the extension (editable), install dev deps
52
+ uv run pytest # run tests; edits to C++ sources trigger a rebuild automatically
53
+ uv build # build sdist + wheel into dist/
54
+
55
+ # Native C ABI test
56
+ cmake -S . -B build/native -DNEXTOK_BUILD_TESTS=ON && cmake --build build/native && ctest --test-dir build/native
57
+
58
+ # Benchmarks
59
+ cmake -S . -B build/bench -DNEXTOK_BUILD_BENCH=ON && cmake --build build/bench
60
+ ./build/bench/bench_sampler --reference # C++ latency per preset, vs. scalar reference
61
+ ./build/bench/bench_sampler --backend SSE4 # force a kernel backend
62
+ ./build/bench/bench_sampler --threads 4 --batch 16 # parallel batch
63
+ uv run python bench/bench.py # Python latency vs. a NumPy sampler
64
+ ```
65
+
66
+ Highway is downloaded at configure time (pinned release). To build offline, point
67
+ `FETCHCONTENT_SOURCE_DIR_HIGHWAY` at a local checkout, or use an installed copy with
68
+ `-DNEXTOK_USE_SYSTEM_HWY=ON`. `-DNEXTOK_SIMD=OFF` builds the scalar backend only.
69
+
70
+ ### Releasing
71
+
72
+ `.github/workflows/wheels.yml` builds the sdist and wheels with cibuildwheel. That means
73
+ abi3 wheels for CPython >= 3.12, per-version wheels for 3.10–3.11, manylinux_2_28 and
74
+ musllinux for x86_64 and aarch64, macOS arm64 and x86_64, and Windows AMD64. Every wheel
75
+ runs the test suite before it is uploaded. Pushing a `v*` tag publishes to PyPI through
76
+ Trusted Publishing. Before the first release:
77
+
78
+ 1. Create a GitHub environment named `pypi`. Adding required reviewers is recommended.
79
+ 2. Register this repository as a trusted publisher of project `nextok` on pypi.org
80
+ (workflow `wheels.yml`, environment `pypi`).
81
+ 3. Bump `NEXTOK_VERSION_MAJOR/MINOR/PATCH` in `include/nextok.h`, the single source of the
82
+ version (CMake and `pyproject.toml` read it from there). Then tag:
83
+ `git tag v0.1.0 && git push --tags`.
84
+
85
+ The Python version is pinned in `.python-version`; uv prefers its managed interpreters,
86
+ which include the headers needed to build the extension.
@@ -0,0 +1,246 @@
1
+ # Performance
2
+
3
+ How fast nextok samples one token, and what it is compared against. All numbers
4
+ below were measured on a single machine (see [Setup](#setup)); rerun the
5
+ commands in [Reproducing](#reproducing) on your own hardware.
6
+
7
+ ## Summary
8
+
9
+ Latency per call with vocabulary 128,256 (Llama 3), batch 1, AVX2, one thread:
10
+
11
+ | preset | nextok | scalar reference | NumPy sampler |
12
+ |---|---:|---:|---:|
13
+ | greedy | **8.6 µs** | 158 µs | 13 µs (`np.argmax`) |
14
+ | temperature 0.8 | **49 µs** | 16,342 µs | 13,146 µs |
15
+ | top_k=40 | **12 µs** | 9,548 µs | 931 µs |
16
+ | top_p=0.95 | **51 µs** | 9,554 µs | 11,659 µs |
17
+ | min_p=0.05 | **10 µs** | 10,273 µs | 11,474 µs |
18
+ | top_k + top_p + min_p | **13 µs** | 9,863 µs | 882 µs |
19
+ | top_p + penalties + bias | **66 µs** | 9,388 µs | — |
20
+
21
+ - **140–1000× faster than a sort-based sampler written in C++.** The
22
+ reference is nextok's own float64 implementation of the spec. It sorts the
23
+ whole vocabulary, which is the textbook approach.
24
+ - **65–1150× faster than a typical NumPy sampler** in the Python benchmark.
25
+ Greedy is 1.4× faster than `np.argmax`.
26
+ - **Greedy, top-k and min-p run at 40–60 GB/s.** That is close to the speed of
27
+ simply reading the logits from L2 cache once. Temperature and top-p must
28
+ evaluate `exp` for every token and run at ~10 GB/s.
29
+ - **Python call overhead is 0.6 µs** (argument parsing, array validation,
30
+ releasing the GIL).
31
+
32
+ ## Setup
33
+
34
+ | | |
35
+ |---|---|
36
+ | CPU | Intel Core i5-14400F (6 P-cores + 4 E-cores; AVX2, no AVX-512) |
37
+ | OS | Ubuntu 24.04 under WSL2, Linux 6.18 |
38
+ | Compiler | GCC 13.3, `-O3`, no `-march=native` (runtime dispatch picks AVX2) |
39
+ | Kernels | Highway 1.3.0 |
40
+ | Python | CPython 3.12, NumPy 2.x |
41
+ | Date | 2026-10-31, commit after `62d3df0` |
42
+
43
+ **Input.** Each row is roughly LLM-shaped: a Gaussian floor (σ = 2) plus 8
44
+ strongly boosted tokens (+8 to +14). The same logits are reused for every call,
45
+ so they stay in L2/L3 cache. That models a sampler that runs right after the
46
+ model wrote its logits. The `top_p+penalties` preset adds a 1,024-token prompt,
47
+ 1,024 generated tokens, repetition/frequency/presence penalties and 16
48
+ `logit_bias` entries.
49
+
50
+ **Method.** `bench/bench_sampler.cpp` calls `nt_sample_batch` back to back and
51
+ reports the median of 15 repetitions of ~10 ms each. "GB/s" means
52
+ `vocab × 4 bytes × batch / time`, i.e. how fast the input is consumed.
53
+
54
+ ## nextok vs. a sort-based sampler
55
+
56
+ The scalar reference (`src/reference/`) implements the spec directly: apply
57
+ penalties, sort all tokens, then truncate and draw, in float64. It is what a
58
+ straightforward C++ sampler looks like, and nextok is tested against it for
59
+ correctness.
60
+
61
+ Speedup of nextok (AVX2, 1 thread) over the reference, batch 1:
62
+
63
+ | preset | V=32,000 | V=128,256 | V=151,936 | V=256,000 |
64
+ |---|---:|---:|---:|---:|
65
+ | greedy | 19× | 18× | 18× | 20× |
66
+ | temperature | 303× | 335× | 307× | 347× |
67
+ | top_k=40 | 442× | 801× | 824× | 924× |
68
+ | top_p=0.95 | 154× | 189× | 190× | 40× |
69
+ | min_p=0.05 | 929× | 1001× | 1063× | 1248× |
70
+ | top_k + top_p + min_p | 417× | 784× | 795× | 913× |
71
+ | top_p + penalties | 76× | 141× | 151× | 31× |
72
+
73
+ Absolute latency of nextok in µs:
74
+
75
+ | preset | V=32,000 | V=128,256 | V=151,936 | V=256,000 |
76
+ |---|---:|---:|---:|---:|
77
+ | greedy | 2.4 | 8.6 | 11.2 | 16.8 |
78
+ | temperature | 11.8 | 48.8 | 63.9 | 94.6 |
79
+ | top_k=40 | 4.6 | 11.9 | 13.0 | 22.5 |
80
+ | top_p=0.95 | 12.9 | 50.6 | 61.9 | 578 |
81
+ | min_p=0.05 | 2.3 | 10.3 | 10.2 | 18.3 |
82
+ | top_k + top_p + min_p | 4.4 | 12.6 | 12.8 | 23.0 |
83
+ | top_p + penalties | 26.3 | 66.4 | 74.7 | 693 |
84
+
85
+ Where the speed comes from:
86
+
87
+ - **No sorting.** Each parameter combination takes its own path. top-k uses a
88
+ streaming heap seeded with a safe threshold, min-p uses a threshold
89
+ compaction, and top-p first guesses a small candidate set. Only the few
90
+ tokens that can matter are ever ordered.
91
+ - **No full copy.** Penalties and `logit_bias` become a sparse patch list that
92
+ is applied block by block on the fly.
93
+ - **SIMD kernels** for max, `exp`, compaction and search, chosen at runtime
94
+ for the CPU (see below).
95
+
96
+ The top-p numbers at V=256,000 are a deliberately bad case; see
97
+ [Known slow case](#known-slow-case).
98
+
99
+ ## nextok vs. NumPy (Python)
100
+
101
+ `bench/bench.py` compares `Sampler.sample` with a common NumPy implementation:
102
+ `argpartition` for top-k, `argsort` plus `cumsum` for top-p, and
103
+ `rng.choice` for the draw. Microseconds per call:
104
+
105
+ | preset | V=32,000 | | V=128,256 | | V=256,000 | |
106
+ |---|---:|---:|---:|---:|---:|---:|
107
+ | | nextok | NumPy | nextok | NumPy | nextok | NumPy |
108
+ | greedy | 2.9 | 4.3 | 9.2 | 13.1 | 19.3 | 25.5 |
109
+ | temperature | 14.1 | 2,405 | 59.2 | 13,146 | 103 | 25,145 |
110
+ | top_k=40 | 4.8 | 225 | 13.2 | 931 | 22.5 | 1,538 |
111
+ | top_p=0.95 | 13.1 | 2,429 | 53.5 | 11,659 | 102 | 25,437 |
112
+ | min_p=0.05 | 2.9 | 2,437 | 10.0 | 11,474 | 18.7 | 23,849 |
113
+ | top_k + top_p + min_p | 6.4 | 234 | 13.3 | 882 | 24.3 | 1,564 |
114
+
115
+ The Python numbers include the 0.6 µs binding overhead. Otherwise they match
116
+ the C++ numbers above.
117
+
118
+ ## SIMD backends
119
+
120
+ The same binary contains kernels for SSE2 through AVX-512/AVX10 (x86) and
121
+ NEON/SVE (ARM), plus a portable scalar fallback. The best one is picked at
122
+ runtime, and `NEXTOK_BACKEND` forces one. µs per call at V=128,256, batch 1:
123
+
124
+ | preset | scalar | SSE4 | AVX2 |
125
+ |---|---:|---:|---:|
126
+ | greedy | 106 | 10.4 | 8.6 |
127
+ | temperature | 597 | 135 | 48.8 |
128
+ | top_k=40 | 135 | 16.3 | 11.9 |
129
+ | top_p=0.95 | 634 | 141 | 50.6 |
130
+ | min_p=0.05 | 123 | 11.8 | 10.3 |
131
+ | top_p + penalties | 569 | 150 | 66.4 |
132
+
133
+ - **Greedy, top-k and min-p** are mostly comparisons and are already limited by
134
+ memory speed with SSE4.
135
+ - **Temperature and top-p** are dominated by `exp`. They gain 2.8× from SSE4
136
+ to AVX2: twice the vector width, plus fused multiply-add.
137
+ - **AVX-512 and NEON** kernels are compiled in but have not been measured yet.
138
+ This CPU has no AVX-512.
139
+
140
+ ## Threads
141
+
142
+ ### One row split across threads
143
+
144
+ With an explicit `num_threads > 1`, a row with vocab ≥ 64k is split into parts
145
+ that run in parallel. µs per call, batch 1, back-to-back calls:
146
+
147
+ | preset | V=128,256, 1 thread | V=128,256, 4 threads | V=256,000, 1 thread | V=256,000, 4 threads |
148
+ |---|---:|---:|---:|---:|
149
+ | greedy | 8.6 | 4.0 | 16.8 | 5.7 |
150
+ | temperature | 48.8 | 23.1 | 94.6 | 29.5 |
151
+ | top_k=40 | 11.9 | 7.4 | 22.5 | 8.2 |
152
+ | top_p=0.95 | 50.6 | 22.7 | 578 | 260 |
153
+ | min_p=0.05 | 10.3 | 5.1 | 18.3 | 6.3 |
154
+ | top_p + penalties | 66.4 | 45.3 | 693 | 300 |
155
+
156
+ At V=32,000 a row is not split (each part needs ≥ 32k tokens), so the numbers
157
+ match one thread. Results are bit-identical for every thread count.
158
+
159
+ ### Batches
160
+
161
+ Rows of a batch are spread across threads. That is what the default
162
+ `num_threads=0` does. µs per call for the whole batch, back-to-back calls:
163
+
164
+ | preset (V=128,256) | batch 16, 1 thread | batch 16, 6 threads | batch 64, 1 thread | batch 64, 6 threads |
165
+ |---|---:|---:|---:|---:|
166
+ | greedy | 136 | 29 | 586 | 108 |
167
+ | temperature | 762 | 161 | 3,187 | 591 |
168
+ | top_k=40 | 177 | 41 | 804 | 141 |
169
+ | top_p=0.95 | 857 | 169 | 3,557 | 646 |
170
+ | min_p=0.05 | 135 | 30 | 550 | 109 |
171
+ | top_p + penalties | 1,145 | 238 | 4,864 | 1,000 |
172
+
173
+ With 6 threads the batch is 4.7–5.7× faster. Throughput reaches 280–300 GB/s for
174
+ greedy and min-p, which is past single-core cache bandwidth.
175
+
176
+ ### Calls with idle time in between
177
+
178
+ The numbers above are for back-to-back calls. In a real decode loop the
179
+ sampler waits ~10–50 ms for the model between calls. Caches cool down, cores
180
+ clock down, and idle worker threads go to sleep. With `--gap-us 20000` (20 ms
181
+ of sleep before each call), V=128,256, median µs per call:
182
+
183
+ | preset | batch 1, 1 thread | batch 1, 4 threads | batch 16, 1 thread | batch 16, 6 threads |
184
+ |---|---:|---:|---:|---:|
185
+ | greedy | 53 | 130 | 195 | 366 |
186
+ | temperature | 159 | 245 | 916 | 598 |
187
+ | top_k=40 | 73 | 260 | 308 | 360 |
188
+ | top_p=0.95 | 154 | 244 | 1,034 | 731 |
189
+ | top_p + penalties | 266 | 261 | 1,974 | 1,124 |
190
+
191
+ - **Every call is 3–6× slower after an idle gap, even with one thread.** This
192
+ comes from the machine (cold caches, frequency scaling), not from the
193
+ sampler. Absolute numbers here are noisy, especially under WSL2.
194
+ - **Splitting a single row does not pay off with idle gaps.** Waking sleeping
195
+ workers costs more than the parallel speedup. That is why the default
196
+ (`num_threads=0`) keeps single rows on the calling thread.
197
+ - **Batches still gain** for exp-heavy presets (temperature, top-p), but not
198
+ for the cheap ones (greedy, top-k).
199
+
200
+ ## Known slow case
201
+
202
+ top-p at V=256,000 takes 578 µs, ten times the V=151,936 number. The synthetic
203
+ input makes the cause visible. With 256k Gaussian tokens, the tail holds over
204
+ 5% of the probability mass, so `top_p=0.95` has to keep tens of thousands of
205
+ tail tokens. nextok's quick path (guessing a small candidate set) gives up
206
+ above 8,192 candidates and falls back to a weighted histogram over the whole
207
+ row. That costs a second `exp` pass plus a scattered histogram update per
208
+ token.
209
+
210
+ Real LLM distributions are usually far more peaked, so they stay on the quick
211
+ path. Even this case is 40× faster than the sort-based reference, and 4
212
+ threads bring it to 260 µs.
213
+
214
+ ## Not measured yet
215
+
216
+ - **llama.cpp's sampler chain, PyTorch/vLLM samplers.** These are the most
217
+ relevant baselines, but they could not be built in this environment.
218
+ `bench/bench_sampler.cpp` shows the input data and presets to mirror.
219
+ - llama.cpp applies temperature *after* top-k/top-p, whereas nextok follows
220
+ vLLM's order (see [semantics.md](semantics.md)). The work is comparable,
221
+ but outputs are not identical.
222
+ - **AVX-512, ARM NEON/SVE, macOS and Windows.** These will be covered by the CI
223
+ matrix (P4/P5).
224
+ - **Real model logits** instead of synthetic rows.
225
+
226
+ ## Reproducing
227
+
228
+ ```bash
229
+ # C++: build the benchmark
230
+ cmake -S . -B build/bench -DNEXTOK_BUILD_BENCH=ON && cmake --build build/bench
231
+ B=./build/bench/bench_sampler
232
+
233
+ $B --reference # vs. the sort-based reference (all vocab sizes)
234
+ $B --backend scalar # force a backend (scalar, SSE4, AVX2, AVX3, NEON, ...)
235
+ $B --threads 4 # split single rows across 4 threads
236
+ $B --batch 16 --threads 6 # batch parallelism
237
+ $B --gap-us 20000 # 20 ms idle before each call (decode-loop conditions)
238
+ $B --quick ... # shorter runs, for any of the above
239
+
240
+ # Python: nextok vs. NumPy, plus call overhead
241
+ uv run python bench/bench.py --vocab 128256
242
+ ```
243
+
244
+ Expect different absolute numbers on other machines. On a quiet bare-metal
245
+ Linux system, the idle-gap and threading numbers in particular should look
246
+ better than under WSL2.