Niko1221/Strata
1
0 Fork 0
mirror of https://github.com/Niko1221/Strata.git synced 2026-11-07 06:15:58 -10:00
Code Issues Packages Projects Releases 1 Wiki Activity

Strata: Qwen3.8-Flash-Next on a 12-24 GB NVIDIA GPU + 64 GB RAM, one-click install

Browse Source 
Inference engine, OpenAI/Anthropic API on localhost, optional image input, one-file setup
(START-HERE.bat / setup.sh). Model files are downloaded from Hugging Face by the setup.
This commit is contained in:
Niko1221 committed 2026-09-24 18:42:20 +02:00
commit 01529aaa82
219 files changed
+50165
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

.gitattributes
+6
View File
@@ -0,0 +1,6 @@
# Frozen prompts are hashed byte for byte (manifest sha256): never convert their line endings.
bench/prompts/** -text
# launchers keep their platform line endings wherever the repository is cloned
*.sh text eol=lf
*.bat text eol=crlf
.gitignore
+48
View File
@@ -0,0 +1,48 @@
# Generated, not source: deterministic build products of tools/ (see docs/).
# tiny-qwen4exp.gguf is ~113 MB, over GitHub's 100 MB file limit -- regenerate with
# python tools/make_tiny_model.py
bench/tiny-*.gguf
bench/wide-*.gguf
bench/micro/*.exe
bench/micro/*.obj
bench/micro/*.bin
bench/micro/*.dll
bench/micro/*.txt
*.log
build*/
pack/
logs/
__pycache__/
.pytest_cache/
*.pyc
# plan v0.3 P1 (repository hygiene, review E1): profiler captures and large diagnostic dumps stay out of git.
*.nsys-rep
*.sqlite
*.npz
*.occ1
td.log
# created by setup.py on the user's machine
/.venv/
/build/
/third_party/llama.cpp/
/models/
/packs/
/mtp/
/run-*.bat
/run-*.sh
/strata-*.json
/strata-*.log
/build-strata.bat
# image encoder (setup.py --vision): build output and the generated build script
build-vision*/
build-vision.bat
# setup.py outputs (engine, downloads) and release builds
/engine/
/build-release*/
/dist/
*.done
/third_party/*.zip
CMakeLists.txt
+536
View File
@@ -0,0 +1,536 @@
# Strata Engine - top-level build (P1.S1).
#
# Targets (phase-1 §P1.S1): `strata_artifact` (no CUDA dependency), `tests`. `strata_plan` and the C++
# `strata-pack` are added when their sources exist - a target with no sources would make the build graph
# claim something the tree does not have, which is worse than an absent target.
#
# CUDA targets arrive in Phase 2 and this build REJECTS any architecture other than 120: the engine is
# written against sm_120 (RTX 5000 / Blackwell) and silently building for an older arch would produce a
# binary that runs and is wrong about its own performance, which is the worst of both outcomes.
cmake_minimum_required(VERSION 3.24)
project(strata VERSION 0.1.0 LANGUAGES CXX)
# The prebuilt release (setup.py downloads it, so users need no compiler): ggml's CPU backend for any AVX2 CPU
# instead of this machine's (Strata's own AVX-512 kernels are per-file and chosen at run time either way), and
# the MSVC runtime linked in, so no Visual C++ redistributable is needed.
option(STRATA_PORTABLE "build for other PCs: AVX2 baseline, static MSVC runtime" OFF)
if(STRATA_PORTABLE AND MSVC)
set(CMAKE_MSVC_RUNTIME_LIBRARY "MultiThreaded")
endif()
set(CMAKE_CXX_STANDARD 20)
set(CMAKE_CXX_STANDARD_REQUIRED ON)
set(CMAKE_CXX_EXTENSIONS OFF)
if(NOT CMAKE_BUILD_TYPE AND NOT CMAKE_CONFIGURATION_TYPES)
set(CMAKE_BUILD_TYPE Release CACHE STRING "" FORCE)
endif()
option(STRATA_ENABLE_CUDA "Build the CUDA targets (Phase 2+)" OFF)
# The test suite and the micro benchmarks live in tests/ and bench/micro/, which the published source leaves out.
set(_strata_tests_default OFF)
if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/tests/CMakeLists.txt AND EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/bench/micro)
set(_strata_tests_default ON)
endif()
option(STRATA_BUILD_TESTS "Build the test suite" ${_strata_tests_default})
option(STRATA_WERROR "Treat compiler warnings as errors" OFF)
if(STRATA_ENABLE_CUDA)
enable_language(CUDA)
# `enable_language(CUDA)` gives the .cu compile rules but NOT the imported targets, so `CUDA::cudart` has to
# be brought in explicitly for the host-side tests that call cudaMalloc themselves.
find_package(CUDAToolkit REQUIRED)
if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
set(CMAKE_CUDA_ARCHITECTURES 120 CACHE STRING "CUDA architectures" FORCE)
endif()
foreach(_a IN LISTS CMAKE_CUDA_ARCHITECTURES)
string(REGEX REPLACE "-.*$" "" _base "${_a}") # 120-real / 120-virtual are still 120
# Plan v0.3 P6: developed and measured on sm_120 (RTX 50); the kernels need sm_80 or newer (tf32 mma in
# the attention scorer, bf16 math), so RTX 30 (86) and RTX 40 (89) build too. Older cards are refused.
if(_base STREQUAL "native" OR _base STREQUAL "all" OR _base STREQUAL "all-major")
continue()
endif()
if(_base LESS 80)
message(FATAL_ERROR
"Strata needs an NVIDIA GPU of compute capability 8.0 or newer (RTX 30 / 40 / 50 series), but "
"CMAKE_CUDA_ARCHITECTURES is '${CMAKE_CUDA_ARCHITECTURES}'.")
endif()
endforeach()
message(STATUS "Strata: CUDA enabled, arch ${CMAKE_CUDA_ARCHITECTURES}")
endif()
# Warnings live on an INTERFACE target the project's own code links, NOT in add_compile_options: applied
# globally they would also hit FetchContent'd Catch2, whose -W4 warnings are not this project's to fix and
# whose /WX failures would be indistinguishable from our own.
#
# Every flag is guarded by COMPILE_LANGUAGE:CXX. Without that guard nvcc receives bare `/W4 /utf-8` and dies
# with "A single input file is required for a non-link phase" - it parses them as inputs rather than host
# flags. Host flags for CUDA must go through -Xcompiler, which is what COMPILE_LANGUAGE:CUDA does below.
add_library(strata_warnings INTERFACE)
target_compile_options(strata_warnings INTERFACE
$<$<AND:$<COMPILE_LANGUAGE:CXX>,$<CXX_COMPILER_ID:MSVC>>:/W4>
$<$<AND:$<COMPILE_LANGUAGE:CXX>,$<CXX_COMPILER_ID:MSVC>>:/utf-8>
$<$<AND:$<COMPILE_LANGUAGE:CXX>,$<NOT:$<CXX_COMPILER_ID:MSVC>>>:-Wall>
$<$<AND:$<COMPILE_LANGUAGE:CXX>,$<NOT:$<CXX_COMPILER_ID:MSVC>>>:-Wextra>
$<$<AND:$<COMPILE_LANGUAGE:CUDA>,$<CXX_COMPILER_ID:MSVC>>:-Xcompiler=/W4>)
if(STRATA_WERROR)
target_compile_options(strata_warnings INTERFACE
$<$<AND:$<COMPILE_LANGUAGE:CXX>,$<CXX_COMPILER_ID:MSVC>>:/WX>
$<$<AND:$<COMPILE_LANGUAGE:CXX>,$<NOT:$<CXX_COMPILER_ID:MSVC>>>:-Werror>)
endif()
# ---- strata_artifact -------------------------------------------------------
# Header-only, and that is a deliberate property rather than an unfinished split: the GGUF reader and the
# scalar dequantizers hold no state worth hiding, and a header-only boundary cannot introduce a
# duplicate-symbol or missing-declaration bug into code that is already validated against the real artifact.
add_library(strata_artifact INTERFACE)
target_include_directories(strata_artifact INTERFACE ${CMAKE_CURRENT_SOURCE_DIR}/include)
target_link_libraries(strata_artifact INTERFACE strata_warnings)
add_executable(strata-gguf src/artifact/gguf_reader.cpp)
target_link_libraries(strata-gguf PRIVATE strata_artifact)
add_executable(strata-dequant src/artifact/dequant.cpp)
target_link_libraries(strata-dequant PRIVATE strata_artifact)
# this file is now nothing but the CLI entry point, and its inherited guard spells that as a "selftest"
target_compile_definitions(strata-dequant PRIVATE STRATA_DEQUANT_SELFTEST)
# ---- strata_plan ----------------------------------------------------------
# Also header-only. A "planner" that hides state would be one whose decision is not inspectable, and the
# whole point of P1.S9 is that the decision is arithmetic you can check by hand.
add_library(strata_plan INTERFACE)
target_include_directories(strata_plan INTERFACE ${CMAKE_CURRENT_SOURCE_DIR}/include)
target_link_libraries(strata_plan INTERFACE strata_warnings)
add_executable(strata-plan src/plan/plan_main.cpp)
target_link_libraries(strata-plan PRIVATE strata_plan)
# ---- strata_core (CUDA, Phase 2) ------------------------------------------
# Only built when CUDA is on, so the Phase 1 targets keep working on a machine with no toolkit. The arch
# guard above has already refused anything but sm_120 by the time this runs.
if(STRATA_ENABLE_CUDA)
add_library(strata_core STATIC src/core/device.cu src/core/pinned.cu src/platform/memory.cpp src/core/graph.cpp src/core/weights.cpp
src/core/layout.cpp)
target_include_directories(strata_core PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include)
# CUDA::cudart is PUBLIC because graph.cpp is a HOST translation unit that calls cudaGraphLaunch and
# cudaEventQuery: the .cu sources get cudart through the CUDA language rules, a .cpp does not, and the
# linker error ("unresolved external symbol cudaGraphInstantiate") names the function rather than the cause.
target_link_libraries(strata_core PUBLIC strata_plan strata_warnings CUDA::cudart)
add_executable(strata-device src/core/device_main.cpp)
target_link_libraries(strata-device PRIVATE strata_core)
add_executable(strata-load src/core/load_main.cpp)
target_link_libraries(strata-load PRIVATE strata_core strata_kernels)
add_executable(strata-concurrent src/core/concurrent_main.cpp)
target_include_directories(strata-concurrent PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
target_link_libraries(strata-concurrent PRIVATE strata_core strata_kernels CUDA::cudart)
add_executable(strata-overlap src/core/overlap_main.cpp)
# a .cpp that calls cudaMemcpyAsync itself needs the toolkit headers and cudart explicitly
target_include_directories(strata-overlap PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
target_link_libraries(strata-overlap PRIVATE strata_core strata_kernels CUDA::cudart)
# The load test needs the 34 GB artifact, so it is registered ONLY where that artifact exists. A test that
# silently skips is a test that stops running; an absent registration is visible in the configure output.
if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/pack/full/experts.bin)
add_test(NAME expert_arena_load
COMMAND strata-load --file ${CMAKE_CURRENT_SOURCE_DIR}/pack/full/experts.bin --layers 1)
set_tests_properties(expert_arena_load PROPERTIES PASS_REGULAR_EXPRESSION "per-layer checksums")
else()
message(STATUS "strata: pack/full/experts.bin absent - the expert arena load test is NOT registered")
endif()
# ---- strata_kernels (P2.S2) ----
# Naive first, each with a parity test against a reference that was itself validated against ggml. The
# kernels link `strata_artifact` for the scalar dequantizers, because using the VALIDATED scalar path as the
# reference is the whole point - a kernel checked against another copy of itself proves nothing.
add_library(strata_kernels STATIC src/kernels/cuda/dequant_s2.cu src/kernels/cuda/s2_gemv.cu
src/kernels/cuda/s_gemv.cu
src/kernels/cuda/s2_gemv_quads.cu
src/kernels/cuda/s2_gemv_fast.cu
src/kernels/cuda/router_top10.cu
src/kernels/cuda/quantize_act.cu
src/kernels/cuda/rope.cu
src/kernels/cuda/sampler.cu
src/kernels/cuda/s2_gemv_q8.cu
src/kernels/cuda/s2_expert_grouped.cu
src/kernels/cuda/shared_expert.cu
src/kernels/cuda/gr.cu
src/kernels/cuda/gdn.cu
src/kernels/cuda/qsa.cu
src/kernels/cuda/elementwise.cu
src/kernels/cuda/bf16_gemv.cu
src/kernels/cuda/native_mmvq.cu
src/kernels/cuda/native_bf16.cu
src/kernels/cuda/native_gr_norm.cu
src/kernels/cuda/native_moe.cu
src/kernels/cuda/native_gdn.cu
src/kernels/cuda/native_gr_postops.cu
src/kernels/cuda/native_router.cu
src/kernels/cuda/native_gdn_preprocess.cu
src/kernels/cuda/native_qsa.cu
src/kernels/cuda/native_rope.cu
src/kernels/cuda/native_ple_postops.cu
src/kernels/cuda/native_flash_attn.cu
src/kernels/cuda/native_qsa_indexer.cu
src/kernels/cuda/native_qsa_score.cu
src/kernels/cuda/ple.cu
src/kernels/cuda/kv_q8.cu
src/kernels/cuda/fused_gr.cu
src/kernels/cuda/qsa_decode_attn.cu
src/kernels/cuda/fused_gdn.cu
src/kernels/cuda/dequant_bf16.cu
src/kernels/cuda/qsa_select.cu
src/kernels/cuda/verify_kernels.cu
src/kernels/cuda/iq_kernels.cu
src/kernels/ngram.cpp
src/ngram/ple_reader.cpp
src/platform/direct_file.cpp)
target_link_libraries(strata_kernels PUBLIC strata_core strata_artifact)
# plan v0.3 P6: llama.cpp's block layouts and codebook grids for the i-quant formats (MIT, third_party/ggml)
target_include_directories(strata_kernels PRIVATE third_party/ggml)
# Match the pinned llama.cpp CUDA arithmetic contract for this adapter only.
set_source_files_properties(src/kernels/cuda/native_mmvq.cu src/kernels/cuda/native_bf16.cu
src/kernels/cuda/native_gr_norm.cu
src/kernels/cuda/native_moe.cu
src/kernels/cuda/native_gdn.cu
src/kernels/cuda/native_gr_postops.cu
src/kernels/cuda/native_router.cu
src/kernels/cuda/native_gdn_preprocess.cu
src/kernels/cuda/native_qsa.cu
src/kernels/cuda/native_rope.cu
src/kernels/cuda/native_ple_postops.cu
src/kernels/cuda/native_flash_attn.cu
src/kernels/cuda/native_qsa_indexer.cu
src/kernels/cuda/native_qsa_score.cu
PROPERTIES COMPILE_OPTIONS "--use_fast_math")
set(STRATA_ORACLE_SOURCE_DIR "" CACHE PATH "Pinned llama.cpp source for independent CUDA parity tests")
set(STRATA_ORACLE_BUILD_DIR "" CACHE PATH "Matching existing shared llama.cpp build")
if(STRATA_ORACLE_SOURCE_DIR AND STRATA_ORACLE_BUILD_DIR)
foreach(parity IN ITEMS native_mmvq native_bf16 native_shared native_gr_norm native_moe native_gdn native_gr_postops native_router native_gdn_preprocess native_gdn_pipeline native_qsa native_rope native_ple_postops native_flash_attn native_qsa_indexer native_qsa_score)
add_executable(${parity}_parity bench/micro/${parity}_parity.cpp)
target_include_directories(${parity}_parity PRIVATE
"${STRATA_ORACLE_SOURCE_DIR}/ggml/include" ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
target_compile_definitions(${parity}_parity PRIVATE GGML_SHARED GGML_BACKEND_SHARED)
target_link_libraries(${parity}_parity PRIVATE strata_kernels CUDA::cudart)
foreach(component IN ITEMS ggml ggml-base ggml-cpu ggml-cuda)
string(REPLACE "-" "_" variable "STRATA_ORACLE_${component}")
find_library(${variable} NAMES ${component} PATHS
"${STRATA_ORACLE_BUILD_DIR}/ggml/src" "${STRATA_ORACLE_BUILD_DIR}/ggml/src/${component}"
"${STRATA_ORACLE_BUILD_DIR}/bin" NO_DEFAULT_PATH REQUIRED)
target_link_libraries(${parity}_parity PRIVATE "${${variable}}")
endforeach()
add_test(NAME ${parity}_parity COMMAND ${parity}_parity --selftest)
if(WIN32)
set_tests_properties(${parity}_parity PROPERTIES
ENVIRONMENT_MODIFICATION "PATH=path_list_prepend:${STRATA_ORACLE_BUILD_DIR}/bin")
endif()
endforeach()
endif()
# ---- strata_engine (P2.S5) -----------------------------------------------------
# The COMPOSED layer lives here and not in `strata_core`, because it calls kernels and `strata_kernels`
# already links `strata_core` - putting `gdn_layer` in core would make the two mutually dependent. A third
# library is the honest shape of the dependency graph: core knows about devices and weights, kernels know
# about arithmetic, and this one knows the ORDER the arithmetic happens in.
add_library(strata_engine STATIC src/core/layer.cpp src/core/session.cpp src/core/expert_source.cpp
src/core/expert_cache.cpp src/core/native_head.cpp src/core/native_dense.cpp src/core/verify.cpp src/core/mtp.cpp)
target_include_directories(strata_engine PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include)
target_link_libraries(strata_engine PUBLIC strata_core strata_kernels strata_kernels_cpu CUDA::cudart)
# `strata generate` (P2.S6). The FIRST user-facing program: everything under it is a component, and this is
# what composes them into a token. No tokenizer - the prompt arrives as ids - because Gate C1 compares logits
# against llama.cpp on the SAME ids and a tokenizer on one side of that comparison is a second variable.
# plan v0.3 P5: batched prompt processing (tensor-core GEMMs through cuBLAS).
add_library(strata_prefill STATIC src/prefill/gemm.cu src/prefill/kernels.cu src/prefill/prefill.cpp)
target_include_directories(strata_prefill PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
target_link_libraries(strata_prefill PUBLIC strata_engine CUDA::cublas CUDA::cudart)
add_executable(strata src/program/generate.cpp)
target_include_directories(strata PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
target_link_libraries(strata PRIVATE strata_engine strata_prefill CUDA::cudart)
add_executable(dequant_s2_parity src/kernels/dequant_s2_parity.cpp)
# The tests are HOST C++ files that call cudaMalloc/cudaMemcpy themselves, so they need the toolkit headers
# and cudart explicitly: enabling the CUDA language gives those to .cu sources automatically, but a .cpp
# target is just a C++ target and gets neither.
target_include_directories(dequant_s2_parity PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
target_link_libraries(dequant_s2_parity PRIVATE strata_kernels CUDA::cudart)
add_test(NAME dequant_s2_parity COMMAND dequant_s2_parity --selftest)
add_executable(s2_gemv_parity src/kernels/s2_gemv_parity.cpp)
target_include_directories(s2_gemv_parity PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
target_link_libraries(s2_gemv_parity PRIVATE strata_kernels CUDA::cudart)
add_test(NAME s2_gemv_parity COMMAND s2_gemv_parity --selftest)
add_executable(iq_parity src/kernels/iq_parity.cpp)
target_include_directories(iq_parity PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
target_link_libraries(iq_parity PRIVATE strata_kernels CUDA::cudart)
add_executable(shared_expert_parity src/kernels/shared_expert_parity.cpp)
target_include_directories(shared_expert_parity PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
target_link_libraries(shared_expert_parity PRIVATE strata_kernels CUDA::cudart)
add_test(NAME shared_expert_parity COMMAND shared_expert_parity --selftest)
add_executable(gr_parity src/kernels/gr_parity.cpp)
target_include_directories(gr_parity PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
target_link_libraries(gr_parity PRIVATE strata_kernels CUDA::cudart)
add_test(NAME gr_parity COMMAND gr_parity --selftest)
add_executable(gdn_parity src/kernels/gdn_parity.cpp)
target_include_directories(gdn_parity PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
target_link_libraries(gdn_parity PRIVATE strata_kernels CUDA::cudart)
add_test(NAME gdn_parity COMMAND gdn_parity --selftest)
add_executable(s2_gemv_q8_parity src/kernels/s2_gemv_q8_parity.cpp)
target_include_directories(s2_gemv_q8_parity PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
target_link_libraries(s2_gemv_q8_parity PRIVATE strata_kernels CUDA::cudart)
add_test(NAME s2_gemv_q8_parity COMMAND s2_gemv_q8_parity --selftest)
add_executable(sampler_parity src/kernels/sampler_parity.cpp)
target_include_directories(sampler_parity PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
target_link_libraries(sampler_parity PRIVATE strata_kernels CUDA::cudart)
add_test(NAME sampler_parity COMMAND sampler_parity --selftest)
add_executable(rope_parity src/kernels/rope_parity.cpp)
target_include_directories(rope_parity PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
target_link_libraries(rope_parity PRIVATE strata_kernels CUDA::cudart)
add_test(NAME rope_parity COMMAND rope_parity --selftest)
add_executable(quantize_act_parity src/kernels/quantize_act_parity.cpp)
target_include_directories(quantize_act_parity PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
target_link_libraries(quantize_act_parity PRIVATE strata_kernels CUDA::cudart)
add_test(NAME quantize_act_parity COMMAND quantize_act_parity --selftest)
add_executable(router_top10_parity src/kernels/router_top10_parity.cpp)
target_include_directories(router_top10_parity PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
target_link_libraries(router_top10_parity PRIVATE strata_kernels CUDA::cudart)
add_test(NAME router_top10_parity COMMAND router_top10_parity --selftest)
add_executable(s_gemv_parity src/kernels/s_gemv_parity.cpp)
target_include_directories(s_gemv_parity PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
target_link_libraries(s_gemv_parity PRIVATE strata_kernels CUDA::cudart)
add_test(NAME s_gemv_parity COMMAND s_gemv_parity --selftest)
add_executable(elementwise_parity src/kernels/elementwise_parity.cpp)
target_include_directories(elementwise_parity PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
target_link_libraries(elementwise_parity PRIVATE strata_kernels CUDA::cudart)
add_test(NAME elementwise_parity COMMAND elementwise_parity --selftest)
add_executable(bf16_gemv_parity src/kernels/bf16_gemv_parity.cpp)
target_include_directories(bf16_gemv_parity PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
target_link_libraries(bf16_gemv_parity PRIVATE strata_kernels CUDA::cudart)
add_test(NAME bf16_gemv_parity COMMAND bf16_gemv_parity --selftest)
add_executable(s_gemv_q8k_parity src/kernels/s_gemv_q8k_parity.cpp)
target_include_directories(s_gemv_q8k_parity PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
target_link_libraries(s_gemv_q8k_parity PRIVATE strata_kernels CUDA::cudart)
add_test(NAME s_gemv_q8k_parity COMMAND s_gemv_q8k_parity --selftest)
add_executable(qsa_parity src/kernels/qsa_parity.cpp)
target_include_directories(qsa_parity PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
target_link_libraries(qsa_parity PRIVATE strata_kernels CUDA::cudart)
add_test(NAME qsa_parity COMMAND qsa_parity --selftest)
# ---- P2.S4: the PLE n-gram hash, the IQ4_NL table read, and the PLE block ----
# `ngram.cpp` is host-only but lives in `strata_kernels` because `ple.cu` shares its geometry constants, and
# splitting them across two libraries would only add a layer.
#
# The pack, the original GGUF and the ggml capture are all REQUIRED inputs: this test's oracles are real
# artifacts, so the working directory is pinned to the source tree and a missing input is a hard error
# INSIDE the test. A test that passes by finding nothing to compare is the failure mode this project keeps
# re-learning.
add_executable(ple_parity src/kernels/ple_parity.cpp)
target_include_directories(ple_parity PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
target_link_libraries(ple_parity PRIVATE strata_kernels CUDA::cudart)
set(STRATA_PLE_FIXTURE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/bench/micro" CACHE PATH
"Directory containing required ple_in.bin and ple_out.bin parity fixtures")
add_test(NAME ple_parity COMMAND ple_parity --selftest
--in "${STRATA_PLE_FIXTURE_DIR}/ple_in.bin" --out "${STRATA_PLE_FIXTURE_DIR}/ple_out.bin"
WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR})
# Plan v0.3 P7: INT8 KV storage vs a host reference and vs the FP16 path (GPU, synthetic, no model).
add_executable(kv_q8_parity src/kernels/kv_q8_parity.cpp)
target_include_directories(kv_q8_parity PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
target_link_libraries(kv_q8_parity PRIVATE strata_kernels CUDA::cudart)
add_test(NAME kv_q8_parity COMMAND kv_q8_parity)
# Plan v0.3 P3: native MMVQ for 1..8 tokens (bitwise selftest + --bench timing); synthetic data, no model.
if(STRATA_BUILD_TESTS)
add_executable(native_mmvq_multi bench/micro/native_mmvq_multi.cpp)
target_include_directories(native_mmvq_multi PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
target_link_libraries(native_mmvq_multi PRIVATE strata_kernels CUDA::cudart)
add_test(NAME native_mmvq_multi COMMAND native_mmvq_multi --selftest)
endif()
# Plan v0.3 P0.1: working-set lock of a host region (CPU only).
add_executable(platform_memory_test src/platform/memory_test.cpp)
target_link_libraries(platform_memory_test PRIVATE strata_core CUDA::cudart)
add_test(NAME platform_memory_test COMMAND platform_memory_test)
# Plan v0.3 P2: the n-gram table read straight from the SSD. The selftest needs only the CPU and a temp file.
add_executable(ple_reader_test src/ngram/ple_reader_test.cpp)
target_include_directories(ple_reader_test PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
target_link_libraries(ple_reader_test PRIVATE strata_kernels CUDA::cudart)
add_test(NAME ple_reader_selftest COMMAND ple_reader_test --selftest)
# The GPU test is registered only when CUDA is on, and it runs ON THE TARGET PC. `--selftest` allocates
# and writes real VRAM, so a machine without a device fails this test rather than skipping it silently -
# which is correct: an engine that targets one GPU should not report a pass for having no GPU.
add_test(NAME cuda_device_selftest COMMAND strata-device --selftest)
# ---- P2.S5: the GraphRegistry --------------------------------------------------
# A .cu because the test needs real kernels to capture; it links `strata_core` and not `strata_kernels`,
# because what it exercises is the capture machinery and not any kernel in it.
if(STRATA_BUILD_TESTS)
add_executable(graph_registry_test bench/micro/graph_registry_test.cu)
target_include_directories(graph_registry_test PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
target_link_libraries(graph_registry_test PRIVATE strata_core CUDA::cudart)
add_test(NAME graph_registry_test COMMAND graph_registry_test)
# ---- P2.S5's GATE: is a device<->pinned-memory handoff graph-capturable? --------------------
# The doorbell protocol puts the D2H of `x_f` and the miss list INSIDE G[l] - that is the whole overlap. The
# phase doc says to verify this "before building on it", and this project has already lost two rounds to
# building first (a memory-only spin that DEADLOCKS on Windows, and a `cudaEventRecord` inside a capture that
# is SILENTLY DROPPED). It links only cudart: it is about the capture machinery, not about any kernel.
add_executable(pinned_capture bench/micro/pinned_capture.cu)
target_include_directories(pinned_capture PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
target_link_libraries(pinned_capture PRIVATE CUDA::cudart)
add_test(NAME pinned_capture COMMAND pinned_capture)
endif()
# Tells tests/CMakeLists.txt to add the CUDA-dependent Catch2 target. It cannot be declared there
# unconditionally: tests/ is added at the end of this file, but the flag has to be known before it runs.
set(STRATA_CUDA_TESTS ON)
endif()
# ---- strata_kernels_cpu (P2.S3) ----
#
# A SEPARATE LIBRARY from `strata_kernels`, and `/arch:AVX512` is set on the SOURCE rather than on the target
# so that the flag applies only to the TU that needs it. Both parts matter:
#
# * the CPU expert kernel uses `_mm256_dpbusd_epi32` (AVX512-VNNI) and `_mm256_multishift_epi64_epi8`
# (AVX512-VBMI); without the flag MSVC reports undeclared intrinsics and never mentions AVX-512;
# * a TU compiled WITH `/arch:AVX512` may use AVX-512 in ANY of its code, so the scalar fallback and the
# test drivers must NOT be in it - otherwise the fallback is not a fallback, it is the same bet with
# worse opt-out. `cpu_require_expert_support()` therefore refuses at startup on a CPU without the
# features, and the fallback exists for tests, exactly as P2.S3 specifies.
# Plan v0.3 P6: speculation components that need no GPU.
add_library(strata_spec STATIC src/spec/suffix_drafter.cpp src/spec/controller.cpp)
target_include_directories(strata_spec PUBLIC include)
if(STRATA_BUILD_TESTS)
add_executable(suffix_drafter_test src/spec/suffix_drafter_test.cpp)
target_link_libraries(suffix_drafter_test PRIVATE strata_spec)
add_test(NAME suffix_drafter_test COMMAND suffix_drafter_test)
add_executable(controller_test src/spec/controller_test.cpp)
target_link_libraries(controller_test PRIVATE strata_spec)
add_test(NAME controller_test COMMAND controller_test)
endif()
add_library(strata_kernels_cpu STATIC src/kernels/cpu/expert.cpp src/kernels/cpu/pool.cpp
src/kernels/cpu/expert_layout.cpp src/kernels/cpu/q2_avx2.cpp)
target_include_directories(strata_kernels_cpu PUBLIC include)
# Plan v0.3 P6: only the kernel files get the AVX-512 / AVX2 flags; the pool and the probe run on any x86-64,
# and the engine picks the kernels at run time (cpu_avx512_ok).
if(MSVC)
set_source_files_properties(src/kernels/cpu/expert.cpp PROPERTIES COMPILE_OPTIONS "/arch:AVX512")
set_source_files_properties(src/kernels/cpu/q2_avx2.cpp PROPERTIES COMPILE_OPTIONS "/arch:AVX2")
else()
# GCC/Clang spell it differently, and `-mavx512vnni` alone is not enough: VBMI is a separate feature bit.
set_source_files_properties(src/kernels/cpu/expert.cpp
PROPERTIES COMPILE_OPTIONS "-mavx512f;-mavx512bw;-mavx512vl;-mavx512dq;-mavx512vnni;-mavx512vbmi;-mf16c;-mfma")
set_source_files_properties(src/kernels/cpu/q2_avx2.cpp PROPERTIES COMPILE_OPTIONS "-mavx2;-mfma;-mf16c")
endif()
find_package(Threads REQUIRED)
target_link_libraries(strata_kernels_cpu PUBLIC Threads::Threads)
# ---- plan v0.3 P6: the i-quant models' CPU experts (IQ2_XS / IQ3_XXS files) ----
#
# The CPU computes a native expert with ggml-cpu's own dot products (`ggml_get_type_traits_cpu`), so an IQ
# weight means exactly what it means in llama.cpp. ggml is built static from a llama.cpp checkout: a local one
# with -DSTRATA_GGML_DIR=..., otherwise the pinned upstream commit (third_party/ggml/VERSION.txt).
option(STRATA_NATIVE_EXPERTS "Native (IQ) expert support through ggml-cpu" ON)
set(STRATA_GGML_DIR "" CACHE PATH "llama.cpp checkout to take ggml from (empty: fetch the pinned commit)")
if(STRATA_NATIVE_EXPERTS)
enable_language(C)
set(BUILD_SHARED_LIBS OFF CACHE BOOL "" FORCE)
set(GGML_STATIC ON CACHE BOOL "" FORCE)
set(GGML_CUDA OFF CACHE BOOL "" FORCE)
set(GGML_OPENMP OFF CACHE BOOL "" FORCE)
set(GGML_BACKEND_DL OFF CACHE BOOL "" FORCE)
set(GGML_BUILD_TESTS OFF CACHE BOOL "" FORCE)
set(GGML_BUILD_EXAMPLES OFF CACHE BOOL "" FORCE)
if(STRATA_PORTABLE)
set(GGML_NATIVE OFF CACHE BOOL "" FORCE)
foreach(_f GGML_AVX GGML_AVX2 GGML_FMA GGML_F16C GGML_BMI2)
set(${_f} ON CACHE BOOL "" FORCE)
endforeach()
else()
set(GGML_NATIVE ON CACHE BOOL "" FORCE)
endif()
if(STRATA_GGML_DIR)
add_subdirectory(${STRATA_GGML_DIR}/ggml ${CMAKE_BINARY_DIR}/ggml EXCLUDE_FROM_ALL)
else()
include(FetchContent)
FetchContent_Declare(strata_llamacpp
GIT_REPOSITORY https://github.com/ggml-org/llama.cpp.git
GIT_TAG 3cf03257f219afbe7334045ff7c6a06ac68c627d
GIT_SHALLOW FALSE)
FetchContent_GetProperties(strata_llamacpp)
if(NOT strata_llamacpp_POPULATED)
FetchContent_Populate(strata_llamacpp)
endif()
add_subdirectory(${strata_llamacpp_SOURCE_DIR}/ggml ${CMAKE_BINARY_DIR}/ggml EXCLUDE_FROM_ALL)
endif()
# part of strata_kernels_cpu (the pool calls these and they call its CPU probe: one library, so GNU ld's
# one-pass static linking has no cycle to resolve)
target_sources(strata_kernels_cpu PRIVATE src/kernels/cpu/native_expert.cpp src/kernels/cpu/iq_avx512.cpp)
target_include_directories(strata_kernels_cpu PRIVATE third_party/ggml)
if(MSVC)
set_source_files_properties(src/kernels/cpu/iq_avx512.cpp PROPERTIES COMPILE_OPTIONS "/arch:AVX512")
else()
set_source_files_properties(src/kernels/cpu/iq_avx512.cpp PROPERTIES COMPILE_OPTIONS
"-mavx512f;-mavx512bw;-mavx512vl;-mavx512dq;-mavx512vnni;-mavx512vbmi;-mf16c;-mfma")
endif()
target_link_libraries(strata_kernels_cpu PUBLIC ggml-cpu ggml-base)
target_compile_definitions(strata_kernels_cpu PUBLIC STRATA_NATIVE_EXPERTS=1)
if(STRATA_ENABLE_CUDA)
add_executable(native_expert_parity src/kernels/native_expert_parity.cpp)
target_include_directories(native_expert_parity PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
target_link_libraries(native_expert_parity PRIVATE strata_kernels strata_kernels_cpu strata_artifact CUDA::cudart)
endif()
endif()
if(STRATA_ENABLE_CUDA AND STRATA_BUILD_TESTS)
add_executable(hit_cpu_order_parity bench/micro/hit_cpu_order_parity.cu)
target_link_libraries(hit_cpu_order_parity PRIVATE strata_kernels strata_kernels_cpu CUDA::cudart)
endif()
if(STRATA_BUILD_TESTS)
# Plan v0.3 P6: one expert applied to several tokens, bitwise vs the single-token kernel (synthetic, CPU only).
# plan v0.3 P5: the device dequantizers vs the artifact's CPU dequantizers on the real shards (not a ctest: needs
# the 38 GB model; run it by hand or from tools/gpu_session.py)
add_executable(dequant_bf16_test src/kernels/dequant_bf16_test.cpp)
target_link_libraries(dequant_bf16_test PRIVATE strata_kernels strata_artifact CUDA::cudart)
add_executable(expert_multi_test src/kernels/cpu/expert_multi_test.cpp)
target_link_libraries(expert_multi_test PRIVATE strata_kernels_cpu)
add_test(NAME expert_multi_test COMMAND expert_multi_test)
add_executable(expert_parity src/kernels/cpu/expert_parity.cpp)
target_link_libraries(expert_parity PRIVATE strata_kernels_cpu)
# The pack is a REQUIRED input, not an optional one: this test compares against a real expert blob, and a
# missing file would otherwise be reported as a pass by a test that skipped everything.
add_test(NAME expert_parity COMMAND expert_parity --selftest
WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR})
add_executable(pool_test src/kernels/cpu/pool_test.cpp)
target_link_libraries(pool_test PRIVATE strata_kernels_cpu)
add_test(NAME pool_test COMMAND pool_test --selftest WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR})
endif()
if(STRATA_BUILD_TESTS)
enable_testing()
add_subdirectory(tests)
endif()
README.md
+313
View File
@@ -0,0 +1,313 @@
<h1 align="center">Qwen3.8-Flash-Next (125B MoE) on a 12-24 GB NVIDIA GPU + 64 GB RAM: one-click install for Windows / Linux</h1>
<p align="center"><b>Strata</b>: an inference engine built for exactly one model and exactly this kind of PC</p>
A 125-billion-parameter mixture-of-experts model (6B active per token) on a **normal gaming PC**: one NVIDIA card with
12-24 GB of VRAM, 64 GB of RAM and a desktop CPU. Strata keeps the always-used part of the model on the GPU, all 24,576
experts in RAM, computes the experts the GPU does not hold on the CPU at the same time, and checks several tokens per
pass with the model's own multi-token-prediction (MTP) layer. The text it produces is identical to plain greedy decoding.
It serves an **OpenAI-** and **Anthropic-compatible API on localhost**, so any chat app, coding agent or script that
talks to those APIs can use it. Optionally it also **reads images** (screenshots, photos, scanned pages) with the
model's own vision encoder.
> **Jump to:** [Speed](#speed-measured) · [Other GPUs](#other-gpus-estimated) · [Which model?](#which-model) ·
> [Windows](#windows) · [Linux](#linux) · [Using it](#using-it) · [Images](#images-vision) ·
> [Troubleshooting](#troubleshooting) · [How it works](#how-it-works)
---
## Speed (measured)
RTX 5070 **12 GB**, Ryzen 5 7600 (6 cores), 64 GB DDR5-5200, Windows. One code-agent prompt per length, 256 generated
tokens, MTP speculative decoding on. "262K" is the model's full context window (a 259,943-token prompt).
### Prompt processing (tokens/s)
| Model | 1K | 4K | 32K | 64K | 128K | 262K |
| --- | ---: | ---: | ---: | ---: | ---: | ---: |
| **Q2_0** | 389 | 539 | 571 | 561 | 543 | 496 |
| **IQ2_XS** | 332 | 463 | 495 | 486 | 472 | 437 |
| **IQ3_XXS** | 285 | 410 | 435 | 427 | 414 | - |
### Output (tokens/s)
| Model | 1K | 4K | 32K | 64K | 128K | 262K |
| --- | ---: | ---: | ---: | ---: | ---: | ---: |
| **Q2_0** | 88.7 | 94.6 | 87.5 | 76.4 | 65.1 | 56.3 |
| **IQ2_XS** | 82.0 | 78.0 | 65.3 | 63.7 | 52.0 | 48.0 |
| **IQ3_XXS** | 64.6 | 65.6 | 57.3 | 54.4 | 44.8 | - |
IQ3_XXS at 262K is not measured: with its 43 GB of experts, a 260K-token context brings a 64 GB PC to its memory
limit. Use up to 128K with IQ3_XXS on 64 GB.
Time to first token is prompt length / prompt speed: about 7 s at 4K, 55 s at 32K, 4 minutes at 128K and 9 minutes at
262K. The raw numbers: `bench/results/`. The [paper](docs/paper/Strata-Paper.pdf) explains every number.
## Other GPUs (estimated)
Not measured - estimated from the runs above (same CPU and 64 GB RAM): the GPU part scaled by memory bandwidth, the CPU
part by how many more experts the card's VRAM holds. Treat as **±20%**. Numbers are *prompt / output* tokens/s.
| GPU | Model | 1K | 4K | 32K | 64K | 128K | 262K |
| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: |
| RTX 5060 Ti 16GB | Q2_0 | ~341 / ~80 | ~472 / ~87 | ~501 / ~81 | ~492 / ~72 | ~476 / ~62 | ~435 / ~53 |
| | IQ2_XS | ~291 / ~80 | ~406 / ~77 | ~434 / ~63 | ~426 / ~62 | ~413 / ~51 | ~383 / ~47 |
| | IQ3_XXS | ~249 / ~66 | ~359 / ~65 | ~381 / ~56 | ~374 / ~54 | ~363 / ~45 | - |
| RTX 3090 24GB | Q2_0 | ~355 / ~128 | ~491 / ~140 | ~521 / ~130 | ~512 / ~115 | ~495 / ~100 | ~453 / ~85 |
| | IQ2_XS | ~303 / ~131 | ~422 / ~128 | ~451 / ~103 | ~444 / ~102 | ~430 / ~85 | ~398 / ~78 |
| | IQ3_XXS | ~260 / ~106 | ~374 / ~103 | ~396 / ~89 | ~390 / ~85 | ~378 / ~71 | - |
More VRAM matters more than a faster GPU: every extra GB holds ~700 more experts, and every expert on the GPU is one the
CPU does not have to compute. A 3090's 24 GB takes most of the CPU work away.
## Which model?
All three are [ISTA-DASLab's GSQ-RCO quantizations](https://huggingface.co/ISTA-DASLab/Qwen3.8-Flash-Next-GSQ-RCO-GGUF)
of [Qwen3.8-Flash-Next](https://huggingface.co/Qwen/Qwen3.8-Flash-Next).
| Model | Download | RAM it uses | Speed | Quality |
| --- | ---: | ---: | --- | --- |
| **Q2_0** | 66 GB | ~34 GB experts + ~6 GB | fastest | good |
| **IQ2_XS** | 68 GB | ~36 GB experts + ~6 GB | close to Q2_0 | a bit better |
| **IQ3_XXS** | 76 GB | ~43 GB experts + ~6 GB | slower (more CPU work) | best |
With 64 GB of RAM all three fit (close the browser for IQ3_XXS, and keep its context at 128K or less). With 48 GB only Q2_0 / IQ2_XS may fit. 32 GB is not enough.
## Before you start
You need **only an NVIDIA driver** (version 580 or newer; update it with the NVIDIA App or from
[nvidia.com/drivers](https://www.nvidia.com/drivers)). Everything else is installed for you the first time.
| | |
| --- | --- |
| GPU | NVIDIA **RTX 30, 40 or 50 series**, **12 GB VRAM or more** (8 GB runs, slowly). Measured on an RTX 5070; RTX 30/40 are untested. |
| RAM | **64 GB** recommended (see the table above). |
| CPU | x86-64 with AVX2 (any Intel/AMD desktop CPU from the last ~8 years). AVX-512 (Ryzen 7000/9000) is a bit faster. |
| Disk | ~70-80 GB free for the model, ~6 GB for the MTP layer (+1 GB with images). **Q2_0 on an AVX-512 CPU** also writes a one-time ~40 GB copy of its experts for the fast CPU kernel. An NVMe SSD is strongly recommended. |
| OS | Windows 10/11, or Linux (Ubuntu 22.04/24.04 get everything installed automatically). |
What the first start installs, all inside this folder (`.venv/`, `engine/`, `third_party/`, `models/`, `packs/`, `mtp/`):
Python 3.12 if you have none (for your user account, no admin), a private Python environment, NVIDIA's CUDA libraries
(from pip, ~0.4 GB), the ready-made Strata engine for RTX 30/40/50, the model and the MTP draft layer. If no
ready-made engine fits your PC, it offers to install the build tools (Visual Studio Build Tools + CUDA Toolkit on
Windows, `build-essential` + CUDA on Ubuntu) and compiles the engine for your GPU (asks first; 20-40 minutes once).
---
## Windows
### Double-click `START-HERE.bat`
**The first time** it asks three questions and does the rest:
1. **Which model?** Q2_0, IQ2_XS or IQ3_XXS (it recommends one for your RAM).
2. **How much context?** 8K to 256K tokens (it recommends one for your VRAM).
3. **Images?** yes / no (see [Images](#images-vision)).
Then it downloads and prepares everything (the model is 66-76 GB, so the first start takes a while; an interrupted
download continues where it stopped) and **starts the model**: your browser opens `http://127.0.0.1:8080`, a small
page that shows it is running and lets you chat. The API is at `http://127.0.0.1:8080/v1` for your apps.
**Every time after that**, `START-HERE.bat` just starts the model (30-90 s to load 34-43 GB into RAM). Nothing is
downloaded again. Closing the window stops the model.
```
START-HERE.bat --setup install another model, or change context / images
START-HERE.bat --model IQ2_XS --context 32768 --vision yes --yes no questions
START-HERE.bat --gguf-dir D:\models\IQ2_XS use GGUF files you already have
START-HERE.bat --port 8081 another port
```
With more than one model installed, it asks which one to start. `run-<model>.bat` starts a model directly.
### Chat in the terminal (optional)
```
.venv\Scripts\python chat.py
```
---
## Linux
```bash
./setup.sh
```
The same three questions, the same automatic install (it uses `sudo apt` for Python and, only if it has to compile,
for the build tools), and the same start: `http://127.0.0.1:8080`. Later runs of `./setup.sh` (or `./run-<model>.sh`)
start the model directly. Options as on Windows (`./setup.sh --setup`, `--model Q2_0 --yes`, `--gguf-dir /data/Q2_0`).
Terminal chat: `.venv/bin/python chat.py`.
---
## Using it
The server listens on `http://127.0.0.1:8080` (change with `--port` in setup, or edit the run script).
| API | Endpoint |
| --- | --- |
| OpenAI Chat Completions (stream and non-stream, tools) | `POST /v1/chat/completions` |
| Anthropic Messages (stream and non-stream, tools) | `POST /v1/messages` |
| Model list / health | `GET /v1/models`, `GET /health` |
```bash
curl http://127.0.0.1:8080/v1/chat/completions -H "Content-Type: application/json" -d '{
"model": "strata", "messages": [{"role": "user", "content": "Write a haiku about GPUs."}], "max_tokens": 512 }'
```
```python
from openai import OpenAI
client = OpenAI(base_url="http://127.0.0.1:8080/v1", api_key="none")
r = client.chat.completions.create(model="strata", messages=[{"role": "user", "content": "Hello!"}])
print(r.choices[0].message.content)
```
- **Thinking levels: none, low, medium, high.** The model thinks before it answers (streamed as
`reasoning_content`, Anthropic: `thinking` blocks). Choose how much per request - in the chat page (the "Thinking"
menu), in `chat.py` (`/think low`), or over the API:
| API | how |
| --- | --- |
| OpenAI | `"reasoning_effort": "none" \| "low" \| "medium" \| "high"` (also `"reasoning": {"effort": ...}`, or `"chat_template_kwargs": {"enable_thinking": false}`) |
| Anthropic | `"output_config": {"effort": "low" \| "medium" \| "high"}`, `"thinking": {"type": "disabled"}`, or `"thinking": {"type": "enabled", "budget_tokens": N}` (under 2K = low, under 8K = medium, more = high) |
Without a setting the model uses its own default, **high**. `none` answers at once (fastest); `low` keeps the thinking
short. The levels are instructions the model was trained with, not a hard token limit: on easy questions all three
think briefly, on hard ones `high` thinks longest and is most accurate.
- **Chat apps.** Any app with an "OpenAI-compatible" provider works: base URL `http://127.0.0.1:8080/v1`, any API key.
- **Context.** Chosen in setup (8K-262K). Requests longer than that are refused, never silently cut.
- **From other devices / the internet.** The server listens on your PC only (`127.0.0.1`). To reach it from elsewhere,
put a tunnel in front of it, for example [cloudflared](https://developers.cloudflare.com/cloudflare-one/connections/connect-networks/do-more-with-tunnels/trycloudflare/):
`cloudflared tunnel --url http://127.0.0.1:8080`. **Set a key first**, or anyone with the link can use your PC:
add `"api_key": "some-long-secret"` to `strata-<model>.json` (or set the `STRATA_API_KEY` environment variable);
clients then send it as their API key.
**Current limits (v1):** one request at a time; greedy decoding (temperature is ignored); every request processes its
whole prompt again (no conversation cache yet, so long chats have a long time-to-first-token); images only when set up
with them (below); no video.
---
## Images (vision)
The model has a vision encoder: [`mmproj-Qwen3.8-Flash-Next-BF16.gguf`](https://huggingface.co/ISTA-DASLab/Qwen3.8-Flash-Next-GSQ-RCO-GGUF)
(0.9 GB, a 27-layer ViT plus the projector into the language model). It is **optional**: say yes when the setup asks
"Images?", or run it again with `--vision gpu` (or `--vision cpu`). The setup downloads the encoder, builds a small
helper (`strata-vision`, from llama.cpp's `mtmd` library) and adds it to your start script. Nothing else changes.
| Encoder on | Time per picture | Cost |
| --- | --- | --- |
| **GPU** (recommended) | **0.1-0.5 s** (up to 1,024 image tokens) | ~1.4 GB of VRAM is kept free for it, so the expert cache is smaller: text output is a few % slower (table below) |
| CPU | 10-30 s (pictures are scaled down to ~300 image tokens) | nothing on the GPU |
A picture becomes up to 1,024 tokens of the context (a 640x480 photo: 300). The same picture sent again, as chat apps
do on every turn, is encoded only once.
### Sending a picture
**Terminal chat:** type `/image <path to a picture>`, press Enter, then type your question.
```
you> /image C:\Users\me\Pictures\receipt.jpg
(picture attached: receipt.jpg - now type your question)
you> What is the total on this receipt?
```
**OpenAI API** (an `image_url` part: a `data:` URL, an `http(s)://` URL or a local file path):
```python
import base64
from openai import OpenAI
client = OpenAI(base_url="http://127.0.0.1:8080/v1", api_key="none")
img = base64.b64encode(open("photo.jpg", "rb").read()).decode()
r = client.chat.completions.create(model="strata", messages=[{"role": "user", "content": [
{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{img}"}},
{"type": "text", "text": "What is in this picture?"}]}])
print(r.choices[0].message.content)
```
**Anthropic API:** an `image` block with a `base64` (or `url`) source, as usual.
JPEG, PNG, BMP and GIF work (WebP does not). Chat apps with image upload work the same way.
### Speed with images on (4K context, measured)
| Model | Prompt tok/s, images off | Prompt tok/s, images on | Output tok/s, images off | Output tok/s, images on |
| --- | ---: | ---: | ---: | ---: |
| **Q2_0** | 541 | 531 | 90.7 | 86.9 |
| **IQ2_XS** | 468 | 458 | 77.1 | 74.0 |
| **IQ3_XXS** | 411 | 401 | 64.7 | 59.5 |
"Off" is the published setup; "on" runs with the encoder loaded on the GPU and its VRAM kept free, so ~1,000 fewer
experts fit in VRAM and the CPU computes a few more per token: 2-8% slower. For text, turning images on changes
nothing else (the output is bit-identical when the VRAM is the same).
A question about a picture (a 640x480 newspaper page = 300 image tokens, a 328-token prompt; the whole request, encoder
on the GPU, measured through the API):
| Model | Answer ("MEN WALK ON MOON") | Picture encoded | Prompt | Output |
| --- | ---: | ---: | ---: | ---: |
| **Q2_0** | 1.9 s | ~0.1 s | 207 tok/s | 65-73 tok/s |
| **IQ2_XS** | 2.3 s | ~0.1 s | 173 tok/s | 50-60 tok/s |
| **IQ3_XXS** | 2.7 s | ~0.1 s | 144 tok/s | 41-50 tok/s |
(Short prompts run below the 4K prompt speed: a 2,048-token chunk is where the prompt path is efficient. Short answers
run below the long-output speed: the first rounds have no draft yet.)
**How it works inside:** the encoder turns the picture into rows of the same width as the model's word embeddings;
Strata puts them where the prompt has `<|image_pad|>` tokens and gives each one its 2-D position (row and column in the
picture; the model uses interleaved M-RoPE). Answers match llama.cpp's multimodal implementation token for token on our
test images.
---
## Troubleshooting
| Symptom | What to do |
| --- | --- |
| `the NVIDIA driver is too old` | Update the driver (NVIDIA App or nvidia.com/drivers), restart, run `START-HERE.bat` again. |
| Python or the build tools could not be installed | Install what it names (links are printed), then run it again. Everything already done is kept. |
| `port 8080 is already in use` | Strata is already running (look for its window), or another program uses the port: `START-HERE.bat --port 8081`. |
| `cudaHostRegister ... out of memory` in the log | Normal on Windows: the engine pins the experts in per-layer slices instead. Only a problem if the load then fails. |
| The first start takes minutes | It is reading 34-43 GB into RAM; the second start is faster while the files are in the OS cache. |
| Slow output, disk light busy | Not enough free RAM: close other programs, or choose Q2_0 / IQ2_XS. |
| `prompt ... exceeds the context` | The request is longer than the context you chose: run setup again with a bigger `--context`. |
| Slower than the tables | The monitor plugged into the GPU and other GPU programs take VRAM from the expert cache; RAM running below its rated speed (enable EXPO/XMP in the BIOS) slows the CPU half. |
| `this server was started without the vision encoder` | The model was set up for text only: run setup again with `--vision gpu`. |
| A picture is refused or `cannot read the image` | Use JPEG, PNG, BMP or GIF (not WebP). |
| Pictures are slow (10-30 s) | The encoder runs on the CPU: run setup again with `--vision gpu` (needs ~1.4 GB of VRAM). |
| Anything else | The engine log is `strata-<model>.log` in this folder. |
---
## How it works
<p align="center"><img src="docs/paper/tiers.svg" width="760" alt="memory tiers"></p>
- **GPU (VRAM):** attention and DeltaNet mixers, the gated-residual weights, routers, shared experts, output head, the MTP
draft layer, the KV cache, and an **expert cache** that fills the rest of VRAM with the most-used experts (it adapts to
the conversation while you chat).
- **RAM:** all 24,576 experts, pinned. The CPU computes the experts that are not on the GPU **in place**, at the same time
as the GPU works on the cached ones (AVX-512 / AVX2 kernels, ggml's for the i-quants).
- **SSD:** the 28.8 GB n-gram table, read a few rows per token through the OS cache.
- **Speculation:** the model's own MTP layer drafts up to 3 tokens; one pass over all 48 layers checks them. 2.4-3.2
tokens per pass on average.
- **Prompts** are processed in 2,048-token chunks with the experts streamed to the GPU over PCIe.
The full story, with measurements, bottlenecks and what comes next: **[docs/paper/Strata-Paper.pdf](docs/paper/Strata-Paper.pdf)**.
---
## Credits and licenses
- Model: [Qwen/Qwen3.8-Flash-Next](https://huggingface.co/Qwen/Qwen3.8-Flash-Next) by the Qwen team; quantizations:
[ISTA-DASLab/Qwen3.8-Flash-Next-GSQ-RCO-GGUF](https://huggingface.co/ISTA-DASLab/Qwen3.8-Flash-Next-GSQ-RCO-GGUF).
Their licenses apply to the weights.
- [llama.cpp / ggml](https://github.com/ggml-org/llama.cpp) (MIT): the i-quant formats, the GPU dot products and
dequantizers transcribed in `src/kernels/cuda/iq_kernels.cu`, the CPU backend linked for the i-quant experts, the
`mtmd` library behind the image encoder (`tools/vision/`), and `gguf-py` used by the tools. See
`third_party/ggml/LICENSE`.
- Ideas from [Splash](https://github.com/incoai/splash), [ninfer](https://github.com/Neroued/ninfer) and
[HyperQwen](https://github.com/syv-ai/HyperQwen); references in the paper.
START-HERE.bat
+52
View File
@@ -0,0 +1,52 @@
@echo off
rem Strata for Windows: the first run installs everything and starts the model; later runs just start it.
rem Needs only an NVIDIA driver. Python is installed for your user account if it is missing (no admin needed).
setlocal
title Strata
cd /d "%~dp0"
if exist ".venv\Scripts\python.exe" goto run
call :findpy
if defined PY goto venv
echo.
echo Python 3.10 or newer is not installed. Installing Python 3.12 for your user account ...
where winget >nul 2>nul
if errorlevel 1 goto pyorg
winget install -e --id Python.Python.3.12 --scope user --silent --source winget --accept-package-agreements --accept-source-agreements --disable-interactivity
call :findpy
if defined PY goto venv
:pyorg
echo Downloading the Python installer from python.org ...
powershell -NoProfile -ExecutionPolicy Bypass -Command "[Net.ServicePointManager]::SecurityProtocol='Tls12'; Invoke-WebRequest -UseBasicParsing https://www.python.org/ftp/python/3.12.10/python-3.12.10-amd64.exe -OutFile \"$env:TEMP\strata-python-setup.exe\""
if exist "%TEMP%\strata-python-setup.exe" "%TEMP%\strata-python-setup.exe" /quiet InstallAllUsers=0 PrependPath=1 Include_launcher=1 Include_test=0
call :findpy
if defined PY goto venv
echo.
echo Python could not be installed automatically.
echo Install 64-bit Python 3.12 from https://www.python.org/downloads/ ("Add python.exe to PATH"),
echo then double-click START-HERE.bat again.
pause
exit /b 1
:venv
rem a private environment inside this folder, so nothing is installed into the system Python
%PY% -m venv .venv
if exist ".venv\Scripts\python.exe" goto run
echo Could not create the Python environment in .venv
pause
exit /b 1
:run
".venv\Scripts\python.exe" setup.py %*
if errorlevel 1 pause
exit /b
:findpy
rem the py launcher first, then python on PATH (not the Microsoft Store stub), then the usual per-user folders
set "PY="
py -3 -c "import sys; sys.exit(0 if sys.version_info >= (3, 10) and sys.maxsize > 2**32 else 1)" >nul 2>nul
if not errorlevel 1 set "PY=py -3" & goto :eof
python -c "import sys; sys.exit(0 if sys.version_info >= (3, 10) and sys.maxsize > 2**32 else 1)" >nul 2>nul
if not errorlevel 1 set "PY=python" & goto :eof
for %%V in (313 312 311 310) do if exist "%LOCALAPPDATA%\Programs\Python\Python%%V\python.exe" set "PY="%LOCALAPPDATA%\Programs\Python\Python%%V\python.exe"" & goto :eof
goto :eof
bench/results/2026-09-24-final/estimates.json
+461
View File
@@ -0,0 +1,461 @@
[
{
"card": "RTX 5070 12GB (measured)",
"model": "Q2_0",
"tier": "1k",
"slots": 4517,
"h": 0.747,
"decode": 88.7,
"prefill": 389.0
},
{
"card": "RTX 5060 Ti 16GB",
"model": "Q2_0",
"tier": "1k",
"slots": 7179,
"h": 0.826,
"decode": 80.5,
"prefill": 341.0
},
{
"card": "RTX 3090 24GB",
"model": "Q2_0",
"tier": "1k",
"slots": 12503,
"h": 0.93,
"decode": 127.5,
"prefill": 355.0
},
{
"card": "RTX 5070 12GB (measured)",
"model": "Q2_0",
"tier": "4k",
"slots": 4464,
"h": 0.722,
"decode": 94.6,
"prefill": 539.0
},
{
"card": "RTX 5060 Ti 16GB",
"model": "Q2_0",
"tier": "4k",
"slots": 7126,
"h": 0.809,
"decode": 87.4,
"prefill": 472.0
},
{
"card": "RTX 3090 24GB",
"model": "Q2_0",
"tier": "4k",
"slots": 12450,
"h": 0.923,
"decode": 140.0,
"prefill": 491.0
},
{
"card": "RTX 5070 12GB (measured)",
"model": "Q2_0",
"tier": "32k",
"slots": 4174,
"h": 0.674,
"decode": 87.5,
"prefill": 571.0
},
{
"card": "RTX 5060 Ti 16GB",
"model": "Q2_0",
"tier": "32k",
"slots": 6836,
"h": 0.774,
"decode": 81.3,
"prefill": 501.0
},
{
"card": "RTX 3090 24GB",
"model": "Q2_0",
"tier": "32k",
"slots": 12160,
"h": 0.907,
"decode": 130.0,
"prefill": 521.0
},
{
"card": "RTX 5070 12GB (measured)",
"model": "Q2_0",
"tier": "64k",
"slots": 3809,
"h": 0.625,
"decode": 76.4,
"prefill": 561.0
},
{
"card": "RTX 5060 Ti 16GB",
"model": "Q2_0",
"tier": "64k",
"slots": 6471,
"h": 0.74,
"decode": 71.6,
"prefill": 492.0
},
{
"card": "RTX 3090 24GB",
"model": "Q2_0",
"tier": "64k",
"slots": 11795,
"h": 0.892,
"decode": 115.3,
"prefill": 512.0
},
{
"card": "RTX 5070 12GB (measured)",
"model": "Q2_0",
"tier": "128k",
"slots": 3079,
"h": 0.571,
"decode": 65.1,
"prefill": 543.0
},
{
"card": "RTX 5060 Ti 16GB",
"model": "Q2_0",
"tier": "128k",
"slots": 5741,
"h": 0.702,
"decode": 61.8,
"prefill": 476.0
},
{
"card": "RTX 3090 24GB",
"model": "Q2_0",
"tier": "128k",
"slots": 11065,
"h": 0.873,
"decode": 100.0,
"prefill": 495.0
},
{
"card": "RTX 5070 12GB (measured)",
"model": "Q2_0",
"tier": "262k",
"slots": 1588,
"h": 0.526,
"decode": 56.3,
"prefill": 496.0
},
{
"card": "RTX 5060 Ti 16GB",
"model": "Q2_0",
"tier": "262k",
"slots": 4250,
"h": 0.668,
"decode": 52.9,
"prefill": 435.0
},
{
"card": "RTX 3090 24GB",
"model": "Q2_0",
"tier": "262k",
"slots": 9574,
"h": 0.849,
"decode": 84.6,
"prefill": 453.0
},
{
"card": "RTX 5070 12GB (measured)",
"model": "IQ2_XS",
"tier": "1k",
"slots": 4621,
"h": 0.799,
"decode": 82.0,
"prefill": 332.0
},
{
"card": "RTX 5060 Ti 16GB",
"model": "IQ2_XS",
"tier": "1k",
"slots": 7176,
"h": 0.86,
"decode": 79.6,
"prefill": 291.0
},
{
"card": "RTX 3090 24GB",
"model": "IQ2_XS",
"tier": "1k",
"slots": 12287,
"h": 0.941,
"decode": 131.1,
"prefill": 303.0
},
{
"card": "RTX 5070 12GB (measured)",
"model": "IQ2_XS",
"tier": "4k",
"slots": 4568,
"h": 0.785,
"decode": 78.1,
"prefill": 463.0
},
{
"card": "RTX 5060 Ti 16GB",
"model": "IQ2_XS",
"tier": "4k",
"slots": 7123,
"h": 0.85,
"decode": 76.9,
"prefill": 406.0
},
{
"card": "RTX 3090 24GB",
"model": "IQ2_XS",
"tier": "4k",
"slots": 12234,
"h": 0.937,
"decode": 127.7,
"prefill": 422.0
},
{
"card": "RTX 5070 12GB (measured)",
"model": "IQ2_XS",
"tier": "32k",
"slots": 4289,
"h": 0.783,
"decode": 65.3,
"prefill": 495.0
},
{
"card": "RTX 5060 Ti 16GB",
"model": "IQ2_XS",
"tier": "32k",
"slots": 6844,
"h": 0.848,
"decode": 62.9,
"prefill": 434.0
},
{
"card": "RTX 3090 24GB",
"model": "IQ2_XS",
"tier": "32k",
"slots": 11955,
"h": 0.935,
"decode": 102.9,
"prefill": 451.0
},
{
"card": "RTX 5070 12GB (measured)",
"model": "IQ2_XS",
"tier": "64k",
"slots": 3936,
"h": 0.759,
"decode": 63.7,
"prefill": 486.0
},
{
"card": "RTX 5060 Ti 16GB",
"model": "IQ2_XS",
"tier": "64k",
"slots": 6491,
"h": 0.83,
"decode": 62.0,
"prefill": 426.0
},
{
"card": "RTX 3090 24GB",
"model": "IQ2_XS",
"tier": "64k",
"slots": 11602,
"h": 0.927,
"decode": 102.2,
"prefill": 444.0
},
{
"card": "RTX 5070 12GB (measured)",
"model": "IQ2_XS",
"tier": "128k",
"slots": 3229,
"h": 0.695,
"decode": 52.0,
"prefill": 472.0
},
{
"card": "RTX 5060 Ti 16GB",
"model": "IQ2_XS",
"tier": "128k",
"slots": 5784,
"h": 0.784,
"decode": 51.3,
"prefill": 413.0
},
{
"card": "RTX 3090 24GB",
"model": "IQ2_XS",
"tier": "128k",
"slots": 10895,
"h": 0.905,
"decode": 85.0,
"prefill": 430.0
},
{
"card": "RTX 5070 12GB (measured)",
"model": "IQ2_XS",
"tier": "262k",
"slots": 1796,
"h": 0.627,
"decode": 48.1,
"prefill": 437.0
},
{
"card": "RTX 5060 Ti 16GB",
"model": "IQ2_XS",
"tier": "262k",
"slots": 4351,
"h": 0.735,
"decode": 47.3,
"prefill": 383.0
},
{
"card": "RTX 3090 24GB",
"model": "IQ2_XS",
"tier": "262k",
"slots": 9462,
"h": 0.876,
"decode": 77.5,
"prefill": 398.0
},
{
"card": "RTX 5070 12GB (measured)",
"model": "IQ3_XXS",
"tier": "1k",
"slots": 3556,
"h": 0.638,
"decode": 64.6,
"prefill": 285.0
},
{
"card": "RTX 5060 Ti 16GB",
"model": "IQ3_XXS",
"tier": "1k",
"slots": 5658,
"h": 0.728,
"decode": 65.8,
"prefill": 249.0
},
{
"card": "RTX 3090 24GB",
"model": "IQ3_XXS",
"tier": "1k",
"slots": 9864,
"h": 0.858,
"decode": 106.5,
"prefill": 260.0
},
{
"card": "RTX 5070 12GB (measured)",
"model": "IQ3_XXS",
"tier": "4k",
"slots": 3511,
"h": 0.712,
"decode": 65.6,
"prefill": 410.0
},
{
"card": "RTX 5060 Ti 16GB",
"model": "IQ3_XXS",
"tier": "4k",
"slots": 5613,
"h": 0.784,
"decode": 64.7,
"prefill": 359.0
},
{
"card": "RTX 3090 24GB",
"model": "IQ3_XXS",
"tier": "4k",
"slots": 9819,
"h": 0.887,
"decode": 103.0,
"prefill": 374.0
},
{
"card": "RTX 5070 12GB (measured)",
"model": "IQ3_XXS",
"tier": "32k",
"slots": 3279,
"h": 0.696,
"decode": 57.3,
"prefill": 435.0
},
{
"card": "RTX 5060 Ti 16GB",
"model": "IQ3_XXS",
"tier": "32k",
"slots": 5381,
"h": 0.772,
"decode": 56.3,
"prefill": 381.0
},
{
"card": "RTX 3090 24GB",
"model": "IQ3_XXS",
"tier": "32k",
"slots": 9587,
"h": 0.879,
"decode": 89.0,
"prefill": 396.0
},
{
"card": "RTX 5070 12GB (measured)",
"model": "IQ3_XXS",
"tier": "64k",
"slots": 2990,
"h": 0.677,
"decode": 54.4,
"prefill": 427.0
},
{
"card": "RTX 5060 Ti 16GB",
"model": "IQ3_XXS",
"tier": "64k",
"slots": 5092,
"h": 0.757,
"decode": 53.7,
"prefill": 374.0
},
{
"card": "RTX 3090 24GB",
"model": "IQ3_XXS",
"tier": "64k",
"slots": 9298,
"h": 0.87,
"decode": 85.1,
"prefill": 390.0
},
{
"card": "RTX 5070 12GB (measured)",
"model": "IQ3_XXS",
"tier": "128k",
"slots": 2407,
"h": 0.57,
"decode": 44.8,
"prefill": 414.0
},
{
"card": "RTX 5060 Ti 16GB",
"model": "IQ3_XXS",
"tier": "128k",
"slots": 4509,
"h": 0.675,
"decode": 44.6,
"prefill": 363.0
},
{
"card": "RTX 3090 24GB",
"model": "IQ3_XXS",
"tier": "128k",
"slots": 8715,
"h": 0.824,
"decode": 70.8,
"prefill": 378.0
}
]
bench/results/2026-09-24-final/estimates.md
+27
View File
@@ -0,0 +1,27 @@
### Output (decode) tok/s - estimated
| Card | Model | 1k | 4k | 32k | 64k | 128k | 262k |
|---|---|---:|---:|---:|---:|---:|---:|
| RTX 5070 12GB (measured) | Q2_0 | ~89 | ~95 | ~88 | ~76 | ~65 | ~56 |
| RTX 5070 12GB (measured) | IQ2_XS | ~82 | ~78 | ~65 | ~64 | ~52 | ~48 |
| RTX 5070 12GB (measured) | IQ3_XXS | ~65 | ~66 | ~57 | ~54 | ~45 | - |
| RTX 5060 Ti 16GB | Q2_0 | ~80 | ~87 | ~81 | ~72 | ~62 | ~53 |
| RTX 5060 Ti 16GB | IQ2_XS | ~80 | ~77 | ~63 | ~62 | ~51 | ~47 |
| RTX 5060 Ti 16GB | IQ3_XXS | ~66 | ~65 | ~56 | ~54 | ~45 | - |
| RTX 3090 24GB | Q2_0 | ~128 | ~140 | ~130 | ~115 | ~100 | ~85 |
| RTX 3090 24GB | IQ2_XS | ~131 | ~128 | ~103 | ~102 | ~85 | ~78 |
| RTX 3090 24GB | IQ3_XXS | ~106 | ~103 | ~89 | ~85 | ~71 | - |
### Prompt processing tok/s - estimated
| Card | Model | 1k | 4k | 32k | 64k | 128k | 262k |
|---|---|---:|---:|---:|---:|---:|---:|
| RTX 5070 12GB (measured) | Q2_0 | ~389 | ~539 | ~571 | ~561 | ~543 | ~496 |
| RTX 5070 12GB (measured) | IQ2_XS | ~332 | ~463 | ~495 | ~486 | ~472 | ~437 |
| RTX 5070 12GB (measured) | IQ3_XXS | ~285 | ~410 | ~435 | ~427 | ~414 | - |
| RTX 5060 Ti 16GB | Q2_0 | ~341 | ~472 | ~501 | ~492 | ~476 | ~435 |
| RTX 5060 Ti 16GB | IQ2_XS | ~291 | ~406 | ~434 | ~426 | ~413 | ~383 |
| RTX 5060 Ti 16GB | IQ3_XXS | ~249 | ~359 | ~381 | ~374 | ~363 | - |
| RTX 3090 24GB | Q2_0 | ~355 | ~491 | ~521 | ~512 | ~495 | ~453 |
| RTX 3090 24GB | IQ2_XS | ~303 | ~422 | ~451 | ~444 | ~430 | ~398 |
| RTX 3090 24GB | IQ3_XXS | ~260 | ~374 | ~396 | ~390 | ~378 | - |
bench/results/2026-09-24-final/matrix.json
+240
View File
@@ -0,0 +1,240 @@
[
{
"model": "Q2_0",
"tier": "1k",
"prompt_tokens": 1017,
"exit": 0,
"wall_s": 40.2,
"prefill_tok_s": 388.83,
"ttft_s": 2.68,
"decode_tok_s": 88.67,
"decoded": 256,
"spec_accept": 0.739,
"tokens_per_round": 2.71,
"vram_slots": 4517
},
{
"model": "Q2_0",
"tier": "4k",
"prompt_tokens": 3562,
"exit": 0,
"wall_s": 44.6,
"prefill_tok_s": 538.75,
"ttft_s": 6.67,
"decode_tok_s": 94.59,
"decoded": 256,
"spec_accept": 0.846,
"tokens_per_round": 3.23,
"vram_slots": 4464
},
{
"model": "Q2_0",
"tier": "32k",
"prompt_tokens": 31566,
"exit": 0,
"wall_s": 91.9,
"prefill_tok_s": 571.03,
"ttft_s": 55.34,
"decode_tok_s": 87.48,
"decoded": 256,
"spec_accept": 0.865,
"tokens_per_round": 3.11,
"vram_slots": 4174
},
{
"model": "Q2_0",
"tier": "64k",
"prompt_tokens": 64162,
"exit": 0,
"wall_s": 151.5,
"prefill_tok_s": 561.16,
"ttft_s": 114.4,
"decode_tok_s": 76.38,
"decoded": 256,
"spec_accept": 0.781,
"tokens_per_round": 2.8,
"vram_slots": 3809
},
{
"model": "Q2_0",
"tier": "128k",
"prompt_tokens": 128478,
"exit": 0,
"wall_s": 273.5,
"prefill_tok_s": 542.95,
"ttft_s": 236.7,
"decode_tok_s": 65.09,
"decoded": 256,
"spec_accept": 0.687,
"tokens_per_round": 2.65,
"vram_slots": 3079
},
{
"model": "Q2_0",
"tier": "262k",
"prompt_tokens": 259943,
"exit": 0,
"wall_s": 562.5,
"prefill_tok_s": 496.33,
"ttft_s": 523.8,
"decode_tok_s": 56.26,
"decoded": 256,
"spec_accept": 0.659,
"tokens_per_round": 2.46,
"vram_slots": 1588
},
{
"model": "IQ2_XS",
"tier": "1k",
"prompt_tokens": 1017,
"exit": 0,
"wall_s": 42.3,
"prefill_tok_s": 332.19,
"ttft_s": 3.14,
"decode_tok_s": 81.97,
"decoded": 256,
"spec_accept": 0.84,
"tokens_per_round": 3.25,
"vram_slots": 4621
},
{
"model": "IQ2_XS",
"tier": "4k",
"prompt_tokens": 3562,
"exit": 0,
"wall_s": 47.4,
"prefill_tok_s": 462.95,
"ttft_s": 7.78,
"decode_tok_s": 78.05,
"decoded": 256,
"spec_accept": 0.856,
"tokens_per_round": 3.28,
"vram_slots": 4568
},
{
"model": "IQ2_XS",
"tier": "32k",
"prompt_tokens": 31566,
"exit": 0,
"wall_s": 102.3,
"prefill_tok_s": 495.14,
"ttft_s": 63.83,
"decode_tok_s": 65.32,
"decoded": 256,
"spec_accept": 0.706,
"tokens_per_round": 2.53,
"vram_slots": 4289
},
{
"model": "IQ2_XS",
"tier": "64k",
"prompt_tokens": 64162,
"exit": 0,
"wall_s": 170.8,
"prefill_tok_s": 486.42,
"ttft_s": 131.99,
"decode_tok_s": 63.66,
"decoded": 256,
"spec_accept": 0.701,
"tokens_per_round": 2.72,
"vram_slots": 3936
},
{
"model": "IQ2_XS",
"tier": "128k",
"prompt_tokens": 128478,
"exit": 0,
"wall_s": 312.7,
"prefill_tok_s": 471.61,
"ttft_s": 272.52,
"decode_tok_s": 52.04,
"decoded": 256,
"spec_accept": 0.602,
"tokens_per_round": 2.44,
"vram_slots": 3229
},
{
"model": "IQ2_XS",
"tier": "262k",
"prompt_tokens": 259943,
"exit": 0,
"wall_s": 637.4,
"prefill_tok_s": 436.86,
"ttft_s": 595.11,
"decode_tok_s": 48.05,
"decoded": 256,
"spec_accept": 0.628,
"tokens_per_round": 2.42,
"vram_slots": 1796
},
{
"model": "IQ3_XXS",
"tier": "1k",
"prompt_tokens": 1017,
"exit": 0,
"wall_s": 43.8,
"prefill_tok_s": 284.64,
"ttft_s": 3.67,
"decode_tok_s": 64.58,
"decoded": 256,
"spec_accept": 0.912,
"tokens_per_round": 3.63,
"vram_slots": 3556
},
{
"model": "IQ3_XXS",
"tier": "4k",
"prompt_tokens": 3562,
"exit": 0,
"wall_s": 48.5,
"prefill_tok_s": 409.76,
"ttft_s": 8.78,
"decode_tok_s": 65.57,
"decoded": 256,
"spec_accept": 0.863,
"tokens_per_round": 3.24,
"vram_slots": 3511
},
{
"model": "IQ3_XXS",
"tier": "32k",
"prompt_tokens": 31566,
"exit": 0,
"wall_s": 112.6,
"prefill_tok_s": 434.75,
"ttft_s": 72.69,
"decode_tok_s": 57.3,
"decoded": 256,
"spec_accept": 0.752,
"tokens_per_round": 2.78,
"vram_slots": 3279
},
{
"model": "IQ3_XXS",
"tier": "64k",
"prompt_tokens": 64162,
"exit": 0,
"wall_s": 190.1,
"prefill_tok_s": 427.26,
"ttft_s": 150.25,
"decode_tok_s": 54.4,
"decoded": 256,
"spec_accept": 0.723,
"tokens_per_round": 2.88,
"vram_slots": 2990
},
{
"model": "IQ3_XXS",
"tier": "128k",
"prompt_tokens": 128478,
"exit": 0,
"wall_s": 351.1,
"prefill_tok_s": 414.43,
"ttft_s": 310.1,
"decode_tok_s": 44.81,
"decoded": 256,
"spec_accept": 0.648,
"tokens_per_round": 2.42,
"vram_slots": 2407
}
]
bench/results/2026-09-24-final/matrix.md
+15
View File
@@ -0,0 +1,15 @@
### Prompt processing (tok/s)
| Model | 1K | 4K | 32K | 64K | 128K | 262K |
|---|---:|---:|---:|---:|---:|---:|
| Q2_0 | 389 | 539 | 571 | 561 | 543 | 496 |
| IQ2_XS | 332 | 463 | 495 | 486 | 472 | 437 |
| IQ3_XXS | 285 | 410 | 435 | 427 | 414 | - |
### Output (tok/s)
| Model | 1K | 4K | 32K | 64K | 128K | 262K |
|---|---:|---:|---:|---:|---:|---:|
| Q2_0 | 88.7 | 94.6 | 87.5 | 76.4 | 65.1 | 56.3 |
| IQ2_XS | 82.0 | 78.0 | 65.3 | 63.7 | 52.0 | 48.0 |
| IQ3_XXS | 64.6 | 65.6 | 57.3 | 54.4 | 44.8 | - |
bench/results/2026-09-24-vision/vision.json
+53
View File
@@ -0,0 +1,53 @@
[
{
"model": "Q2_0",
"off": {
"prefill_tok_s": 541.31,
"decode_tok_s": 90.74,
"vram_slots": 4464,
"spec_accept": 0.846
},
"on": {
"prefill_tok_s": 530.64,
"decode_tok_s": 86.92,
"vram_slots": 3402,
"spec_accept": 0.858
},
"flag_identical": true,
"on_identical": false
},
{
"model": "IQ2_XS",
"off": {
"prefill_tok_s": 468.17,
"decode_tok_s": 77.14,
"vram_slots": 4568,
"spec_accept": 0.856
},
"on": {
"prefill_tok_s": 457.61,
"decode_tok_s": 73.96,
"vram_slots": 3542,
"spec_accept": 0.866
},
"flag_identical": true,
"on_identical": false
},
{
"model": "IQ3_XXS",
"off": {
"prefill_tok_s": 411.3,
"decode_tok_s": 64.7,
"vram_slots": 3511,
"spec_accept": 0.863
},
"on": {
"prefill_tok_s": 401.48,
"decode_tok_s": 59.46,
"vram_slots": 2662,
"spec_accept": 0.868
},
"flag_identical": true,
"on_identical": true
}
]
bench/results/2026-09-24-vision/vision.md
+5
View File
@@ -0,0 +1,5 @@
| Model | Prompt tok/s (vision off) | Prompt tok/s (vision on) | Output tok/s (vision off) | Output tok/s (vision on) | VRAM expert slots off / on | --vision alone: same tokens |
|---|---:|---:|---:|---:|---:|---|
| Q2_0 | 541 | 531 | 90.7 | 86.9 | 4464 / 3402 | yes |
| IQ2_XS | 468 | 458 | 77.1 | 74.0 | 4568 / 3542 | yes |
| IQ3_XXS | 411 | 401 | 64.7 | 59.5 | 3511 / 2662 | yes |
chat.py
+116
View File
@@ -0,0 +1,116 @@
#!/usr/bin/env python3
"""A tiny terminal chat for a running Strata server (start it with run-<model>.bat / run-<model>.sh first).
python chat.py [--port 8080] [--think none|low|medium|high]
Type a message and press Enter. /image <path> attaches a picture to your next message (when the server was set up
with images), /think <none|low|medium|high> sets how long the model thinks first, /reset starts a new conversation,
/quit leaves. Standard library only.
"""
from __future__ import annotations
import argparse
import base64
import json
import mimetypes
import os
import sys
import time
import urllib.request
def stream(url, messages, think, max_tokens):
body = {"model": "strata", "messages": messages, "stream": True, "max_tokens": max_tokens,
"reasoning_effort": think}
req = urllib.request.Request(url, data=json.dumps(body).encode(), headers={"Content-Type": "application/json"})
with urllib.request.urlopen(req, timeout=3600) as r:
for raw in r:
line = raw.decode("utf-8", "replace").strip()
if not line.startswith("data: ") or line == "data: [DONE]":
continue
d = json.loads(line[6:])["choices"][0]["delta"]
yield d.get("reasoning_content") or "", d.get("content") or ""
def main() -> int:
ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
ap.add_argument("--port", type=int, default=8080)
ap.add_argument("--host", default="127.0.0.1")
ap.add_argument("--think", default="medium", choices=["none", "low", "medium", "high"],
help="how long the model thinks before answering (none = answer directly)")
ap.add_argument("--no-think", action="store_true", help="same as --think none")
ap.add_argument("--max-tokens", type=int, default=4096)
a = ap.parse_args()
url = f"http://{a.host}:{a.port}/v1/chat/completions"
gray, reset = ("\033[90m", "\033[0m") if sys.stdout.isatty() else ("", "")
messages, pending = [], []
think = "none" if a.no_think else a.think
print(f"Strata chat ({url}). /image <path> = attach a picture, /think none|low|medium|high (now: {think}), "
f"/reset = new conversation, /quit = leave.")
while True:
try:
user = input("\nyou> ").strip()
except (EOFError, KeyboardInterrupt):
print()
return 0
if user in ("/quit", "/exit"):
return 0
if user == "/reset":
messages, pending = [], []
print("(new conversation)")
continue
if user.startswith("/think"):
level = user[6:].strip().lower()
if level in ("none", "low", "medium", "high"):
think = level
print(f"(thinking: {think})")
else:
print("(usage: /think none|low|medium|high)")
continue
if user.startswith("/image"):
path = user[6:].strip().strip('"').strip("'")
if not os.path.isfile(path):
print(f"(no such file: {path or '?'} - usage: /image <path to a picture>)")
continue
mime = mimetypes.guess_type(path)[0] or "image/jpeg"
data = base64.b64encode(open(path, "rb").read()).decode()
pending.append({"type": "image_url", "image_url": {"url": f"data:{mime};base64,{data}"}})
print(f"(picture attached: {os.path.basename(path)} - now type your question)")
continue
if not user:
continue
if pending:
messages.append({"role": "user", "content": [*pending, {"type": "text", "text": user}]})
pending = []
else:
messages.append({"role": "user", "content": user})
answer, n, t0 = [], 0, time.time()
print("model> ", end="", flush=True)
try:
in_think = False
for reasoning, content in stream(url, messages, think, a.max_tokens):
if reasoning:
if not in_think:
print(gray, end="")
in_think = True
print(reasoning, end="", flush=True)
if content:
if in_think:
print(reset + "\n", end="")
in_think = False
print(content, end="", flush=True)
answer.append(content)
n += 1
if in_think:
print(reset, end="")
except OSError as e:
print(f"\n(could not reach the server at {url}: {e})")
messages.pop()
continue
dt = time.time() - t0
print(f"\n{gray}[{n} chunks in {dt:.1f} s]{reset}")
messages.append({"role": "assistant", "content": "".join(answer)})
if __name__ == "__main__":
sys.exit(main())
data/draft_vocab.bin
BIN
View File
Binary file not shown.
data/expert-profile.bin
BIN
View File
Binary file not shown.
docs/paper/Strata-Paper.pdf
BIN
View File
Binary file not shown.
docs/paper/tiers.svg
+18
View File
@@ -0,0 +1,18 @@
<svg viewBox="0 0 700 250" width="100%" xmlns="http://www.w3.org/2000/svg">
<rect x="0" y="0" width="700" height="72" rx="6" fill="#eaf1f9" stroke="#2f6fb3"/>
<text x="10" y="18" font-size="12" font-weight="700" fill="#2f6fb3">GPU VRAM 12 GB &#183; 672 GB/s</text>
<rect x="10" y="26" width="150" height="38" rx="3" fill="#2f6fb3"/><text x="85" y="42" font-size="10" fill="#fff" text-anchor="middle">dense weights ~3.3 GB</text><text x="85" y="56" font-size="9" fill="#fff" text-anchor="middle">mixers &#183; GR &#183; routers &#183; head</text>
<rect x="166" y="26" width="330" height="38" rx="3" fill="#6d9fd4"/><text x="331" y="42" font-size="10" fill="#fff" text-anchor="middle">expert cache: everything left (~6 GB)</text><text x="331" y="56" font-size="9" fill="#fff" text-anchor="middle">~4,500 of 24,576 experts, adaptive</text>
<rect x="502" y="26" width="90" height="38" rx="3" fill="#8fb6e0"/><text x="547" y="42" font-size="10" fill="#fff" text-anchor="middle">MTP layer</text><text x="547" y="56" font-size="9" fill="#fff" text-anchor="middle">0.8 GB</text>
<rect x="598" y="26" width="94" height="38" rx="3" fill="#b3cde9"/><text x="645" y="42" font-size="10" fill="#1b1d21" text-anchor="middle">KV + state</text><text x="645" y="56" font-size="9" fill="#1b1d21" text-anchor="middle">0.3-3 GB</text>
<text x="350" y="88" font-size="10" fill="#5a606b" text-anchor="middle">PCIe 4.0 x16 &#183; ~26 GB/s &#183; experts by DMA (prompt, 20% of misses, cache swaps)</text>
<rect x="0" y="96" width="700" height="72" rx="6" fill="#fbf0e6" stroke="#c7702a"/>
<text x="10" y="114" font-size="12" font-weight="700" fill="#c7702a">System RAM 64 GB DDR5 &#183; ~41-52 GB/s &#183; 6 CPU cores (AVX-512)</text>
<rect x="10" y="122" width="440" height="38" rx="3" fill="#c7702a"/><text x="230" y="138" font-size="10" fill="#fff" text-anchor="middle">pinned expert arena: all 24,576 experts</text><text x="230" y="152" font-size="9" fill="#fff" text-anchor="middle">34 GB (Q2_0) &#183; 35.5 GB (IQ2_XS) &#183; 43 GB (IQ3_XXS); the CPU computes misses in place</text>
<rect x="456" y="122" width="236" height="38" rx="3" fill="#e0a877"/><text x="574" y="138" font-size="10" fill="#fff" text-anchor="middle">OS page cache: hot n-gram rows</text><text x="574" y="152" font-size="9" fill="#fff" text-anchor="middle">+ OS, embedding table, buffers</text>
<text x="350" y="184" font-size="10" fill="#5a606b" text-anchor="middle">NVMe &#183; 16 row reads per token, issued a step ahead</text>
<rect x="0" y="192" width="700" height="54" rx="6" fill="#f1f3e8" stroke="#6b7b3a"/>
<text x="10" y="210" font-size="12" font-weight="700" fill="#6b7b3a">SSD</text>
<rect x="60" y="200" width="300" height="38" rx="3" fill="#6b7b3a"/><text x="210" y="223" font-size="10" fill="#fff" text-anchor="middle">GGUF shard 1 (37.6-47 GB): read once at start</text>
<rect x="366" y="200" width="326" height="38" rx="3" fill="#95a15f"/><text x="529" y="223" font-size="10" fill="#fff" text-anchor="middle">GGUF shard 2: n-gram table 28.8 GB, memory-mapped</text>
</svg>
Side by Side

After

Width:  |  Height:  |  Size: 3.0 KiB

include/strata/artifact/dequant.hpp
+292
View File
@@ -0,0 +1,292 @@
#pragma once
// src/artifact/dequant.cpp - P1.S3: scalar reference dequantizers, transcribed from ggml.
//
// P1.S3 asks for every encoding present in the Q2_0 file. This starts with the ones on the critical
// path (the routed experts are all Q2_0) plus the trivial element types, and is structured so the
// rest drop in. Source of truth per encoding is quoted in each function.
//
// Q2_0 is the one that matters, and its contract is settled from source in docs/q2_0-contract.md:
// block = { fp16 d; uint8 qs[16] } = 64 elements, 18 bytes
// code {0,1,2,3} -> symbol {-1,0,+1,+2}; v = (code - 1) * d
// 4 codes per byte, LSB-first: byte j/4, bits (j%4)*2
// QK2_0 = 64 is PROVEN from the artifact's own offset brackets, not assumed (the llama.cpp PR that
// introduced the type used 128; this artifact is 64).
//
// Validation: `--check <file.gguf>` dequantizes real blocks and asserts the STRUCTURAL invariant that
// every output value is one of {-d, 0, +d, +2d} for that block's own scale. That catches bit-order,
// sign and stride bugs, which are the failure modes a dequantizer actually has. It is deliberately
// not a correctness-vs-ggml test - that needs ggml linked, which is P1.S3's next step.
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <cmath>
#include <string>
#include <vector>
#define STRATA_GGUF_MAIN_DISABLED 1
#include "strata/artifact/gguf_reader.hpp"
namespace strata {
// ---- fp16/bf16 -> fp32. Both handle subnormals and infinities; naive shifts do not.
inline float fp16_to_fp32(uint16_t h) {
const uint32_t sign = (uint32_t)(h >> 15) & 1u;
uint32_t exp = (h >> 10) & 0x1Fu, man = h & 0x3FFu, f;
if (exp == 0) {
if (man == 0)
f = sign << 31;
else {
exp = 127 - 15 + 1;
while (!(man & 0x400u)) {
man <<= 1;
--exp;
}
man &= 0x3FFu;
f = (sign << 31) | (exp << 23) | (man << 13);
}
} else if (exp == 31)
f = (sign << 31) | 0x7F800000u | (man << 13);
else
f = (sign << 31) | ((exp - 15 + 127) << 23) | (man << 13);
float out;
std::memcpy(&out, &f, 4);
return out;
}
inline float bf16_to_fp32(uint16_t h) {
const uint32_t f = (uint32_t)h << 16;
float out;
std::memcpy(&out, &f, 4);
return out;
}
inline uint16_t read_u16(const uint8_t* p) {
return (uint16_t)(p[0] | (p[1] << 8));
}
// ---- Q2_0: 64 elements from an 18-byte block. Mirrors dequantize_row_q2_0 (ggml master 3cf03257).
inline void dequantize_q2_0(const uint8_t* block, float* out) {
const float d = fp16_to_fp32(read_u16(block));
const uint8_t* qs = block + 2;
for (int j = 0; j < 64; ++j) {
const int byte_index = j / 4;
const int bit_offset = (j % 4) * 2;
const int code = (qs[byte_index] >> bit_offset) & 0x03;
out[j] = (float)(code - 1) * d; // code {0,1,2,3} -> symbol {-1,0,+1,+2}
}
}
// ---- Q8_0: 32 elements from a 34-byte block.
inline void dequantize_q8_0(const uint8_t* block, float* out) {
const float d = fp16_to_fp32(read_u16(block));
const int8_t* qs = (const int8_t*)(block + 2);
for (int j = 0; j < 32; ++j) out[j] = (float)qs[j] * d;
}
// ---- Q4_0: 32 elements from an 18-byte block. Low nibble is element j, high nibble is j+16.
inline void dequantize_q4_0(const uint8_t* block, float* out) {
const float d = fp16_to_fp32(read_u16(block));
const uint8_t* qs = block + 2;
for (int j = 0; j < 16; ++j) {
out[j] = (float)((int)(qs[j] & 0x0F) - 8) * d;
out[j + 16] = (float)((int)(qs[j] >> 4) - 8) * d;
}
}
// ---- Q5_0: 32 elements from a 22-byte block. Same nibble layout as Q4_0 plus a 5th bit per element
// packed in qh[4]; the 32 bits of qh are the high bits of elements 0..31 in order.
inline void dequantize_q5_0(const uint8_t* block, float* out) {
const float d = fp16_to_fp32(read_u16(block));
const uint8_t* qh = block + 2;
const uint8_t* qs = block + 6;
const uint32_t h =
(uint32_t)qh[0] | ((uint32_t)qh[1] << 8) | ((uint32_t)qh[2] << 16) | ((uint32_t)qh[3] << 24);
for (int j = 0; j < 16; ++j) {
const int x0 = (int)(qs[j] & 0x0F) | (int)(((h >> j) & 1u) << 4);
const int x1 = (int)(qs[j] >> 4) | (int)(((h >> (j + 16)) & 1u) << 4);
out[j] = (float)(x0 - 16) * d;
out[j + 16] = (float)(x1 - 16) * d;
}
}
// ---- IQ4_NL: 32 elements from an 18-byte block. Non-linear 4-bit codebook, NOT a linear grid, so
// the table must be exact. Copied verbatim from ggml-common.h `kvalues_iq4nl` at 3cf03257:
// -127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113
// Layout matches dequantize_row_iq4_nl (ggml-quants.c): low nibble is element j, high nibble j+16.
static const int8_t kvalues_iq4nl[16] = {-127, -104, -83, -65, -49, -35, -22, -10,
1, 13, 25, 38, 53, 69, 89, 113};
inline void dequantize_iq4_nl(const uint8_t* block, float* out) {
const float d = fp16_to_fp32(read_u16(block));
const uint8_t* qs = block + 2;
for (int j = 0; j < 16; ++j) {
out[j] = d * (float)kvalues_iq4nl[qs[j] & 0x0F];
out[j + 16] = d * (float)kvalues_iq4nl[qs[j] >> 4];
}
}
// ---- Q6_K: 256 elements from a 210-byte super-block, transcribed from dequantize_row_q6_K
// (ggml-quants.c at 3cf03257). Layout: ql[128] low nibbles, qh[64] high 2 bits, scales[16] int8,
// fp16 d LAST. Processed in two 128-element halves, each advancing ql by 64, qh by 32, sc by 8.
// The `sc` indices are is+{0,2,4,6} with is = l/16 - so each 32-element group draws from 8 of the 16
// scales, and the four output positions within a 32-group use different scales. Easy to get subtly
// wrong, which is why it is transcribed rather than recalled.
inline void dequantize_q6_K(const uint8_t* block, float* out) {
const uint8_t* ql = block; // 128 bytes
const uint8_t* qh = block + 128; // 64
const int8_t* sc = (const int8_t*)(block + 192); // 16
const float d = fp16_to_fp32(read_u16(block + 208));
float* y = out;
for (int n = 0; n < 256; n += 128) {
for (int l = 0; l < 32; ++l) {
const int is = l / 16;
const int q1 = (int)((ql[l + 0] & 0x0F) | (((qh[l] >> 0) & 3) << 4)) - 32;
const int q2 = (int)((ql[l + 32] & 0x0F) | (((qh[l] >> 2) & 3) << 4)) - 32;
const int q3 = (int)((ql[l + 0] >> 4) | (((qh[l] >> 4) & 3) << 4)) - 32;
const int q4 = (int)((ql[l + 32] >> 4) | (((qh[l] >> 6) & 3) << 4)) - 32;
y[l + 0] = d * (float)sc[is + 0] * (float)q1;
y[l + 32] = d * (float)sc[is + 2] * (float)q2;
y[l + 64] = d * (float)sc[is + 4] * (float)q3;
y[l + 96] = d * (float)sc[is + 6] * (float)q4;
}
y += 128;
ql += 64;
qh += 32;
sc += 8;
}
}
// ---- Q4_K: 256 elements from a 144-byte super-block. Layout is `fp16 d`, `fp16 dmin` (so the FIRST
// four bytes are the scale pair, unlike Q6_K where d is last), `scales[12]`, `qs[128]`.
// Both functions transcribed from ggml-quants.c at 3cf03257. The 6-bit scale/min unpacking in
// get_scale_min_k4 is the part that is impossible to recall reliably: the low 4 bits of the second
// half come from q[j+4], and the high 2 bits from the low halves' top bits (q[j-4] and q[j]).
static inline void get_scale_min_k4(int j, const uint8_t* q, uint8_t& d, uint8_t& m) {
if (j < 4) {
d = q[j] & 63;
m = q[j + 4] & 63;
} else {
d = (q[j + 4] & 0x0F) | ((q[j - 4] >> 6) << 4);
m = (q[j + 4] >> 4) | ((q[j - 0] >> 6) << 4);
}
}
inline void dequantize_q4_K(const uint8_t* block, float* out) {
const float d = fp16_to_fp32(read_u16(block)); // d
const float mn = fp16_to_fp32(read_u16(block + 2)); // dmin
const uint8_t* scales = block + 4; // 12 bytes
const uint8_t* q = block + 16; // 128 bytes
float* y = out;
int is = 0;
for (int j = 0; j < 256; j += 64) {
uint8_t sc, m;
get_scale_min_k4(is + 0, scales, sc, m);
const float d1 = d * (float)sc, m1 = mn * (float)m;
get_scale_min_k4(is + 1, scales, sc, m);
const float d2 = d * (float)sc, m2 = mn * (float)m;
for (int l = 0; l < 32; ++l) *y++ = d1 * (float)(q[l] & 0x0F) - m1;
for (int l = 0; l < 32; ++l) *y++ = d2 * (float)(q[l] >> 4) - m2;
q += 32;
is += 2;
}
}
// ---- Q5_K: 256 elements from a 176-byte super-block: `fp16 d`, `fp16 dmin`, `scales[12]`,
// `qh[32]`, `qs[128]`. Same shape as Q4_K plus the 5th bit plane. Transcribed from
// dequantize_row_q5_K (ggml-quants.c, 3cf03257). The 5th bits are selected by a mask that SHIFTS LEFT
// BY 2 each 64-element group (u1 = 1,4,16,64; u2 = 2,8,32,128) - a recalled version would very likely
// have used a fixed mask or a shift by 1.
inline void dequantize_q5_K(const uint8_t* block, float* out) {
const float d = fp16_to_fp32(read_u16(block));
const float mn = fp16_to_fp32(read_u16(block + 2));
const uint8_t* scales = block + 4; // 12
const uint8_t* qh = block + 16; // 32
const uint8_t* ql = block + 48; // 128
float* y = out;
int is = 0;
uint8_t u1 = 1, u2 = 2;
for (int j = 0; j < 256; j += 64) {
uint8_t sc, m;
get_scale_min_k4(is + 0, scales, sc, m);
const float d1 = d * (float)sc, m1 = mn * (float)m;
get_scale_min_k4(is + 1, scales, sc, m);
const float d2 = d * (float)sc, m2 = mn * (float)m;
for (int l = 0; l < 32; ++l) *y++ = d1 * (float)((ql[l] & 0x0F) + ((qh[l] & u1) ? 16 : 0)) - m1;
for (int l = 0; l < 32; ++l) *y++ = d2 * (float)((ql[l] >> 4) + ((qh[l] & u2) ? 16 : 0)) - m2;
ql += 32;
is += 2;
u1 = (uint8_t)(u1 << 2);
u2 = (uint8_t)(u2 << 2);
}
}
// ---- Q3_K: 256 elements from a 110-byte super-block: `hmask[32]`, `qs[64]`, `scales[12]`, `fp16 d`
// LAST. Transcribed from dequantize_row_q3_K (ggml-quants.c, 3cf03257). This type has NO `dmin` - its
// scales are SIGNED 6-bit with a `- 32` bias - so it does not follow the Q4_K/Q5_K pattern at all.
// Two details that cannot be recalled:
// * the 12 scale bytes are bit-repacked through three 32-bit masks before use (kmask1/2 below),
// because 16 six-bit scales do not fit in 12 bytes linearly;
// * the low 2 bits are offset by -4 when the high bit is CLEAR, i.e. the correction is NEGATIVE.
inline void dequantize_q3_K(const uint8_t* block, float* out) {
const uint8_t* hm = block; // 32
const uint8_t* q = block + 32; // 64
const float d_all = fp16_to_fp32(read_u16(block + 108));
const uint32_t kmask1 = 0x03030303u, kmask2 = 0x0f0f0f0fu;
uint32_t aux[4] = {0, 0, 0, 0};
std::memcpy(aux, block + 96, 12);
const uint32_t tmp = aux[2];
aux[2] = ((aux[0] >> 4) & kmask2) | (((tmp >> 4) & kmask1) << 4);
aux[3] = ((aux[1] >> 4) & kmask2) | (((tmp >> 6) & kmask1) << 4);
aux[0] = (aux[0] & kmask2) | (((tmp >> 0) & kmask1) << 4);
aux[1] = (aux[1] & kmask2) | (((tmp >> 2) & kmask1) << 4);
const int8_t* scales = (const int8_t*)aux;
float* y = out;
int is = 0;
uint8_t m = 1; // hoisted: ggml declares this BEFORE the n loop, so it does not reset between halves
for (int n = 0; n < 256; n += 128) {
int shift = 0;
for (int j = 0; j < 4; ++j) {
float dl = d_all * (float)(scales[is++] - 32);
for (int l = 0; l < 16; ++l)
*y++ = dl * (float)((int8_t)((q[l + 0] >> shift) & 3) - ((hm[l + 0] & m) ? 0 : 4));
dl = d_all * (float)(scales[is++] - 32);
for (int l = 0; l < 16; ++l)
*y++ = dl * (float)((int8_t)((q[l + 16] >> shift) & 3) - ((hm[l + 16] & m) ? 0 : 4));
shift += 2;
m = (uint8_t)(m << 1);
}
q += 32;
}
}
// ---- IQ4_XS: 256 elements from a 136-byte super-block: `fp16 d`, `uint16 scales_h`, `scales_l[4]`,
// `qs[128]`. Transcribed from dequantize_row_iq4_xs (ggml-quants.c, 3cf03257). Reuses the IQ4_NL
// codebook. Eight 32-element groups; each group's scale is 6 bits assembled from the LOW nibble of
// `scales_l[ib/2]` (selected by ib%2) and TWO bits of `scales_h` at shift 2*ib - so the high bits come
// from a 16-bit field indexed by group, not from a byte array. Then `- 32`, as in Q3_K.
inline void dequantize_iq4_xs(const uint8_t* block, float* out) {
const float d = fp16_to_fp32(read_u16(block));
const uint16_t scales_h = read_u16(block + 2);
const uint8_t* scales_l = block + 4; // 4
const uint8_t* qs = block + 8; // 128
float* y = out;
for (int ib = 0; ib < 8; ++ib) {
const int ls = ((scales_l[ib / 2] >> (4 * (ib % 2))) & 0x0F) | (((scales_h >> (2 * ib)) & 3) << 4);
const float dl = d * (float)(ls - 32);
for (int j = 0; j < 16; ++j) {
y[j + 0] = dl * (float)kvalues_iq4nl[qs[j] & 0x0F];
y[j + 16] = dl * (float)kvalues_iq4nl[qs[j] >> 4];
}
y += 32;
qs += 16;
}
}
// ---- element types
inline void dequantize_f32(const uint8_t* p, float* out, int n) {
std::memcpy(out, p, (size_t)n * 4);
}
inline void dequantize_f16(const uint8_t* p, float* out, int n) {
for (int i = 0; i < n; ++i) out[i] = fp16_to_fp32(read_u16(p + 2 * i));
}
inline void dequantize_bf16(const uint8_t* p, float* out, int n) {
for (int i = 0; i < n; ++i) out[i] = bf16_to_fp32(read_u16(p + 2 * i));
}
} // namespace strata
include/strata/artifact/gguf_reader.hpp
+486
View File
@@ -0,0 +1,486 @@
// include/strata/artifact/gguf_reader.hpp - generated from src/artifact/gguf_reader.cpp by
// scripts/split_artifact.py. Header-only on purpose: the reader is one translation unit's worth of
// code with no state to hide, and a header-only split cannot introduce a duplicate-symbol or
// missing-declaration bug in code that is already validated.
#pragma once
// src/artifact/gguf_reader.cpp - P1.S2: GGUF v3 reader, mmap, no ggml dependency.
//
// The C++ counterpart of tools/gguf_reader.py, which has been the reference since P0.S6 (written
// because gguf-py cannot represent type 42 / Q2_0). Per P1.S2 it must:
// * mmap the file and parse header, metadata KV (ALL value types incl. arrays), tensor directory
// * be multi-shard aware (split.* keys)
// * carry an architecture guard: general.architecture == "qwen4exp" and the compiled-in constants
// must match, refusing with a precise error otherwise
// * have no ggml dependency
//
// This file is the reader plus a `--check` mode that validates it the way the Python one is
// validated: parse the tiny model AND both real shards, and report counts that other harnesses have
// already established independently (1,223 / 1 tensors; 202 Q2_0 in shard 1). Agreement with those
// numbers is the test.
//
// Build: scripts/build_artifact.bat Run: gguf_reader.exe <file.gguf> [--check]
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <string>
#include <vector>
#include <map>
#include <stdexcept>
#include <algorithm>
#ifdef _WIN32
#define WIN32_LEAN_AND_MEAN
#include <windows.h>
#else
#include <fcntl.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <unistd.h>
#endif
namespace strata {
// ---- ggml type ids we care about. 42 = Q2_0, the PrismML ternary 2-bit encoding this engine targets.
inline const char* ggml_type_name(uint32_t t) {
switch (t) {
case 0:
return "F32";
case 1:
return "F16";
case 2:
return "Q4_0";
case 3:
return "Q4_1";
case 6:
return "Q5_0";
case 7:
return "Q5_1";
case 8:
return "Q8_0";
case 9:
return "Q8_1";
case 10:
return "Q2_K";
case 11:
return "Q3_K";
case 12:
return "Q4_K";
case 13:
return "Q5_K";
case 14:
return "Q6_K";
case 15:
return "Q8_K";
case 16:
return "IQ2_XXS";
case 17:
return "IQ2_XS";
case 18:
return "IQ3_XXS";
case 19:
return "IQ1_S";
case 20:
return "IQ4_NL";
case 21:
return "IQ3_S";
case 22:
return "IQ2_S";
case 23:
return "IQ4_XS";
case 30:
return "BF16";
case 34:
return "TQ1_0";
case 35:
return "TQ2_0";
case 39:
return "MXFP4";
case 40:
return "NVFP4";
case 41:
return "Q1_0";
case 42:
return "Q2_0";
default:
return "?";
}
}
// Block geometry: (elements per block, bytes per block). Q2_0 is 64/18 - proven from this artifact's
// own offset brackets in P0.S6 and recorded in docs/q2_0-contract.md.
inline bool block_geometry(uint32_t t, int& elems, int& bytes) {
switch (t) {
case 0:
elems = 1;
bytes = 4;
return true;
case 1:
case 30:
elems = 1;
bytes = 2;
return true;
case 2:
elems = 32;
bytes = 18;
return true;
case 3:
elems = 32;
bytes = 20;
return true;
case 6:
elems = 32;
bytes = 22;
return true;
case 7:
elems = 32;
bytes = 24;
return true;
case 8:
elems = 32;
bytes = 34;
return true;
case 9:
elems = 32;
bytes = 36;
return true;
case 10:
elems = 256;
bytes = 84;
return true;
case 11:
elems = 256;
bytes = 110;
return true;
case 12:
elems = 256;
bytes = 144;
return true;
case 13:
elems = 256;
bytes = 176;
return true;
case 14:
elems = 256;
bytes = 210;
return true;
case 16:
elems = 256;
bytes = 66;
return true;
case 17:
elems = 256;
bytes = 74;
return true;
case 18:
elems = 256;
bytes = 98;
return true;
case 20:
elems = 32;
bytes = 18;
return true;
case 21: // IQ3_S
elems = 256;
bytes = 110;
return true;
case 22: // IQ2_S
elems = 256;
bytes = 82;
return true;
case 23:
elems = 256;
bytes = 136;
return true;
case 29: // IQ1_M
elems = 256;
bytes = 56;
return true;
case 42:
elems = 64;
bytes = 18;
return true;
default:
return false;
}
}
struct TensorInfo {
std::string name;
std::vector<uint64_t> shape; // GGUF order: dim 0 varies fastest
uint32_t type = 0;
uint64_t offset = 0;
uint64_t elements() const {
uint64_t n = 1;
for (auto d : shape) n *= d;
return n;
}
const char* type_name() const { return ggml_type_name(type); }
};
// A bounds-checked cursor over the mmapped header region. Every read is checked, so a truncated or
// corrupt file produces a precise error rather than a segfault.
class Cursor {
public:
Cursor(const uint8_t* base, size_t size) : base_(base), size_(size) {}
void need(size_t n) const {
if (pos_ + n > size_) throw std::runtime_error("GGUF: unexpected end of file in header");
}
template <class T> T read() {
need(sizeof(T));
T v;
std::memcpy(&v, base_ + pos_, sizeof(T));
pos_ += sizeof(T);
return v;
}
std::string str() {
const uint64_t n = read<uint64_t>();
need((size_t)n);
std::string s(reinterpret_cast<const char*>(base_ + pos_), (size_t)n);
pos_ += (size_t)n;
return s;
}
size_t pos() const { return pos_; }
private:
const uint8_t* base_;
size_t size_;
size_t pos_ = 0;
};
enum class MetaType : uint32_t { U8 = 0, I8, U16, I16, U32, I32, F32, BOOL, STRING, ARRAY, U64, I64, F64 };
struct MetaValue {
MetaType type = MetaType::U32;
uint64_t u = 0; // integer payloads
double f = 0; // float payloads
std::string s; // string payloads
MetaType elem = MetaType::U32; // arrays
uint64_t count = 0;
std::vector<MetaValue> items;
bool is_num() const { return type != MetaType::STRING && type != MetaType::ARRAY; }
double num() const { return (type == MetaType::F32 || type == MetaType::F64) ? f : (double)u; }
};
inline MetaValue read_value(Cursor& c, MetaType t, int depth = 0) {
if (depth > 2) throw std::runtime_error("GGUF: array nesting too deep");
MetaValue v;
v.type = t;
switch (t) {
case MetaType::U8:
v.u = c.read<uint8_t>();
break;
case MetaType::I8:
v.u = (uint64_t)(int64_t)c.read<int8_t>();
break;
case MetaType::U16:
v.u = c.read<uint16_t>();
break;
case MetaType::I16:
v.u = (uint64_t)(int64_t)c.read<int16_t>();
break;
case MetaType::U32:
v.u = c.read<uint32_t>();
break;
case MetaType::I32:
v.u = (uint64_t)(int64_t)c.read<int32_t>();
break;
case MetaType::F32: {
float x = c.read<float>();
v.f = x;
v.u = 0;
break;
}
case MetaType::BOOL:
v.u = c.read<uint8_t>() ? 1 : 0;
break;
case MetaType::STRING:
v.s = c.str();
break;
case MetaType::U64:
v.u = c.read<uint64_t>();
break;
case MetaType::I64:
v.u = (uint64_t)c.read<int64_t>();
break;
case MetaType::F64:
v.f = c.read<double>();
break;
case MetaType::ARRAY: {
v.elem = (MetaType)c.read<uint32_t>();
v.count = c.read<uint64_t>();
if (v.count > (1u << 24)) throw std::runtime_error("GGUF: implausible array length");
v.items.reserve((size_t)std::min<uint64_t>(v.count, 64));
for (uint64_t i = 0; i < v.count; ++i) {
MetaValue e = read_value(c, v.elem, depth + 1);
if (i < 64) v.items.push_back(std::move(e)); // keep a sample; the count is what matters
}
break;
}
default:
throw std::runtime_error("GGUF: unknown metadata value type");
}
return v;
}
class GgufFile {
public:
explicit GgufFile(const std::string& path) : path_(path) {
try { open(); }
catch (...) { close(); throw; }
}
~GgufFile() { close(); }
GgufFile(const GgufFile&) = delete;
GgufFile& operator=(const GgufFile&) = delete;
const std::vector<TensorInfo>& tensors() const { return tensors_; }
const std::map<std::string, MetaValue>& metadata() const { return meta_; }
uint32_t version() const { return version_; }
uint64_t data_start() const { return data_start_; }
uint64_t file_size() const { return size_; }
const TensorInfo* find(const std::string& name) const {
for (const auto& t : tensors_)
if (t.name == name) return &t;
return nullptr;
}
uint64_t count_type(const char* tn) const {
uint64_t n = 0;
for (const auto& t : tensors_)
if (std::strcmp(t.type_name(), tn) == 0) ++n;
return n;
}
const MetaValue* get(const std::string& key) const {
auto it = meta_.find(key);
return it == meta_.end() ? nullptr : &it->second;
}
const uint8_t* tensor_data(const TensorInfo& t) const { return base_ + data_start_ + t.offset; }
private:
void open() {
#ifdef _WIN32
HANDLE h = CreateFileA(path_.c_str(), GENERIC_READ, FILE_SHARE_READ, nullptr, OPEN_EXISTING,
FILE_ATTRIBUTE_NORMAL, nullptr);
if (h == INVALID_HANDLE_VALUE) throw std::runtime_error("cannot open " + path_);
LARGE_INTEGER li{};
GetFileSizeEx(h, &li);
size_ = (uint64_t)li.QuadPart;
HANDLE m = CreateFileMappingA(h, nullptr, PAGE_READONLY, 0, 0, nullptr);
if (!m) {
CloseHandle(h);
throw std::runtime_error("CreateFileMapping failed");
}
base_ = (const uint8_t*)MapViewOfFile(m, FILE_MAP_READ, 0, 0, 0);
if (!base_) {
CloseHandle(m);
CloseHandle(h);
throw std::runtime_error("MapViewOfFile failed");
}
map_ = m;
file_ = h;
#else
int fd = ::open(path_.c_str(), O_RDONLY);
if (fd < 0) throw std::runtime_error("cannot open " + path_);
fd_ = fd;
struct stat st{};
if (fstat(fd, &st) != 0) throw std::runtime_error("fstat failed");
size_ = (uint64_t)st.st_size;
void* p = mmap(nullptr, size_, PROT_READ, MAP_PRIVATE, fd, 0);
if (p == MAP_FAILED) throw std::runtime_error("mmap failed");
base_ = (const uint8_t*)p;
#endif
parse();
}
void close() {
#ifdef _WIN32
if (base_) UnmapViewOfFile(base_);
if (map_) CloseHandle((HANDLE)map_);
if (file_) CloseHandle((HANDLE)file_);
#else
if (base_) munmap((void*)base_, size_);
if (fd_ >= 0) ::close(fd_);
#endif
}
void parse() {
Cursor c(base_, size_);
const uint32_t magic = c.read<uint32_t>();
if (magic != 0x46554747u) throw std::runtime_error("not a GGUF file (bad magic)");
version_ = c.read<uint32_t>();
if (version_ != 3)
throw std::runtime_error("GGUF v" + std::to_string(version_) + ", this reader handles v3");
const uint64_t n_tensors = c.read<uint64_t>();
const uint64_t n_kv = c.read<uint64_t>();
for (uint64_t i = 0; i < n_kv; ++i) {
std::string key = c.str();
MetaType t = (MetaType)c.read<uint32_t>();
meta_.emplace(std::move(key), read_value(c, t));
}
tensors_.reserve((size_t)n_tensors);
for (uint64_t i = 0; i < n_tensors; ++i) {
TensorInfo t;
t.name = c.str();
const uint32_t nd = c.read<uint32_t>();
if (nd == 0 || nd > 4) throw std::runtime_error("GGUF: bad n_dims for " + t.name);
t.shape.resize(nd);
for (uint32_t d = 0; d < nd; ++d) t.shape[d] = c.read<uint64_t>();
t.type = c.read<uint32_t>();
t.offset = c.read<uint64_t>();
tensors_.push_back(std::move(t));
}
uint64_t align = 32;
if (const MetaValue* a = get("general.alignment"))
if (a->u) align = a->u;
alignment_ = align;
data_start_ = (c.pos() + align - 1) / align * align;
if (data_start_ > size_) throw std::runtime_error("GGUF: data section starts past EOF");
}
std::string path_;
const uint8_t* base_ = nullptr;
uint64_t size_ = 0, data_start_ = 0, alignment_ = 32;
uint32_t version_ = 0;
#ifdef _WIN32
void* file_ = nullptr;
void* map_ = nullptr;
#else
int fd_ = -1;
#endif
std::vector<TensorInfo> tensors_;
std::map<std::string, MetaValue> meta_;
};
// ---- architecture guard (P1.S2). The engine is specialised to ONE model; anything else must be
// refused with a precise error rather than silently mis-run.
struct Qwen4ExpGuard {
uint32_t block_count = 48, hidden = 2560, experts = 512, experts_used = 10, head_count = 24,
head_count_kv = 2;
};
inline std::string check_architecture(const GgufFile& g, const Qwen4ExpGuard& want = {}) {
const MetaValue* arch = g.get("general.architecture");
if (!arch) return "missing general.architecture";
if (arch->s != "qwen4exp") return "architecture is '" + arch->s + "', this engine requires 'qwen4exp'";
struct Req {
const char* key;
uint64_t want;
};
const Req reqs[] = {
{"qwen4exp.block_count", want.block_count},
{"qwen4exp.embedding_length", want.hidden},
{"qwen4exp.expert_count", want.experts},
{"qwen4exp.expert_used_count", want.experts_used},
{"qwen4exp.attention.head_count", want.head_count},
{"qwen4exp.attention.head_count_kv", want.head_count_kv},
};
for (const auto& r : reqs) {
const MetaValue* v = g.get(r.key);
if (!v) return std::string("missing ") + r.key;
if (v->u != r.want)
return std::string(r.key) + " = " + std::to_string(v->u) + ", expected " + std::to_string(r.want);
}
return {}; // empty == ok
}
} // namespace strata
include/strata/core/device.hpp
+71
View File
@@ -0,0 +1,71 @@
// include/strata/core/device.hpp - P2.S1: the device arena and the runtime's device facts.
//
// `DeviceArena` is ONE cudaMalloc per planner region with bump sub-allocation below it and no frees. That is
// not a simplification for the first version: the memory plan from P1.S9 is fixed at startup, so the set of
// regions and their sizes is known before anything is allocated, and an allocator that can free would be
// solving a problem the engine does not have while adding fragmentation and failure modes it does.
//
// The reason to get this in early is that the VRAM budget is the binding constraint of the whole design
// (5.95 GB pooled between KV and the expert cache, 33.97 GB of experts in DRAM). A runtime that discovers at
// token 4000 that it has overcommitted has already lost; the plan is printed against `cudaMemGetInfo` at
// startup so the discrepancy is visible immediately.
#pragma once
#include <cstdint>
#include <stdexcept>
#include <string>
#include <vector>
namespace strata::core {
struct DeviceInfo {
int ordinal = -1;
std::string name;
int cc_major = 0, cc_minor = 0;
uint64_t total_bytes = 0; // as reported by cudaMemGetInfo at query time
uint64_t free_bytes = 0;
int driver_version = 0, runtime_version = 0;
int multi_processor_count = 0;
};
// Throws when there is no CUDA device. The engine targets sm_120 specifically and must say so rather than
// run slowly on something else: `CMakeLists.txt` already refuses to COMPILE for another architecture, and
// this is the matching check at run time (a binary can be carried to a different machine).
DeviceInfo device_info(int ordinal = 0);
class CudaError : public std::runtime_error {
public:
CudaError(const std::string& what, int code) : std::runtime_error(what), code_(code) {}
int code() const { return code_; }
private:
int code_;
};
// One cudaMalloc, bump-allocated below. `poison` fills new allocations with a NaN-ish pattern in a debug
// build so that reading uninitialised VRAM gives a NaN rather than a plausible number - the same reasoning as
// the harness work in Phase 1: a wrong value that looks right is the expensive kind.
class DeviceArena {
public:
explicit DeviceArena(uint64_t bytes, int ordinal = 0, bool poison = false);
~DeviceArena();
DeviceArena(const DeviceArena&) = delete;
DeviceArena& operator=(const DeviceArena&) = delete;
// `align` must be a power of two; 256 keeps every sub-allocation at a sector boundary.
void* alloc(uint64_t bytes, uint64_t align = 256);
uint64_t capacity() const { return capacity_; }
uint64_t used() const { return used_; }
uint64_t peak() const { return used_; } // no frees, so used IS the peak
int ordinal() const { return ordinal_; }
void* base() const { return base_; }
private:
void* base_ = nullptr;
uint64_t capacity_ = 0, used_ = 0;
int ordinal_ = 0;
bool poison_ = false;
};
} // namespace strata::core
include/strata/core/expert_cache.hpp
+155
View File
@@ -0,0 +1,155 @@
// include/strata/core/expert_cache.hpp - R4: the VRAM-resident expert tier.
//
// **WHY THIS EXISTS, IN ONE PARAGRAPH.** The CPU expert pool is the engine's largest single cost: measured,
// **663.6 MB of expert bytes per token at ~40 GB/s = 16.2 ms**, on a token of ~53 ms, and `test_expert_pool`
// says the pipeline hides **1.055 ms of 19.035** - the CPU work is 96% exposed because the residual chain is
// strictly serial. The only way past it is to stop reading those bytes: put some of the 48 x 512 experts in
// VRAM and let the GPU compute them, so the CPU reads only the misses.
//
// **THE DESIGN IS SIZED AGAINST MEASUREMENTS, NOT AGAINST THE PLAN'S ASSUMPTIONS.**
//
// * `h_expert` at 4,105 slots, **ten-fold leave-one-out: 0.6447 [0.6110, 0.6750]** (`tools/routing_kfold.py`).
// The 0.6573 the README published is in-sample and 0.4720 measured a single-prompt CORPUS, not the metric.
// * VRAM, measured on a live run rather than at startup: **6414 MiB free** with the engine at steady state,
// against **5413 MiB** for 4,105 slots. It fits, with ~1 GB spare. `strata-device`'s planner agrees and says
// FITS, but the planner runs before the dense arena exists, so the live figure is the one that counts.
// * `h_layer` held-out is **0.0456** - only 4.6% of (layer, token) pairs have all ten experts resident - so
// this cannot be a per-layer grouped kernel and the three-way split is not an optimisation.
//
// **WHAT THIS FILE IS AND IS NOT, TODAY.** It is the SLOT STORAGE and the RESIDENCY TABLE: it allocates the
// VRAM, fills it from the host arena, and answers `(layer, expert) -> slot or -1`. It does **not** yet compute
// anything - `moe_hit_grouped_s2` does not exist and the hit/miss split is not wired into the graph - so with
// the cache on and nothing consuming it, **the engine is slower by the fill cost and faster by nothing.**
// That is stated here rather than discovered from a benchmark, and it is why `--expert-cache` defaults to 0.
//
// The order of work is `Memory/R4-design-note.md` §7: slots and residency first, then the split, then the
// kernel. This is that first step, and the step it unblocks is the one that can be measured.
#pragma once
#include <cstdint>
#include <string>
#include <vector>
namespace strata::core {
/// Where one expert's VRAM lives, or -1. `int32_t` and not a bool because the slot index is what the kernel
/// needs, and a table that has to be re-derived from a bitmask is a table that will be re-derived differently
/// in the two places that use it.
inline constexpr int32_t kNotResident = -1;
/// **THE STATIC RESIDENCY PLAN, READ BACK FROM THE FILE `tools/make_profile.py` WROTE.**
///
/// `profile.bin` is `STRP`, then `version, n_layers, n_expert, slots, n_ranked`, then `n_ranked` `(layer,
/// expert)` pairs ranked by descending routing frequency, then a `n_layers x n_expert` frequency table.
///
/// **WHY A PROFILE AND NOT THE COMPULSORY-MISS POLICY THIS ENGINE ALREADY HAD.** Compulsory-miss fills with
/// the first `S` DISTINCT experts the run happens to route, which on this workload is whatever the prompt
/// touched first - it measured `h = 0.4864`. A profile ranked by frequency over a whole trace is what the
/// plan's `0.6447` refers to, and the CPU's half of the drain is linear in `h`, so the gap is worth about
/// 4 ms of drain. The point of reading it from a file is that a profile built on ONE prompt can be scored on
/// ANOTHER, which is the only way to know whether it transfers.
///
/// The frequency table is not returned: it is the profile's own working, and the engine measures `h` by
/// running rather than by re-deriving what the file claims.
bool read_expert_profile(const std::string& path, int64_t n_layers, int64_t n_expert,
std::vector<std::pair<int32_t, int32_t>>& ranked, int64_t& slots, std::string& err);
class ExpertCache {
public:
ExpertCache() = default;
~ExpertCache();
ExpertCache(const ExpertCache&) = delete;
ExpertCache& operator=(const ExpertCache&) = delete;
/// Allocates `n_slots * blob_bytes` of DEVICE memory and a `n_layers * n_expert` residency table, and
/// **CHECKS THE ALLOCATION AGAINST WHAT THE CARD ACTUALLY HAS** rather than assuming the planner was right.
/// A cache that silently allocates less than it was asked for would report a hit rate for slots it does
/// not have, which is the shape of error this project keeps paying for.
bool open(int64_t n_slots, int64_t n_layers, int64_t n_expert, int64_t blob_bytes, std::string& err);
/// Plan v0.3 P6: slots of the given sizes, back to back (a native pack's blobs differ per layer, and a
/// profile-filled tier never moves an expert to another layer's slot, so each slot keeps its first size).
bool open_sized(const std::vector<int64_t>& slot_bytes, int64_t n_layers, int64_t n_expert, std::string& err);
/// Byte offset of each slot in the arena (null for uniform slots).
const uint64_t* slot_offsets() const { return off_.empty() ? nullptr : off_.data(); }
void close();
bool valid() const { return base_ != nullptr; }
int64_t slots() const { return slots_; }
/// Slots actually claimed. Not the same as `slots()` - the cache does not evict, so a run that routes
/// fewer distinct experts than there are slots leaves the rest empty.
int64_t resident() const { return per_layer_ ? admitted_ : next_free_; }
int64_t bytes() const { return off_.empty() ? slots_ * blob_ : (int64_t) off_.back(); }
double gib() const { return (double) bytes() / 1073741824.0; }
/// `(layer, expert)` -> slot index, or `kNotResident`. Bounds-checked: a bad layer or expert returns
/// `kNotResident` rather than reading whatever is adjacent in the table.
int32_t slot_of(int64_t layer, int64_t expert) const;
/// Claims the next free slot for `(layer, expert)`. Returns the slot, or `kNotResident` when the cache is
/// full - **it never evicts**, because eviction policy is a measured question (`R4.1`'s LFU-decay vs LRU
/// sweep) and a placeholder policy would set the hit rate that everything downstream is then sized against.
int32_t admit(int64_t layer, int64_t expert);
/// **R4.2g: GIVE EACH LAYER ITS OWN SLOTS. ROUND 328 MEASURED WHY THE GLOBAL FORM CANNOT WORK.**
///
/// `admit`'s original policy hands out slots in ARRIVAL ORDER from one counter shared by every layer
/// (`if (next_free_ >= slots_) return kNotResident;`). A position looks at 10 experts in each of 48
/// layers, so the first `n_slots` distinct (layer, expert) pairs are **about 26 layers of POSITION 0** -
/// the cache fills entirely inside the first position it ever sees and never changes again, so hits are
/// confined to those few layers for the rest of the run.
///
/// Measured, `--expert-cache 256`: **1781 of 60000 = 2.97%**. `Memory/cache_allocation.py` scores the
/// same run's routing and says **8 slots per layer is worth 21.4%** and 64 is worth **70.4%**. The cache
/// was never short of capacity - it was handing all of it to two layers.
///
/// With this on, layer `l` may only use slots `[l*q, (l+1)*q)` where `q = slots()/n_layers`, the last
/// layer taking the remainder. **EVICTION IS STILL ABSENT** - the property the comment above calls
/// deliberate is preserved exactly; only WHICH slot a layer may take changes. Default off, so every
/// caller that predates this flag is bit-for-bit unaffected.
void set_per_layer_admission(bool on) { per_layer_ = on; }
bool per_layer_admission() const { return per_layer_; }
/// The slot range layer `l` may admit into under per-layer admission. Exposed so a test can check it.
void layer_slot_range(int64_t layer, int64_t& lo, int64_t& hi) const;
/// The device address of one slot.
uint8_t* device_slot(int32_t slot);
const uint8_t* device_slot(int32_t slot) const;
/// Copies `(layer, expert)`'s blob from `host_blob` into `slot` on `stream`. Asynchronous: the caller
/// orders it. Returns false if the indices are out of range rather than reading past the arena.
/// `bytes` (plan v0.3 P6): the blob's own size when it is smaller than the slot (a native pack); 0 = the slot.
bool fill_slot(int32_t slot, const uint8_t* host_blob, void* stream, std::string& err, int64_t bytes = 0);
/// **THE SAME COPY, BUT BLOCKING, AND THE PROFILE FILL NEEDS IT.** `fill_slot` is asynchronous because on
/// the token path the fill and the kernel that reads it are ordered by one stream and waiting would be a
/// synchronisation per layer. At STARTUP there is no such ordering to lean on: the source is pageable host
/// memory, `verify_slot` reads the slot back on the LEGACY stream, and a legacy-stream copy is not ordered
/// against a `cudaStreamNonBlocking` one. The first version used the async form and `verify_slot` refused
/// the whole run with "slot 0 differs from the arena at byte 0" - which is the check doing its job.
bool fill_slot_blocking(int32_t slot, const uint8_t* host_blob, std::string& err, int64_t bytes = 0);
/// Reads `slot` back to the host and compares it to `host_blob`, byte for byte. **THE ONLY THING THAT SAYS
/// THE CACHE HOLDS THE EXPERT IT CLAIMS TO.** A slot table that is right about indices and wrong about
/// bytes produces a plausible token, which is exactly the failure this project has paid for most often.
bool verify_slot(int32_t slot, const uint8_t* host_blob, std::string& err, int64_t bytes = 0);
/// Slots filled so far, for the startup report.
int64_t fills() const { return fills_; }
private:
uint8_t* base_ = nullptr;
std::vector<int32_t> residency_; ///< [n_layers * n_expert] -> slot or kNotResident
int64_t slots_ = 0;
int64_t n_layers_ = 0;
int64_t n_expert_ = 0;
int64_t blob_ = 0;
int64_t next_free_ = 0;
int64_t fills_ = 0;
/// R4.2g. `per_layer_` off (the default) leaves `next_free_` as the only admission counter, so the
/// pre-existing path is untouched.
bool per_layer_ = false;
std::vector<int32_t> layer_next_; ///< [n_layers] -> that layer's next free slot
std::vector<uint64_t> off_; ///< plan v0.3 P6: slot offsets (slots + 1 entries) when sized
int64_t admitted_ = 0;
};
} // namespace strata::core
include/strata/core/expert_source.hpp
+354
View File
@@ -0,0 +1,354 @@
// include/strata/core/expert_source.hpp - P2.S3/P2.S5: the expert pool's adapter to the host loop.
//
// `session_loop` publishes, per layer, the NORMED activation `x_f`, the ten routed expert ids and their router
// weights, and expects `k x n_embd` floats back. This is the piece that turns those four things into an
// `ExpertPool::run` call. Nothing here is clever, and that is the point: the pool, the kernel and the loop are
// each already verified, so this file has exactly one job - get the CONTRACT between them right.
//
// THE CONTRACT, and each clause is a way to be wrong:
//
// 1. **`x_f` IS THE SAME PINNED BUFFER EVERY LAYER.** Its ADDRESS never changes, so the `ActQ` cannot be
// cached by pointer - a cache keyed on `x_f` would quantize layer 0 and reuse it for all 47 remaining
// layers, which is a finite, plausible, completely wrong token. There is no cache here at all: the
// conversion is one 2560-element pass against a 0.3 ms/layer budget, and a correct answer is worth more
// than the microseconds.
// 2. **THE POOL DOES NOT APPLY THE ROUTER WEIGHT.** `moe_combine` (`src/core/layer.cpp:479`) sums
// `w[i] * parts[i]` on the device. `ExpertJob::weight` is a diagnostic field; setting it here would
// apply the weight twice, which is invisible in a single layer and compounds over 48.
// 3. **THE KERNEL'S GEOMETRY IS FIXED BY THE ARTIFACT** (`H = 2560`, `FF = 640`, `BLOB = 1,382,400`). A
// different `n_embd` or a different expert width is REFUSED rather than mis-indexed, because the blob's
// internal offsets are compile-time constants and reading a 640-wide expert as a 2560-wide one walks off
// the end into the next expert's bytes without faulting.
#pragma once
#include "strata/core/expert_cache.hpp"
#include "strata/core/hit_hook.hpp"
#include "strata/kernels/cpu/pool.hpp"
#include <cstdint>
#include <string>
#include <vector>
namespace strata::core {
/// Where one routed expert's bytes come from.
///
/// Phase 2 has NO cache (`phase-2-correct-engine.md`: hit rate `h = 0`), so the only implementation is a
/// file-backed reader. The interface exists anyway because Phase 3 replaces exactly this object with the VRAM
/// cache, and because a test can supply an in-memory source without a 34 GB artifact.
class ExpertSource {
public:
virtual ~ExpertSource() = default;
/// The 1,382,400-byte blob for `(layer, expert)`, or nullptr if it cannot be produced.
///
/// The pointer only has to stay valid until the next `blob()` call: with `h = 0` every expert is computed
/// immediately and nothing is retained. A CACHING source must return pointers into the cache, not into a
/// reused staging buffer - otherwise the pool would read the next expert's bytes while computing this one.
virtual const uint8_t* blob(int64_t layer, int64_t expert) = 0;
/// Blobs touched, for the driver to report. A source that does not count returns 0.
virtual int64_t reads() const { return 0; }
/// Plan v0.3 P5: whether `blob(layer, expert)` lies in page-locked, CUDA-registered memory, so an
/// asynchronous host-to-device copy can DMA it directly (no staging copy on the CPU).
virtual bool pinned(int64_t layer, int64_t expert) const { (void) layer; (void) expert; return false; }
/// Called once before the first expert of a layer. A source that reads from disk wants to start the read
/// here so it overlaps the quantisation, and a prefetching source in Phase 3 wants the ids.
virtual void begin_layer(int64_t layer, const int32_t* ids, int64_t k) { (void) layer; (void) ids; (void) k; }
/// Plan v0.3 P6: the DEVICE address of a pinned, mapped blob (the GPU can read it over PCIe), or null.
virtual const uint8_t* device_alias(int64_t layer, int64_t expert) const { (void) layer; (void) expert; return nullptr; }
};
/// Plan v0.3 P6: what the GPU computes in a verify window's layer, written by the pool (mapped host memory) right
/// after the ring and published before the CPU starts its own share. Groups of entries that share a blob; the
/// blob is a VRAM slot or a pinned host blob read over PCIe.
struct GpuPlanSink {
int32_t* counts = nullptr; ///< [0] groups, [1] entries
int32_t* start = nullptr; ///< cap + 1
int32_t* dst = nullptr; ///< cap: the entry's row of parts (token * k + j)
int32_t* tok = nullptr; ///< cap: the entry's token
unsigned long long* ptr = nullptr; ///< cap: the group's blob, device address
/// Plan v0.3 P6 (DMA): the PCIe share is a second list of groups - `ptr2[q]` is staging slot q, `start2` indexes
/// the same dst/tok entries - computed by the GPU after `fetch` has copied their blobs into staging with the
/// copy engine. counts[0] = VRAM groups, counts[1] = all entries, counts[2] = PCIe groups.
unsigned long long* ptr2 = nullptr;
int32_t* start2 = nullptr;
unsigned long long staging = 0;
int64_t staging_cap = 0;
int64_t cap = 0;
void (*publish)(void* ctx) = nullptr;
/// Starts the DMA copies of `n` host blobs (pinned) into staging slots 0..n-1 and signals the GPU when they land.
void (*fetch)(void* ctx, const uint8_t* const* src, int n, size_t bytes) = nullptr;
void* ctx = nullptr;
/// Plan v0.3 P6: how a PCIe group reaches the GPU. 0 = the copy engine stages it (DMA, `fetch`); 1 = the grouped
/// kernel reads the mapped arena directly; 2 = a copy kernel stages it inside the graph. For 1 and 2 `ptr2`
/// holds the arena's device alias.
int pcie_mode = 0;
};
/// The adapter's own state. One per session, reused every layer so the token path allocates nothing (P2.T10).
struct ExpertDispatch {
strata::kernels::cpu::ExpertPool* pool = nullptr;
ExpertSource* src = nullptr;
int64_t n_expert = strata::kernels::cpu::NE;
/// Counters, for the driver to report rather than for control flow.
int64_t layers = 0;
int64_t experts = 0;
int64_t missing = 0;
// ================================ R4: THE VRAM TIER, MEASURED BEFORE IT IS USED ================================
//
// **THE CACHE IS CONSULTED AND FILLED HERE, AND NOTHING IS COMPUTED FROM IT YET.** That is deliberate and
// it is `R4-design-note.md` §7 step 2: the *dispatch* is measured on its own before any kernel is written,
// because a hit rate measured after the kernel exists cannot say whether a disappointing result was the
// policy, the split or the kernel.
//
// So every expert is still computed by the CPU, exactly as before, and the run is numerically identical
// with the cache on or off. What changes is that the engine now reports **h on its own routing, on a real
// workload** - which is the number `R4.1` asks for and which until now existed only from offline traces.
//
// The policy is compulsory-miss: the first time `(layer, expert)` is routed, if a slot is free it is
// admitted and filled from the arena. No eviction, because eviction policy is the measured question
// (`R4.1`'s LFU-decay vs LRU sweep) and a placeholder would set the hit rate everything is sized against.
ExpertCache* cache = nullptr;
int64_t cache_hits = 0; ///< lookups already resident
int64_t cache_admitted = 0; ///< lookups that took a slot
int64_t cache_refused = 0; ///< lookups with no slot free (the cache is full)
void* cache_stream = nullptr;
const char* cache_fail = nullptr;
// ================================ R4.2c: THE HITS GO TO THE GPU ================================
//
// **THE SPLIT IS BY ROUTER INDEX, AND THAT IS THE WHOLE TRICK.** A layer routes ten experts; the resident
// ones are computed on the GPU and the rest on the CPU, and both answers have to end up in `parts` at the
// index the ROUTER gave them, because that is the order `moe_combine` weights against. So the CPU's `out`
// rows are zeroed for the hits, and each hit carries its routed index to the kernel as `dst`.
//
// Everything below is per-session state rather than per-call, so the token path allocates nothing (P2.T10).
const uint8_t* cache_base = nullptr; ///< the slot arena on the DEVICE
int64_t cache_blob = 0; ///< bytes per slot
const uint64_t* cache_slot_off = nullptr; ///< plan v0.3 P6: per-slot offsets when the slots differ in size
void* hit_scratch = nullptr; ///< `moe_hit_grouped_scratch_bytes(K, ...)`
float* parts_out = nullptr; ///< the graph's `parts` buffer, on the device
/// Where the GPU's hits land, `K x n_embd`, DEVICE and separate from `parts_out` on purpose: see
/// `HitPhase`. Zeroed by the hit path each layer before the kernel writes it.
float* hit_out = nullptr;
int64_t parts_elems = 0; ///< `K * n_embd`, the length of both buffers
const float* mixed = nullptr; ///< the layer's normed activation, for the hit kernel's Q8_0
uint8_t* x_q8_0_hit = nullptr; ///< `(n_embd/32) * 34` bytes, its own buffer
/// **R4.2h: THE CPU's fp32 ACTIVATION SCALES, `(n_embd/32)` FLOATS.** Without them the GPU's hits are
/// computed with the `block_q8_0`'s fp16 `d` while the CPU's misses use the fp32 `ActQ::scale`
/// (`cpu/expert.cpp:92`) - **4.761e-04 relative on 80 of 80 chunks**, measured with both real
/// implementations linked in `bench/micro/act_quant_parity.cu`. That disagreement is why turning the
/// cache on changed the generated tokens. Required whenever the hit path runs.
float* x_q8_0_hit_scale = nullptr;
int32_t* d_slot = nullptr; ///< device, K entries
int32_t* d_dst = nullptr; ///< device, K entries
std::vector<int32_t> h_slot, h_dst; ///< host staging, sized at session setup
/// **PER ROUTER INDEX, DECIDED IN `Launch` AND CONSUMED BY THE POOL.** The two callbacks share it
/// so the decision is made exactly once, on this layer's ids, and neither side can re-decide it.
std::vector<uint8_t> is_hit;
bool decided = false;
/// Plan v0.3 P4 token graph: the STATIC residency table (`n_layers x n_expert`, slot or -1), the host's copy
/// of what the device hit path reads. When set, the pool leaves a resident expert's row at zero (the GPU
/// computes it) without any `Launch` callback.
const int32_t* host_res = nullptr;
/// Plan v0.3 P4: split every expert by rows across the pool's threads (default on; A/B `--no-split-rows`).
bool split_rows = true;
// ================================ R4.2d: DID THE GPU ACTUALLY START? ================================
//
// **THE HIT PATH IS 0.91 ms SLOWER AND THE ONLY EXPLANATION LEFT IS THAT IT NEVER OVERLAPS.** The kernel
// is 159 GB/s against the CPU's 35, the CPU's half of the drain falls 6 ms, and none of it reaches the
// token - which is what it would look like if the enqueued hit kernels did not BEGIN until the host next
// entered the driver, i.e. after `pool()` returned. That is the third sighting of this driver behaving
// that way (rounds 195/287 on the doorbell, round 36 on the head and sampler) and it decides whether R4 can
// pay at all, so it gets measured rather than argued.
//
// `hit_done` is recorded on the stream straight after the hit kernel. `Combine` - which runs after the
// pool - queries it: SUCCESS means the GPU finished while the CPU was working, NOT-READY means it had not.
// Same shape as the doorbell's `rings_mid_graph`, and for the same reason.
void* hit_done = nullptr; ///< cudaEvent_t, created at session setup
int64_t hit_ready = 0; ///< layers where the hit work was DONE by the time the pool returned
int64_t hit_late = 0; ///< layers where it was not
/// The candidate fix, as an A/B arm: enter the driver once right after the hit launch. If submission is
/// lazy, this starts the GPU work before the pool instead of after it.
bool hit_poke = false;
bool hit_cpu_order = false; ///< experimental CPU-order GPU expert arithmetic; opt-in only
int64_t n_hits = 0; ///< this layer's hits
/// Set by `Launch` and consumed by `Combine`, so a `Combine` with no `Launch` in front of it cannot
/// add a stale buffer into `parts`.
bool hit_pending = false;
const char* hit_fail = nullptr;
/// Whether the hit path is wired up. All of it or none of it: a half-configured hit path would compute
/// some experts twice and others not at all, which is a wrong token rather than an error.
bool hits_ready() const {
return cache != nullptr && cache_base != nullptr && parts_out != nullptr && hit_out != nullptr &&
mixed != nullptr &&
hit_scratch != nullptr && x_q8_0_hit != nullptr && x_q8_0_hit_scale != nullptr && d_slot != nullptr && d_dst != nullptr;
}
std::vector<strata::kernels::cpu::ExpertJob> jobs;
strata::kernels::cpu::ActQ act;
/// Plan v0.3 P6 verify window: one quantized activation per token, the multi-token jobs, and the expert ->
/// job map (reset after every layer).
std::vector<strata::kernels::cpu::ActQ> act_multi;
/// Plan v0.3 P6: a native pack's per-token activations (MAXT x kNativeActBytes).
std::vector<uint8_t> nact_multi;
std::vector<strata::kernels::cpu::ExpertJobMulti> jobs_multi;
std::vector<int16_t> job_of;
/// Plan v0.3 P6: the verify window's GPU plan (VRAM hits + the PCIe share of the misses); `pcie_num`/256 of
/// each layer's distinct missed experts (the last ones in routing order) are read by the GPU over PCIe.
GpuPlanSink* plan = nullptr;
int pcie_num = 0;
int64_t pcie_experts = 0; ///< distinct experts the GPU read over PCIe in verify windows
double ms_plan = 0, ms_actq = 0, ms_jobs = 0, ms_run = 0; ///< verify-window dispatch sections
/// Plan v0.3 P6: decayed routing counts per (layer, expert) during decode (sized by the caller; empty = off),
/// which the driver uses to swap the most-routed missing experts into the VRAM tier between rounds.
std::vector<float> usage;
int64_t multi_misses = 0; ///< distinct (layer, expert) pairs the CPU computed in verify windows
int64_t multi_entries = 0; ///< routed (token, expert) entries the CPU served in verify windows
/// Set when `dispatch` could not produce an answer. The loop itself has no error channel, so this is
/// where a source failure surfaces: the driver checks it after `session_loop` returns rather than the
/// engine computing from a half-filled `parts`.
///
/// **`failed` LATCHES AND `fail` IS ONLY THE MESSAGE.** Clearing the message must not re-arm the adapter:
/// a session that failed at layer 5 has already fed `moe_combine` whatever `parts` held, so every layer
/// after it is built on a hole - and resuming into a plausible-looking token is the exact outcome this
/// whole mechanism exists to prevent. `failed` is what the short-circuit reads.
bool failed = false;
const char* fail = nullptr;
int64_t fail_layer = -1;
int64_t fail_expert = -1;
};
/// `strata::core::PoolFn`, exactly. Silent on failure BY SIGNATURE - see `ExpertDispatch::fail`.
void expert_pool_dispatch(void* user, const float* x_f, const int32_t* ids, const float* weights, int64_t n_embd,
int64_t k, float* out);
/// Plan v0.3 P6: the pool for a verify window of `n_tok` tokens. `x_f` is (n_tok, n_embd), `ids` (n_tok, k) and
/// `out` (n_tok * k, n_embd). Each distinct missed expert is computed once for all the tokens routed to it;
/// resident experts' rows are zeroed (the GPU adds them). Requires `host_res` (the token-graph residency).
void expert_pool_dispatch_multi(ExpertDispatch& d, const float* x_f, const int32_t* ids, int64_t n_tok, int64_t k,
float* out);
/// **THE HITS, LAUNCHED AFTER THE MISSES ARE STAGED AND BEFORE `post[l]`.** Same shape as `PoolFn` and for the
/// same reason: `session_loop` owns the ORDER and this owns the work, so the loop needs to know nothing about
/// expert caches. Returns immediately when there are no hits, which is every layer until the cache is warm.
///
/// It must run AFTER the host has copied the misses into `parts_dev` (it writes into the same buffer, on rows
/// the CPU zeroed) and BEFORE `post[l]` (which reads it). Both are stream-ordered on the loop's own stream.
void expert_hit_run(void* user, void* stream, HitPhase phase, const int32_t* ids, int64_t k);
/// The pool half of the same decision; see `ExpertDispatch::is_hit`.
/// **PHASE 2'S ONLY SOURCE: `experts.bin`, memory-mapped, no cache.**
///
/// `experts.bin` is 33,973,862,400 B and `BLOB` is 1,382,400, so it holds exactly `48 x 512 = 24,576` blobs and
/// the index is `layer * 512 + expert` with NO padding. A mapping is therefore the whole implementation: the
/// blob pointer is base plus a multiply, and the page fault that follows is the read.
///
/// **AND THAT IS NOT AS SLOW AS IT SOUNDS, WHICH IS THE POINT.** The expert set is 34 GB and this machine has
/// 64 GB of DDR5, so a warm OS page cache holds ALL of it - the second token onward is a DRAM read at the
/// measured 44.14 GB/s, not a disk read. The first token's 663.6 MB comes off the disk and is the cold-path
/// number; `benches` should report the two separately rather than blending them.
///
/// IT IS NOT AN LRU, AN LFU OR A PREFETCHER. Phase 3 replaces this object with those. Phase 2 is hit rate
/// `h = 0` on purpose, so that the CPU path and the host round trip are exercised on every layer of every
/// token - which is the only way they get debugged before speed matters.
class FileExpertSource : public ExpertSource {
public:
FileExpertSource() = default;
~FileExpertSource() override;
FileExpertSource(const FileExpertSource&) = delete;
FileExpertSource& operator=(const FileExpertSource&) = delete;
/// Maps `<pack_dir>/experts.bin` and checks its size against `n_layers * n_expert * BLOB`.
///
/// The size check is not a formality: a short file would fault at the END of a long sequence, and an
/// over-long one means the pack is not the one the geometry came from. Refuses with the two numbers.
bool open(const std::string& pack_dir, int64_t n_layers, int64_t n_expert, std::string& err);
void close();
bool mapped() const { return base_ != nullptr; }
int64_t blobs() const { return blobs_; }
const uint8_t* blob(int64_t layer, int64_t expert) override;
/// Blobs touched, for the driver to report. With `h = 0` this is `48 * k` per token and the number is only
/// interesting once Phase 3 makes it not so.
int64_t reads() const override { return reads_; }
private:
const uint8_t* base_ = nullptr;
int64_t blobs_ = 0;
int64_t n_expert_ = 0;
int64_t reads_ = 0;
#if defined(_WIN32)
void* file_ = nullptr;
void* mapping_ = nullptr;
#else
int fd_ = -1;
#endif
};
// ================================ THE RESIDENT ARENA (R2.1) ================================
//
// **THE MMAP ABOVE IS THE REVIEW'S FINDING C1 AND IT IS WORTH 2.2x.** Read the comment on `FileExpertSource`
// about the page cache holding all 34 GB: that is true of the expert file ALONE, and it is not what this engine
// does. The PLE/n-gram shard is a 26.8 GB file that the engine also maps, so 34 GB of experts plus 26.8 GB of
// n-gram is 60.8 GB of mapped, file-backed pages on a 63 GB machine - and file-backed pages are exactly the ones
// the OS drops from the standby list when it wants memory. The expert stream then re-faults from disk.
//
// Measured, this round: the pool runs at **~19 GB/s in the engine against 42.8 GB/s in `bench/micro/cpu_s2.cpp`
// on the same machine, reading the same 34 GB**. The micro does `std::fread` into a heap arena and reads
// ordinary (anonymous, resident) memory; the engine reads `MapViewOfFile`. That is the whole difference, and it
// is why this class exists.
//
// It reads `experts.bin` into a `PinnedArena` once at startup, so the expert stream comes from anonymous memory
// the OS has no cheaper reason to evict. `PinnedArena` also tries `cudaHostRegister`, which the GPU needs for
// Phase 3's cache fills and the CPU/PCIe miss split - but registration is best-effort and reported, not assumed.
class ArenaExpertSource : public ExpertSource {
public:
ArenaExpertSource() = default;
~ArenaExpertSource() override;
ArenaExpertSource(const ArenaExpertSource&) = delete;
ArenaExpertSource& operator=(const ArenaExpertSource&) = delete;
/// Allocates and loads `<pack_dir>/experts.bin`. Prints nothing; the caller reports `note()` and the load
/// rate, because those are the two numbers that say whether the arena is the one that was asked for.
bool open(const std::string& pack_dir, int64_t n_layers, int64_t n_expert, int threads, std::string& err);
/// Plan v0.3 P6: a native pack without experts.bin takes its experts from the model's shard 1.
void set_gguf(const std::string& shard1) { gguf_ = shard1; }
void close();
bool mapped() const { return base_ != nullptr; }
int64_t blobs() const { return blobs_; }
const uint8_t* blob(int64_t layer, int64_t expert) override;
int64_t reads() const { return reads_; }
bool pinned(int64_t layer, int64_t expert) const override;
const uint8_t* device_alias(int64_t layer, int64_t expert) const override;
/// What backing was obtained and why, for the startup print. "The engine adapts to the machine it is on" is
/// only true if the engine says what it got.
const std::string& note() const { return note_; }
double load_gib_per_second() const { return gib_per_s_; }
private:
void* arena_ = nullptr; ///< the PinnedArena, owned
std::vector<const uint8_t*> dev_slice_; ///< device alias of each registered slice (or of the whole range)
uint64_t slice_bytes_ = 0;
const uint8_t* base_ = nullptr;
int64_t blobs_ = 0;
int64_t n_expert_ = 0;
int64_t reads_ = 0;
std::string note_;
double gib_per_s_ = 0.0;
uint64_t pinned_bytes_ = 0;
std::string gguf_;
};
} // namespace strata::core
include/strata/core/graph.hpp
+117
View File
@@ -0,0 +1,117 @@
// include/strata/core/graph.hpp - P2.S5: the GraphRegistry.
//
// `GraphRegistry`: capture/replay wrappers keyed by `(layer_type, n_tokens)`, with static input/output
// buffers at fixed addresses.
//
// THREE THINGS ARE BAKED INTO THIS INTERFACE, and each one is a measurement rather than a preference. All
// three come from `bench/micro/graph_capture.cu`, which was written to answer the phase's own instruction -
// "verify on the installed CUDA version with a minimal test before building on it".
//
// 1. **THE WAIT MUST POLL THE DRIVER, NOT MEMORY.** The phase's host loop says "spin on
// doorbell_seq == expected ... no cudaStreamSynchronize anywhere". A spin that only reads memory NEVER
// RUNS THE KERNEL: 5,907,703 spins over 500 ms, and the datum flips to 1 the instant anything calls into
// the driver and never before. On Windows the display driver model BATCHES command submission, and a
// thread that only reads memory gives it no reason to flush. Hence `wait_ms`, which polls an event -
// a QUERY and not a blocking sync, so P2.X3's "zero synchronization calls in the layer loop" still holds,
// and it costs 0.021 ms.
//
// 2. **A GRAPH RE-READS ITS INPUT BUFFERS AT REPLAY, BUT NOT ITS KERNEL ARGUMENTS.** Verified: eight
// replays with a different pinned input each time all tracked correctly, so the data path is live; but a
// kernel argument is copied into the node at capture and is never re-read. **Anything that changes per
// step - a position, a page-table base, a token id - must be DATA in a device buffer, never an argument.**
// This is the phase's "graph capture with pointers that change between steps silently replays stale
// addresses" hazard, and it now has a measurement behind it.
//
// 3. **FIXED ADDRESSES ARE THE CALLER'S OBLIGATION.** The registry cannot check them. Every buffer a
// captured body touches must be allocated before `record` and never reallocated; the registry holds the
// graph, not the memory.
//
// What the registry is FOR: a layer's ops are the same every token, so recording them once and replaying
// removes the per-call host cost of submitting them individually. Measured, that cost is 10-22 us PER CALL
// into the driver on this machine, and a layer with 8 submissions is ~0.1 ms of round trip.
#pragma once
#include <cstdint>
#include <cstring>
#include <functional>
#include <map>
#include <string>
#include <vector>
#include <cuda_runtime.h>
namespace strata::core {
/// The two attention families. `docs/semantics.md`: 36 GDN layers and 12 QSA, at indices 3, 7, ... 47.
enum class LayerType : int { GDN = 0, QSA = 1 };
inline const char* to_string(LayerType t) { return t == LayerType::GDN ? "GDN" : "QSA"; }
/// One recorded graph. Move-only: it owns a `cudaGraphExec_t` and an event, and two owners would
/// double-destroy.
class CapturedGraph {
public:
CapturedGraph() = default;
~CapturedGraph() { reset(); }
CapturedGraph(const CapturedGraph&) = delete;
CapturedGraph& operator=(const CapturedGraph&) = delete;
CapturedGraph(CapturedGraph&& o) noexcept { *this = std::move(o); }
CapturedGraph& operator=(CapturedGraph&& o) noexcept;
/// Start recording everything the caller launches into `stream`.
bool begin(void* stream, std::string& err);
/// Stop recording, instantiate, and record the completion event. A capture that produced ZERO nodes is
/// reported as an error: it means the body did nothing, which is a wiring mistake and NOT an empty graph
/// to be replayed happily forever.
bool end(void* stream, std::string& err);
/// Replay. Does not synchronise.
bool launch(void* stream, std::string& err) const;
/// The host loop's wait: poll `done_` until it completes or `timeout_ms` elapses. Returns false on
/// timeout, which the caller must treat as a hang and not as "probably fine". See NOTE 1 above for why
/// this exists instead of a memory spin.
bool wait_ms(int timeout_ms) const;
size_t nodes() const { return nodes_; }
bool valid() const { return exec_ != nullptr; }
private:
void reset();
cudaGraph_t graph_ = nullptr;
cudaGraphExec_t exec_ = nullptr;
cudaEvent_t done_ = nullptr;
size_t nodes_ = 0;
};
/// Capture-once, replay-many, keyed by `(layer_type, n_tokens)`.
///
/// `record` is given a CALLABLE that launches the layer's ops. On the first call for a key it captures;
/// afterwards it returns the cached graph and does NOT run the callable - so a caller that mutates state
/// inside the body will be surprised, which is why the body must only launch.
class GraphRegistry {
public:
explicit GraphRegistry(void* stream) : stream_(stream) {}
/// Capture `body` for `key` if it is not already recorded. Returns false and fills `err` on failure.
/// Recording the same key twice is a no-op and is NOT an error - that is the point of a registry, and a
/// regression test asserts that the body runs exactly once across many calls.
bool record(LayerType type, int n_tokens, const std::function<void()>& body, std::string& err);
const CapturedGraph* find(LayerType type, int n_tokens) const;
/// Launch and wait, in that order, the two things the host loop does per layer.
bool launch(LayerType type, int n_tokens, int timeout_ms, std::string& err) const;
size_t size() const { return graphs_.size(); }
/// How many times a BODY was actually captured. With a working cache this equals `size()`.
size_t captures() const { return captures_; }
private:
using Key = std::pair<int, int>;
void* stream_ = nullptr;
std::map<Key, CapturedGraph> graphs_;
size_t captures_ = 0;
};
} // namespace strata::core
include/strata/core/hit_hook.hpp
+41
View File
@@ -0,0 +1,41 @@
// include/strata/core/hit_hook.hpp - R4.2c: the seam between the host loop and the VRAM expert tier.
//
// **IT IS ITS OWN HEADER BECAUSE BOTH SIDES NEED IT AND NEITHER SHOULD OWN IT.** `session.hpp` declares the
// loop that calls the hook; `expert_source.hpp` declares the implementation that answers it. Putting the type
// in either one makes the other include a header it has no other use for, and the loop knowing what an expert
// cache is (or the adapter knowing what a graph is) is exactly the coupling this seam exists to avoid.
#pragma once
#include <cstdint>
namespace strata::core {
/// **TWO PHASES, AND THE SPLIT IS THE WHOLE POINT.** `Launch` runs BEFORE the CPU pool and `Combine` runs
/// AFTER the misses have been copied into `parts`. That ordering is what lets the GPU's resident experts run
/// at the same time as the CPU's misses instead of after them:
///
/// ring -> hits(Launch) [GPU: quantize + grouped experts into `hit_out`]
/// -> pool() [CPU: the misses, CONCURRENT with the GPU]
/// -> H2D misses into `parts`
/// -> hits(Combine) [GPU: parts += hit_out]
/// -> post[l]
///
/// **MEASURED, AND WHY IT IS NOT THE OBVIOUS ORDER.** The first version called the hit path once, after the
/// pool, writing straight into `parts`. The drain duly fell **18.2 -> 10.2 ms** - the CPU really was doing
/// half the experts - and **the token did not move at all**, because the GPU's half had simply been moved onto
/// the critical path behind the CPU's instead of beside it. A second buffer and one `add_inplace` buy the
/// overlap back.
enum class HitPhase { Launch, Combine };
/// `ids` AND `k` ARE THE ROUTER'S OWN OUTPUT, AND THEY ARE WHY THIS HOOK TAKES ARGUMENTS AT ALL.
///
/// The first version took only `(user, stream, phase)` and left the hit/miss decision inside the pool
/// callback - which runs AFTER `Launch`. So `Launch` read the hit list from the PREVIOUS layer and computed
/// the previous layer's experts into the current layer's rows. It produced a token, the timings looked
/// plausible, and **C1 moved from mean KL 9.69e-02 to 1.03e+00 and top-1 from 0.867 to 0.333.** The decision
/// has to happen in `Launch`, on the ids the doorbell published for THIS layer, and the pool then consumes it.
///
/// A null `HitFn` means "no VRAM tier" - the default, and what every caller had before R4.
using HitFn = void (*)(void* user, void* stream, HitPhase phase, const int32_t* ids, int64_t k);
} // namespace strata::core
include/strata/core/layer.hpp
+566
View File
@@ -0,0 +1,566 @@
// include/strata/core/layer.hpp - one GDN LAYER, composed. P2.S5's mixer.
//
// This is the first place in the engine where the kernels are COMPOSED rather than tested individually, so
// the thing it has to get right is the ORDER and the buffers, not the arithmetic. Every kernel it calls is
// parity-tested on its own; what can go wrong here is a buffer sized for the wrong shape, a step done in the
// wrong order, or a value that should have been converted and was not.
//
// `ref/gdn.py::gdn_decode` is the specification, and the order below is its docstring verbatim:
//
// qkv = wqkv @ x (10240 = q|k|v, each a contiguous run of 128-wide heads)
// conv -> SiLU (SiLU is applied to the WHOLE conv output before the split)
// q,k = l2_norm (v is NOT normalised)
// beta = sigmoid(wbeta @ x)
// gate = softplus(walpha @ x + dt) * ssm_a
// recurrent step
// y = rms_norm(o, ssm_norm) * sigmoid(z) <- SIGMOID for qwen4exp; qwen3.5 used SiLU
// out = ssm_out @ y
//
// FOUR THINGS IN THAT ORDER ARE EASY TO READ PAST, and each produces a plausible number:
//
// * **SiLU is applied BEFORE the split**, to the whole 10240-wide conv output. Applying it only to q and
// k, or after the reshape, is the same arithmetic on different data.
// * **v is NOT normalised.** `l2_norm` applies to q and k only - two rows of 16 heads, not the 48 of v.
// * **the `/hc`-style scale on q is `1/sqrt(S)`**, applied by the CALLER because `ref/gdn.py` says the two
// llama.cpp paths apply it in different places. It goes on AFTER the norm.
// * **the state is transposed** relative to the reference: (S, h_v, S) here, (S, S, h_v) there. That is
// `gdn.hpp`'s choice and it is the reason `gdn_step`'s arguments are laid out as they are.
//
// THE ACTIVATIONS FOLLOW THE CONTRACT, which is why three different GEMVs appear below: `attn_qkv`,
// `attn_gate` and `ssm_out` are K-quants (Q8_K), while `ssm_alpha` and `ssm_beta` are BF16 (bf16). Using one
// activation format for all five would be a 0.66-1.41% error on the three that matter - measured, in
// `s_gemv_q8k_parity` and `bf16_gemv_parity`.
#pragma once
#include "strata/core/layout.hpp"
#include "strata/core/weights.hpp"
#include "strata/kernels/gr.hpp"
#include "strata/kernels/ngram.hpp"
#include "strata/kernels/ple.hpp"
#include <cstdint>
#include <string>
namespace strata::core {
/// Select pinned CUDA BF16/F32 projections for SSM gates, routing and the sparse indexer.
/// Configure before session capture; captured graphs retain the selected implementation.
/// Does not change GR, PLE or the shared expert's scalar gate.
void layer_set_native_bf16(bool enabled);
// Diagnostic vector-attention adapter; configure before capture, context <=256.
void layer_set_native_flash_attn_short(bool enabled);
/// Every buffer one GDN layer needs, carved from an arena ONCE. Sizes are functions of `ModelGeometry`, and
/// `gdn_buffers_bytes` returns the total so a caller can size the arena before committing.
struct GdnBuffers {
// the activation, in the formats the contract asks for. BOTH quantized images are produced for every
// layer, because WHICH ONE A WEIGHT WANTS IS A PROPERTY OF THE WEIGHT and the pack mixes them by layer -
// `attn_qkv` alone is IQ4_XS on 13 layers, Q3_K on 18, Q4_K on 4 and S2 on 1.
uint8_t* x_q8k = nullptr; ///< n_embd, block_q8_K
uint8_t* x_q8_0 = nullptr; ///< n_embd, block_q8_0
uint16_t* x_bf16 = nullptr; ///< n_embd, bf16 bits
// the qkv path
float* qkv = nullptr; ///< conv_channels
float* conv_out = nullptr; ///< conv_channels
float* h = nullptr; ///< conv_channels = silu(conv_out), and it IS q|k|v in place
// the per-head scalars and the recurrence
float* alpha = nullptr; ///< h_v
float* beta = nullptr; ///< h_v
float* gate = nullptr; ///< h_v
float* o = nullptr; ///< h_v * S
float* z = nullptr; ///< value_dim
float* y = nullptr; ///< value_dim
uint8_t* y_q8k = nullptr; ///< value_dim, block_q8_K
uint8_t* y_q8_0 = nullptr; ///< value_dim, block_q8_0
// the recurrent state, which outlives the call
float* state = nullptr; ///< S * h_v * S
float* conv_state = nullptr; ///< conv_channels * (d_conv-1)
};
/// Bytes needed for `b`, 16-byte aligned between regions.
uint64_t gdn_buffers_bytes(const ModelGeometry& g);
/// Carves `base` (which must be DEVICE memory) into `b`. Returns the bytes used.
uint64_t gdn_buffers_init(const ModelGeometry& g, void* base, GdnBuffers& b);
/// Zeroes the recurrent and conv state, so a fresh sequence starts from the reference's own `zeros()`.
void gdn_buffers_zero_state(const GdnBuffers& b, const ModelGeometry& g, void* stream);
/// `Mixed` (n_embd f32) through one GDN layer to `out` (n_embd f32). `state` and `conv_state` are updated.
///
/// `layer` selects the weights; `tables` must outlive the call. Returns false and fills `err` if a tensor is
/// missing - which is a wiring mistake here rather than a shape problem, because `check_layer` already
/// asserted the shapes at load time.
bool gdn_layer(const WeightTable& tables, const ModelGeometry& g, int64_t layer, const GdnBuffers& b,
const float* mixed, float* out, void* stream, std::string& err);
// ================================ the MoE block ================================
struct Doorbell; // defined with the rest of the handoff, below; only a pointer is needed here
/// Buffers for the MoE block, on EVERY layer.
struct MoEBuffers {
uint16_t* x_bf16 = nullptr; ///< n_embd, for the ROUTER's weight, which is BF16
uint16_t* x_f16 = nullptr; ///< n_embd, for the SHARED EXPERT's weights (see the note in `moe_layer`)
float* logits = nullptr; ///< n_expert
int* ids = nullptr; ///< k
float* weights = nullptr; ///< k
float* shared = nullptr; ///< n_embd
float* sh_scratch = nullptr; ///< shared_expert_scratch_bytes(n_ff): the shared expert's own scratch
uint8_t* x_q8_0 = nullptr; ///< n_embd, block_q8_0 - the shared expert's weights want Q8_0 on the layers
///< whose types are Q2_0/Q4_0/Q5_0/Q8_0/IQ4_NL, which is most of them
uint8_t* x_q8k = nullptr; ///< n_embd, block_q8_K - and Q8_K on the K-quant ones
};
uint64_t moe_buffers_bytes(const ModelGeometry& g, int64_t k);
uint64_t moe_buffers_init(const ModelGeometry& g, int64_t k, void* base, MoEBuffers& b);
/// The MoE block: router -> (the caller's expert outputs) -> combine.
///
/// `parts` is (k, n_embd) device f32 and is **filled by the CALLER** - by the CPU pool and the VRAM-resident
/// experts. That split is the architecture's, not a shortcut: the routed experts are the only part of the
/// model that does not fit in VRAM, and every design decision in this engine exists to keep them out of the
/// graph. This function therefore owns the two ends - which experts, and how their outputs are combined -
/// and deliberately not the middle.
///
/// TWO ACTIVATION FORMATS AGAIN, and the router's is the one that matters most:
///
/// * `ffn_gate_inp.weight` is BF16, so the router sees a BF16 activation. `ref/moe.py` L88-92 calls this
/// "the highest-consequence rounding in the model": the logits were measured **8.100e-03** off without it,
/// which is enough to FLIP a top-10 selection - and a flipped selection is not a small numeric error, it
/// is a different expert reading a different weight.
/// * the shared expert's weights are Q3_K/IQ4_XS/Q5_0 IQ-quants whose contract is Q8_K, which
/// `shared_expert` does not implement; it takes fp16. That is the known gap its own header documents and
/// `docs/activation-contract.md` names as outstanding.
bool moe_layer(const WeightTable& tables, const ModelGeometry& g, int64_t layer, int64_t k, const MoEBuffers& b,
const float* x, const float* parts, float* out, void* stream, std::string& err,
const Doorbell* db = nullptr);
// ================================ THE MOE, SPLIT AT THE ROUTER ================================
//
// **THE SPLIT EXISTS BECAUSE THE ONE-LAYER DELAY WAS NOT A TIMING CHANGE.** `moe_layer` is
// `route -> publish -> shared expert -> combine`, and `combine` multiplies `parts` by THIS layer's
// `b.weights`. So `parts` must be THIS layer's expert outputs. A host loop that hands layer `l`'s pool
// result to layer `l+1` computes `sum_j w_{l+1}[j] * expert_{ids_l,j}(x_l)` while the model is
// `sum_j w_l[j] * expert_{ids_l,j}(x_l)` - **both the expert selection and the input one layer stale while the
// weights are current** - and `session.hpp` states the contract it violates: the doorbell is "a change to WHEN
// the host runs, not to what this function computes".
//
// The one-layer delay was introduced for a real reason - it is the only way to overlap a CPU pool with a GPU
// that has already queued the combine - and the fix is to split the GRAPH, not to delay the answer:
//
// block_layer_pre[l] attn -> mixer -> gr_write -> ffn gr_read -> moe_route [rings the doorbell]
// ... the host runs the pool here, on the ids `moe_route` just published ...
// block_layer_post[l] moe_finish -> gr_write
//
// `post[l]` is launched AFTER the pool returns, so it sees layer `l`'s experts with layer `l`'s weights, and it
// is still a captured graph. What it costs is the overlap: the window is now whatever GPU work follows the
// ring inside `pre[l]` - the shared expert and nothing else, because everything after the combine depends on
// `parts`. **A per-layer CPU pool cannot be hidden behind a strictly serial residual chain**, which is why
// Phase 3's VRAM expert cache is the answer to the CPU term and the doorbell is not.
/// The routing half: bf16 activation, the router GEMV, top-k, and the doorbell publish. Ends with the ring
/// rung, so a host that sees it knows `h_x_f`/`h_ids`/`h_weights` are in place.
bool moe_route(const WeightTable& tables, const ModelGeometry& g, int64_t layer, int64_t k, const MoEBuffers& b,
const float* x, void* stream, std::string& err, const Doorbell* db = nullptr);
/// The finishing half: the shared expert and the combination.
///
/// **`parts` MUST BE THE EXPERTS THIS LAYER'S OWN ROUTER SELECTED.** It is combined with `b.weights`, which
/// `moe_route` wrote for THIS layer. Passing the previous layer's vectors is the L100 bug: it produces a
/// finite, fluent, deterministic token that is not the model's.
bool moe_finish(const WeightTable& tables, const ModelGeometry& g, int64_t layer, int64_t k, const MoEBuffers& b,
const float* x, const float* parts, float* out, void* stream, std::string& err);
/// Plan v0.3 P3: the two parts of `moe_finish` separately. `moe_shared` computes the shared expert into
/// `b.shared` and needs only `x`; `moe_combine_parts` combines `parts` with `b.weights` and `b.shared`.
/// `moe_finish` is exactly `moe_shared` followed by `moe_combine_parts`.
bool moe_shared(const WeightTable& tables, const ModelGeometry& g, int64_t layer, const MoEBuffers& b,
const float* x, void* stream, std::string& err);
bool moe_combine_parts(const ModelGeometry& g, int64_t layer, int64_t k, const MoEBuffers& b, const float* parts,
float* out, void* stream, std::string& err);
/// Plan v0.3 P3 (default ON): the shared expert runs at the END OF `pre[l]`, after the doorbell has rung, so
/// the GPU computes it while the host runs the CPU pool; `post[l]` then only combines. Same kernels on the
/// same inputs in the same stream order relative to their consumers, so the result is bitwise unchanged.
/// Must be set before capture. `false` restores the shared expert inside `post[l]` (A/B arm).
void layer_set_shared_early(bool enabled);
bool layer_shared_early();
// ================================ the QSA mixer ================================
/// THE PERSISTENT STATE OF ONE QSA LAYER FOR ONE SEQUENCE: the KV cache, the indexer's pooled keys, the rope
/// tables, and the per-token step buffer. None of it is scratch - it survives every token.
///
/// `max_cells` is the sequence capacity the arena was sized for, and it fixes the rope table's length: a table
/// shorter than the sequence would have `rope_neox_apply` read past it, which is a wrong rotation rather than a
/// fault.
struct QsaState {
uint16_t* k_pool = nullptr; ///< [page][kv_head][page_size][head_dim] fp16 (null in INT8 mode)
uint16_t* v_pool = nullptr;
/// Plan v0.3 P7: INT8 KV (qsa_set_kv_int8). Codes [page][kv_head][page_size][head_dim], one FP16 scale per 64.
bool kv_int8 = false;
int8_t* k_q = nullptr;
int8_t* v_q = nullptr;
uint16_t* k_scale = nullptr;
uint16_t* v_scale = nullptr;
int32_t* page_table = nullptr; ///< (n_pages,) logical page -> physical page
int64_t n_pages = 0;
int64_t max_cells = 0;
float* idx_tail = nullptr; ///< (idx_block - 1, idx_dim): the raw tail of the block being filled
float* idx_dead = nullptr; ///< (idx_dim,): the spare slot's key, CONSTANT for the sequence
float* idx_pooled = nullptr; ///< (max_cells/idx_block + 2, idx_dim)
int32_t* idx_block_pos = nullptr;
float* cos_tab = nullptr; ///< (max_cells, n_rot/2), built on the HOST in float64
float* sin_tab = nullptr;
/// THE PER-TOKEN COUNTS, IN DEVICE MEMORY - the whole reason this layer can be a graph. `qsa_step_fill`
/// computes all four from the position, so one H2D per token updates every count every kernel needs.
int32_t* step = nullptr;
int32_t* attention_status = nullptr; ///< per-layer device status, captured D2H to host_step[kStepCount]
/// The rope positions, (n_head,) int32, all equal to the current position. `rope_neox_apply` reads its
/// positions from the DEVICE, so a host scalar here would be an illegal access (and uncapturable).
int32_t* pos_dev = nullptr;
/// PINNED HOST STAGING AT FIXED ADDRESSES, and they are not an optimisation. The uploads in `qsa_layer`
/// copy FROM these, so a CUDA graph captures the SOURCE POINTER - and the first version of the layer filled
/// a stack array and a local `std::vector` instead. Those addresses are dead by the time the graph is
/// replayed, which is this project's own recorded pitfall ("graph inputs live at fixed addresses") walked
/// straight into: the replay would read whatever the stack held, and the step counts would be wrong in a way
/// that looks like a working token.
///
/// Pinned as well as fixed, because a pageable source would make `cudaMemcpyAsync` a synchronous staged
/// copy and give back the overlap the doorbell protocol exists for.
int32_t* host_step = nullptr; ///< (kStepCount + 1,) pinned; final element is attention status
int32_t* host_pos = nullptr; ///< (n_head,) pinned
};
/// Plan v0.3 P7: the RoPE cos/sin table (max_cells x n_rot/2 x 2 floats, 64 MiB at 262K) is identical in every
/// QSA layer. `with_rope = false` sizes a state that borrows it; `share_rope` points `st` at another state's table
/// instead of building a copy (the session builds it once, in the first QSA layer).
uint64_t qsa_state_bytes(const ModelGeometry& g, int64_t max_cells, bool with_rope = true);
/// Plan v0.3 P7: store K/V as INT8 with FP16 scales per 64 values (half the VRAM of FP16). Set before sizing and
/// initializing the session; default off until gate G-C accepts it.
void qsa_set_kv_int8(bool enabled);
bool qsa_kv_int8();
uint64_t qsa_state_init(const ModelGeometry& g, int64_t max_cells, void* base, QsaState& st,
const QsaState* share_rope = nullptr);
/// Zeroes the pools AND the indexer, so a fresh sequence matches the reference's own `zeros()`. The KV pool
/// matters even for cells that are never attended, because `kv_gather` reads whatever the selection names.
void qsa_state_zero(const QsaState& st, const ModelGeometry& g, void* stream);
// ================================ PER-STAGE TIMING, DEBUG ONLY ================================
//
// Event pairs around the pieces of `block_layer_pre`/`_post`, recorded per (layer, stage) and read once at the
// end so that **NOTHING IS SYNCHRONISED DURING THE RUN** - a per-stage sync would serialise the pipeline and
// report the serialised time. **CAPTURED GRAPHS CANNOT USE IT**: an event record inside a stream capture is
// silently dropped, so this is for the `--no-capture` path only.
//
// It exists because the engine spends ~1.047 ms per layer with the experts off while every cost model in
// `bench/` predicts less than half that, and the gap has to be attributed before it can be attacked.
// slot 0 gr_read(attn) 1 the attention block 2 gr_write(attn) 3 gr_read(ffn) 4 router
// slot 5 moe_finish 6 gr_write(ffn)
bool stage_timing_enable();
void stage_timing_name(int slot, const char* name);
void stage_timing_report(int64_t n_layers);
/// The same event pairs, reachable from `session.cpp`, for marks that sit OUTSIDE the captured graphs - the
/// two `cudaGraphLaunch` calls per layer. Those are ordinary stream operations, so an event around them is
/// legal and is not captured, which is what lets the captured path's 1.047 ms/layer be split into the two
/// graphs' execution and the gap the host spends between them. Slots 16 and 17 are reserved for it; slots
/// 0-15 are the in-layer marks and are only meaningful with `--no-capture`.
void stage_mark_begin(int64_t layer, int slot, void* stream);
void stage_mark_end(int64_t layer, int slot, void* stream);
/// Scratch for one QSA layer. Every buffer here is overwritten by every token.
struct QsaBuffers {
uint8_t* x_q8_0 = nullptr; ///< n_embd, block_q8_0 - `attn_q` is Q2_0, whose vec_dot_type is Q8_0
uint8_t* x_q8k = nullptr; ///< n_embd, block_q8_K - the three K-quant projections
uint16_t* x_bf16 = nullptr; ///< n_embd - the two BF16 indexer projections
float* q_full = nullptr; ///< n_head * 2 * head_dim: the raw `attn_q` output, q | gate per head
float* qcur = nullptr; ///< n_head * head_dim: the FIRST half of each head, contiguously
float* kcur = nullptr; ///< n_head_kv * head_dim
float* vcur = nullptr; ///< n_head_kv * head_dim
float* idx_raw = nullptr; ///< idx_key_dim: NEVER normed and NEVER rotated
float* q_idx = nullptr; ///< idx_q_heads * idx_key_dim
float* cell_scores = nullptr; ///< max_cells
int32_t* ids = nullptr; ///< cap = qsa_selection_width(kTopkMaxCells, s)
uint16_t* k_scratch = nullptr; ///< cap * n_head_kv * head_dim
uint16_t* v_scratch = nullptr;
float* attn = nullptr; ///< n_head * head_dim
uint16_t* attn16 = nullptr; ///< n_head * head_dim, for a fp16-activation consumer
/// `attn * sigmoid(gate)` in f32 and then in Q8_K - what `attn_output` ACTUALLY wants. Its weights are
/// Q4_K/Q5_K/Q6_K, whose `vec_dot_type` is Q8_K, and the fp16 path is worth 0.66-1.41% against the
/// reference. `docs/activation-contract.md`.
float* attn32 = nullptr;
uint8_t* attn_q8k = nullptr;
float* attn_scratch = nullptr; ///< plan v0.3 P3: split-K decode attention partials
};
uint64_t qsa_buffers_bytes(const ModelGeometry& g, int64_t max_cells);
uint64_t qsa_buffers_init(const ModelGeometry& g, int64_t max_cells, void* base, QsaBuffers& b);
/// `x` (n_embd f32) through one QSA layer to `out` (n_embd f32) at sequence position `pos`.
///
/// `pos_base` is the sequence's FIRST cell's position, which the indexer needs and cannot derive - equating a
/// cell's position with its index is only right for a sequence starting at 0.
///
/// THE OP ORDER IS `ref/model.py::_qsa`, and two things in it are easy to read past:
///
/// * the indexer's raw key is appended BEFORE any norm or rotation, and it is the raw `indexer.k_proj`
/// output that is pooled and rotated LATER, once per block and at the block's FIRST cell;
/// * `q_full` is split into q and gate by HALVES of each head's 2*head_dim block, contiguous, not
/// interleaved - `qsa_gate_apply` reads the second half and `qcur` is copied out of the first.
///
/// THE THREE ACTIVATION FORMATS FOLLOW `docs/activation-contract.md`, and they are not interchangeable:
/// `attn_q` is Q2_0 (Q8_0 activation), `attn_k`/`attn_v`/`attn_output` are K-quants (Q8_K), and the two
/// indexer projections are BF16 (bf16).
bool qsa_layer(const WeightTable& tables, const ModelGeometry& g, int64_t layer, int64_t pos, int32_t pos_base,
const QsaState& st, const QsaBuffers& b, const float* x, float* out, void* stream,
std::string& err, float* dump = nullptr);
// ================================ THE DOORBELL ================================
/// THE HANDOFF: `x_f` and the miss list reach the CPU WHILE THE LAYER IS STILL RUNNING.
///
/// Everything about the shape of this struct comes from two measurements, and neither was a guess:
///
/// * **THE WAIT MUST POLL THE DRIVER.** Round 195: a host spin that only reads a memory location never runs
/// the kernel - 5,907,703 spins over 500 ms and the datum flips the instant `cudaStreamSynchronize` is
/// called and never before, because on Windows the driver BATCHES command submission and a memory-only
/// spin gives it no reason to flush. So `h_seq` is read AND `cudaEventQuery` is called, which is a query
/// and not a blocking sync - P2.X3's "zero synchronization calls in the layer loop" still holds.
/// * **THE WRITE MUST BE CAPTURABLE, AND A KERNEL IS.** Round 199: `cudaEventRecord` inside a capture is
/// SILENTLY DROPPED - 2 nodes for two kernels plus an event record, and the event never completed across a
/// 41 ms graph. A mid-graph doorbell KERNEL is captured normally, and round 209 measured the host seeing a
/// pinned handoff **0.050 ms** into a 39.8 ms graph.
///
/// The memory is MAPPED PINNED, so it has one address for the device and another for the host and both refer
/// to the same bytes. Pinned as well as mapped, because a pageable destination would turn the device's write
/// into a staged copy and give back the overlap this exists for.
struct Doorbell {
int64_t n_embd = 0;
int64_t k = 0;
// the HOST's view
float* h_x_f = nullptr; ///< n_embd
int32_t* h_ids = nullptr; ///< k
float* h_weights = nullptr; ///< k
uint32_t* h_seq = nullptr; ///< 1
/// Plan v0.3 P3 token graph: the HOST's answer. The host writes the ring value it has served (after the
/// pool's output is in place); the device's `doorbell_wait` spins until the flag equals the ring.
uint32_t* h_flag = nullptr; ///< 1
// the DEVICE's view of the same bytes
float* d_x_f = nullptr;
int32_t* d_ids = nullptr;
float* d_weights = nullptr;
uint32_t* d_seq = nullptr;
uint32_t* d_flag = nullptr;
};
/// Allocates the four mapped-pinned regions. Returns the bytes, or 0 on failure.
uint64_t doorbell_init(const ModelGeometry& g, int64_t k, Doorbell& db);
void doorbell_free(Doorbell& db);
/// Zeroes `h_seq` (and `h_flag`), so a caller can wait for it to become `expected` from a known state.
void doorbell_reset(const Doorbell& db);
// ================================ ONE WHOLE BLOCK ================================
/// The gated-residual stack plus the two buffers the halves hand to each other.
///
/// `R` is (hc, n_embd) and is BOTH the input and the output: `hc` residual streams, each n_embd wide, which is
/// what `gr_read`/`gr_write` exist for. It is not one vector with a gate.
struct BlockBuffers {
float* R = nullptr; ///< (hc, n_embd)
float* mixed = nullptr; ///< n_embd: what the mixer consumes
float* block_out = nullptr; ///< n_embd: what a half produces
float* inject = nullptr; ///< hc: `gr_read`'s per-stream injection, reused per half
/// Plan v0.3 P3 fused GR: the FFN half's injection lives here (the attention half's stays in `inject`),
/// because the fused read of one half reads the other half's injection while writing its own.
float* inject2 = nullptr; ///< hc
float* gr_rs = nullptr; ///< hc: the fused read's per-stream 1/rms
uint8_t* head_q8k = nullptr; ///< `q8k_bytes(n_embd)`: the head's Q8_K activation image
strata::kernels::GrWorkspace gr; ///< `gr_read`'s scratch
/// THE HALF-LEVEL C1 ORACLE, and it is a HOST buffer of `(2 * n_embd + 2 * hc + n_head * head_dim)` floats
/// PER LAYER: `[0, n_embd)` the attention half's `block_out`, `[n_embd, 2*n_embd)` the MoE half's, the two
/// `hc`-vectors of `inject`, and finally the QSA layers' PRE-GATE attention (`n_head * head_dim`, unused by
/// GDN). The residual ladder (`session_loop`'s `dump_layers`) says WHICH LAYER diverges; this says which
/// stage, and it is the only way to separate the reference's `linear_attn_out-<l>` from its `ffn_out-<l>`
/// without re-deriving one of them. The copies are issued from inside `block_layer_pre`/`_post`, so they
/// are recorded as nodes of the per-layer captured graph and replay with the layer they belong to. Null
/// means no dumping, which is the default and costs nothing.
float* dump = nullptr;
};
uint64_t block_buffers_bytes(const ModelGeometry& g);
/// Plan v0.3 P3: the hyper-connection read as two fused kernels with the previous half's write folded in
/// (`strata/kernels/fused_gr.hpp`). Requires the native MMVF contract (FP32 activations). The residual `R`
/// then lags by one pending write between halves; it is materialised before the layer-1 PLE and after the last
/// layer. Set before capture; not for the per-stage or dump measurements.
void layer_set_fused_gr(bool enabled);
bool layer_fused_gr();
/// Plan v0.3 P3/P7: QSA attention by the split-K kernel reading the KV pools directly (default ON; `false` keeps
/// the gather + one-block-per-head kernel). Set before capture.
void layer_set_fast_attn(bool enabled);
/// Plan v0.3 P3: doorbell payload + ring as one kernel, QSA step staging read by a kernel from mapped memory
/// (default ON; `false` keeps the memcpy nodes). Set before capture.
void layer_set_publish_kernel(bool enabled);
/// Plan v0.3 P3: GDN step + output norm as one coalesced kernel under the native contract (default ON).
void layer_set_fused_gdn(bool enabled);
/// Plan v0.3 P7: QSA block scores in FP32 and a radix top-k over blocks (default ON; `false` = the FP64 row scores
/// and the bit-serial cell top-k). Set before capture.
void layer_set_fast_select(bool enabled);
/// Plan v0.3 P6: whether the current decode configuration is the one the speculative verify window reproduces
/// bit for bit (native projections, fused GR and GDN, split-K attention, block selection, native indexer).
bool layer_verify_compatible(std::string& why);
uint64_t block_buffers_init(const ModelGeometry& g, void* base, BlockBuffers& b);
// ================================ THE TWO ENDS OF A TOKEN ================================
/// THE EMBEDDING LOOKUP, done as a COLUMN DEQUANT rather than a GEMV.
///
/// `token_embd.weight` is `[n_embd, n_vocab]`, so `out[i] = W[i][token]` - one column, `n_embd` values. Running
/// `output`-style arithmetic on a one-hot activation would produce the same numbers with 248320x the work.
///
/// **THE DECODE IS TRANSCRIBED FROM `s_gemv_q8k_split` AND MUST STAY BIT-IDENTICAL TO IT**: the code is the
/// `CODE_BITS`-wide field, least-significant first, at bit `(i % PER_BYTE) * CODE_BITS` of byte `i / PER_BYTE`;
/// the value is `(code + code_bias) * scales[i / group_elems] + offsets[i / group_elems]`; and the column's
/// planes start at `token * n_in / PER_BYTE` and `token * n_in / group_elems`. A second decoder that drifts
/// from the kernel's would make the prompt embed differently from the tokens generated after it.
bool embed_row(const WeightTable& tables, const ModelGeometry& g, int64_t token, float* out_dev, void* stream,
std::string& err);
/// `ref/model.py` L520: `fin = GR.gr_read(R, output_hc_*..., inject=None, rms_eps); logits = output @ fin`.
///
/// `R` is read from `bb.R` and the result goes to `logits`, which the caller sizes at `output.weight`'s `ne1`
/// (the vocabulary). `bb.head_q8k` and `bb.mixed` are scratch.
bool lm_head(const WeightTable& tables, const ModelGeometry& g, const BlockBuffers& bb, float* logits,
void* stream, std::string& err);
/// Compute the final residual mix in bb.mixed, before the output projection. This is shared by the
/// canonical and experimental native head so only the projection's arithmetic changes in their A/B test.
bool lm_head_mix(const WeightTable& tables, const ModelGeometry& g, const BlockBuffers& bb,
void* stream, std::string& err);
/// THE PER-LAYER OP ORDER of `phases/phase-2-correct-engine.md` P2.S5, architecture §6.1, for ONE layer:
///
/// gr_read (attn) -> mixer -> gr_write (attn)
/// gr_read (ffn) -> MoE -> gr_write (ffn)
///
/// `layer` picks the mixer: `is_qsa_layer` selects `qsa_layer`, otherwise `gdn_layer`. `pos`/`pos_base` reach
/// the QSA layer and are ignored by GDN, which has no position argument at all - its whole per-token input is
/// `x` and what advances between tokens is its recurrent state.
///
/// **`parts` IS THE CALLER'S**, and that is the architecture's split rather than an omission: the routed experts
/// are the only part of the model that does not fit in VRAM, so the CPU pool and the VRAM-resident experts fill
/// `parts` and this function owns the two ends. In P2.S5 everything is a miss, so a caller passes whatever the
/// pool produced; the doorbell protocol is what lets it be produced DURING this layer rather than before it.
// ================================ THE PLE (LAYER 1'S PER-LAYER EMBEDDING) ================================
//
// **THIS MODULE EXISTED, WAS PARITY-TESTED, AND WAS NEVER CALLED - WHICH IS WHY GATE C1 FAILED.** Layer 1 is
// the only layer with `blk.1.ple_*` tensors (six of them) and `block_layer` ran GDN/QSA + GR + MoE without it,
// so the engine computed a different model. `ple_parity` verifies all three pieces against things that are
// not this project's code - the hash against `ref/ngram.py` with the artifact's own constants, the IQ4_NL row
// read against numpy over the original GGUF, and the block against ggml's own graph - and none of that helps
// if nothing calls it.
//
// THE TABLE IS NOT IN THE PACK. `per_layer_token_embd.weight` is 51.2e9 elements (28.8 GB) of IQ4_NL in the
// ORIGINAL second GGUF shard, and it is the only tensor this engine reads from the GGUF rather than from the
// canonical pack. It is also why the PLE is cheap: sixteen 90-byte rows per token is 1440 bytes of traffic,
// against the experts' 664 MB.
struct PleRun {
/// Null when the PLE is not wired (no table, or a pack without `blk.1.ple_*`). Then `block_layer_pre`
/// SKIPS the module entirely, which is what every measurement before this change did.
strata::kernels::PleTable* table = nullptr;
strata::kernels::PleWeights w;
/// The session's PLE state, by pointer so `block_layer_pre` needs no `SessionState`.
const int32_t* token = nullptr; ///< the token being decoded
const int32_t* prev = nullptr; ///< (2,) predecessors, OLDEST FIRST, `-1` where there is none
float* hist = nullptr; ///< NG_HIST x hc_dim NORMALIZED history, row-fastest
/// The gathered 2560-wide row vector on the DEVICE, and the host staging it is copied from: `gather`
/// writes HOST memory (the table is a host mapping) and the block wants the device.
float* emb_dev = nullptr;
float* emb_host = nullptr;
/// Internal workspace plus a disjoint normalized export for the next history row.
/// Allocate ple_run_scratch_bytes(), not merely the lower-level ple_block_scratch_bytes().
float* scratch = nullptr;
strata::kernels::PleConsts consts;
bool ready() const {
return table != nullptr && table->is_open() && hist != nullptr && emb_dev != nullptr &&
emb_host != nullptr && scratch != nullptr && token != nullptr && prev != nullptr &&
((w.key_codes != nullptr && w.key_scales != nullptr) || w.key_bf16 != nullptr ||
w.key_native_data != nullptr) &&
w.value_bf16 != nullptr &&
w.norm_key != nullptr && w.norm_query != nullptr && w.norm_conv != nullptr &&
w.conv1d_f16 != nullptr;
}
};
/// **CALLED ONCE PER TOKEN, FROM THE DRIVER, *OUTSIDE* ANY CAPTURE.** The n-gram hash and the table gather
/// are host operations, and `block_layer_pre` is captured - so doing them there would run them once, at capture
/// time, and replay the same rows for every token forever. This is the doorbell's cloned-literal bug in a
/// different coat: silent, finite, and stable.
bool ple_stage_token(const PleRun& p, void* stream, std::string& err);
/// Plan v0.3 P2: `ple_stage_token` in two halves, so the SSD reads overlap other host work.
/// `ple_issue_token` hashes the token and starts its 16 row reads; `ple_finish_token` waits for them,
/// dequantizes and queues the upload. Both are driver-side, outside any capture, like `ple_stage_token`.
bool ple_issue_token(const PleRun& p, std::string& err);
bool ple_finish_token(const PleRun& p, void* stream, std::string& err);
/// Internal PLE workspace plus one separate hc_dim normalized row for advancing history.
uint64_t ple_run_scratch_bytes();
/// `ple` is the layer-1 PLE, or null to skip it (which is what every measurement before LEDGER L123 did).
/// It is applied INSIDE `pre[1]`, before layer 1's own `gr_read`, because the block's output is
/// `hidden + gated + silu(conv)` on the residual itself - it is a residual UPDATE, and the architecture puts
/// the gathered rows in place "before layer 2".
bool block_layer(const WeightTable& tables, const ModelGeometry& g, int64_t layer, int64_t pos, int32_t pos_base,
const GdnBuffers& gb, const QsaState& qst, const QsaBuffers& qb, const MoEBuffers& mb,
int64_t k, const BlockBuffers& bb, const float* parts, void* stream, std::string& err,
const Doorbell* db = nullptr, const PleRun* ple = nullptr);
/// **THE BLOCK, SPLIT AT THE ROUTER, SO A HOST LOOP CAN PUT THE CPU POOL IN THE MIDDLE.**
///
/// pre[l] gr_read(attn) -> mixer -> gr_write(attn) -> gr_read(ffn) -> moe_route [rings the doorbell]
/// post[l] moe_finish -> gr_write(ffn)
///
/// `parts` reaches only `post`, and it must be THIS layer's experts. The two calls are NOT independent: `post`
/// reads `bb.mixed` and `bb.inject` exactly as `pre` left them, so nothing may run between them that touches the
/// block buffers - and in particular `pre[l+1]` may NOT be launched between them, because its own `gr_read`
/// overwrites both. That constraint is why the doorbell's window is the shared expert and nothing more, and
/// why the CPU expert term is answered by Phase 3's VRAM cache rather than by this pipeline.
///
/// **`half` EXISTS SO THE CAPTURED GRAPH CAN BE SPLIT FOR MEASUREMENT, AND IT IS THE ONLY REASON.**
///
/// The plan's R3 items - GR, GDN, the router - are each worth several milliseconds and there has never been a
/// per-stage breakdown of the CAPTURED graph to choose between them. The `--no-capture` table cannot supply
/// one: round 312 showed its `moe_finish` figure was a host-gap artefact, four times the captured cost of the
/// same work. An event recorded BETWEEN two `cudaGraphLaunch` calls IS valid, so the answer is to capture the
/// layer as more than one graph and time them from outside.
///
/// `half` selects a prefix of the stages, and the seam is the one this file already documents:
///
/// 1 = stages 0..2 gr_read -> attention -> gr_write the MIXER
/// 2 = stages 3..4 gr_read -> moe_route the FFN front and the ROUTER
/// 0 = both, which is what every caller before this used and what `session_loop` still uses
///
/// **A half is not a free-standing call.** `2` reads what `1` left in `bb.mixed`/`bb.inject`, exactly as
/// `block_layer_post` does; the two are only meaningful replayed in order on one stream, which is what
/// `session_replay_stages` does.
bool block_layer_pre(const WeightTable& tables, const ModelGeometry& g, int64_t layer, int64_t pos,
int32_t pos_base, const GdnBuffers& gb, const QsaState& qst, const QsaBuffers& qb,
const MoEBuffers& mb, int64_t k, const BlockBuffers& bb, void* stream, std::string& err,
const Doorbell* db = nullptr, const PleRun* ple = nullptr, int half = 0,
int stage_prefix = 0);
bool block_layer_post(const WeightTable& tables, const ModelGeometry& g, int64_t layer, int64_t k,
const MoEBuffers& mb, const BlockBuffers& bb, const float* parts, void* stream,
std::string& err);
} // namespace strata::core
include/strata/core/layout.hpp
+97
View File
@@ -0,0 +1,97 @@
// include/strata/core/layout.hpp - the pack's LAYOUT, resolved by name and checked against the kernels.
//
// `WeightTable` answers "where are the bytes". This answers "is this the tensor I think it is", and it is
// the layer where a pack and a kernel disagree - which is the failure this project keeps paying for. Two
// examples, both real:
//
// * round 194 sized the indexer key store from `indexer.head_count = 4`, which counts QUERY heads. The
// cached key is ONE shared head of 128 (`indexer.k_proj` is [2560, 128]), so the term was 4x too big and
// the error survived a round because it moved in the direction that TIGHTENS the budget.
// * `docs/semantics.md` and the tensor manifest agree on every dimension here, but nothing CHECKED that a
// named tensor had the shape the kernel reading it assumes.
//
// So this header does two things: it names the tensors per layer type, and it asserts every shape the
// kernels depend on. A mismatch is reported with the tensor name, the shape found and the shape required,
// at LOAD time - not as a wrong number inside a GEMV at token 4000.
#pragma once
#include "strata/core/weights.hpp"
#include <cstdint>
#include <string>
namespace strata::core {
/// The model's geometry, taken from `docs/semantics.md` and the artifact's own metadata. Every field here
/// is a number a kernel depends on, so a change is a change to a kernel contract and not a tuning knob.
struct ModelGeometry {
int64_t n_embd = 2560;
int64_t n_layers = 48;
int64_t qsa_interval = 4; ///< every 4th layer is full attention: layers 3, 7, ... 47
// GDN (36 layers)
int64_t ssm_state_size = 128;
int64_t ssm_k_heads = 16;
int64_t ssm_v_heads = 48;
int64_t ssm_d_conv = 4;
int64_t ssm_conv_channels = 10240; ///< 2*128*16 + 128*48
int64_t ssm_value_dim = 6144; ///< 128 * 48
// QSA (12 layers)
int64_t n_head = 24;
int64_t n_head_kv = 2;
int64_t head_dim = 256;
int64_t idx_q_heads = 4;
int64_t idx_key_dim = 128;
// gated residual, on every layer
int64_t hc = 4;
int64_t hc_lr = 320;
// MoE, on every layer
int64_t n_expert = 512;
int64_t n_ff = 640;
int64_t hc_dim() const { return hc * n_embd; }
/// `layer % qsa_interval == qsa_interval - 1` is full attention. Derived, not a second list.
int64_t n_qsa_layers() const { return n_layers / qsa_interval; }
int64_t n_gdn_layers() const { return n_layers - n_qsa_layers(); }
};
/// True for the full-attention layers. `docs/semantics.md` gives this twice over - `full_attention_interval
/// = 4` and an explicit `attention.compress_ratios` array - and this is the first of the two.
inline bool is_qsa_layer(const ModelGeometry& g, int64_t layer) {
return layer % g.qsa_interval == g.qsa_interval - 1;
}
/// One layer's tensors, resolved by NAME. `get("attn_qkv.weight")` looks up `blk.<layer>.attn_qkv.weight`
/// and returns null if the pack does not have it - a null is information, because a GDN layer has no
/// `attn_q` and a QSA layer has no `attn_qkv`.
///
/// Holds a reference to the table; the table must outlive it.
class LayerView {
public:
LayerView(const WeightTable& table, int64_t layer) : table_(&table), layer_(layer) {}
int64_t layer() const { return layer_; }
std::string name(const char* suffix) const;
const WeightRef* get(const char* suffix) const { return table_->find(name(suffix)); }
private:
const WeightTable* table_;
int64_t layer_;
};
/// Every shape the kernels depend on, asserted for ONE layer. Returns false and fills `err` with the first
/// mismatch, naming the tensor, what it has and what is required.
///
/// This is deliberately a separate function from `LayerView`: a caller that only wants the pointers should
/// not pay for the checks, and a caller that wants the checks should get ALL of them rather than the ones
/// its own call site happens to touch.
bool check_layer(const WeightTable& table, const ModelGeometry& g, int64_t layer, std::string& err);
/// `check_layer` over every layer, plus the cross-layer properties: exactly 12 QSA layers at the right
/// indices, every GDN layer having the GDN set and no QSA tensor, and vice versa. Returns the first failure.
bool check_all(const WeightTable& table, const ModelGeometry& g, std::string& err);
} // namespace strata::core
include/strata/core/mtp.hpp
+135
View File
@@ -0,0 +1,135 @@
// include/strata/core/mtp.hpp - plan v0.3 P6: the MTP draft layer on the GPU.
//
// The model ships one multi-token-prediction layer (vLLM 0.30.0 `qwen4_exp` MTP; transcribed and validated in
// `tools/mtp_probe.py`, where it accepts 0.89 / 0.86 / 0.85 of greedy drafts at steps 1-3). Cell i of its
// sequence pairs the main model's final multi-stream residual at position i with the token at position i+1, at
// rope position i (vLLM's convention); its output predicts the token at position i+2, and its own residual feeds
// the next draft step.
//
// e = fc_embedding(rms(embed(tok)) * (1+w)) h = fc_hidden per stream(rms_10240(R) * (1+w))
// R = h + e -> attention hyper-connection -> QSA attention (own K/V) -> MLP hyper-connection ->
// MoE (512 experts, top-10, shared expert) -> R += y * 2 sigmoid(inject / hc) -> final mixer -> head
//
// Runtime choices, each a trade of exactness for simplicity that only affects DRAFT quality, never the output
// (the verify window decides every emitted token):
// * the large projections run as Q8_0 through the multi-column MMVQ (`tools/mtp_rt.py` quantizes them);
// * the attention is DENSE over every cell the layer has seen - identical to the model's sparse selection
// below 2,051 cells - so speculative cells (draft steps, rejected window rows) never touch indexer state and
// are simply overwritten when their positions are processed again;
// * all 512 routed experts live in VRAM (708 MB) and run through the grouped hit kernels.
#pragma once
#include "strata/core/layer.hpp"
#include "strata/core/session.hpp"
#include <cuda_runtime.h>
#include <cstdint>
#include <string>
#include <vector>
namespace strata::core {
class NativeHead;
class MtpDrafter {
public:
MtpDrafter() = default;
~MtpDrafter();
MtpDrafter(const MtpDrafter&) = delete;
MtpDrafter& operator=(const MtpDrafter&) = delete;
/// Loads `rt_dir` (from tools/mtp_rt.py) and allocates the layer's K/V and buffers for up to `max_t` rows.
/// Call before the VRAM expert tier is sized: this takes ~0.9 GB.
bool load(const std::string& rt_dir, const ModelGeometry& g, SessionState& ss, int max_t, std::string& err,
int64_t window = 32768);
/// The prompt's length: prefill() skips the cells the attention window can never reach again.
void set_prompt_len(int64_t n) { prompt_len_ = n; }
uint64_t vram_bytes() const { return vram_; }
/// The main model's embedding and head, and the verify window's final residuals (T rows, hc*n_embd each).
bool bind(const WeightTable& wt, const NativeHead* head, const float* window_R, std::string& err);
/// Prompt cells [cell0, cell0 + n): residual rows `R_rows` (device, hc*n_embd each) and `next_tokens` (host,
/// the token at position cell+1). Runs in batches of up to max_t rows.
bool prefill(const float* R_rows, const int32_t* next_tokens, int64_t n, int64_t cell0, std::string& err);
/// One round: catch-up over T cells from `p` (rows = the window's final residuals, `tokens` = the window's
/// argmaxes: row t pairs R_{p+t} with the token at p+t+1), then the draft chain from row `a` (the last
/// accepted row) for T-1 drafts at cells p+a+1 ... `drafts` gets T-1 tokens.
bool draft(int T, const int32_t* tokens, int64_t p, int a, int32_t* drafts, std::string& err,
float* probs = nullptr, float min_p = 0.0f, int* n_drafts = nullptr);
/// The first round: one cell (`cell`) from `R_row` (device) and `token` -> T-1 drafts.
bool draft_first(int T, const float* R_row, int32_t token, int64_t cell, int32_t* drafts, std::string& err,
float* probs = nullptr, float min_p = 0.0f, int* n_drafts = nullptr);
double ms_draft = 0, ms_prefill = 0;
int64_t rounds = 0;
private:
bool record_forward(int T, int step_row0, cudaStream_t cs, std::string& err);
bool capture_prefill(int T, std::string& err);
bool capture_round(int T, std::string& err);
bool capture_step(int j, std::string& err);
cudaGraphExec_t step_exec_[9] = {};
const float* f32(const char* name) const;
const uint16_t* bf16(const char* name) const;
const void* q8(const char* name) const;
const ModelGeometry* g_ = nullptr;
SessionState* ss_ = nullptr;
const WeightTable* wt_ = nullptr;
const NativeHead* head_ = nullptr;
const float* window_R_ = nullptr;
int max_t_ = 0;
int64_t n_vocab_ = 0;
uint64_t vram_ = 0;
cudaStream_t cs_ = nullptr;
cudaGraphExec_t prefill_exec_[9] = {};
cudaGraphExec_t round_exec_[9] = {};
struct Tensor { std::string name, kind; int64_t rows = 0, cols = 0; uint64_t off = 0, bytes = 0; };
std::vector<Tensor> tensors_;
uint8_t* dense_ = nullptr;
uint8_t* experts_ = nullptr;
void* state_arena_ = nullptr;
QsaState st_;
void* arena_ = nullptr;
// mapped staging: tokens, step records (2*max_t rows), positions per head (2*max_t rows), the selected row,
// the drafts out
int32_t *h_tok_ = nullptr, *m_tok_ = nullptr, *h_step_ = nullptr, *m_step_ = nullptr;
int32_t *h_pos_ = nullptr, *m_pos_ = nullptr, *h_row_ = nullptr, *m_row_ = nullptr;
int32_t *h_out_ = nullptr, *m_out_ = nullptr;
float *h_prob_ = nullptr, *m_prob_ = nullptr; // each draft's probability under the draft layer
// the draft head: the main head's rows for a token subset (rt/draft_vocab.bin), or the whole head
uint8_t* dhead_ = nullptr;
int32_t* dvocab_ = nullptr;
int64_t n_dvocab_ = 0;
std::string rt_dir_;
int64_t window_ = 0; // attention over the last window_ cells (0 = every cell)
int64_t prompt_len_ = 0;
float* probs_ = nullptr;
// device
int32_t *tok_ = nullptr, *step_ = nullptr, *pos_ = nullptr, *row_ = nullptr, *ident_ = nullptr;
float *Rin_ = nullptr, *R_ = nullptr, *emb_ = nullptr, *en_ = nullptr, *e2_ = nullptr, *hn_ = nullptr, *h2_ = nullptr;
float *mixed_ = nullptr, *inj_ = nullptr, *inj2_ = nullptr, *lo_ = nullptr, *rs_ = nullptr, *bo_ = nullptr;
float* xn_ = nullptr;
uint8_t* xq_ = nullptr;
float *qfull_ = nullptr, *qcur_ = nullptr, *kcur_ = nullptr, *vcur_ = nullptr, *attn_ = nullptr, *attn32_ = nullptr;
float* attn_scratch_ = nullptr;
float *logits_ = nullptr, *w_ = nullptr, *shared_ = nullptr, *parts_ = nullptr, *y_ = nullptr, *sample_ = nullptr;
int32_t *ids_ = nullptr, *hit_slot_ = nullptr, *hit_dst_ = nullptr, *hit_count_ = nullptr, *out_ids_ = nullptr;
uint8_t* hit_xq_ = nullptr;
unsigned long long* grp_ptr_ = nullptr;
int32_t *grp_start_ = nullptr, *grp_counts_ = nullptr;
float* hit_xs_ = nullptr;
void* hit_scratch_ = nullptr;
float* sh_scratch_ = nullptr;
uint16_t* x_bf16_ = nullptr;
float* head_logits_ = nullptr;
float* dummy_inj_ = nullptr;
int64_t cap_ = 0, attn_scratch_floats_ = 0;
};
} // namespace strata::core
include/strata/core/native_dense.hpp
+36
View File
@@ -0,0 +1,36 @@
#pragma once
#include <cstdint>
#include <set>
#include <string>
#include <vector>
namespace strata::core {
class WeightTable;
// Experimental GDN/QSA/shared-expert projection overrides. Upload unchanged native GGUF
// blocks once, then attach them to the matching canonical WeightRef. Unsupported
// types retain their canonical paths. Owns one Q8_1 scratch vector shared by all
// these projections, so use one ordered session stream and keep this object
// alive until all graphs that reference it have been destroyed and synchronized.
class NativeDense {
public:
NativeDense() = default;
~NativeDense();
NativeDense(const NativeDense&) = delete;
NativeDense& operator=(const NativeDense&) = delete;
bool load(const std::vector<std::string>& shards, WeightTable& table, std::string& err,
bool include_ple_key = false);
/// Plan v0.3 P1: the canonical tensor names `load` would serve natively from these shards (eligible name,
/// supported type, 2-D), read from the GGUF headers only - so the canonical arena can skip them.
static bool served_names(const std::vector<std::string>& shards, bool include_ple_key,
std::set<std::string>& out, std::string& err);
uint64_t weight_bytes() const { return bytes_; }
size_t tensor_count() const { return weights_.size(); }
private:
std::vector<void*> weights_;
void* scratch_ = nullptr;
uint64_t bytes_ = 0;
};
} // namespace strata::core
include/strata/core/native_head.hpp
+65
View File
@@ -0,0 +1,65 @@
#pragma once
#include <cstddef>
#include <cstdint>
#include <string>
namespace strata::core {
/// Opt-in native output projection. Uploads unchanged GGUF blocks once at startup and owns its
/// activation scratch. The normal canonical head remains available for numerical A/B comparisons.
/// Plan v0.3 P6: any type the native MMVQ takes (Q5_K in the Q2_0 / IQ3_XXS files, IQ4_XS in IQ2_XS).
class NativeHead {
public:
NativeHead() = default;
~NativeHead();
NativeHead(const NativeHead&) = delete;
NativeHead& operator=(const NativeHead&) = delete;
bool load(const std::string& gguf, int64_t n_in, int64_t n_out, std::string& err);
bool run(const float* mixed, float* logits, void* stream, std::string& err) const;
uint64_t weight_bytes() const { return bytes_; }
bool loaded() const { return weights_ != nullptr; }
/// Plan v0.3 P6: the GGUF blocks, for the verify window's multi-column head.
const void* weights() const { return weights_; }
int type() const { return type_; }
/// Bytes of one vocabulary row.
size_t row_bytes() const { return n_out_ > 0 ? (size_t) (bytes_ / (uint64_t) n_out_) : 0; }
private:
void* weights_ = nullptr;
void* scratch_ = nullptr;
uint64_t bytes_ = 0;
int n_in_ = 0, n_out_ = 0;
int type_ = -1;
};
/// Plan v0.3 P6: `token_embd.weight` in its GGUF form (the IQ model files), in mapped pinned host memory: a row
/// is read over PCIe per token, so the table costs no VRAM. `embed_row`, the verify window and the MTP drafter
/// use it instead of the canonical table when it is set (`set_native_embed`).
class NativeEmbed {
public:
NativeEmbed() = default;
~NativeEmbed();
NativeEmbed(const NativeEmbed&) = delete;
NativeEmbed& operator=(const NativeEmbed&) = delete;
bool load(const std::string& gguf, int64_t n_embd, int64_t n_vocab, std::string& err);
/// Rows for device token ids.
void gather_dev(const int32_t* tokens, int64_t n_tok, float* out, void* stream) const;
/// One row for a host token id.
void gather_one(int64_t token, float* out, void* stream) const;
uint64_t bytes() const { return bytes_; }
int type() const { return type_; }
private:
void* host_ = nullptr;
const void* dev_ = nullptr;
uint64_t bytes_ = 0;
size_t row_ = 0;
int64_t n_embd_ = 0, n_vocab_ = 0;
int type_ = -1;
};
void set_native_embed(const NativeEmbed* e);
const NativeEmbed* native_embed();
} // namespace strata::core
include/strata/core/pinned.hpp
+95
View File
@@ -0,0 +1,95 @@
// include/strata/core/pinned.hpp - P2.S1: the host arena the expert weights live in.
//
// 33.97 GB of expert weights cannot fit in a 12 GB card, so they stay in host memory and are streamed. That
// makes this arena the engine's real working set: it must be PAGE-LOCKED for the copy engine to reach full
// bandwidth, and it must say which backing it got, because the two options differ by more than twice in TLB
// reach:
//
// * LARGE PAGES (2 MB) - `mmap(MAP_HUGETLB)` / hugetlbfs on Linux, `VirtualAlloc(MEM_LARGE_PAGES)` on
// Windows. 33.97 GB at 4 KB pages is 8.3 million TLB entries, which does not fit in any TLB, so every
// block of every expert matvec takes TLB misses.
// * NORMAL PAGES - the fallback. Correct, slower, and it must be REPORTED rather than silently accepted:
// "the engine adapts to the machine it is on" is only true if the engine says what it got. On Windows
// large pages additionally need SeLockMemoryPrivilege, which a normal user account does not have, so this
// fallback is the common case and not an error path.
#pragma once
#include <cstdint>
#include <string>
#include <vector>
namespace strata::core {
enum class PageBacking { LargePages, NormalPages, PinnedByCuda };
struct PinnedArena {
void* base = nullptr;
uint64_t capacity = 0;
PageBacking backing = PageBacking::NormalPages;
std::string note; // why the backing is what it is, for the startup print
uint64_t locked_bytes = 0; // resident via the working-set lock when CUDA could not pin it
/// Plan v0.3 P5: when the whole arena cannot be registered, it is registered in `slice`-byte pieces from the
/// start; this is the pinned prefix (a copy that stays inside one slice can then DMA straight from the arena).
uint64_t registered_bytes = 0;
uint64_t slice_bytes = 0;
int registered_slices = 0;
PinnedArena() = default;
/// `slice`: the piece size for the per-slice registration fallback (0 = none).
explicit PinnedArena(uint64_t bytes, uint64_t slice = 0);
/// Plan v0.3 P6: slices of different sizes (one per layer of a native pack), given as their start offsets
/// followed by the end of the last one. `slice_starts` holds the registered ones.
PinnedArena(uint64_t bytes, const std::vector<uint64_t>& bounds);
std::vector<uint64_t> slice_starts;
~PinnedArena();
PinnedArena(const PinnedArena&) = delete;
PinnedArena& operator=(const PinnedArena&) = delete;
bool valid() const { return base != nullptr; }
uint8_t* data() const { return (uint8_t*) base; }
};
struct LoadStats {
double seconds = 0.0;
uint64_t bytes = 0;
uint64_t layers = 0;
std::vector<uint64_t> layer_checksums; // one FNV-1a per layer
double gib_per_second() const { return seconds > 0 ? (double) bytes / (1024.0 * 1024 * 1024) / seconds : 0.0; }
};
// Load `layers` layers of the expert arena into `dst` with `threads` readers, `chunk` bytes at a time.
// Each thread opens its OWN handle and seeks, which is the portable form of parallel pread: a shared handle
// needs a lock around the seek and defeats the parallelism on Windows.
LoadStats load_experts(const std::string& path, uint8_t* dst, uint64_t blob_bytes, uint64_t blobs_per_layer,
uint64_t layers, int threads, uint64_t chunk);
/// Plan v0.3 P6: the same with one byte range per layer (`layer_off[L]`, `layer_bytes[L]`).
LoadStats load_experts_ranges(const std::string& path, uint8_t* dst, const std::vector<uint64_t>& layer_off,
const std::vector<uint64_t>& layer_bytes, int threads, uint64_t chunk);
// FNV-1a 64. Per layer, so a corrupt or short read names WHICH layer rather than just failing a whole-file
// comparison - the same reason the Phase 1 tools report the first differing element.
uint64_t fnv1a64(const uint8_t* p, uint64_t n, uint64_t seed = 1469598103934665603ull);
// ---- the expert stream -----------------------------------------------------
//
// THIS IS THE NUMBER THE OBJECTIVE DEPENDS ON. With one expert per stream, the design's own arithmetic is:
//
// per token: E = 663.6 MB of expert weights, of which only the MISSES cross the bus
// miss fraction (1 - h) x 663.6 MB / measured bandwidth = milliseconds per token
//
// So `h` (the VRAM hit rate, ~0.167 by design) is only worth anything if a per-expert transfer can actually
// reach the bus's bandwidth. A 1,382,400-byte copy is SMALL, and if small copies run at a fraction of large
// ones then the architecture's central constant is wrong and no kernel optimisation can fix it. That is what
// this measures: the same total bytes at several granularities, so the difference is visible.
struct StreamStats {
double seconds = 0.0;
uint64_t bytes = 0;
uint64_t chunk = 0;
double gib_per_second() const { return seconds > 0 ? (double) bytes / (1024.0 * 1024 * 1024) / seconds : 0.0; }
};
// Copy `bytes` from pinned host memory at `src` to a device buffer of `chunk` bytes, in `chunk`-sized
// transfers, `iters` times. The destination is REUSED, so this measures the bus and not allocation.
StreamStats stream_bandwidth(const uint8_t* src, uint64_t bytes, uint64_t chunk, int iters);
} // namespace strata::core
include/strata/core/session.hpp
+411
View File
@@ -0,0 +1,411 @@
// include/strata/core/session.hpp - ONE TOKEN through all 48 layers, and the doorbell protocol between them.
//
// P2.S5's host loop, and the first place the engine runs a whole forward pass rather than one block. What it
// owns is the ORDER across layers and the state that outlives a layer:
//
// * the gated residual stack `R` - (hc, n_embd), which every block reads and writes;
// * 12 QSA states (KV cache + indexer) and 36 GDN states (recurrence + conv history), because `is_qsa_layer`
// splits the 48 layers 12/36 and the two kinds of state have nothing in common;
// * the per-token expert handoff, which is where the CPU pool meets the GPU.
//
// WHY THE LAYER GRAPHS ARE CAPTURED AND THE LOOP IS NOT. A graph bakes in its arguments, so the graph unit is
// ONE LAYER: everything per-token (the position, the step counts) arrives through fixed-address device buffers
// (rounds 211-214), and the HOST decides when each layer runs. That decision is the whole point - it is what
// lets the CPU expert pool start while the GPU is still working on the next layer.
#pragma once
#include "strata/core/hit_hook.hpp"
#include "strata/core/layer.hpp"
#include <cuda_runtime.h>
#include <cstdint>
#include <string>
#include <vector>
namespace strata::core {
/// Everything a sequence needs that is NOT a weight: the per-layer state, the block scratch, and the pinned
/// handoff between the GPU and the CPU expert pool.
struct SessionState {
int64_t max_cells = 0;
GdnBuffers gdn; ///< the 36 GDN layers share one set of scratch; their STATE is per layer
float* gdn_state = nullptr; ///< (n_gdn_layers, gdn_state_floats)
QsaState* qsa_states = nullptr; ///< one per QSA layer
void* qsa_state_arena = nullptr;
QsaBuffers qsa_bufs; ///< scratch, shared across the 12 (they never run concurrently)
void* qsa_buf_arena = nullptr;
MoEBuffers moe;
void* moe_arena = nullptr;
BlockBuffers block;
void* block_arena = nullptr;
float* R = nullptr; ///< alias of `block.R`, named for what it means at this level
int64_t k = 10; ///< experts per token
/// The doorbell the host loop polls. Null until a caller provides one - the engine runs without it, and
/// neither `session_token` nor `session_replay` looks at it.
const Doorbell* db = nullptr;
// ================================ THE PLE, AND IT IS SEQUENCE STATE ================================
//
// **THE PLE IS A LAYER-1 MODULE AND THIS ENGINE DID NOT RUN IT, WHICH IS WHY C1 FAILED** (LEDGER L123).
// Every part of it was already built and parity-tested - the hash against `ref/ngram.py` with the
// artifact's own constants, the IQ4_NL row read against numpy over the original GGUF, and the block
// against ggml's own graph - and none of it was ever called from the model path. These fields are what
// was missing on the state side.
//
// `ple_prev` is the last two tokens OLDEST FIRST with `-1` where there is none, which is exactly
// `ngram_rows`' `prev` argument; `ple_token` is the token being decoded. **They are advanced by the
// DRIVER, once per token**, because the hash is a function of the sequence and not of a layer.
int32_t ple_prev[2] = {-1, -1};
int32_t ple_token = -1;
/// `NG_HIST` rows of `hc_dim` NORMALIZED history, **ROW-FASTEST**: `ple_hist[row + NG_HIST*channel]`.
/// That is `ggml_reshape_3d(state, d_conv-1, conv_channels, n_seqs)`, not a layout choice - the transposed
/// reading fails ggml's own assert (`ple_layer_xcheck.cpp` L79-83).
float* ple_hist = nullptr;
/// **THE WHOLE PLE RUN LIVES HERE, AND THAT IS NOT A CONVENIENCE.** Its history and its token window are
/// sequence state, so the session is where they belong - and putting the rest of it here too means
/// `session_capture`/`session_loop`/`session_token` need no new parameter to reach it. A `PleRun` whose
/// `ready()` is false is skipped, which is what every caller had before LEDGER L123.
PleRun ple;
};
/// Bytes for a whole session at `max_cells` of context. Every layer's state is sized at once, because P2.T10
/// requires ZERO token-path allocations - a `cudaMalloc` that happened on the first token of a longer sequence
/// would satisfy every test here and fail that one.
uint64_t session_bytes(const ModelGeometry& g, int64_t max_cells, int64_t k);
/// Carves `base` (DEVICE memory) into `s`. Returns the bytes used.
uint64_t session_init(const ModelGeometry& g, int64_t max_cells, int64_t k, void* base, SessionState& s);
/// Zeroes every layer's state - the residual to `R_init`, everything else to zero, so a fresh sequence starts
/// from the reference's own `zeros()`.
void session_zero(SessionState& s, const ModelGeometry& g, const float* R_init, void* stream);
/// One token: layers 0..47 in order, each a `block_layer`, and the residual is updated in place.
///
/// `parts` is (k, n_embd) DEVICE memory, filled by the caller - by the CPU expert pool and the VRAM-resident
/// experts. In this phase `moe_layer` consumes it directly; the doorbell protocol's job is to let it be
/// produced DURING the layer rather than before it, which is a change to *when* the host runs, not to what this
/// function computes.
///
/// `sync_every_layer` inserts a device synchronisation and an error check after EVERY layer. That is the
/// `--sync-every-layer` debug mode, and it is deliberately a parameter rather than a compile-time choice: the
/// errors it catches (a kernel reading a buffer that a later layer writes) are invisible otherwise, and the
/// synchronisation it costs is exactly what P2.X3 forbids in the fast path.
bool session_token(const WeightTable& tables, const ModelGeometry& g, int64_t pos, int32_t pos_base,
SessionState& s, const float* parts, void* stream, bool sync_every_layer,
std::string& err);
// ================================ the captured form ================================
/// ONE GRAPH PER LAYER, because a graph bakes in its POINTERS and every layer has different weights AND
/// different state. That is not a limitation to work around - it is what makes the per-layer graph the right
/// unit: the host decides when each layer runs, which is what lets the CPU expert pool start while the GPU is
/// still working on the next one.
///
/// WHY THIS IS WORTH DOING, MEASURED: **the engine is launch-bound, not bandwidth-bound.** Every kernel in
/// `gdn_layer` measured ~0.028 ms - including `gdn_gate`, which processes FORTY-EIGHT elements, and a kernel
/// cannot execute for 28 us. That is what a null-stream launch costs on this driver, and `gdn_layer` makes ~19
/// of them while a block makes ~45. Replaying the 43-node block graph measured **1.585 ms against 2.393 ms**
/// for the same work as direct launches - a 1.51x that is pure launch overhead.
struct SessionGraphs {
cudaGraphExec_t* execs = nullptr; ///< one per layer, in layer order
int64_t n = 0;
/// **THE SECOND GRAPH PER LAYER: `moe_finish` + `gr_write`, launched AFTER the host has run the pool.**
///
/// Without it the loop has to hand layer `l`'s expert vectors to layer `l+1`, which multiplies them by
/// layer `l+1`'s router weights - see `layer.hpp`'s note on `moe_route`. With it, `post[l]` is launched
/// after the pool returns, so layer `l`'s experts are combined with layer `l`'s weights.
///
/// The cost is the overlap, and it is worth stating: the window is now whatever GPU work follows the ring
/// inside `pre[l]` - the shared expert, and nothing else, because everything after the combine depends on
/// `parts`. **A per-layer CPU pool cannot be hidden behind a strictly serial residual chain**, which is
/// why the CPU term is answered by Phase 3's VRAM cache and not by this pipeline.
cudaGraphExec_t* posts = nullptr;
bool captured = false;
/// THE DEVICE ADDRESS THE GRAPHS WERE CAPTURED WITH. The host loop copies each layer's expert outputs
/// here before launching the next layer, and a graph bakes the POINTER, so it has to be this one.
float* parts_dev = nullptr;
/// **DIAGNOSTIC, AND IT IS THE ONE THAT SAYS WHETHER THE DOORBELL IS DECORATIVE.**
///
/// Counts the layers where the ring was observed while the graph was STILL RUNNING - i.e. the layers where
/// the pool actually had a window to work in. A nonzero `h_seq` read at the same moment
/// `cudaEventQuery` returns `cudaSuccess` means the layer was already over, and the pool then runs with the
/// GPU IDLE, which makes the whole pipeline equivalent to a serialised one.
///
/// That distinction is invisible in the loop's RETURN VALUE - both paths produce a correct token - and it
/// is invisible in a per-layer timing, because the layer still takes the same time either way. It shows up
/// only as "the overlap hides nothing", which is why it is counted here.
int64_t rings_mid_graph = 0;
/// **THE WINDOW, MEASURED RATHER THAN ASSUMED.** Milliseconds from `cudaGraphLaunch` to the moment the
/// ring was observed, summed over the loop. The pool's window is `layer period - this`, and that is the
/// entire budget the overlap has: if the CPU's per-layer work exceeds it, the GPU goes idle for the
/// difference on every layer and a "pipelined" loop measures the same as a serialised one.
///
/// It is NOT the same as the doorbell's position inside the graph. It includes the launch itself, the
/// driver's submission, and however long the host took to notice - all of which are real and none of which
/// the pool gets to use.
double ms_to_ring = 0;
// ================================ R0.9: THE LAYER AS THREE GRAPHS ================================
//
// **THE CAPTURED PER-STAGE TABLE THE PROTOCOL HAS DEMANDED SINCE R0 AND NOBODY HAS HAD.** Round 309's
// came from `--no-capture`, and round 312 showed its `moe_finish` figure was a host-gap artefact four
// times the captured cost of the same work - so the plan has been choosing between R3.3, R3.4 and R3.5
// without knowing what any of them is worth.
//
// An event recorded BETWEEN two `cudaGraphLaunch` calls is not inside a capture and is therefore valid,
// so the layer is captured at its own documented seam (`block_layer_pre`'s `half`) and timed from outside.
/// **R0.11: FIVE PREFIX GRAPHS, ADDED BESIDE `preA`/`preB` RATHER THAN REPLACING THEM.**
/// `preP[k]` (k = 1..5) runs stages `0..k-1`. Consecutive differences are the per-stage times, which is
/// what prices the two `gr_read`s and the `gr_write` (R3.3) against GDN's interior (R3.4). A prefix and not
/// a range, because the stages are a chain: stage 1 reads what stage 0 wrote and stage 2 mutates the
/// residual in place.
///
/// **THE FIFTH PREFIX USED TO BE `execs[l]`, AND THAT WAS THE BUG ROUND 324 COULD NOT FIND.** Prefix 5
/// has ~7 more nodes than prefix 4, so the two differ in launch setup as well as in execution - and the
/// difference attributed ALL of it to stage 4, the router. Capturing prefix 5 the same way as the others
/// makes the two comparable and the difference purely stage 4's.
cudaGraphExec_t* preP[5] = {nullptr, nullptr, nullptr, nullptr, nullptr};
cudaGraphExec_t* preA = nullptr; ///< `half == 1`: gr_read -> attention -> gr_write (the MIXER)
cudaGraphExec_t* preB = nullptr; ///< `half == 2`: gr_read -> moe_route (FFN front + ROUTER)
bool split_captured = false;
/// How many times `session_loop` has run. `rings_mid_graph` and `ms_to_ring` are CUMULATIVE and the loop
/// is called once per token, so anything reporting a ratio needs this - and prefill runs the loop too, so
/// "per generated token" is the wrong denominator.
int64_t calls_total = 0;
/// **WHERE THE PER-LAYER ROUND TRIP GOES, AND `ms_to_ring` ALONE CANNOT SAY.**
///
/// `--no-pool` measures **38.73 ms/token against a 26.32 ms pure-GPU floor**, so ~12.4 ms/token - a quarter
/// of the token - is spent in the loop with no expert work at all: 0.258 ms per layer. `ms_to_ring` covers
/// only the first half of that; it stops when the ring is seen, and everything the host then does before
/// the next `pre` is launched is invisible to it.
///
/// So the round trip is split at the ring, and **`ms_to_ring` above IS the first half** - do not add a
/// second counter for it:
///
/// * `ms_to_ring` - launch -> ring observed. The GPU is working (or already finished) here.
/// * `ms_host` - ring observed -> `pre[l+1]` launched. **The GPU has nothing queued and is IDLE for
/// all of this unless `pre[l]`'s tail still covers it.** This is
/// `cudaMemcpyAsync(parts)` + `cudaGraphLaunch(post[l])` + `cudaGraphLaunch(pre[l+1])` +
/// `cudaEventRecord`, and if it is the large half then the fix is fewer driver calls
/// (R2.4), not a faster kernel.
///
/// Cumulative, like `ms_to_ring`; divide by `calls_total * n_layers`.
double ms_host = 0;
};
/// Captures all 48 layer graphs. Idempotent: a second call is a no-op.
///
/// `pos`/`pos_base` are baked in ONLY through the fixed-address pinned staging in each `QsaState`, which the
/// replay re-reads every token - so the position is data, not an argument. That is the whole reason rounds
/// 211-214 moved the per-token counts into device buffers.
bool session_capture(const WeightTable& tables, const ModelGeometry& g, SessionState& s, const float* parts,
SessionGraphs& gr, std::string& err, bool split = false);
/// One token by REPLAYING the captured graphs. Identical arithmetic to `session_token`; the only difference is
/// that ~2,000 kernel launches become 48 graph launches.
bool session_replay(const ModelGeometry& g, int64_t pos, int32_t pos_base, SessionState& s, SessionGraphs& gr,
void* stream, std::string& err);
/// One token by replaying BOTH halves of every layer - `pre[l]` and `post[l]` - with no pool running.
///
/// **THIS IS THE TRUE PER-TOKEN GPU FLOOR, AND `session_replay` IS NOT.** `session_replay` launches only
/// `gr.execs[l]`, the `pre` graphs, so it omits the 48 `post` graphs (shared expert, combine, second `gr_write`)
/// entirely - and everything published as "39.8 ms pure GPU" came from it while being described as the whole
/// GPU. See `Memory/ERRORS.md` A4.
///
/// `parts` is left at whatever the buffer holds, because no pool writes it; that is deliberate and makes the
/// measurement pure GPU work. The LM head is NOT included here - the caller times it separately, since it runs
/// on the host-side path rather than inside a captured graph.
bool session_replay_full(const ModelGeometry& g, int64_t pos, int32_t pos_base, SessionState& s,
SessionGraphs& gr, void* stream, std::string& err);
/// **ONE TOKEN THROUGH THE SPLIT GRAPHS, WITH EVERY STAGE TIMED FROM OUTSIDE THE CAPTURE.** Requires
/// `session_capture(..., split = true)`. Replays `preA[l]`, `preB[l]`, `post[l]` in order on one stream -
/// the same order `session_loop` uses - and reports the three totals in milliseconds for the whole token.
///
/// The stages are ordered by the stream, so an event pair around each launch measures that launch's GPU work
/// plus the gap the driver left before it. That gap is real and is why the three totals sum slightly above
/// `session_replay_full`; the comparison is printed by the caller rather than hidden here.
/// **R0.11: THE FIVE STAGES SEPARATELY, BY DIFFERENCING PREFIXES.** `stage_ms` is resized to 5 and holds
/// stages 0..4 in milliseconds for a whole token; `mixer_per_layer` holds prefix 3's time for each layer, which
/// is what R0.10's GDN-versus-QSA split needs. Requires `session_capture(..., split = true)`.
///
/// **THE RESIDUAL IS SAVED AND RESTORED BETWEEN PREFIXES.** Stage 2 is `gr_write`, which mutates `R` in
/// place, so without the restore every prefix after the first reads a residual the earlier ones already
/// advanced - and every number after the first would be measuring a different tensor while looking correct.
bool session_replay_stage_prefixes(const ModelGeometry& g, int64_t pos, int32_t pos_base, SessionState& s,
SessionGraphs& gr, void* stream, std::vector<double>& stage_ms,
std::vector<double>& mixer_per_layer, std::string& err);
/// **R3.5c: ONE PREFIX, ALL 48 LAYERS, BACK TO BACK.** `k` in 1..5 replays `preP[k-1][l]` (or `execs[l]` for
/// k = 5) for every layer with no other graph in between, and reports the whole sweep in milliseconds.
///
/// **WHY THIS EXISTS: TO RULE OUT THE PREFIX TABLE'S OWN BIAS.** The five-stage table differences prefix
/// times measured with FIVE DIFFERENT GRAPHS LAUNCHED PER LAYER. Stage 4's interval therefore contains the
/// launch of `execs[l]` - the full five-stage graph - while stage 3's contains `preP[3][l]`, a four-stage one.
/// A launch cost that scales with graph size would inflate stage 4 specifically, and stage 4 is the router,
/// which is the thing under investigation. This sweep launches ONE graph type per layer, so there is nothing
/// to switch between and nothing to attribute wrongly.
///
/// Five separate sweeps give five totals; their consecutive differences are the stage times on the same
/// arithmetic as before, and **if the two methods disagree the difference is the bias and the sweep is right.**
bool session_replay_stage_sweep(const ModelGeometry& g, int64_t pos, int32_t pos_base, SessionState& s,
SessionGraphs& gr, void* stream, int k, double& ms, std::string& err);
bool session_replay_stages(const ModelGeometry& g, int64_t pos, int32_t pos_base, SessionState& s,
SessionGraphs& gr, void* stream, double& ms_mixer, double& ms_ffn, double& ms_post,
std::string& err);
/// **THE SAME REPLAY, KEEPING EVERY LAYER INSTEAD OF SUMMING THEM.**
///
/// The 48 layers are not 48 of the same thing: 36 are GDN and 12 are QSA, and the mixer - 55.8% of the GPU
/// floor - contains both. A total cannot separate them; a per-layer vector can, and it is already being
/// measured. `out` is resized to `n_layers` and holds each layer's mixer time in milliseconds.
///
/// **THIS IS THE CHEAPEST HALF OF R0.10 AND IT NEEDS NO NEW SPLIT.** Round 309's uncaptured table put GDN at
/// 10.88 ms for 36 layers and QSA at 4.56 for 12, which would make the recurrence the largest single R3
/// target; that table also had `moe_finish` wrong by 4x, so the ratio is worth re-deriving from the captured
/// graph rather than inherited.
bool session_replay_stages_per_layer(const ModelGeometry& g, int64_t pos, int32_t pos_base, SessionState& s,
SessionGraphs& gr, void* stream, std::vector<double>& mixer_per_layer,
double& ms_ffn, double& ms_post, std::string& err);
/// Releases the graph execs. The `SessionGraphs` struct holds no device memory of its own.
void session_graphs_free(SessionGraphs& gr);
// ================================ THE HOST LOOP ================================
/// WHAT THE POOL HAS TO DO, as the loop sees it: given the `x_f` a layer published and the miss list it
/// selected, produce the `k` expert outputs into `out` (k x n_embd, HOST memory). A function pointer rather
/// than a virtual class because the loop is the thing under test and the pool is the thing being faked.
using PoolFn = void (*)(void* user, const float* x_f, const int32_t* ids, const float* weights, int64_t n_embd,
int64_t k, float* out);
/// **R4.2c: THE VRAM-RESIDENT HALF OF THE LAYER'S EXPERTS.** Called by the loop after the misses have been
/// staged into `parts` and before `post[l]` is launched, so its kernels are stream-ordered between them: the
/// misses are already in `parts`, the hits overwrite the rows the CPU zeroed, and `post[l]`'s `moe_combine`
/// reads the sum and cannot tell which engine produced which row.
///
/// **WHY IT IS A CALLBACK RATHER THAN MORE PARAMETERS.** `session_loop` owns the ORDER across layers and has
/// no business knowing what an expert cache is. Everything the hit path needs lives on the caller's `user`,
/// and this has the same shape as `PoolFn` for the same reason.
///
/// A null pointer means "no VRAM tier" - the default, and what every caller had before R4.
// The hit hook's type lives in its own header - see `hit_hook.hpp` for why and for the ordering it wants.
/// `phases/phase-2-correct-engine.md` P2.S5's host loop, which is the point of everything above it:
///
/// launch G[l] -> poll the doorbell -> read the miss list -> dispatch the pool -> wait for it ->
/// write y_miss -> launch G[l+1]
///
/// **NO `cudaStreamSynchronize` IN THE LOOP**, which is P2.X3. Waiting is `cudaEventQuery` on an event recorded
/// right after the launch - a QUERY, and it is not optional: round 195 measured a host spin that only read
/// memory never running the kernel at all, 5,907,703 spins over 500 ms with the datum flipping only when
/// `cudaStreamSynchronize` was called.
///
/// **THE OVERLAP IS THE ONE-LAYER DELAY.** A layer publishes `x_f` at ~71% of its own duration, so the pool has
/// the remaining 29% of that layer to work in; its result is consumed by the NEXT layer, which is why the
/// combine is not waiting on it. `overlap = false` runs the same sequence with a synchronisation after every
/// layer, which is the comparison that says whether the pipeline is worth its complexity - and it is a
/// parameter rather than a second function so that the two cannot drift.
///
/// **`dump_layers`, WHEN NOT NULL, IS THE C1 ORACLE AND IT IS A HOST BUFFER OF `(n_layers + 1) * hc * n_embd`
/// FLOATS.** Slot 0 is the residual the function was ENTERED with - the token's embedding broadcast to every
/// stream, which is the reference's `hc_init` node - and slot `l + 1` is `s.R` after layer `l`, which is the
/// reference's `l_last-<l>`. Each is enqueued as a device-to-host copy, and the final `cudaStreamSynchronize`
/// at the end of the loop is what makes every one of them complete. **NO EXTRA SYNCHRONISATION IS INSERTED**,
/// which is the whole reason this is written as an enqueue rather than a read: `pre[l+1]` does not write `R` and
/// `post[l+1]` is launched after the copy on the same stream, so FIFO order alone makes slot `l + 1` a true
/// snapshot. It costs 40 KB per layer per token, and it is the only way to bisect the engine against the
/// reference layer by layer instead of against one end-to-end perplexity.
/// Per-token scratch that MUST NOT be reallocated per token (P2.T10, review finding H3).
///
/// **`session_loop` IS CALLED ONCE PER TOKEN, SO ANYTHING IT ALLOCATES IS A PER-TOKEN ALLOCATION.** It used to
/// allocate three things on every token of every run:
///
/// * `cudaHostAlloc` for the pool's pinned staging, freed by an RAII destructor at the end of the call - and
/// `cudaFreeHost` IMPLICITLY SYNCHRONISES THE DEVICE, so every token ended with a device sync;
/// * `cudaEventCreate` / `cudaEventDestroy` for the doorbell probe;
/// * a host-affinity change (`physical_cores` + pin + restore).
///
/// A comment in the old code said "allocated once", but once per token is not once.
///
/// Own one of these at session setup and pass it in. **Passing `nullptr` keeps the old behaviour** - the loop
/// allocates and frees locally - so every existing caller and test is unaffected, but the allocation is then
/// visible at the call site instead of hidden inside a hot function.
struct SessionLoopScratch {
float* y_miss = nullptr; ///< pinned host staging for the pool's answer, `parts_bytes` long
size_t parts_bytes = 0;
cudaEvent_t probe = nullptr;
long long pinned_core = -1; ///< the affinity to restore, or -1 if the host was never pinned
bool pinned = false;
/// Allocates the buffers and pins the host thread. Call ONCE, at session setup.
bool init(size_t parts_bytes_in, std::string& err);
/// Frees everything and restores the affinity. Safe to call twice, or on a default-constructed object.
void free();
};
bool session_loop(const ModelGeometry& g, int64_t pos, int32_t pos_base, SessionState& s, SessionGraphs& gr,
PoolFn pool, HitFn hits, void* user, bool overlap, void* stream, std::string& err,
float* dump_layers = nullptr, SessionLoopScratch* scratch = nullptr);
// ================================ plan v0.3 P3: THE WHOLE TOKEN AS ONE GRAPH ================================
//
// Under WDDM every graph launch costs ~0.3-0.4 ms of submission latency, and the per-layer loop above makes 96
// of them per token: measured 48.6 ms/token for 12 ms of GPU kernels and 27 ms of pool. The token graph holds
// all 48 layers. Between `pre[l]` (which rings the doorbell) and the parts copy, a one-thread kernel waits ON
// THE DEVICE until the host writes the ring value into `db.h_flag`; the host only polls rings, runs the pool
// and writes flags, with no driver call in the loop (bench/micro/device_wait.cu: 5.2 us per handoff).
//
// for l: pre[l] (.. router, doorbell ring, shared expert) -> doorbell_wait -> parts <- y_miss (H2D)
// -> post[l] (combine, gr_write)
//
// Same kernels, same order per layer as `session_loop`, so the token is bitwise the same. No VRAM expert tier
// yet (the hit decision is a host step between ring and post); the caller falls back to `session_loop` then.
struct TokenGraph {
cudaGraphExec_t exec = nullptr;
bool captured = false;
int64_t n_layers = 0;
const float* y_src = nullptr; ///< the pinned host staging the graph's H2D copies read (baked in)
size_t parts_bytes = 0;
int64_t calls = 0;
double ms_wait = 0; ///< host waiting for rings (the GPU is running)
double ms_pool = 0; ///< host inside the pool
int64_t flushes = 0; ///< cudaStreamQuery calls made because a ring was slow to appear
};
/// Plan v0.3 P4: the VRAM expert tier inside the token graph. Residency is STATIC during a token (a
/// `n_layers x n_expert` table, slot or -1, on the device; the host pool reads its own copy), so the hit list
/// is built on the device right after the ring and the hits run while the CPU computes the misses.
struct TokenHits {
const int32_t* d_res = nullptr; ///< device [n_layers * n_expert]
int64_t n_expert = 0;
const uint8_t* cache_base = nullptr; ///< slot 0 of the VRAM expert arena
int64_t blob = 0; ///< bytes per slot
int32_t* d_slot = nullptr; ///< device, k entries
int32_t* d_dst = nullptr; ///< device, k entries
int32_t* d_count = nullptr; ///< device, 1
uint8_t* x_q8 = nullptr; ///< the activation's q8_0 blocks, (n_embd/32)*34 B
float* x_scale = nullptr; ///< and its fp32 scales (the CPU pool's contract)
void* scratch = nullptr; ///< moe_hit_grouped_scratch_bytes(k, ...)
float* hit_out = nullptr; ///< device, k * n_embd
bool on() const { return d_res && cache_base && d_slot && d_dst && d_count && x_q8 && x_scale && scratch && hit_out; }
};
bool session_capture_token(const WeightTable& tables, const ModelGeometry& g, SessionState& s, float* parts_dev,
const float* y_miss_host, size_t parts_bytes, TokenGraph& tg, std::string& err,
const TokenHits* hits = nullptr);
bool session_run_token(const ModelGeometry& g, int64_t pos, int32_t pos_base, SessionState& s, TokenGraph& tg,
PoolFn pool, void* user, float* y_miss_host, void* stream, std::string& err);
void token_graph_free(TokenGraph& tg);
} // namespace strata::core
include/strata/core/verify.hpp
+160
View File
@@ -0,0 +1,160 @@
// include/strata/core/verify.hpp - plan v0.3 P6: the speculative VERIFY window.
//
// T tokens at consecutive positions p0 .. p0+T-1 - the last accepted token and T-1 drafts - go through all 48
// layers in ONE captured graph, and the head's argmax is produced for every one of them. Token t's argmax is
// what plain greedy decode would produce after token t, BIT FOR BIT: every kernel here is either the single-token
// kernel applied per token, or a multi-token kernel whose per-token arithmetic is the single-token kernel's
// (multi-column MMVQ in exact mode, the T-token GDN kernels, the per-token hit activation, the multi-token CPU
// expert rows). So a draft is accepted exactly when greedy decode would have produced it.
//
// What the window costs is the dense weights read ONCE for T tokens and the union of the T tokens' missed
// experts on the CPU (measured on decode traces: 1.75x one token's misses for T=2, 2.4x for 3, 3.05x for 4).
//
// STATE. The window appends K/V and indexer keys for all T positions and leaves the GDN state untouched.
// `commit(n_keep)` then makes the first `n_keep` tokens permanent: the GDN conv history and recurrent state are
// advanced by replaying those tokens from inputs the window stored, the indexer's key tail is restored from a
// snapshot and the accepted keys re-appended (a rejected key can land in a slot the current block still needs),
// and the PLE history is set to its snapshot after token n_keep-1. K/V cells past the accepted prefix are simply
// overwritten when those positions are processed again, before any query can read them.
//
// Requires the default native decode configuration (native projections, fused GR, fused GDN, fast attention and
// selection, native indexer) and a profile-filled VRAM expert tier with its residency table on the device.
#pragma once
#include "strata/core/expert_source.hpp"
#include "strata/core/layer.hpp"
#include "strata/core/session.hpp"
#include <cuda_runtime.h>
#include <cstdint>
#include <string>
namespace strata::core {
class NativeHead;
/// The CPU pool for a window: x_f (n_tok, n_embd), ids (n_tok, k) -> out (n_tok * k, n_embd), hit rows zeroed.
using PoolMultiFn = void (*)(void* user, const float* x_f, const int32_t* ids, int64_t n_tok, int64_t k, float* out,
int64_t layer);
struct VerifyHits {
const int32_t* d_res = nullptr; ///< device [n_layers * n_expert] slot or -1
const uint8_t* cache_base = nullptr; ///< slot 0 of the VRAM expert arena
int64_t blob = 0;
};
class Verifier {
public:
Verifier() = default;
~Verifier();
Verifier(const Verifier&) = delete;
Verifier& operator=(const Verifier&) = delete;
/// `max_t` <= kVerifyMaxT. `head` may be null (the canonical head is then run per token).
bool init(const WeightTable& wt, const ModelGeometry& g, SessionState& ss, const VerifyHits& hits,
const NativeHead* head, int max_t, std::string& err);
/// One window: `tokens[0..T)` at positions pos0.., the pool served per layer; `out[t]` = argmax after token t.
/// The PLE rows are gathered here from `ss.ple_prev` and the tokens. Captures the T-token graph on first use.
bool run(int T, const int32_t* tokens, int64_t pos0, PoolMultiFn pool, void* user, int32_t* out, std::string& err);
/// Keep the first `n_keep` (1..T) tokens of the last window; advances `ss.ple_prev` by them.
bool commit(int n_keep, std::string& err);
/// Token t's residual after the last layer, (hc, n_embd) on the device, valid until the next `run`.
const float* final_R(int t) const;
const float* final_R_all() const { return R_; }
/// The GPU plan the pool writes each layer (VRAM hits + the PCIe share of the misses); give it to the
/// dispatch (`ExpertDispatch::plan`) before the first `run`.
GpuPlanSink* plan_sink() { return &sink_; }
/// Plan v0.3 P6: split the window into two token groups and pipeline the CPU experts of one with the GPU work
/// of the other (default on). Set before the first `run`.
void set_split(bool on) { split_ = on; }
/// Plan v0.3 P6: how the PCIe share of the misses reaches the GPU: 0 = DMA into staging (the copy engine works
/// beside the CPU; best when the CPU is compute-bound, the i-quants), 1 = the grouped kernel reads the mapped
/// arena directly, 2 = a copy kernel stages it inside the graph (no API calls on the pool's thread; best when
/// the CPU is RAM-bound, Q2_0). Set before the first `run`.
void set_pcie_mode(int mode) { sink_.pcie_mode = mode; }
double ms_wait = 0, ms_pool = 0, ms_host = 0, ms_commit = 0;
int64_t windows = 0;
private:
bool capture(int T, std::string& err);
bool capture_commit(std::string& err);
bool record_window(int T, cudaStream_t cs, std::string& err);
const WeightTable* wt_ = nullptr;
const ModelGeometry* g_ = nullptr;
SessionState* ss_ = nullptr;
VerifyHits hits_;
const NativeHead* head_ = nullptr;
int max_t_ = 0;
int last_t_ = 0;
int64_t last_pos0_ = 0;
int32_t last_tokens_[8] = {};
int64_t n_vocab_ = 0;
cudaStream_t cs_ = nullptr;
cudaGraphExec_t exec_[9] = {};
cudaGraphExec_t commit_exec_ = nullptr;
// mapped staging (host pointer, device alias)
int32_t* h_tok_ = nullptr; int32_t* m_tok_ = nullptr; // T
int32_t* h_step_ = nullptr; int32_t* m_step_ = nullptr; // T * kStepCount
int32_t* h_pos_ = nullptr; int32_t* m_pos_ = nullptr; // T * n_head
int32_t* h_commit_ = nullptr; int32_t* m_commit_ = nullptr; // [n_keep, n_keep-1, pos_0 .. pos_{T-1}]
float* h_ple_ = nullptr; float* m_ple_ = nullptr; // T * n_embd
int32_t* h_out_ = nullptr; int32_t* m_out_ = nullptr; // T argmax ids
float* h_x_ = nullptr; float* m_x_ = nullptr; // doorbell payload: T * n_embd
int32_t* h_ids_ = nullptr; int32_t* m_ids_ = nullptr; // T * k
float* h_w_ = nullptr; float* m_w_ = nullptr; // T * k
uint32_t* h_seq_ = nullptr; uint32_t* m_seq_ = nullptr;
uint32_t* h_flag_ = nullptr; uint32_t* m_flag_ = nullptr;
uint32_t* h_flagA_ = nullptr; uint32_t* m_flagA_ = nullptr; // the GPU plan is in place
uint32_t* h_flagB_ = nullptr; uint32_t* m_flagB_ = nullptr; // the PCIe share's DMA copies have landed
cudaStream_t copy_ = nullptr; // the copy engine's stream (DMA of missed experts)
struct FlagSet { uint32_t* flag; uint32_t value; };
FlagSet flag_sets_[2 * 64 * 2] = {}; // host-function arguments, one per (layer, group)
static void fetch_dma(void* ctx, const uint8_t* const* src, int n, size_t bytes);
static void raise_flag(uint32_t* flag, uint32_t value);
int32_t* h_plan_ = nullptr; int32_t* m_plan_ = nullptr; // counts | start | dst | tok | ptr (as int32 pairs)
int64_t plan_i32_ = 0; // int32 words in the plan block
GpuPlanSink sink_;
uint32_t cur_layer_ = 0;
static void publish_plan(void* ctx);
void set_plan_slot(int grp);
bool split_ = false; // opt-in (--spec-split): exact but slower, see the overlap study
int groups_[9] = {};
float* h_ymiss_ = nullptr; float* m_ymiss_ = nullptr; // T * k * n_embd
// device
void* arena_ = nullptr;
int32_t *tok_ = nullptr, *step_ = nullptr, *pos_ = nullptr, *commit_ = nullptr;
float *ple_ = nullptr, *emb_ = nullptr, *R_ = nullptr, *mixed_ = nullptr, *bo_ = nullptr;
float *inj_ = nullptr, *inj2_ = nullptr, *lo_ = nullptr, *rs_ = nullptr, *xn_ = nullptr;
uint8_t* xq_ = nullptr; // T columns of q8_1
float *qkv_L_ = nullptr, *h_L_ = nullptr, *gate_L_ = nullptr, *beta_L_ = nullptr; // per GDN layer
float *z_ = nullptr, *y_ = nullptr, *y_dummy_ = nullptr;
float *qfull_ = nullptr, *qcur_ = nullptr, *kcur_ = nullptr, *vcur_ = nullptr, *idx_raw_L_ = nullptr;
float *qidx_ = nullptr, *scores_ = nullptr, *attn_ = nullptr, *attn32_ = nullptr, *attn_scratch_ = nullptr;
float* tail_snap_ = nullptr; // per QSA layer
int32_t* sel_ = nullptr;
float *logits_ = nullptr, *w_ = nullptr, *shared_ = nullptr, *parts_ = nullptr, *hit_out_ = nullptr;
int32_t *ids_ = nullptr, *hit_slot_ = nullptr, *hit_dst_ = nullptr, *hit_count_ = nullptr;
int32_t* plan_ = nullptr; // device copy of the plan block
uint8_t* staging_ = nullptr; // VRAM slots for the PCIe share of the misses
static constexpr int64_t kStagingBlobs = 16;
uint8_t* hit_xq_ = nullptr;
uint8_t* nat_xq_ = nullptr; // plan v0.3 P6: q8_1 activations for a native pack's grouped experts
float* hit_xs_ = nullptr;
void* hit_scratch_ = nullptr;
float *head_mixed_ = nullptr, *head_inj_ = nullptr, *head_logits_ = nullptr;
uint16_t* sh_bf16_ = nullptr;
float *sh_gate_ = nullptr, *sh_up_ = nullptr, *sh_g_ = nullptr;
float* hist_snap_ = nullptr; // T * NG_HIST * NG_HC_DIM
int64_t cap_ = 0, max_blocks_ = 0, attn_scratch_floats_ = 0;
};
} // namespace strata::core
include/strata/core/weights.hpp
+135
View File
@@ -0,0 +1,135 @@
// include/strata/core/weights.hpp - the dense weights, loaded into VRAM in ENGINE form.
//
// The pack is the ARTIFACT's business; this is the ENGINE's. The differences are deliberate, and they run in
// BOTH directions:
//
// * SHRINKING. A tensor whose source type is BF16 is stored in the pack promoted to 32 bits, because that is
// lossless and it keeps the canonical decoder simple, but a running engine re-rounds it to 16 - which costs
// nothing at all, because the value CAME from 16 bits. Per `tools/pack_budget.py` that saving is the
// difference between the dense weights fitting in VRAM beside the expert cache and not fitting at all.
//
// * GROWING. 90 of the 303 quantized tensors carry fp16 SCALES in the pack (all 58 Q2_0 expert tensors, plus
// Q4_0, Q5_0, Q8_0 and IQ4_NL) because that is what the source block holds. Every kernel here takes
// `const float* scales`, and fp16 is a subset of f32, so this loader WIDENS them - exactly, changing no
// value - for 11,407,360 B. See `docs/pack-format.md` §2.4; assuming 4 bytes per scale instead of reading
// `scales_fp16` is what made the plane arithmetic disagree with the pack for a whole round.
//
// pack 5.305 GiB dense + 0.444 GiB embd
// engine 4.067 GiB dense + 0.444 GiB embd + 10.88 MiB widened scales = 4.522 GiB
//
// WHAT THE LOADER DOES NOT DO: it does not know what any tensor MEANS. It reads a flat index
// (`<pack>/index.txt`, written by `tools/pack_index.py`), copies each tensor into one arena at a recorded
// offset - as a list of PLANE SEGMENTS, because a scale plane can change width and the rest cannot - and hands
// back a name -> WeightRef map. Geometry, forms and semantics belong to the caller.
//
// WHY AN INDEX FILE AND NOT A JSON PARSER: a JSON parser in the engine would be a new, unaudited component
// whose failure mode is a wrong byte offset - which decodes to a plausible weight and produces plausible
// logits. Python wrote the manifest and already parses it correctly.
#pragma once
#include <cstdint>
#include <map>
#include <set>
#include <string>
#include <vector>
namespace strata::core {
/// What the loader had to DO for a tensor, decided by `tools/pack_index.py` and written into the index
/// rather than inferred here. The four cases are distinguishable in the manifest only by combining
/// `source_type` with `values_fp16`, and getting that combination wrong yields a plausible tensor.
enum class WeightKind : int {
Verbatim = 0, ///< quantized planes, or F16 already stored as 2 bytes: copy straight through
Bf16InF32 = 1, ///< the pack holds the bf16 value promoted to f32: take the HIGH 16 BITS (exact)
F32 = 2, ///< copy 4 bytes per element
F16InF32 = 3, ///< an f16 value promoted to f32: a real f32->f16 conversion, not a truncation
};
struct WeightRef {
const void* data = nullptr; ///< device pointer inside the arena
uint64_t bytes = 0;
int64_t ne0 = 0; ///< the manifest's own shape, ne0 the CONTIGUOUS axis
int64_t ne1 = 0;
int64_t elements = 0;
WeightKind kind = WeightKind::Verbatim;
/// The canonical-form attributes, valid only for a quantized VERBATIM tensor (`code_bits != 0`).
int code_bits = 0;
int code_bias = 0;
int group_elems = 0;
bool codebook_iq4nl = false;
bool has_offset = false;
/// THE THREE PLANE SIZES, IN ENGINE FORM, in the order they sit inside `data`. Recorded by the loader
/// from what the index TOLD it rather than re-derived by each caller: the one time this layout was
/// re-derived independently it assumed 4-byte scales for tensors whose pack scales are fp16, and the
/// only reason that was not a silent wrong answer is that the re-derivation checked itself and refused.
/// Two places computing one layout is two places to disagree.
uint64_t codes_bytes = 0;
uint64_t scales_bytes = 0; ///< always 4 B per scale in the arena, whatever the pack held
uint64_t offset_bytes = 0; ///< 0 when the form has no offset plane
/// WHERE THIS TENSOR CAME FROM, so a caller can check the loader's work against the pack's own bytes
/// instead of against the loader's opinion of them. Without `src_off` a test can only compare two
/// derived numbers; with it, the widening of an fp16 scale plane can be verified value by value against
/// `<file>.bin` - and the pack files are named, so `file_id` says which one to open.
uint64_t src_off = 0;
uint64_t src_bytes = 0;
int file_id = 0;
bool scales_fp16 = false; ///< the PACK's width, before the loader widened it
/// WHICH ACTIVATION THIS WEIGHT WANTS - 0 for Q8_0, 1 for Q8_K. Decided by ools/pack_index.py from
/// source_type and carried here, because **the S-form CANNOT express it**: Q5_0 and Q5_K are both 8-bit
/// with bias -16, and IQ4_NL and IQ4_XS differ in nothing the S-form holds. docs/activation-contract.md.
int act_kind = 0;
bool wants_q8k() const { return act_kind == 1; }
bool quantized() const { return code_bits != 0; }
// Optional native GGUF projection, owned by NativeDense. All references in
// one table share its scratch and must execute on one ordered session stream.
const void* native_data = nullptr;
void* native_q8_1 = nullptr;
int native_type = -1;
/// Plan v0.3 P1: false when the loader SKIPPED this tensor's canonical bytes because another form serves it
/// (native GGUF projections, the native head). The metadata above stays valid; `data` is null.
bool resident = true;
};
struct LoadReport {
size_t tensors = 0;
uint64_t arena_bytes = 0;
size_t re_rounded = 0; ///< tensors written as 16 bits out of a 32-bit container
uint64_t bytes_saved = 0; ///< what those conversions saved against the pack
double read_ms = 0.0;
double upload_ms = 0.0;
};
class WeightTable {
public:
/// The arena size the index asks for, readable WITHOUT loading anything - so a caller can size its
/// `DeviceArena` before committing to the load, and a plan that cannot fit is refused at startup rather
/// than halfway through a 5 GB upload.
/// With `skip`, the size of the compacted arena that holds every tensor EXCEPT the named ones.
static bool pool_bytes(const std::string& pack_dir, uint64_t& out, std::string& err,
const std::set<std::string>* skip = nullptr);
/// Load every tensor in `<pack_dir>/index.txt` into `arena_base`.
///
/// `arena_bytes` MUST be at least `pool_bytes`; the loader checks rather than trusting the caller,
/// because the failure mode otherwise is a device write past the end of the arena.
/// With `skip`, the named tensors are not read: their rows keep their metadata with `data == nullptr` and
/// `resident == false`, and the other tensors are packed into the compacted arena `pool_bytes` sized.
bool load(const std::string& pack_dir, void* arena_base, uint64_t arena_bytes, std::string& err,
const std::set<std::string>* skip = nullptr);
const WeightRef* find(const std::string& name) const;
const std::map<std::string, WeightRef>& all() const { return table_; }
const LoadReport& report() const { return report_; }
private:
friend class NativeDense;
std::map<std::string, WeightRef> table_;
LoadReport report_;
};
} // namespace strata::core
include/strata/kernels/bf16_bits.hpp
+46
View File
@@ -0,0 +1,46 @@
// include/strata/kernels/bf16_bits.hpp - the bf16 conversions, written out, as bits.
//
// bf16 is the TOP HALF of an f32, so `f32_from_bf16` is a shift and needs no rounding at all. The forward
// direction does: ggml's `ggml_compute_fp32_to_bf16` adds `((i >> 16) & 1) + 0x7FFF` and truncates 16 bits.
//
// **THE `& 1` IS ON BIT 16**, the lowest KEPT bit, which is what makes an exact tie round to EVEN. Using bit
// 15 there is a plausible mistake: it differs only on exact ties, i.e. almost never on real data, which is
// exactly why it would survive any amount of end-to-end testing. `ref/quant.py` documents the same rule for
// the same reason.
//
// WHY THIS IS SEPARATE FROM `f16_bits.hpp`: fp16 and bf16 are not two spellings of one idea. fp16 has 5
// exponent bits and 10 mantissa bits; bf16 has 8 and 7. A conversion routine written for one and reused for
// the other is wrong in a way that produces plausible numbers - bf16's 2^-9 relative precision looks like a
// small fp16 error and is not small at all next to fp16's 2^-11.
//
// `gr.cu` carried a private copy of `f32_to_bf16_bits` until this header existed. It now includes this.
#pragma once
#include <cstdint>
#include <cstring>
#if defined(__CUDACC__)
#define STRATA_BF16_HD __host__ __device__
#else
#define STRATA_BF16_HD
#endif
namespace strata::kernels {
/// Round-to-nearest-even f32 -> bf16, returned as raw bits in the LOW half of the uint16.
STRATA_BF16_HD inline uint16_t bf16_from_f32(float f) {
uint32_t i;
std::memcpy(&i, &f, 4);
i = (i + ((i >> 16) & 1u) + 0x7FFFu) & 0xFFFF0000u;
return (uint16_t) (i >> 16);
}
/// bf16 bits -> f32. Exact, and no rounding: the low 16 bits of an f32 are simply zero.
STRATA_BF16_HD inline float f32_from_bf16(uint16_t h) {
const uint32_t i = (uint32_t) h << 16;
float f;
std::memcpy(&f, &i, 4);
return f;
}
} // namespace strata::kernels
include/strata/kernels/bf16_gemv.hpp
+48
View File
@@ -0,0 +1,48 @@
// include/strata/kernels/bf16_gemv.hpp - the GEMV for a BF16 weight, P2.S5.
//
// WHY THIS EXISTS WHEN `s_gemv` ALREADY DOES. `s_gemv` decodes the S-family canonical forms (S2/S4/S8) from
// `codes` + `scales`, and a BF16 weight is not one of them - it has no codes and no scales, it is a plain
// 16-bit matrix. Six tensors per layer are in that state:
//
// ssm_alpha.weight, ssm_beta.weight GDN, [n_embd, 48]
// indexer.q_proj.weight, indexer.k_proj.weight QSA, [n_embd, 512] and [n_embd, 128]
// ple_value.weight layer 1, [n_embd, n_embd]
// ffn_gate_inp.weight handled inside `router_top10`
//
// The original entry points take BF16 activation bits for the historical CPU-derived contract. The explicit
// FP32 entry point below follows the pinned CUDA single-token MMVF contract. See docs/activation-contract.md
// for why activation precision depends on backend and batch geometry rather than weight type alone.
#pragma once
#include <cstdint>
namespace strata::kernels {
/// `y[o] = sum_i f32(x[i]) * f32(w[o*n_in + i])`, with both sides bf16-valued so every product is exact in
/// f32 and only the summation order differs.
///
/// THE WEIGHT LAYOUT IS THE MANIFEST'S: shape [n_in, n_out] with ne0 = n_in CONTIGUOUS, so row `o` is a
/// contiguous run of `n_in` bf16 values. That is the same convention `s_gemv` uses, and it is what
/// `strata::core::WeightRef::ne0/ne1` already report.
void bf16_gemv(const uint16_t* x, const uint16_t* w, float* y, int64_t n_in, int64_t n_out, void* stream);
/// The row-split variant, for the same reason `s_gemv_split` exists: with one thread per output row the
/// parallelism IS the output width, and `ssm_alpha` has 48 of them over 48 SMs. One warp per row, lanes
/// striding the reduction axis - which is also what makes the access pattern coalesced in this layout.
///
/// NOT bit-identical to `bf16_gemv`: the partial sums are added in a different order. Both are kept and both
/// are tested, because the naive one is the reference the split one is checked against.
void bf16_gemv_split(const uint16_t* x, const uint16_t* w, float* y, int64_t n_in, int64_t n_out,
int threads_per_row, void* stream);
/// Single-token, contiguous 2D BF16-weight/F32-activation MMVF from the pinned llama.cpp CUDA path.
/// Preserves pair accumulation, adaptive block size and warp-XOR reduction order. This is a separate
/// opt-in entry point; the legacy BF16-activation kernels above retain their behavior.
/// `w[o*n_in+i]` is BF16 bits; x/y are F32. n_in must be positive and even, n_out positive, both <= INT_MAX.
/// x must be 8-byte aligned, w 4-byte aligned, and y 4-byte aligned. Throws on invalid geometry/pointers.
/// The device buffers must not overlap; allocation bounds remain the caller's responsibility.
/// No allocations or synchronization, including when stream is null (the CUDA default stream).
void bf16_gemv_fp32_mmvf(const float* x, const uint16_t* w, float* y,
int64_t n_in, int64_t n_out, void* stream);
} // namespace strata::kernels
include/strata/kernels/cpu/expert.hpp
+179
View File
@@ -0,0 +1,179 @@
// include/strata/kernels/cpu/expert.hpp - P2.S3: the CPU expert path.
//
// This is the Q2_0 expert kernel promoted out of `bench/micro/cpu_s2.cpp`, which is the instrument that
// measured `c` (Memory/LEDGER.md L9/L10: 15.03 ms for one token's full routed set on 6 cores, DRAM-bound at
// 44.14 GB/s) and whose P0.T2 parity passes at 1.461e-06.
//
// IT IS THE SAME CODE, NOT A COPY. Keeping a second implementation next to a validated one is how a project
// ends up with two answers and no way to say which is right, so `cpu_s2.cpp` is expected to become a thin
// benchmark over these functions. (Until it is, the two exist side by side and this header says so.)
//
// WHAT MAKES IT FAST, and each of the three was paid for once:
//
// * `vpdpbusd` (AVX512-VNNI) does four int8 MACs per instruction. The scalar loop it replaces ran at
// 9.89 GB/s on one core; VNNI is the reason the 6-core figure is 42.55 GB/s.
// * `vpmultishiftqb` (AVX512-VBMI) unpacks the 2-bit codes eight at a time. Its ARGUMENT ORDER is
// (control, data) and was determined empirically - the Intel guide states the opposite, and following the
// guide produced 43 of 64 wrong codes.
// * The row accumulator stays a FLOAT VECTOR and is reduced ONCE at the end of the row. Reducing per
// 64-weight block cost more than the dot products themselves: it made the kernel compute-bound at
// 21 GB/s instead of DRAM-bound at 32.
//
// THE CONTRACT THAT MATTERS, because collapsing it is invisible and cannot match ggml's logits: one Q2_0
// WEIGHT block is QK=64 weights with one fp16 scale, while one ACTIVATION chunk is QKA=32 elements with
// its OWN scale. So a weight block spans TWO activation chunks with two different scales. Using one scale
// per weight block is the natural-looking mistake.
#pragma once
#include <cstdint>
#include <cstddef>
namespace strata::kernels::cpu {
// ---- geometry, all of it fixed by the artifact (tools/verify_q2_0_geometry.py: 202/202 tensors at QK=64) ----
inline constexpr int H = 2560; // n_embd
inline constexpr int FF = 640; // expert intermediate width
inline constexpr int NE = 512; // routed experts per layer
inline constexpr int QK = 64; // QK2_0: weights per fp16 scale
inline constexpr int QKA = 32; // QK8_1/QK8_0: ACTIVATION elements per scale
inline constexpr int BB = 18; // bytes per 64-weight Q2_0 block: 2 fp16-scale + 16 codes
inline constexpr int MAXC = H / QKA; // 80 activation chunks across the widest reduction
inline constexpr int ROW_GU = H * 2 / 8; // 640 B of codes per gate/up row
inline constexpr int ROW_D = FF * 2 / 8; // 160 B of codes per down row
inline constexpr int SC_GU = H / QK; // 40 weight blocks per gate/up row
inline constexpr int SC_D = FF / QK; // 10 weight blocks per down row
inline constexpr size_t BLOB = 3ull * (H * FF * BB / QK); // 1,382,400
// ---- the blob's internal layout (one expert) ----
inline constexpr size_t O_GU_CODES = 0;
inline constexpr size_t O_D_CODES = 2ull * FF * ROW_GU;
inline constexpr size_t O_GU_SCALES = O_D_CODES + 1ull * H * ROW_D;
inline constexpr size_t O_D_SCALES = O_GU_SCALES + 2ull * FF * SC_GU * 2;
/// One activation in Strata's planar storage: `QKA` elements per chunk.
/// The default legacy quantizer uses FP32 scales and half-away rounding. It is
/// not the pinned ggml CPU Q8_0 contract; expert_set_oracle_q8_0 selects that
/// experimental contract (FP16-rounded scales and x86 nearest-even codes).
///
/// The historical function name `act_quant_q8_1` does not describe a native
/// ggml block layout. One 64-weight block always spans two 32-element chunks.
///
/// `hx[k] = scale[k] * sum[k]` is the weight-independent correction the Q2_0 identity needs
/// (`sum (c-1) d_w xhat d_x = d_w (d_x*sum(c*xhat) - d_x*sum(xhat))`), precomputed once per layer rather than
/// once per row: 640 rows would otherwise recompute it 640 times.
struct ActQ {
alignas(64) int8_t q[H];
float scale[MAXC];
int32_t sum[MAXC];
float hx[MAXC];
int nchunks;
};
/// Per-worker scratch. Owned by the caller and passed in, so the token path performs NO allocations
/// (P2.T10) - the kernel itself allocates nothing.
struct ExpertScratch {
ActQ a1;
ActQ a2;
alignas(64) float ff[FF];
};
/// Runtime CPU feature check. The kernel uses AVX512-VNNI + AVX512-VBMI + AVX512VL, and code compiled with
/// `/arch:AVX512` can emit AVX-512 anywhere in its translation unit, so a machine without them must be
/// REFUSED rather than silently run. `s2_expert_scalar` is the fallback and exists for tests.
struct CpuFeatures {
bool avx512f = false;
bool avx512bw = false;
bool avx512vl = false;
bool avx512_vnni = false;
bool avx512_vbmi = false;
bool usable() const { return avx512f && avx512bw && avx512vl && avx512_vnni && avx512_vbmi; }
/// A one-line description of what is missing, or "ok".
const char* reason() const;
};
CpuFeatures cpu_features();
/// Exits with a clear message if the CPU cannot run `s2_expert_vnni`. Called once at startup, so a user on an
/// older CPU learns why at second zero instead of seeing an illegal instruction at token 4000.
void cpu_require_expert_support();
/// Select the pinned x86 Q8_0 activation contract at BOTH expert quantization
/// boundaries and the pinned generic CPU Q2_0 dot reduction in all projections.
/// Default false preserves the legacy quantizer and eight-lane FP32 reduction.
/// Configure before starting worker threads and do not change during a request.
/// GPU expert-cache parity is not established for this mode; the driver must
/// reject combining it with that cache. No allocation occurs in either mode.
void expert_set_oracle_q8_0(bool enabled);
/// `x` -> `a` in the configured contract. The function name is historical;
/// `n` must be a multiple of QKA and at most H. No allocation occurs.
void act_quant_q8_1(const float* x, int n, ActQ& a);
/// One expert: `out(H) = down( swiglu( gate(x), up(x) ) )`, all three projections Q2_0,
/// with the configured activation and dot contracts at both quantization boundaries.
///
/// `blob` is one 1,382,400-byte expert in the pack's layout; `x` is the H-wide input; `out` is H-wide.
void s2_expert_vnni(const uint8_t* blob, const float* x, float* out, ExpertScratch& ws);
/// The same, with the activation ALREADY quantized in the configured contract.
/// Workers may share this read-only activation; each must own its scratch/output.
///
/// This split exists for the pool. A layer runs TEN experts against ONE activation, so quantizing inside
/// `s2_expert_vnni` would do the same 2560-element conversion ten times per layer per token - 480 redundant
/// conversions per token. P2.S3's work item is `(blob, x-hat quantized, block sums, weight, out)`, and this
/// is the entry point that matches it.
void s2_expert_vnni_q(const uint8_t* blob, const ActQ& a1, float* out, ExpertScratch& ws);
/// Plan v0.3 P6: ONE expert applied to up to MAXT tokens that were all routed to it (a speculative verify
/// window routes ~1.5 tokens to each distinct expert). The 1.38 MB blob is read and its 2-bit codes unpacked
/// ONCE per block for all tokens; only the VNNI dot and the accumulation run per token. Every token's output
/// is bitwise equal to `s2_expert_vnni_q` on that token (same operation order per token). Legacy activation
/// contract only; with `expert_set_oracle_q8_0(true)` it falls back to one `s2_expert_vnni_q` per token.
inline constexpr int MAXT = 8;
/// Plan v0.3 P4: the expert in row ranges, so several threads can share one expert. `s2_expert_gu_rows` writes
/// ff[r] = silu(gate_r . x) * (up_r . x) for r in [r0, r1) (of FF); `s2_expert_down_rows` writes out[r] for r in
/// [r0, r1) (of H) from the intermediate's quantized image. Together with `act_quant_q8_1(ff, FF, a2)` in
/// between they are exactly `s2_expert_vnni_q`, row for row (same `row_dot`), so the result is bitwise the same.
bool expert_oracle_q8_0_enabled();
void s2_expert_gu_rows(const uint8_t* blob, const ActQ& a1, float* ff, int r0, int r1);
void s2_expert_down_rows(const uint8_t* blob, const ActQ& a2, float* out, int r0, int r1);
/// Plan v0.3 P6: the row-range kernels for up to MAXT tokens routed to the same expert (legacy contract only).
/// Each token's rows are bitwise the single-token `s2_expert_gu_rows` / `s2_expert_down_rows`.
void s2_expert_gu_rows_multi(const uint8_t* blob, const ActQ* const* a1, int n_tokens, float* const* ff, int r0,
int r1);
void s2_expert_down_rows_multi(const uint8_t* blob, const ActQ* const* a2, int n_tokens, float* const* out, int r0,
int r1);
struct ExpertScratchMulti {
ActQ a2[MAXT];
alignas(64) float ff[MAXT][FF];
ExpertScratch single; // the oracle-contract fallback
};
void s2_expert_vnni_multi(const uint8_t* blob, const ActQ* const* a1, int n_tokens, float* const* out,
ExpertScratchMulti& ws);
/// The scalar transcription of ggml's formula - NOT a second opinion, but the ORACLE the VNNI path is checked
/// against, and the only path available on a CPU without AVX-512.
///
/// `quant_acts` makes the oracle consume the SAME INT8 activation values the VNNI path uses, **at BOTH
/// stages**. With it OFF the oracle is exact-FP32-activation and the gap between them IS the activation
/// contract; with it ON the two differ only by FP32 evaluation order, so a small gap there proves the kernel
/// itself correct. That separation is the whole point of P0.T2 and it is why this parameter exists rather
/// than two functions.
///
/// CORRECTION TO THE ORIGINAL: `bench/micro/cpu_s2.cpp`'s comment says "at both stages" but its code
/// quantizes only the intermediate, leaving the gate/up projections on raw f32 `x`. That is invisible when
/// the caller passes an already-quantized input - which the engine does - and shows up as a ~1.2% gap when it
/// does not, which reads as a kernel bug and is not one. This version quantizes both.
/// Plan v0.3 P6: rows [r0, r1) of a Q2_0 matrix in the GGUF block layout (18 bytes per 64 weights, `row_bytes`
/// per row, `nblocks` blocks each) against `nt` activations: out[t][r].
void q2_0_gguf_rows_multi(const uint8_t* w, size_t row_bytes, int nblocks, const ActQ* const* a, int nt,
float* const* out, int r0, int r1);
/// The same two for CPUs without AVX-512 (src/kernels/cpu/q2_avx2.cpp, compiled for AVX2 only).
void q2_0_gguf_rows_multi_avx2(const uint8_t* w, size_t row_bytes, int nblocks, const ActQ* const* a, int nt,
float* const* out, int r0, int r1);
void act_quant_q8_1_avx2(const float* x, int n, ActQ& a);
void s2_expert_scalar(const uint8_t* blob, const float* x, float* out, bool quant_acts);
} // namespace strata::kernels::cpu
include/strata/kernels/cpu/expert_layout.hpp
+51
View File
@@ -0,0 +1,51 @@
// include/strata/kernels/cpu/expert_layout.hpp - plan v0.3 P6: where each routed expert lives in experts.bin.
//
// A Q2_0 pack (tools/strata_pack.py) has one blob size for every layer, `BLOB`, in the Strata expert form. A
// native pack (tools/iq_pack.py, the IQ2_XS / IQ3_XXS files) keeps each expert's raw GGUF slices, so the blob
// size and the formats change from layer to layer; `native_experts.txt` says how. Everything that touches an
// expert blob - the arena, the VRAM tier, the prompt path, the CPU pool, the GPU window - asks this table.
#pragma once
#include "strata/kernels/cpu/expert.hpp"
#include "strata/kernels/cpu/native_expert.hpp"
#include <cstdint>
#include <string>
#include <vector>
namespace strata::kernels::cpu {
struct ExpertLayout {
bool native = false;
int64_t n_layers = 0, n_expert = NE;
std::vector<NativeFmt> fmt; ///< per layer (native packs)
std::vector<uint64_t> offset, bytes; ///< per layer: where its 512 blobs start, bytes per blob
/// Plan v0.3 P6: per layer, the absolute offsets of the gate / up / down tensors in the model's shard 1, so
/// the arena can be filled from the GGUF itself when the pack has no experts.bin (3 x n_layers, 0 = unknown).
std::vector<uint64_t> gguf_off;
uint64_t max_blob = BLOB;
uint64_t total = 0; ///< experts.bin size
uint64_t blob_bytes(int64_t layer) const { return native ? bytes[(size_t) layer] : (uint64_t) BLOB; }
uint64_t layer_offset(int64_t layer) const {
return native ? offset[(size_t) layer] : (uint64_t) layer * (uint64_t) n_expert * (uint64_t) BLOB;
}
uint64_t blob_offset(int64_t layer, int64_t expert) const {
return layer_offset(layer) + (uint64_t) expert * blob_bytes(layer);
}
};
/// Plan v0.3 P6: whether this CPU (and its OS) runs the AVX-512 kernels (F, BW, VL, VNNI, VBMI). Probed in a
/// file compiled without AVX-512, so asking is safe everywhere; STRATA_FORCE_AVX2=1 answers no (for tests).
bool cpu_avx512_ok();
/// Q2_0 GGUF rows / activation quantizer on the kernels this CPU has.
void q2_rows_any(const uint8_t* w, size_t row_bytes, int nblocks, const ActQ* const* a, int nt, float* const* out,
int r0, int r1);
void act_quant_any(const float* x, int n, ActQ& a);
/// The process-wide layout (canonical Q2_0 until `expert_layout_load` finds a native pack).
const ExpertLayout& expert_layout();
/// Reads `<pack_dir>/native_experts.txt` when it exists (a native pack), else sets the canonical layout.
bool expert_layout_load(const std::string& pack_dir, int64_t n_layers, int64_t n_expert, std::string& err);
} // namespace strata::kernels::cpu
include/strata/kernels/cpu/iq_avx512.hpp
+18
View File
@@ -0,0 +1,18 @@
// include/strata/kernels/cpu/iq_avx512.hpp - plan v0.3 P6: AVX-512 multi-token dot products for the i-quant
// expert formats (IQ2_XXS, IQ2_XS, IQ2_S, IQ3_XXS, IQ3_S) against Q8_K activations (ggml's block_q8_K).
#pragma once
#include <cstddef>
#include <cstdint>
namespace strata::kernels::cpu {
bool iq512_supported(int ggml_type) noexcept;
/// ff[t][r] = silu(gate_r . a[t]) * (up_r . a[t]), rows [r0, r1); gate rows at blob, up rows at blob + up_off.
void iq512_gu_rows(int ggml_type, const uint8_t* blob, size_t gu_row, size_t up_off, int n, const void* const* act,
int nt, float* const* ff, int r0, int r1);
/// out[t][r] = w_r . a[t], rows [r0, r1).
void iq512_rows(int ggml_type, const uint8_t* w, size_t row_bytes, int n, const void* const* act, int nt,
float* const* out, int r0, int r1);
} // namespace strata::kernels::cpu
include/strata/kernels/cpu/native_expert.hpp
+49
View File
@@ -0,0 +1,49 @@
// include/strata/kernels/cpu/native_expert.hpp - plan v0.3 P6: one routed expert in its GGUF form on the CPU.
//
// The IQ2_XS / IQ3_XXS model files keep their experts in i-quant formats (IQ1_M, IQ2_XXS, IQ2_XS, IQ2_S, IQ3_XXS,
// IQ3_S gate/up; Q2_0 or IQ4_NL down) whose values cannot be re-expressed in the Q2_0 pack form. A native expert
// blob is the three GGUF slices back to back, [gate rows | up rows | down rows], and the arithmetic is ggml-cpu's
// own (`ggml_get_type_traits_cpu`): the activation is quantized with the weight type's `vec_dot_type` and each row
// is one `vec_dot`, exactly what llama.cpp's CPU backend computes for the same tensor.
#pragma once
#include <cstddef>
#include <cstdint>
#include <string>
namespace strata::kernels::cpu {
/// Bytes of the largest quantized activation any native layer uses (2560 values as Q8_K: 10 x 292).
inline constexpr size_t kNativeActBytes = 4096;
/// Bytes of the largest quantized down activation (640 values as Q8_0: 20 x 34, or Q8_K 3 x 292).
inline constexpr size_t kNativeHBytes = 1024;
/// One layer's native expert geometry.
struct NativeFmt {
int gu_type = -1, d_type = -1; ///< ggml types of gate/up and of down
int gu_act = -1, d_act = -1; ///< their vec_dot_type (the activation formats)
int64_t n_embd = 0, n_ff = 0;
size_t gu_row = 0, d_row = 0; ///< bytes per weight row
size_t up_off = 0, down_off = 0; ///< inside the blob
size_t bytes = 0; ///< the whole blob
size_t act_bytes = 0, h_bytes = 0; ///< quantized activation sizes (n_embd of gu_act, n_ff of d_act)
};
/// Whether this build has the ggml-cpu path.
bool native_experts_available() noexcept;
/// Fills `f` for a layer; false (with a reason) when ggml-cpu has no dot product for a type.
bool native_fmt(int gu_type, int d_type, int64_t n_embd, int64_t n_ff, NativeFmt& f, std::string& err);
/// x (n_embd floats) -> the gate/up activation (act_bytes).
void native_quant_act(const NativeFmt& f, const float* x, void* dst);
/// h (n_ff floats) -> the down activation (h_bytes).
void native_quant_h(const NativeFmt& f, const float* h, void* dst);
/// ff[t][r] = silu(gate_r . a[t]) * (up_r . a[t]) for rows r in [r0, r1), `nt` tokens.
void native_gu_rows(const NativeFmt& f, const uint8_t* blob, const void* const* act, int nt, float* const* ff,
int r0, int r1);
/// out[t][r] = down_r . hq[t] for rows r in [r0, r1).
void native_down_rows(const NativeFmt& f, const uint8_t* blob, const void* const* hq, int nt, float* const* out,
int r0, int r1);
} // namespace strata::kernels::cpu
include/strata/kernels/cpu/pool.hpp
+195
View File
@@ -0,0 +1,195 @@
// include/strata/kernels/cpu/pool.hpp - P2.S3: the CPU expert pool.
//
// A layer runs TEN experts against ONE activation, and the expert kernel is DRAM-bound (Memory/LEDGER.md L9:
// 42.55 GB/s on 6 cores, tracking core count almost exactly). So the pool's job is not to be clever - it is
// to keep every physical core reading expert bytes for the whole layer, and to be cheap enough that ten
// dispatches per layer cost less than one expert.
//
// WHY A FLAT BATCH AND NOT A RING. P2.S3 describes a "lock-free SPMC ring", which is what you need when
// jobs arrive while others are still being computed. Here they cannot: the host must SUM all ten outputs
// before the next layer starts, so a layer is a barrier by construction and the queue never holds more than
// one batch. A ring would add a wrap-around to get wrong and buy nothing. What is kept from the phase is
// the part that matters: `head`/`done` are single fetch_add counters, one claim per worker, no lock.
//
// THE COMPLETION PROTOCOL, because this is where a pool usually goes wrong. `run()` waits for `done == n`
// AND for every worker to PARK. Waiting only for `done` is not enough: a worker can still be inside the
// drain loop after its last `done` increment, and the host resetting `head` underneath it would let that
// worker claim a job from the NEXT batch before the next batch has been published. The
// `done`-then-`parked` pair makes the handover unambiguous, and the second wait costs a few hundred cycles
// against a layer that takes milliseconds.
#pragma once
#include "strata/kernels/cpu/expert.hpp"
#include "strata/kernels/cpu/native_expert.hpp"
#include <atomic>
#include <cstdint>
#include <thread>
#include <vector>
namespace strata::kernels::cpu {
/// One expert evaluation. `act` is SHARED and read-only across the whole batch - that sharing is the point
/// of `s2_expert_vnni_q` and it is what saves 480 redundant activation conversions per token.
struct ExpertJob {
const uint8_t* blob = nullptr; ///< one 1,382,400-byte expert
const ActQ* act = nullptr; ///< the layer's quantized activation, shared
float* out = nullptr; ///< H floats, written by exactly one worker
float weight = 1.0f; ///< the router weight; applied by the HOST, not here
int slot = -1; ///< the job's index, for diagnostics
};
/// Plan v0.3 P6: one expert for the `nt` tokens of a verify window that were routed to it. Every token's output
/// is bitwise the single-token job's.
struct ExpertJobMulti {
const uint8_t* blob = nullptr;
int nt = 0;
const ActQ* act[MAXT] = {};
float* out[MAXT] = {};
/// Plan v0.3 P6: a native pack's activations (the layer's `vec_dot_type`), one per token.
const void* nact[MAXT] = {};
};
/// One logical processor per PHYSICAL core, so a worker is never scheduled onto an SMT sibling of another
/// worker. On the 6-core/12-thread machine this project measures on, `hardware_concurrency()/2` workers on
/// logical processors 0..5 would put every worker on a sibling pair and halve the useful bandwidth - which is
/// exactly the kind of error that shows up as "the CPU path is slower than the model says" with no clue why.
///
/// `skip_first` drops the first core, which P2.S3 reserves for the host loop.
std::vector<int> physical_cores(bool skip_first);
/// **THE RESERVATION IS A FICTION UNLESS THE HOST IS ACTUALLY PUT THERE.**
///
/// `physical_cores(true)` keeps the workers off the first physical core so that the host loop can spin on
/// `cudaEventQuery` without stealing a worker's cycles. Nothing in the pool can enforce the other half of
/// that, so this is it: the host loop calls this on entry and restores on exit.
///
/// MEASURED, and this is why it exists: the pool runs at **36.32 GB/s on 5 workers with nothing else running**
/// - exactly 5/6 of L9's 44.14 on 6 - and at **26.9 GB/s inside the host loop**, where the unpinned spinning
/// host is free to land on a worker's core or its SMT sibling. That 1.35x is not the kernel.
///
/// Returns the PREVIOUS affinity mask, or -1 if the platform refused; pass it to `restore_thread_affinity`.
long long pin_current_thread(int core);
void restore_thread_affinity(long long previous);
class ExpertPool {
public:
/// `n_workers <= 0` means "every physical core except the first". Workers are pinned to physical cores
/// (minus core 0 by default) and each owns one `ExpertScratch`, so nothing in the token path allocates.
///
/// **`host_works` PUTS THE HOST THREAD INTO THE DRAIN (R2.2's FIRST HALF).**
///
/// The pool reserves core 0 for the host loop so the doorbell spin cannot steal a worker's cycles - but
/// during `run()` the host does not spin, it waits, so core 0 is idle for the whole drain. Measured on the
/// 6-core machine this project targets: the engine's pool drains at **33.7 GB/s** (663.6 MB of expert
/// blobs in 19.71 ms/token) where the same kernel on 5 workers should reach 5/6 x 44.14 = 36.8 and the
/// machine measures 44.14 GB/s on all six. So the sixth core is being paid for and not used.
///
/// With `host_works`, `run()` claims jobs itself instead of spinning on `done_`, and the pool is six
/// threads on six cores. `false` is the A/B arm and exists so the change is measurable rather than
/// asserted - the counter it moves is `pool phases ... drain`, which is host-side and needs no profiler.
explicit ExpertPool(int n_workers = 0, bool pin = true, bool host_works = true);
~ExpertPool();
ExpertPool(const ExpertPool&) = delete;
ExpertPool& operator=(const ExpertPool&) = delete;
int workers() const { return n_; }
/// Whether the host thread also drains. Reported at startup, because "the engine adapts to the machine it
/// is on" is only true if the engine says which adaptation it took.
bool host_works() const { return host_works_; }
/// Publish `n` jobs, then block until every one has been claimed AND every worker has parked.
/// `jobs` must outlive the call (it does, and the workers never touch it afterwards).
void run(ExpertJob* jobs, int n);
/// Plan v0.3 P4: the same outputs as `run`, bitwise, with every expert split by rows across all threads
/// (gate/up rows, then the intermediate's quantization, then down rows). With fewer experts than threads -
/// the case once the VRAM tier takes half of them - `run` leaves cores idle and each expert streams at one
/// core's bandwidth; this streams every expert at all of them. At most `kMaxSplit` experts.
void run_split(ExpertJob* jobs, int n);
static constexpr int kMaxSplit = 16;
/// Plan v0.3 P6: `run_split` for multi-token jobs (at most `kMaxSplitMulti`); the rows of each expert are
/// read once for all of its tokens.
void run_split_multi(ExpertJobMulti* jobs, int n);
/// Plan v0.3 P6: the same for a native pack's layer (ggml-cpu arithmetic, `nact` activations).
void run_split_multi_native(const NativeFmt& f, ExpertJobMulti* jobs, int n);
static constexpr int kMaxSplitMulti = 96;
/// run_split_multi's phases, accumulated ms: gate/up rows, the intermediate quantization, down rows.
double ms_multi_gu = 0, ms_multi_q = 0, ms_multi_down = 0;
int64_t multi_bytes = 0;
/// Total `_mm_pause` iterations spent waiting, over all workers, is no longer counted - see the note on the
/// atomics below. It was a LOCKED read-modify-write in the park loop, so measuring the contention added to
/// it.
long long pauses() const { return 0; }
/// **WHERE `run()` SPENDS ITS TIME, in milliseconds accumulated over its lifetime.** Three phases per
/// layer - wait for every worker to be parked, wait for the drain, wait for them to re-park - and until now
/// all three were reported as one number. Without the split there is no way to tell a pool that is slow at
/// the WORK from one that is slow at the SYNCHRONISATION, and those need opposite fixes: the first is a
/// kernel problem and the second is a barrier problem.
///
/// Only the host thread touches these, in `run()`, so they need no atomics.
void phase_ms(double& wait_park, double& drain, double& repark) const {
wait_park = ms_wait_park_;
drain = ms_drain_;
repark = ms_repark_;
}
private:
void worker(int id);
void drain(int id, ExpertScratch& scratch);
void run_phase(int mode, int n_tasks);
int n_ = 0;
bool host_works_ = true;
ExpertJob* jobs_ = nullptr;
int njobs_ = 0;
/// The host's own scratch when `host_works_`. A separate object rather than a share of `scratch_[i]`,
/// because a worker may own any index and the two must not be able to collide.
ExpertScratch host_scratch_;
// `run()`'s three phases, accumulated. Host-thread only; see `phase_ms`.
double ms_wait_park_ = 0.0;
double ms_drain_ = 0.0;
double ms_repark_ = 0.0;
// ---- EACH ATOMIC GETS ITS OWN CACHE LINE, AND THE SPIN COUNTER IS GONE. (Review finding C3.)
//
// These were six adjacent atomics, which put `head_`, `done_`, `parked_` and `epoch_` on ONE cache line -
// the four that workers and the host actually contend on, invalidating each other on every access.
//
// Worse, the parked spin did `pauses_.fetch_add(1)` on EVERY iteration: a LOCKED read-modify-write, five
// workers against one line, at roughly one iteration per `_mm_pause`. So the line the host must WRITE to
// publish work (`epoch_`) and READ to confirm the workers are parked (`parked_`) was being hammered by the
// very threads waiting for it. That is contention the pool imposes on itself.
//
// The counter was diagnostic only - `pauses()` was read in one place, to print a number nothing branched on
// - so it is deleted rather than amortised. `alignas(64)` then stops the remaining four sharing.
alignas(64) std::atomic<uint32_t> head_{0};
alignas(64) std::atomic<uint32_t> done_{0};
alignas(64) std::atomic<uint32_t> parked_{0};
alignas(64) std::atomic<uint32_t> epoch_{0};
alignas(64) std::atomic<bool> stop_{false};
std::vector<std::thread> threads_;
std::vector<ExpertScratch> scratch_; // one per worker: no allocation, no false sharing of the hot data
// run_split state: mode 0 = whole experts, 1 = gate/up row parts, 2 = down row parts
int mode_ = 0;
int parts_a_ = 1, parts_b_ = 1;
struct SplitBuf {
alignas(64) float ff[FF];
ActQ a2;
};
std::vector<SplitBuf> split_;
// run_split_multi state: mode 3 = gate/up row parts, 4 = down row parts
ExpertJobMulti* mjobs_ = nullptr;
int64_t mrows_ = 0; // rows of the current multi phase across all its experts (n * FF, then n * H)
int mtasks_ = 1; // equal row ranges the phase is cut into
struct SplitBufMulti {
alignas(64) float ff[MAXT][FF];
ActQ a2[MAXT];
alignas(64) uint8_t hq[MAXT][kNativeHBytes]; // plan v0.3 P6: native down activations
};
const NativeFmt* nfmt_ = nullptr;
std::vector<SplitBufMulti> split_multi_;
};
} // namespace strata::kernels::cpu
include/strata/kernels/dequant_bf16.hpp
+27
View File
@@ -0,0 +1,27 @@
// include/strata/kernels/dequant_bf16.hpp - plan v0.3 P5: native GGUF blocks -> BF16 on the device, for the
// tensor-core GEMMs of batched prompt processing.
//
// One thread per 32-element group, transcribed from the reference `dequantize_row_*` functions of the pinned
// llama.cpp (`ggml/src/ggml-quants.c`, MIT): Q4_0 (2), Q5_0 (6), Q8_0 (8), Q3_K (11), Q4_K (12), Q5_K (13), Q6_K (14),
// IQ4_NL (20), IQ4_XS (23) and Q2_0 (42). A row-major (n_rows, n_cols) tensor of blocks becomes a row-major BF16
// matrix; `n_cols` must be a multiple of the type's block size.
#pragma once
#include <cstdint>
namespace strata::kernels {
bool dequant_bf16_supported(int ggml_type) noexcept;
/// `rows` rows of `cols` values starting at block row `row0` (so a caller can dequantize a slice of a tensor).
void dequant_bf16(int ggml_type, const void* blocks, int64_t row0, int64_t rows, int64_t cols, uint16_t* out,
void* stream);
/// The same into FP16 bits (the prompt path's quantized-weight GEMMs: Q2_0 values are exact in FP16).
void dequant_f16(int ggml_type, const void* blocks, int64_t row0, int64_t rows, int64_t cols, uint16_t* out,
void* stream);
/// The same into FP32 (tests and small tensors).
void dequant_f32(int ggml_type, const void* blocks, int64_t row0, int64_t rows, int64_t cols, float* out, void* stream);
} // namespace strata::kernels
include/strata/kernels/dequant_s2.hpp
+18
View File
@@ -0,0 +1,18 @@
// include/strata/kernels/dequant_s2.hpp - the S2 decode, host-callable.
//
// S2 is Q2_0's canonical form and 31.64 GiB of the artifact is Q2_0, so this is the hottest decode in the
// engine. Plane layout is docs/pack-format.md's: codes packed 4 per byte LSB-first (element i at byte i/4,
// bits (i%4)*2) and one FP32 scale per group of 64. Decode is `(code - 1) * scale` with the subtraction in
// the INTEGER domain, which is what makes it bit-exact (docs/pack-format.md §3.1).
#pragma once
#include <cstdint>
namespace strata::kernels {
// Decode `n_blocks` S2 blocks. `codes` is n_blocks*16 bytes, `scales` n_blocks floats, `out` n_blocks*64
// floats. All device pointers. Synchronises before returning, so a caller can read `out` immediately -
// correctness first; overlap is a Phase 3 concern.
void dequant_s2(const uint8_t* codes, const float* scales, float* out, int64_t n_blocks);
} // namespace strata::kernels
include/strata/kernels/elementwise.hpp
+116
View File
@@ -0,0 +1,116 @@
// include/strata/kernels/elementwise.hpp - the small ops a LAYER needs between its GEMVs, P2.S5.
//
// These are not interesting kernels and that is why they are in one file with one parity test: they are the
// glue in the per-layer graph, and glue that is written inline at each call site is glue with several
// slightly different float conventions. Three of them exist because a projection returns a DIFFERENT
// quantity from the one the next kernel wants:
//
// * `gdn_gate` turns (alpha, dt, ssm_a) into the recurrence's `gate`. `ref/gdn.py` L265:
// `gate = softplus(w_alpha . x + dt) * ssm_a`, with `softplus(x) = log1p(exp(x))` and the large-x branch
// `ggml_compute_softplus_f32` uses. `ssm_a` is NEGATIVE, which is what puts `exp(gate)` in (0,1).
// * `scale_inplace` applies the delta net's `1/sqrt(S)` to q. It is a separate op because
// `ref/gdn.py` says the two llama.cpp paths apply it in DIFFERENT places, so the caller has to own the
// choice rather than have it buried in a kernel.
// * `f32_to_f16_bulk` is what lets the f32 arithmetic of the mixer feed the fp16-activation GEMVs.
// `docs/activation-contract.md` is explicit that this is NOT the contract for a quantized weight - it is
// the bridge between two stages that both already exist, and the contract decision is a separate one.
//
// `silu_inplace` is here because `shared_expert` has its own copy inline and `gdn` needs one for the conv
// output; a third copy would be two too many.
#pragma once
#include <cstdint>
namespace strata::kernels {
/// Decode one canonical embedding row already selected by the caller. Codes are
/// packed low bits first (2, 4 or 8 bits); scales/offsets are FP32 per group.
/// The optional offset defaults to +0. Multiplication and addition round separately,
/// matching the former CPU decoder. Launches asynchronously on the supplied stream.
/// The caller validates positive, divisible dimensions and row/plane bounds.
void embedding_gather(const uint8_t* codes, const float* scales, const float* offsets,
int64_t n, int code_bits, int code_bias, int group_elems,
float* out, void* stream);
/// `y[t][h] = softplus(alpha[t][h] + dt[h]) * ssm_a[h]`, one thread per (t, h). In place is not possible:
/// `alpha` is read and `gate` is written, and they are different buffers by construction.
///
/// `h_v` is the number of v heads (48) and `n_tokens` the leading dimension; the real call is (1, 48).
void gdn_gate(const float* alpha, const float* dt, const float* ssm_a, float* gate, int64_t n_tokens,
int64_t h_v, void* stream);
/// `x[i] *= s`, in place.
void scale_inplace(float* x, int64_t n, float s, void* stream);
/// `dst[i] += src[i]`. **ADDED FOR R4's HIT/MISS SPLIT, WHERE THE TWO HALVES COME FROM DIFFERENT
/// ENGINES.** The CPU writes the misses into `parts` and the GPU writes the resident experts into a
/// buffer of their own, so that the GPU's work can run CONCURRENTLY with the CPU's instead of after it -
/// and the two are summed here, on the stream, immediately before `moe_combine` reads the result.
///
/// The alternative - having the hit kernel write straight into `parts` - costs the whole overlap, because
/// the copy of the misses into `parts` cannot happen until the pool has produced them, and the kernel is
/// then ordered behind that copy. Measured: the drain fell 18.2 -> 10.2 ms and the token did not move.
void add_inplace(float* dst, const float* src, int64_t n, void* stream);
/// `y[i] = f16(x[i])`, round-to-nearest-even, using the shared conversion in `f16_bits.hpp`.
void f32_to_f16_bulk(const float* x, uint16_t* y, int64_t n, void* stream);
/// `y[i] = bf16(x[i])`, for the weights whose contract is a bf16 activation - the BF16 `hc_*`, `ssm_alpha`,
/// `ssm_beta`, `indexer.*` and `ple_value` tensors.
///
/// A SEPARATE ENTRY POINT FROM `f32_to_f16_bulk` and not a flag on it, because the two conversions are not
/// variants of one idea: fp16 has 5 exponent bits and 10 mantissa bits, bf16 has 8 and 7. A caller that
/// picked the wrong one would get a plausible tensor at the wrong precision - 2^-11 against 2^-9, a factor of
/// four - and nothing downstream could tell which was intended.
void f32_to_bf16_bulk(const float* x, uint16_t* y, int64_t n, void* stream);
/// `x[i] = x[i] / (1 + exp(-x[i]))`, in place.
void silu_inplace(float* x, int64_t n, void* stream);
/// `build_norm`: `y[r][c] = x[r][c] / sqrt(MEAN_c(x[r]^2) + eps) * w[c]`, over the LAST axis.
///
/// QSA's norm, used on `attn_q` (24x256), `attn_k` (2x256) and `indexer.q_proj` (4x128). `w` may be null.
///
/// **IT IS NOT `gdn_l2_norm` AND THE DIFFERENCE IS A FACTOR OF sqrt(cols) = 16 HERE.** `gdn_l2_norm` divides by
/// `sqrt(SUM(x^2) + eps)` and applies no weight; this divides by `sqrt(MEAN(x^2) + eps)` and multiplies by `w`.
/// Two kernels whose names both say "norm" and whose bodies differ in one `/ cols` is exactly the pair that
/// gets swapped, and swapping them produces a well-scaled, plausible tensor - so they are separate entry
/// points and the parity test asserts the two readings are OBSERVABLE apart rather than trusting the names.
///
/// There is NO `(1 + w)` folding: `ref/qsa.py::rms_norm` is `y * w`, and `build_norm` in the source is
/// `ggml_rms_norm` followed by a plain multiply. Some norms in this artifact DO fold (the GDN path records
/// which); this one does not, and the same `sqrt(cols)`-sized argument applies to getting it wrong.
void rms_norm_weighted(float* x, const float* w, int64_t rows, int64_t cols, float eps, void* stream);
/// *seq = v in MAPPED PINNED memory - THE DOORBELL. A KERNEL and not cudaEventRecord, because round 199
/// found that an event record inside a capture is SILENTLY DROPPED while a kernel is captured normally, and
/// round 209 measured the host seeing this one's result 0.050 ms into a 39.8 ms graph.
///
/// It lives here because the caller that needs it - moe_layer - is a HOST translation unit, where
/// __global__ and <<<>>> do not exist.
/// *seq += 1 - AN INCREMENT, NOT A STORE OF A CALLER-SUPPLIED VALUE. A value passed in is evaluated on the
/// HOST at CAPTURE time, so a captured graph would ring the same number on every replay. Reading and
/// incrementing in memory is what makes the count advance under cudaGraphLaunch.
void doorbell_ring(uint32_t* d_seq, void* stream);
/// Plan v0.3 P3 token graph: one thread spins until the host-written `*d_flag` equals the ring value `*d_seq`
/// (both mapped pinned memory), then fences so everything the host wrote before the flag is visible to the
/// nodes that follow. Captured once per layer between `pre[l]` and the parts copy; it is what lets one graph
/// hold the whole token (bench/micro/device_wait.cu: 5.2 us per handoff, no driver call).
void doorbell_wait(const uint32_t* d_flag, const uint32_t* d_seq, void* stream);
/// Plan v0.3 P3 token graph: copy `n` floats from MAPPED pinned host memory (`src` is its device pointer) into
/// device memory with a kernel, so the handoff stays on the compute queue (a memcpy node is a copy-engine
/// operation, which WDDM submits separately and which measured 67 flushes per token).
void copy_from_mapped(float* dst, const float* src, int64_t n, void* stream);
/// Plan v0.3 P3: the doorbell's payload and its ring in ONE kernel. Copies `x` (n floats), `ids` and `weights`
/// (k each) into the mapped host regions, fences, and increments the mapped sequence number - replacing three
/// device-to-host memcpy nodes (copy-engine operations in the middle of the layer chain) and the ring kernel.
void doorbell_publish(const float* x, const int32_t* ids, const float* weights, int64_t n, int64_t k, float* x_out,
int32_t* ids_out, float* weights_out, uint32_t* d_seq, void* stream);
/// Plan v0.3 P3: copy `n` int32 from mapped pinned host memory into device memory with a kernel (the QSA
/// per-token step and positions), instead of a host-to-device memcpy node in the middle of a layer.
void copy_i32_from_mapped(int32_t* dst, const int32_t* src, int64_t n, void* stream);
} // namespace strata::kernels
include/strata/kernels/f16_bits.hpp
+96
View File
@@ -0,0 +1,96 @@
// include/strata/kernels/f16_bits.hpp - the fp16 conversions, written out, as bits.
//
// WHY THIS IS NOT `__float2half`. Round 193 found `__float2half` producing 0x2600 for a value whose correct
// fp16 encoding is 0x26DB, and 0x3C01 for a value that is exactly 1.0 - i.e. wrong bits, not a rounding
// difference. The cause was never isolated (it is either `__float2half` itself or `__half_as_ushort`), so
// `src/kernels/cuda/quantize_act.cu` replaced it with the version below, which then came out byte-exact over
// all 20,480 blocks of the fixture. `shared_expert.cu` carries the same function privately for the same
// reason. Three copies of a converter whose failure mode is silent wrong bits is two copies too many, so the
// copy lives here and the kernel and its TEST both include it - with the test checking it against vectors
// generated by numpy's float16 (an oracle that shares no code with this), which is what neither private copy
// had.
//
// The rounding is ROUND-TO-NEAREST-EVEN, which is what numpy's `.astype(np.float16)` does and what IEEE-754
// specifies for the default mode. The three interesting regimes are all handled explicitly and all three are
// covered by the test's oracle vectors:
//
// * normal - exponent fits, mantissa is rounded from 13 surviving bits with a sticky remainder
// * subnormal - the implicit 1 is restored and the result is shifted right, so the rounding point MOVES
// * overflow - exp >= 31 saturates to inf (0x7C00) and keeps the sign; a NaN input stays a NaN
//
// `f32_from_f16` deliberately does NOT round: it is used to compare against a reference that holds the same
// bits, so it only has to be exact.
#pragma once
#include <cstdint>
#include <cstring>
#if defined(__CUDACC__)
#define STRATA_HD __host__ __device__
#else
#define STRATA_HD
#endif
namespace strata::kernels {
/// Round-to-nearest-even f32 -> fp16, returned as raw bits.
STRATA_HD inline uint16_t f16_from_f32(float f) {
uint32_t x;
std::memcpy(&x, &f, 4);
const uint32_t sign = (x >> 16) & 0x8000u;
const uint32_t rawexp = (x >> 23) & 0xFFu;
int exp = (int) rawexp - 127 + 15;
uint32_t man = x & 0x7FFFFFu;
// TWO different things reach here and they must not be conflated: an f32 inf/NaN (raw exponent 255) stays
// an inf/NaN, and a FINITE value too large for fp16 (1e30) SATURATES TO INF. The version of this function
// in quantize_act.cu and shared_expert.cu tests `if (exp >= 31) return sign | 0x7C00 | (man ? 0x200 : 0)`
// and therefore returns a NaN for every overflowing finite value - which their own fixtures cannot see,
// because all of their values are O(1). Round 198's numpy oracle caught it here on the first run (1e30,
// 65536, 1e45: numpy says inf, the conflated test said NaN).
if (rawexp == 0xFFu) return (uint16_t) (sign | 0x7C00u | (man ? 0x200u : 0u)); // inf, or a quiet/NaN
if (exp >= 31) return (uint16_t) (sign | 0x7C00u); // finite overflow -> inf
if (exp <= 0) {
if (exp < -10) return (uint16_t) sign; // underflows to a signed zero
man |= 0x800000u; // restore the implicit 1
const uint32_t sh = (uint32_t) (14 - exp);
uint32_t h = (man >> sh) & 0x3FFu;
const uint32_t rem = man & ((1u << sh) - 1u);
if (rem > (1u << (sh - 1)) || (rem == (1u << (sh - 1)) && (h & 1u))) ++h;
return (uint16_t) (sign | h);
}
uint16_t h = (uint16_t) (sign | ((uint32_t) exp << 10) | (man >> 13));
const uint32_t rem = man & 0x1FFFu;
if (rem > 0x1000u || (rem == 0x1000u && (h & 1u))) ++h; // the carry may bump the exponent
return h;
}
/// Exact fp16 bits -> f32. No rounding, no exceptions: every one of the 65,536 inputs maps to a value f32
/// represents exactly, including the subnormals (which is why the `ex == 0` case cannot be skipped).
STRATA_HD inline float f32_from_f16(uint16_t h) {
const uint32_t sign = (uint32_t) (h & 0x8000u) << 16;
const uint32_t ex = (h >> 10) & 0x1Fu, man = h & 0x3FFu;
uint32_t out;
if (ex == 0) {
// Subnormal or zero. The value is `man * 2^-24`, so after `s` left shifts the leading 1 reaches bit
// 10 and the value is `(1 + frac) * 2^(-14-s)`: the f32 exponent field is 113 - s and `frac` is the low
// ten bits. Flushing these to zero - which the two private copies of this function in the tree do -
// maps 0x0001 to 0.0 instead of 2^-24, a 100% relative error on any KL-quant scale that lands here.
if (man == 0) {
out = sign;
} else {
uint32_t m = man;
int s = 0;
while ((m & 0x400u) == 0) { m <<= 1; ++s; }
out = sign | ((uint32_t) (113 - s) << 23) | ((m & 0x3FFu) << 13);
}
} else if (ex == 31) {
out = sign | 0x7F800000u | (man << 13);
} else {
out = sign | ((ex - 15 + 127) << 23) | (man << 13);
}
float f;
std::memcpy(&f, &out, 4);
return f;
}
} // namespace strata::kernels
include/strata/kernels/fused_gdn.hpp
+35
View File
@@ -0,0 +1,35 @@
// include/strata/kernels/fused_gdn.hpp - plan v0.3 P3: the gated delta-net step and its output norm in ONE kernel.
//
// The native step (after llama.cpp's gated_delta_net.cu) gives each warp one (head, column) and strides its 32
// lanes over the 128 ROWS of the state, which is stored `[row][head][col]`: every lane's load is 24.5 KB from its
// neighbour's, so each 32 B sector delivers 4 useful bytes - 47.6 us per layer for 6.3 MB of state traffic. Here a
// block owns one whole head, its 512 threads are (row group of 32, column), and a row's 128 columns are one
// contiguous 512 B load. With the head in one block, the per-head RMS norm, the norm weight and the sigmoid(z)
// gate of `gdn_out_norm` run in the same kernel.
//
// S <- g S + k (beta (v - g S^T k))^T (g = exp(gate[head]), per value head; q/k shared by h_v / h_k heads)
// o = (S^T q) / sqrt(128)
// y = rmsnorm(o) * gamma * sigmoid(z)
#pragma once
#include <cstdint>
namespace strata::kernels {
/// The 4-tap causal conv + SiLU over all `channels` (history is `[channel][3]`, updated), then the L2 norm of
/// each of the first `qk_heads` 128-channel heads (q then k), eps as the native `l2_norm` (x / sqrt(sum + eps)).
/// One block per 128-channel head; replaces conv_silu + two l2_norm launches.
void fused_gdn_conv_l2(float* history, const float* qkv, const float* conv_w, float* h, int channels, int qk_heads,
float eps, void* stream);
/// alpha and beta, both `(h_v, n_embd)` BF16 against the FP32 activation, with their epilogues:
/// gate = softplus(alpha + dt) * ssm_a, beta = sigmoid(beta). One warp per row; replaces two MMVF launches and
/// two elementwise kernels.
void fused_gdn_ab(const float* x, const uint16_t* w_alpha, const uint16_t* w_beta, const float* dt, const float* ssm_a,
float* gate, float* beta, int n_embd, int h_v, void* stream);
void fused_gdn_step_norm(float* state, const float* q, const float* k, const float* v, const float* gate,
const float* beta, const float* z, const float* gamma, float eps, float* y, int h_k, int h_v,
void* stream);
} // namespace strata::kernels
include/strata/kernels/fused_gr.hpp
+51
View File
@@ -0,0 +1,51 @@
// include/strata/kernels/fused_gr.hpp - plan v0.3 P3: the hyper-connection read in TWO kernels, with the
// previous half's write folded in.
//
// The native path spends six kernels per `gr_read` (norm, down MMVF, silu, up MMVF, gate+mean, inject MMVF) and
// one per `gr_write`, 96 + 96 times per token, and its up projection (10240 rows of 320) runs one 160-thread
// block per row: 20.6 us for 6.5 MB. Here:
//
// fused_gr_down : R' = R + bo_prev * 2 sigmoid(inj_prev / hc) (only when `apply`, computed on the fly)
// rs[c] = rsqrt(mean(R'[c]^2) + eps), xn = R' * w_norm * rs
// lo[k] = silu((w_down[k] . xn) / hc) k < hc_lr
// inject[c] = w_inject[c] . xn when w_inject is given
// fused_gr_up : R <- R' in place for this block's columns (when `apply`)
// mixed[d] = mean_c xn[c,d] * sigmoid(w_up[c*n_embd + d] . lo)
//
// FP32 activations and BF16 weights, like the native MMVF contract; the summation order differs from it (G-C
// judges the result). Geometry is the artifact's: n_embd 2560, hc 4, hc_lr 320. `inj_prev` and `inject_out`
// must be different buffers (every block reads the former while one block writes the latter).
#pragma once
#include <cstdint>
namespace strata::kernels {
struct FusedGrArgs {
const float* R = nullptr; ///< (hc, n_embd), read by `down`; `up` updates it in place when apply
float* R_out = nullptr; ///< == R for the in-place update
bool apply = false; ///< fold the previous half's gr_write
const float* bo_prev = nullptr; ///< that half's block output, n_embd
const float* inj_prev = nullptr; ///< that half's injection, hc
const float* w_norm = nullptr; ///< (hc * n_embd) f32
const uint16_t* w_down = nullptr; ///< bf16 [hc_lr][hc*n_embd]
const uint16_t* w_up = nullptr; ///< bf16 [hc*n_embd][hc_lr]
const uint16_t* w_inject = nullptr;///< bf16 [hc][hc*n_embd], or null (the final mixer)
float eps = 1e-6f;
float* lo = nullptr; ///< workspace, hc_lr floats
float* rs = nullptr; ///< workspace, hc floats
float* inject_out = nullptr; ///< hc floats (when w_inject)
float* mixed = nullptr; ///< n_embd
};
bool fused_gr_supported(int64_t n_embd, int64_t hc, int64_t hc_lr);
void fused_gr_read(const FusedGrArgs& a, void* stream);
/// Plan v0.3 P6: the same read for up to 8 tokens that share the weights (a verify window): the weights are read
/// once for all of them. `a[t]` is token t's arguments (its own R, pending write, lo, rs, inject, mixed; the four
/// weight pointers and eps must be the same for every t); `xn_scratch` is n_tok * hc * n_embd floats. Every
/// token's outputs are bitwise `fused_gr_read(a[t])`.
constexpr int kFusedGrMaxT = 8;
void fused_gr_read_multi(const FusedGrArgs* a, int n_tok, float* xn_scratch, void* stream);
} // namespace strata::kernels
include/strata/kernels/gdn.hpp
+97
View File
@@ -0,0 +1,97 @@
// include/strata/kernels/gdn.hpp - the gated delta-net (GDN) layer's non-projection parts, P2.S2.
//
// Transcribed from `ref/gdn.py`, which is itself transcribed from `build_layer_attn_linear` and
// `build_delta_net_autoregressive`. The projections (`wqkv`, `wqkv_gate`, `ssm_out`) are the existing
// `s_gemv`/`s2_gemv` kernels; what lives here is everything BETWEEN them.
//
// Geometry of the real artifact (48 layers are GDN, 12 are QSA):
//
// S = state_size = 128 head_k_dim == head_v_dim
// h_k = 16 k/q heads h_v = 48 v heads (idx[h] = h % h_k)
// C = 2*128*16 + 128*48 = 10240 channels d_conv = 4
//
// STATE LAYOUT. This is the one place the kernels deliberately differ from the reference, and the reason is
// coalescing. `ref/gdn.py` carries the state as (S, S, H_v) = (i, j, h); here it is **(S, H_v, S)** =
// (i, h, j), so that for a fixed i a warp walking consecutive j reads CONTIGUOUS memory. The state is a
// runtime buffer, not a weight, so this costs nothing to choose - but it must be said, because the two
// layouts index the same three numbers and a mismatch is silent.
//
// WHAT MAKES THE RECURRENCE PARALLEL, and it is worth stating because it is not obvious:
//
// dec = exp(gate) (h_v,)
// st = state * dec
// sk[h,j] = sum_i st[i,j,h] * k[idx[h]][i]
// d[h,j] = (v[h,j] - sk[h,j]) * beta[h]
// st[i,j,h] += k[idx[h]][i] * d[h,j]
// o[h,j] = sum_i st[i,j,h] * q[idx[h]][i]
//
// Every line touches ONLY the (j, h) column. Nothing couples one column to another, so one thread can own a
// column for the whole step and no barrier is needed anywhere in the recurrence. (The decay must still be
// applied BEFORE the rank-1 update - `ref/gdn.py` PROPERTY 5 pins that, and the test re-pins it here.)
#pragma once
#include <cstdint>
namespace strata::kernels {
/// Geometry of one GDN layer. The real artifact: S = 128, h_k = 16, h_v = 48.
struct GdnShapes {
int64_t S = 0; ///< head dim, and the state's contracted axis
int64_t h_k = 0; ///< key/query heads (q and k are (h_k, S))
int64_t h_v = 0; ///< value heads (v, z, o are (h_v, S)); h_v must be a multiple of h_k
};
/// One step of the delta-rule recurrence, IN PLACE on `state`.
///
/// state (S, h_v, S) f32, j fastest, updated in place
/// q, k (h_k, S) already L2-normalised, and q ALREADY SCALED by 1/sqrt(S) - the two llama.cpp paths
/// apply that scale in different places, so the caller owns it
/// v (h_v, S)
/// gate (h_v,) the PRE-exp value; `dec = exp(gate)` is formed here
/// beta (h_v,)
/// o (h_v, S)
///
/// The head pairing is `h % h_k` (MODULO), not `h / (h_v/h_k)` (INTERLEAVE). `ref/gdn.py` PROPERTY 8 and 11
/// pin this and the test re-pins it, because both conventions produce well-formed output.
void gdn_step(float* state, const float* q, const float* k, const float* v, const float* gate,
const float* beta, float* o, const GdnShapes& s, void* stream);
/// `ggml_ssm_conv` with the carried state prepended. Updates `conv_state` in place.
///
/// conv_state (d_conv-1, channels) the previous d_conv-1 input columns, OLDEST FIRST
/// x (channels,) the new column
/// kW (channels, d_conv) the weight in GGML-NATIVE layout: `kW[c*d_conv + i]` is tap i of
/// channel c. The manifest gives `ssm_conv1d.weight` ne = [4, 10240],
/// i.e. ne0 = d_conv is the FAST axis, which is this and NOT the
/// reference's row-major (d_conv, channels).
/// out (channels,)
///
/// out[c] = sum_{i<d_conv} inp[i][c] * kW[c*d_conv + i], inp = [conv_state | x]
///
/// `kernel[0]` reads the OLDEST state row and the new input lands in the LAST state row.
void gdn_conv_step(float* conv_state, const float* x, const float* kW, float* out, int64_t channels,
int64_t d_conv, void* stream);
/// `build_gdn_l2_norm`: `x / sqrt(sum(x^2) + eps)`, over the LAST axis.
///
/// The `+ eps` is an absolute floor on the SQUARED NORM, not on the mean - so the sum is formed first and
/// there is NO division by the width. The plausible wrong reading (`x / sqrt(mean(x^2) + eps)`) differs by
/// sqrt(S) = 11.3x here, which is why the test asserts the two are distinguishable.
void gdn_l2_norm(float* x, int64_t rows, int64_t cols, float eps, void* stream);
/// `beta = sigmoid(beta)`, in place over `h_v` values.
///
/// **`gdn_step` DOES NOT DO THIS AND THE LAYER MUST.** `build_layer_attn_linear` L889 is
/// `beta = ggml_sigmoid(beta)` and the step kernel's contract is `d = (v - sk) * beta`, so what it receives has
/// to be a fraction. Passing the raw `ssm_beta @ cur` leaves the delta-net write strength unbounded and signed.
/// `ref/gdn.py` and `ref/model.py` both sigmoid it; this call is what makes the C++ agree with them.
void gdn_beta_gate(float* beta, int64_t h_v, void* stream);
/// The layer's closing norm: `y = rms_norm(o, eps) * ssm_norm * sigmoid(z)`.
///
/// Note SIGMOID, not SiLU - qwen3.5 used SiLU and this artifact does not, and `ref/gdn.py` says so outright.
/// The norm is over the LAST axis: one RMS per head, so it is per-ROW here and not over the whole (h_v, S).
void gdn_out_norm(const float* o, const float* z, const float* ssm_norm, float* y, int64_t h_v, int64_t S,
float eps, void* stream);
} // namespace strata::kernels
include/strata/kernels/gr.hpp
+122
View File
@@ -0,0 +1,122 @@
// include/strata/kernels/gr.hpp - the gated residual (hyper-connection), P2.S2.
//
// `ref/gr.py::gr_read` / `gr_write`, which are transcribed from
// `llama_model_qwen4exp::graph::build_hc_mix` / `build_hc_combine` (qwen4exp.cpp L267-350 at 3cf03257).
//
// The residual is a STACK of `hc` streams, each `n_embd` wide - not one vector with a gate. `gr_read`
// collapses the stack to the one vector the mixer consumes and produces a per-stream injection; `gr_write`
// puts the mixer's output back into every stream, weighted per stream.
//
// WEIGHTS ARE BF16, stored as the high 16 bits of an f32 in a `uint16_t`. That is not a shortcut: every one
// of these tensors has source type BF16 in the pack (288 of the shard's 483 BF16 tensors are `hc_*`), so the
// pack's 32-bit promotion is lossless padding and re-rounding to 16 bits costs nothing. It is also worth
// 1.374 GiB of the 5.305 GiB dense file, which is the difference between the dense weights fitting in VRAM
// beside the expert cache and not fitting at all.
//
// The default activation contract is the historical BF16-rounded reference. The opt-in FP32 contract
// matches the pinned llama.cpp CUDA single-token BF16 MMVF input precision: BF16 weights multiply FP32
// activations. The CPU oracle and some batched CUDA paths instead round activations to BF16. Therefore
// activation precision must be pinned with the oracle backend and batch geometry, not inferred from weights.
#pragma once
#include <cstdint>
#include <cstddef>
namespace strata::kernels {
/// Select the activation precision for subsequent gr_read launches (default: false/BF16).
/// Configure before creating/capturing a session, and do not mutate concurrently with gr_read.
/// A captured graph retains the kernel variants chosen at capture; changing this flag cannot alter it.
/// This changes activation precision only; the existing FP32 reduction trees remain unchanged.
void gr_set_fp32_activations(bool enabled);
/// Select pinned single-token BF16/F32 MMVF projections for subsequent gr_read launches (default false).
/// This implies FP32 activations regardless of gr_set_fp32_activations. Configure before session capture;
/// a captured graph retains its selected kernels. Requires even hc*n_embd and hc_lr, as MMVF reads pairs.
/// This also selects pinned weighted F32 RMSNorm and nonlinear/combine arithmetic.
/// Layer reads use fused HC pre; the final mixer (null w_inject) retains its unfused graph order.
/// Subsequent gr_write calls select pinned sigmoid/scatter arithmetic under the same setting.
void gr_set_native_mmvf(bool enabled);
/// Geometry of the hyper-connection. The real model: n_embd = 2560, hc = 4, hc_lr = 320.
struct GrShapes {
int64_t n_embd = 0; ///< width of one residual stream
int64_t hc = 0; ///< number of streams in the stack
int64_t hc_lr = 0; ///< rank of the down/up bottleneck
};
/// Scratch for `gr_read`. EXPLICIT and caller-owned, because P2.T10 requires ZERO token-path allocations:
/// a hidden `cudaMalloc` on first use would satisfy every test here and fail the memory test in P2.S10.
///
/// **IT MUST BE DEVICE MEMORY.** `gr_workspace_init` takes a raw pointer and does no checking, so handing it
/// a host buffer produces device kernels writing to a truncated host address - an illegal access whose
/// reported fault address is the low 32 bits of a host pointer, which reads as a wild pointer rather than as
/// the mistake it is. That cost one round in `gr_parity`;
///
/// xn hc*n_embd floats the unrounded activated streams (the gated mean needs the unrounded value)
/// xq hc*n_embd uint16 bf16(xn), used by the default activation contract
/// lq hc_lr uint16 bf16(lo), used by the default activation contract
/// gated hc*n_embd floats xn * sigmoid(gate), before the mean over streams
/// lo hc_lr floats SiLU bottleneck, used by the FP32 activation contract
///
/// `cudaMalloc(gr_workspace_bytes(s))`, then `gr_workspace_init(s, ptr, ws)` to carve it.
struct GrWorkspace {
float* xn = nullptr;
uint16_t* xq = nullptr;
uint16_t* lq = nullptr;
float* gated = nullptr;
float* lo = nullptr;
/// Set by `gr_workspace_init`. `gr_read` CHECKS it, so an under-sized buffer is a loud failure at the
/// first call instead of silent corruption - which is what a mis-ordered region table used to produce.
size_t bytes = 0;
};
/// Carves `base` (which must be DEVICE memory, or null to size only) into `out` and returns the bytes needed.
/// 16-byte aligned between regions.
size_t gr_workspace_init(const GrShapes& s, void* base, GrWorkspace& out);
inline size_t gr_workspace_bytes(const GrShapes& s) { GrWorkspace w; return gr_workspace_init(s, nullptr, w); }
/// `build_hc_mix`. ONE token: `R` is (hc, n_embd) with n_embd fastest, i.e. stream c starts at c*n_embd.
///
/// w_norm (hc*n_embd) per-stream RMSNorm gamma, STORED AS (1 + w) by the converter
/// w_down (hc_lr, hc*n_embd) MANIFEST LAYOUT - row k is contiguous, `w_down[k*hc*n_embd + i]`
/// w_up (hc*n_embd, hc_lr) MANIFEST LAYOUT - row i is contiguous, `w_up[i*hc_lr + k]`
/// w_inject (hc, hc*n_embd) null for the FINAL mixer, which has no write-back to gate
///
/// mixed (n_embd) = mean over streams of xn * sigmoid(bf16(lo) @ w_up.T)
/// inject (hc) = bf16(xn) @ w_inject.T, or untouched when w_inject is null
///
/// `lo = silu((bf16(xn) @ w_down.T) / hc)` - the division by hc is INSIDE the silu, which is easy to read
/// past and changes the result materially.
/// With gr_set_fp32_activations(true), replace bf16(xn) and bf16(lo) by xn and lo respectively.
///
/// **NO PERMUTATION: both matrices are taken exactly as the manifest stores them**, so the loader does nothing
/// to them. Both are coalesced because each projection assigns ONE WARP PER OUTPUT ROW and strides the
/// reduction axis across its 32 lanes: for a fixed output, consecutive lanes touch consecutive addresses in
/// either orientation. An earlier version assigned one THREAD per row, which is coalesced with itself and 32
/// transactions away from its warp neighbours - see the history below.
///
/// **HISTORY, because this kernel was the engine's worst and the reason is instructive.** The first version
/// was launched `<<<1, 256>>>` with `xn`/`xq` in shared memory, and at the real geometry it cost **2.67 ms per
/// call, 96 calls per token = 262 ms/token, about 3.8 tok/s, 134x off the 643 GB/s memory floor.** The cause
/// was isolated by elimination: NOT memory latency (four independent accumulator chains made it 17% SLOWER),
/// NOT the SM clock (`nvidia-smi` reported 2917 MHz and 99% utilisation throughout), and only partly
/// coalescing (transposing one of the two matrices gave 31% and transposing the other gave nothing). It was
/// ISSUE THROUGHPUT - ~8.4M iterations of 6-8 instructions issued by ONE of 48 SMs.
///
/// This version splits the work the way the shapes want: `hc` blocks for the per-stream norm, one warp per
/// row for each of the two projections, and one block-tile for the mean. The transposes are gone because the
/// warp-per-row mapping makes them unnecessary.
void gr_read(const float* R, const float* w_norm, const uint16_t* w_down, const uint16_t* w_up,
const uint16_t* w_inject, float eps, const GrShapes& s, const GrWorkspace& ws, float* mixed,
float* inject, void* stream);
/// `build_hc_combine`. `R_out[i] = R[i] + block_out[d] * w[c]` with `w[c] = 2*sigmoid(inject[c]/hc)`.
///
/// The `2*sigmoid` is what CENTRES the gate on 1, so a zero injection is a plain residual add - the source
/// comment says so outright and `gr_write`'s test asserts it. Note the block output is added to EVERY stream
/// identically; only the weight is per-stream.
void gr_write(const float* R, const float* block_out, const float* inject, const GrShapes& s, float* R_out,
void* stream);
} // namespace strata::kernels
include/strata/kernels/iq_kernels.hpp
+57
View File
@@ -0,0 +1,57 @@
// include/strata/kernels/iq_kernels.hpp - the i-quant formats (IQ1_M, IQ2_XXS, IQ2_XS, IQ2_S, IQ3_XXS, IQ3_S,
// IQ4_NL) and Q2_0 on the GPU for the IQ2_XS / IQ3_XXS model files.
//
// The block layouts, codebook grids and dot products are llama.cpp's (ggml-common.h, ggml-cuda/vecdotq.cuh,
// ggml-cuda/dequantize.cuh; MIT, see third_party/ggml/LICENSE and VERSION.txt), so a weight means exactly what it
// means in llama.cpp. Activations are q8_1 (32 values, fp16 scale and fp16 sum), the llama.cpp CUDA contract.
#pragma once
#include <cstddef>
#include <cstdint>
namespace strata::kernels {
/// ggml type ids handled here.
bool iq_supported(int ggml_type) noexcept;
/// Bytes of one row of `n` values of `ggml_type` (n a multiple of the type's block).
size_t iq_row_bytes(int ggml_type, int64_t n) noexcept;
/// q8_1 blocks for `n_rows` rows of `n_cols` floats (n_cols a multiple of 32): y is n_rows * n_cols/32 blocks.
void quantize_q8_1_rows(const float* x, int64_t n_rows, int64_t n_cols, void* y, void* stream);
/// y[c][r] = W[r] . x[c] for `ncols` columns of q8_1 activations (x stride n_in/32 blocks per column).
void iq_mmvq(int ggml_type, const void* w, const void* x_q8_1, float* y, int n_in, int n_out, int ncols, void* stream);
/// Dequantize `n` contiguous values (n a multiple of 256) to fp16 / fp32.
void iq_dequant_f16(int ggml_type, const void* src, int64_t n, uint16_t* dst, void* stream);
void iq_dequant_f32(int ggml_type, const void* src, int64_t n, float* dst, void* stream);
/// Rows `tokens[0..n_tok)` (device ids) of a GGUF embedding table (`row_bytes` per row; the table may be mapped
/// host memory) dequantized to fp32, `n_embd` per row (a multiple of 256).
void iq_embed_rows(int ggml_type, const void* table, size_t row_bytes, const int32_t* tokens, int64_t n_tok,
int64_t n_embd, float* out, void* stream);
/// One expert's gate and up matrices (n_ff rows of n_embd each) into the interleaved fp16 layout the prompt path
/// uses: row 2r = gate row r, row 2r+1 = up row r.
void iq_dequant_gu_f16(int ggml_type, const void* gate, const void* up, int64_t n_ff, int64_t n_embd, uint16_t* dst,
void* stream);
/// The layout of one native expert blob: [gate rows | up rows | down rows], raw GGUF blocks.
struct NativeExpertLayout {
int gu_type = -1, d_type = -1;
int64_t n_embd = 0, n_ff = 0;
size_t gu_row = 0, d_row = 0; // bytes per row
size_t up_off = 0, down_off = 0; // byte offsets inside the blob
size_t bytes = 0; // the whole blob
};
NativeExpertLayout native_expert_layout(int gu_type, int d_type, int64_t n_embd, int64_t n_ff);
/// Bytes of scratch `native_expert_grouped` needs for `cap_entries` entries.
size_t native_expert_scratch_bytes(int64_t cap_entries, int64_t n_ff);
/// Grouped experts in the native format: group g's blob at device address grp_ptr[g]; its entries
/// [grp_start[g], grp_start[g+1]) read token ent_tok[e]'s q8_1 activation (n_embd/32 blocks per token in x_q8_1)
/// and write row ent_dst[e] of `out` (n_embd floats). Counts are read on the device.
void native_expert_grouped(const NativeExpertLayout& L, const unsigned long long* grp_ptr, const int32_t* grp_start,
const int32_t* n_groups, const int32_t* ent_dst, const int32_t* ent_tok, int64_t cap_groups,
int64_t cap_entries, const void* x_q8_1, void* scratch, float* out, void* stream);
} // namespace strata::kernels
include/strata/kernels/kv_q8.hpp
+38
View File
@@ -0,0 +1,38 @@
// include/strata/kernels/kv_q8.hpp - plan v0.3 P7: 8-bit KV storage for the QSA layers.
//
// The FP16 KV pools cost 2 x 12,288 bytes per token over the 12 QSA layers; with the FP32 pooled indexer keys the
// engine allocates 29 KB per token, 7.65 GB at 262K (bench/results/2026-09-23-p7-int-audit), which cannot fit
// 12 GB. This stores K and V as int8 codes with one FP16 scale per 64 values (per cell, per KV head):
// scale = fp16(max|x| / 127), code = clamp(rint(x / scale), -127, 127), x' = fp16(code * scale)
// Per layer and token: 2 KV heads x 256 values x (K, V) = 1,024 values -> 1,024 B of codes + 32 B of scales,
// instead of 2,048 B in FP16. Over the 12 QSA layers: 12,672 B/token instead of 24,576 (3.3 GB at 262K, not 6.4).
//
// Only STORAGE changes: `kv_gather_q8_step` dequantizes the selected cells into the same FP16 scratch the
// attention kernels already read, so `qsa_attend` and the native attention adapter are untouched. Layout mirrors
// the FP16 pools: codes `[page][kv_head][page_size][head_dim]`, scales `[page][kv_head][page_size][head_dim/64]`.
#pragma once
#include "strata/kernels/qsa.hpp"
#include <cstdint>
namespace strata::kernels {
inline constexpr int KV_Q8_GROUP = 64;
/// Bytes per cell (one token, one layer): K and V codes plus their scales.
inline uint64_t kv_q8_bytes_per_cell(const QsaShapes& s) {
return (uint64_t) s.n_head_kv * s.head_dim * 2 + (uint64_t) s.n_head_kv * (s.head_dim / KV_Q8_GROUP) * 2 * 2;
}
/// Append the cell at step[kStepPos] (graph-capturable: position and page come from device memory).
void kv_append_q8_step(int8_t* k_q, int8_t* v_q, uint16_t* k_scale, uint16_t* v_scale, const int32_t* page_table,
const int32_t* step, const float* kcur, const float* vcur, const QsaShapes& s, void* stream);
/// Gather step[kStepWidth] cells named by `ids` into FP16 scratch `[id][kv_head][head_dim]`; the grid is sized by
/// `max_ids` (capacity), the kernel reads the real count from `step`.
void kv_gather_q8_step(const int8_t* k_q, const int8_t* v_q, const uint16_t* k_scale, const uint16_t* v_scale,
const int32_t* page_table, const int32_t* ids, const int32_t* step, int64_t max_ids,
const QsaShapes& s, uint16_t* k_scratch, uint16_t* v_scratch, void* stream);
} // namespace strata::kernels
include/strata/kernels/mrope.hpp
+32
View File
@@ -0,0 +1,32 @@
// strata/kernels/mrope.hpp - multimodal rotary positions (the vision path).
//
// qwen4exp rotates with interleaved M-RoPE (llama.cpp LLAMA_ROPE_TYPE_IMROPE, rope.dimension_sections 11/11/10/0):
// each of the 32 rotated pairs takes its position from one of three streams, time t, height h and width w. Text
// has t = h = w, which is why every rope kernel here takes one position per row. An image does not: its tokens
// share t and differ in h and w (mtmd: t = p, h = p + y, w = p + x), and the text after it continues at
// p + max(nx, ny), not at its cell index.
//
// The table: nullptr (the default) keeps every kernel exactly as before - the position a caller passes is the
// rotary position. Set, it is a DEVICE int32 [cells][3] (t, h, w) and the position a caller passes is read as a
// CELL index; pair i rotates by tab[cell * 3 + mrope_sector(i)]. Every caller in the engine passes cell indices
// (pos_base is 0 everywhere), so the table covers the prompt path, the verify window and the MTP drafter at once.
// The pointer must be set before any CUDA graph is captured (kernels take it as an argument); the contents may
// change between requests.
#pragma once
#include <cstdint>
namespace strata::kernels {
void mrope_table_set(const int32_t* device_table);
const int32_t* mrope_table();
#if defined(__CUDACC__)
/// ggml rope_multi, is_imrope, sections {11, 11, 10, 0}: sector = pair % 32; sector % 3 == 1 -> h (sector < 33),
/// == 2 -> w (sector < 30), == 0 -> t (sector < 33). For pairs 0..31 all three bounds hold, so it is pair % 3.
__device__ __forceinline__ int mrope_pos(const int32_t* tab, int pos, int pair) {
return tab ? __ldg(tab + (size_t) pos * 3 + pair % 3) : pos;
}
#endif
} // namespace strata::kernels
include/strata/kernels/native_flash_attn.hpp
+36
View File
@@ -0,0 +1,36 @@
#pragma once
#include "strata/kernels/qsa.hpp"
namespace strata::kernels {
enum NativeFlashAttnStatus : int32_t {
kNativeFlashAttnSuccess = 0,
kNativeFlashAttnUnsupportedStep = 1,
};
// Explicitly selected short-context adapter for pinned CUDA F16/F16 vector
// FlashAttention. Geometry is exactly Q24x256, KV2x256, one query/sequence,
// scale=1/16, no ALiBi/softcap/sink. No allocation, synchronization, or KV rewrite.
// q/output: contiguous [24,256] F32. k/v: existing gathered [capacity,2,256]
// F16 buffers, with only step[kStepWidth] initialized rows required. Optional
// mask: 256 F16 additive logits, broadcast over heads; at least one valid key
// must be unmasked. Finite live q/k/v and finite or -inf live mask entries are
// preconditions; unused k/v/mask padding may contain arbitrary bytes.
// Padding is synthesized as zero K/V and -inf mask without reading it.
//
// capacity >=256; max_context in [1,256] is the HOST session limit. The device
// step must satisfy pos+1==n_kv==width in [1,max_context] and n_bid==n_kv/4.
// An invalid step writes status=UnsupportedStep and NaN output, without reading
// q/k/v/mask. status is one caller-owned device int32, overwritten on every call;
// integration MUST inspect it at its normal completion boundary. Host context
// limits must prevent unsupported steps before decoding. This is not a 20K path.
//
// All buffers (including optional mask and status) must have disjoint spans,
// with four-byte alignment for F32/step/status and two-byte alignment for F16.
// A nonnull explicit stream is mandatory. All pointers remain live for replay.
void native_flash_attn_short_step(const float* q, const uint16_t* k, const uint16_t* v,
const int32_t* step, int64_t capacity, int max_context,
const QsaShapes& shapes, float* output, int32_t* status,
const uint16_t* mask, void* stream);
} // namespace strata::kernels
include/strata/kernels/native_gdn.hpp
+22
View File
@@ -0,0 +1,22 @@
#pragma once
#include "strata/kernels/gdn.hpp"
namespace strata::kernels {
// Configure before creating/capturing a session. Existing captures retain their
// selected kernel, so changing the flag does not change an existing graph.
void native_gdn_set_enabled(bool enabled);
bool native_gdn_enabled();
// Pinned CUDA gated_delta_net recurrence for one token, S=128 and h_v % h_k=0.
// q/k are normalized F32 inputs; q MUST NOT be pre-scaled by 1/sqrt(S).
// Scaling is applied to the completed readout, as in the pinned fused operator.
// State retains Strata's layout: state[(i*h_v+h)*S+j]. The caller owns it and
// output; both must be disjoint from each other and from every read-only input.
// All spans must be aligned to four bytes. A nonnull ordered stream is required.
// No allocation, wait, or change to normalization/gate preprocessing occurs.
void native_gdn_step(float* state, const float* q, const float* k, const float* v,
const float* gate, const float* beta, float* output,
const GdnShapes& shapes, void* stream);
} // namespace strata::kernels
include/strata/kernels/native_gdn_preprocess.hpp
+37
View File
@@ -0,0 +1,37 @@
#pragma once
#include <cstdint>
namespace strata::kernels {
// Optional pinned CUDA GDN preprocessing. The layer selects these functions with
// native_gdn_enabled(); direct callers explicitly select them by calling them.
// Require a nonnull ordered stream, finite inputs, and four-byte aligned spans.
// No function allocates or synchronizes. All pointers are caller-owned.
// Four-tap convolution, then FP32 fast-math SiLU. History is [channels,3], oldest
// first; weights are [channels,4]. Both outputs are required, mutually disjoint,
// and disjoint from input/weights/history. History must not alias other spans.
void native_gdn_conv_silu(float* history, const float* input, const float* weights,
float* raw_output, float* silu_output, int64_t channels,
int64_t d_conv, void* stream);
// In-place [rows,128]: scale(rms_norm(x,epsilon/128),1/sqrt(128)). This includes
// the L2 normalization scale, but NOT the recurrence's readout scale.
// Epsilon is finite/nonnegative; epsilon=0 additionally requires nonzero rows.
void native_gdn_l2_norm(float* input, int64_t rows, int64_t cols, float epsilon, void* stream);
// In-place sigmoid over per-head beta; one token.
void native_gdn_beta_gate(float* beta, int64_t heads, void* stream);
// One token: gate[h] = softplus(alpha[h]+dt[h])*ssm_a[h], using the pinned CUDA
// log(1+exp(x)) expression and threshold 20. Output is disjoint from inputs.
void native_gdn_gate(const float* alpha, const float* dt, const float* ssm_a,
float* gate, int64_t heads, void* stream);
// [heads,128]: rms_norm(output,epsilon)*gamma*sigmoid(z). Gamma contains only
// 128 floats and broadcasts across heads. Destination is disjoint from inputs.
void native_gdn_out_norm(const float* output, const float* z, const float* gamma,
float* destination, int64_t heads, int64_t cols,
float epsilon, void* stream);
} // namespace strata::kernels
include/strata/kernels/native_gr_norm.hpp
+18
View File
@@ -0,0 +1,18 @@
#pragma once
namespace strata::kernels {
// Pinned llama.cpp 3cf03257f219afbe7334045ff7c6a06ac68c627d CUDA weighted F32
// RMSNorm for contiguous [n_cols, n_rows] data and equally shaped F32 gamma.
// Every row has independent norm statistics. Output is (scale * input) * gamma.
// Uses 256 threads below n_cols=1024, otherwise 1024, with the pinned XOR tree.
// Compile the implementation with --use_fast_math; no FP64, allocation, or wait.
// All buffers are caller-owned, nonoverlapping, and at least four-byte aligned.
// Require positive dimensions and finite nonnegative epsilon. Inputs/gamma and
// intermediate sums must be finite; epsilon=0 additionally requires nonzero rows.
// Enqueues on the supplied CUDA stream. A null handle explicitly selects CUDA's
// default stream, preserving gr_read's existing API; this function never syncs.
void native_gr_rms_norm_weighted(const float* input, const float* gamma, float* output,
int n_cols, int n_rows, float epsilon, void* stream);
} // namespace strata::kernels
include/strata/kernels/native_gr_postops.hpp
+28
View File
@@ -0,0 +1,28 @@
#pragma once
namespace strata::kernels {
// Pinned llama.cpp 3cf03257f219afbe7334045ff7c6a06ac68c627d contiguous F32
// single-token GR post-operations. Explicit stream, no allocation or wait; null
// selects the default stream. Positive dimensions with n*hc <= INT_MAX, finite
// values/intermediates, and four-byte alignment are required.
// Buffers must not overlap except the explicitly in-place arguments below.
// Implementation compiles separately with --use_fast_math.
// In-place pinned SCALE(1/hc), then SiLU, applied to the raw down projection.
void native_gr_down_silu(float* lo, int hc_lr, int hc, void* stream);
// xn and raw gate contain [n_embd,hc]. gate is overwritten with xn*sigmoid(gate).
// Layer path: pinned DSV4_HC_PRE_GATED ordered FMA accumulation, then scale.
// Head path (fused_layer=false): individually rounded products, ordered adds
// starting from stream zero, then SCALE. This retains the oracle's il=-1 graph.
void native_gr_pre_gated(const float* xn, float* gate, float* mixed,
int n_embd, int hc, bool fused_layer, void* stream);
// inject[hc] -> SCALE(1/hc), sigmoid, SCALE(2), identity DSV4_HC_POST.
// residual/output contain [n_embd,hc], block_out has n_embd elements.
// output may equal residual exactly; all other overlaps are invalid.
void native_gr_post(const float* residual, const float* block_out, const float* inject,
float* output, int n_embd, int hc, void* stream);
} // namespace strata::kernels
include/strata/kernels/native_mmvq.hpp
+112
View File
@@ -0,0 +1,112 @@
#pragma once
#include <cstddef>
namespace strata::kernels {
// Native GGUF Q2_0, Q4_0, Q5_0, Q8_0, Q3_K, Q4_K, Q5_K, Q6_K, IQ4_NL
// and IQ4_XS / CUDA Q8_1 adapters, pinned to llama.cpp
// 3cf03257f219afbe7334045ff7c6a06ac68c627d, sm_120 generic MMVQ.
// All pointers are device pointers, at least 4-byte aligned, with no overlap.
// All calls enqueue on the explicit non-null CUDA stream; no allocation or wait.
// The translation unit must use --use_fast_math, as the pinned CUDA oracle does.
//
// Shapes use GGUF order: n_in is the contiguous reduction dimension, n_out is
// the weight row count, and ncols is the activation column/token count, 1..8
// (plan v0.3 P3). Columns are contiguous: activation column j is x + j * n_in
// (its Q8_1 blocks at j * n_in / 32) and output column j is y + j * n_out. See
// native_mmvq_set_multi_exact for how ncols > 1 relates to ncols == 1.
// n_in must be a positive multiple of 32 for quantization/Q4_0/Q5_0/Q8_0/IQ4_NL,
// 64 for Q2_0, or 256 for the other formats. n_out must be positive. Weights remain
// unmodified row-major GGUF blocks: Q3_K=110, Q4_K=144, Q5_K=176, Q6_K=210 and
// IQ4_XS=136 bytes per 256 elements; Q2_0 is 18 bytes per 64 elements.
// Q4_0/IQ4_NL=18, Q5_0=22, Q8_0=34 bytes per 32 elements.
// Q8_1 scratch has 36 bytes per 32 elements, with no extra row padding here.
// Input floats must be finite, and their block scales/sums representable in FP16.
std::size_t native_q8_1_bytes(int n_in, int ncols = 1);
// Layout for ncols > 1. false: llama.cpp's generic multi-column table (upstream), equal to ncols == 1 to
// float rounding, speed not yet measured. true (default): the ncols == 1 layout, every column bitwise equal to a
// single-column call. Set before
// graph capture; captured graphs keep the kernels they captured.
void native_mmvq_set_multi_exact(bool exact);
bool native_mmvq_multi_exact();
// One quantization may serve multiple weight matrices sharing the same input.
// Q8_1 stores FP16 scale and FP16 warp sum of the ORIGINAL float inputs; it does
// not reconstruct that sum from the quantized integers.
void native_quantize_q8_1(const float* x, void* x_q8_1, int n_in, int ncols,
void* stream);
void native_q5_k_mmvq(const void* weights, const void* x_q8_1, float* y,
int n_in, int n_out, int ncols, void* stream);
// Convenience composition: caller owns scratch sized by native_q8_1_bytes.
void native_q5_k_f32(const void* weights, const float* x, void* scratch_q8_1,
float* y, int n_in, int n_out, int ncols, void* stream);
void native_q2_0_mmvq(const void* weights, const void* x_q8_1, float* y,
int n_in, int n_out, int ncols, void* stream);
void native_q2_0_f32(const void* weights, const float* x, void* scratch_q8_1,
float* y, int n_in, int n_out, int ncols, void* stream);
void native_q3_k_mmvq(const void* weights, const void* x_q8_1, float* y,
int n_in, int n_out, int ncols, void* stream);
void native_q3_k_f32(const void* weights, const float* x, void* scratch_q8_1,
float* y, int n_in, int n_out, int ncols, void* stream);
void native_iq4_xs_mmvq(const void* weights, const void* x_q8_1, float* y,
int n_in, int n_out, int ncols, void* stream);
void native_iq4_xs_f32(const void* weights, const float* x, void* scratch_q8_1,
float* y, int n_in, int n_out, int ncols, void* stream);
void native_q4_k_mmvq(const void* weights, const void* x_q8_1, float* y,
int n_in, int n_out, int ncols, void* stream);
void native_q4_k_f32(const void* weights, const float* x, void* scratch_q8_1,
float* y, int n_in, int n_out, int ncols, void* stream);
void native_q6_k_mmvq(const void* weights, const void* x_q8_1, float* y,
int n_in, int n_out, int ncols, void* stream);
void native_q6_k_f32(const void* weights, const float* x, void* scratch_q8_1,
float* y, int n_in, int n_out, int ncols, void* stream);
void native_q4_0_mmvq(const void* weights, const void* x_q8_1, float* y,
int n_in, int n_out, int ncols, void* stream);
void native_q4_0_f32(const void* weights, const float* x, void* scratch_q8_1,
float* y, int n_in, int n_out, int ncols, void* stream);
void native_q5_0_mmvq(const void* weights, const void* x_q8_1, float* y,
int n_in, int n_out, int ncols, void* stream);
void native_q5_0_f32(const void* weights, const float* x, void* scratch_q8_1,
float* y, int n_in, int n_out, int ncols, void* stream);
void native_q8_0_mmvq(const void* weights, const void* x_q8_1, float* y,
int n_in, int n_out, int ncols, void* stream);
void native_q8_0_f32(const void* weights, const float* x, void* scratch_q8_1,
float* y, int n_in, int n_out, int ncols, void* stream);
void native_iq4_nl_mmvq(const void* weights, const void* x_q8_1, float* y,
int n_in, int n_out, int ncols, void* stream);
void native_iq4_nl_f32(const void* weights, const float* x, void* scratch_q8_1,
float* y, int n_in, int n_out, int ncols, void* stream);
// Storage/dispatch helpers take stable GGML type IDs, avoiding a ggml runtime
// dependency in the engine: Q4_0=2, Q5_0=6, Q8_0=8, Q3_K=11, Q4_K=12, Q5_K=13,
// Q6_K=14, IQ4_NL=20, IQ4_XS=23, Q2_0=42. Unsupported IDs throw in the byte-count
// and launch helpers; only the
// capability query returns false.
bool native_mmvq_supported(int ggml_type) noexcept;
std::size_t native_mmvq_weight_bytes(int ggml_type, int n_in, int n_out);
void native_mmvq(int ggml_type, const void* weights, const void* x_q8_1, float* y,
int n_in, int n_out, int ncols, void* stream);
} // namespace strata::kernels
include/strata/kernels/native_moe.hpp
+18
View File
@@ -0,0 +1,18 @@
#pragma once
#include <cstdint>
namespace strata::kernels {
// Configure before constructing/capturing sessions. Existing captures keep their
// selected kernels; changing this flag does not rewrite an existing graph.
void native_moe_combine_set_enabled(bool enabled);
bool native_moe_combine_enabled();
// Pinned CUDA weighted-reduction contract for one token, k in [1, 15].
// The backend fuses k=2..15; k=1 is its ordinary multiply/contiguous/add path.
// parts: k contiguous F32 rows of n_embd; weights: k F32 values. The first
// product rounds to F32, following products accumulate with FMA in expert order,
// and optional shared is added once afterward. Shared is not router weighted.
// Requires a nonnull ordered stream and disjoint output. No allocation or sync.
void native_moe_combine(const float* parts, const float* weights, const float* shared,
float* output, int64_t n_embd, int64_t k, void* stream);
}
include/strata/kernels/native_ple_postops.hpp
+30
View File
@@ -0,0 +1,30 @@
#pragma once
#include "strata/kernels/ple.hpp"
namespace strata::kernels {
// Fixed qwen4exp single-token geometry: embd=2560, streams=4, history=9,
// convolution kernel=4/dilation=3. All buffers are caller-owned and capture-safe.
struct NativePlePostopsBuffers {
float* key; // 10240, normalized projected key
float* query; // 10240, temporary normalized hidden; may equal normalized
float* gate; // 4
float* gated; // 10240
float* normalized; // 10240, append this to history AFTER this function
float* conv; // 10240
float* result; // 10240, may equal hidden exactly
};
// Matches pinned llama.cpp 3cf03257f219afbe7334045ff7c6a06ac68c627d CUDA
// build_ple AFTER the two projections. Only norm_key/query/conv and conv1d_f16
// fields of weights are used. Gamma is full [2560,4]; history is row-fastest
// [9,10240], and F16 taps are [4,10240]. Inputs must be finite.
// Require an explicit stream, nonnull aligned pointers and disjoint writable
// spans, with only the two exact aliases documented above allowed. No allocation
// or synchronization; retain every buffer through any captured graph execution.
void native_ple_postops(const float* projected_key, const float* hidden,
const float* value, const float* history,
const PleWeights& weights, const NativePlePostopsBuffers& buffers,
void* stream);
} // namespace strata::kernels
include/strata/kernels/native_qsa.hpp
+30
View File
@@ -0,0 +1,30 @@
#pragma once
namespace strata::kernels {
// Set before session construction/capture; existing CUDA graphs retain their
// selected kernels. Disabled by default. Integration chooses the call sites.
void native_qsa_set_enabled(bool enabled);
bool native_qsa_enabled();
// Pinned single-token CUDA F32 arithmetic from llama.cpp
// 3cf03257f219afbe7334045ff7c6a06ac68c627d. No allocation or synchronization;
// both functions require a nonnull caller-owned stream and four-byte alignment.
// Positive dimensions with n_cols*n_rows (or head_dim*n_head) <= INT_MAX.
// Finite inputs and finite intermediate sums/products are caller preconditions.
// Contiguous input/output [n_cols,n_rows], gamma[n_cols] broadcast over rows.
// Each row has independent RMS statistics. Exact output==input is supported;
// otherwise all three spans must be disjoint. Epsilon must be finite and >= 0;
// epsilon==0 additionally requires a nonzero sum of squares in every row.
// Pinned 256-thread (width<1024) or 1024-thread XOR reduction, (scale*x)*gamma.
void native_qsa_rms_norm_weighted(const float* input, const float* gamma, float* output,
int n_cols, int n_rows, float epsilon, void* stream);
// attn/output [head_dim,n_head]; q_full [2*head_dim,n_head], each row is
// [query channels, gate channels]. Output is attn*sigmoid(second-half gate).
// Exact output==attn is supported; all other spans must be disjoint.
void native_qsa_gate_apply(const float* attn, const float* q_full, float* output,
int n_head, int head_dim, void* stream);
} // namespace strata::kernels
include/strata/kernels/native_qsa_indexer.hpp
+35
View File
@@ -0,0 +1,35 @@
#pragma once
#include "strata/kernels/qsa.hpp"
namespace strata::kernels {
// Configure before session construction/capture. Disabled independently of the
// native QSA norm/gate switch; existing captures retain their selected kernels.
void native_qsa_indexer_set_enabled(bool enabled);
bool native_qsa_indexer_enabled();
// Pinned F16 raw-key cache / F32 pooling / weighted RMSNorm / text IMRoPE.
// Fixed single-sequence geometry: idx_dim=128, idx_block=4, n_rot=64.
// Uses the existing buffers: tail[3,128] stores half-rounded raw values expanded
// to F32; dead[128] is the repeated-cell0 spare key; pooled[max_cells/4+1,128]
// contains completed blocks followed by the spare; block_pos[1] changes only on
// completion and records that block's first absolute position.
//
// relative_pos_device is a cell index, advancing 0..max_cells-1 (use kStepPos,
// not the absolute RoPE position vector). The sequence is text,
// contiguous, with nonnegative pos_base divisible by four. Clear state before
// starting a new sequence. Completed block b rotates at pos_base+4*b; the spare
// always rotates at zero, including nonzero sequence bases, as in the oracle.
//
// Explicit nonnull stream; no allocation or synchronization. All spans must be
// aligned and disjoint, remain valid through replay, and contain finite values;
// raw values must round to finite F16 and epsilon must be finite and positive.
// max_cells is a positive fixed capacity; pos_base+max_cells must fit int32.
// Out-of-range device positions are ignored defensively, not host-validated.
// This adapter does not alter score accumulation, top-k, or finite 1e9 tail bias.
void native_qsa_indexer_append(const float* raw, const int32_t* relative_pos_device,
int32_t pos_base, const float* gamma, float epsilon,
const QsaIndexerBuffers& buffers, const QsaShapes& shapes,
int64_t max_cells, float freq_base, void* stream);
} // namespace strata::kernels
include/strata/kernels/native_qsa_score.hpp
+31
View File
@@ -0,0 +1,31 @@
#pragma once
#include "strata/kernels/qsa.hpp"
namespace strata::kernels {
// Configure before capture; independent of native key pooling and norm/gate.
void native_qsa_score_set_enabled(bool enabled);
bool native_qsa_score_enabled();
// SM120, single contiguous text sequence: F32 pooled[128,max_blocks] and
// query[128,4], four-cell blocks, top-k budget2048. Matches the pinned CUDA MMF
// dispatch with the model's n_kv padded to256 (n_blocks a multiple of64), even
// when only a partial32-row tile is live. This is not the unpadded odd-row
// cuBLAS contract. MMA consumes the same raw F32 register bits as the oracle.
//
// step[kStepCount] is device memory, refreshed before each replay. max_blocks
// must equal max_cells/4+1. Completed pooled rows[0,n_bid) and the spare row
// [n_bid] must be initialized. Writes cell_scores[0,n_kv) only. Applies ReLU
// separately per head, ordered F32 head addition, optional F32 block bias,
// then finite1e9 F32 bias for the incomplete tail; full blocks also receive
// the reference's +0 bias. Does not alter selection, ties, or selected order.
//
// All spans are aligned, disjoint, persistent through replay; inputs finite.
// Explicit nonnull stream. No allocation, host count reads, or synchronization.
// Invalid device counts suppress writes defensively; callers own count validity.
void native_qsa_score(const float* pooled, const float* query, const float* block_bias,
const QsaShapes& shapes, const int32_t* step,
int64_t max_blocks, int64_t max_cells,
float* cell_scores, void* stream);
} // namespace strata::kernels
include/strata/kernels/native_rope.hpp
+14
View File
@@ -0,0 +1,14 @@
#pragma once
namespace strata::kernels {
void native_rope_set_enabled(bool enabled);
bool native_rope_enabled();
// Pinned CUDA text-only IMRoPE: F32 rows, 64 rotated channels, no YaRN/frequency
// factors, equal text positions in all four IMRoPE sections. Each device position
// must be nonnegative. The position buffer remains live through graph replay.
// Supports head_dim 128/256 and exact x==out; partial overlap is rejected.
// Explicit stream required. No allocation or synchronization.
void native_rope_apply(const float* x, float* out, int rows, int head_dim,
int n_rot, float freq_base, const int* positions, void* stream);
}
include/strata/kernels/native_router.hpp
+17
View File
@@ -0,0 +1,17 @@
#pragma once
#include <cstdint>
namespace strata::kernels {
// Configure before session capture. Existing graphs retain their selected kernels.
// This experiment is off by default and does not modify the legacy router.
void native_router_set_enabled(bool enabled);
bool native_router_enabled();
// Pinned CUDA topk-moe contract for ONE token, 512 experts, top 10, softmax,
// no selection bias, lower normalization clamp 2^-14, and scale 1.
// Reads 512 finite F32 logits; writes 10 I32 IDs and 10 F32 weights. Equal
// computed probabilities select the lower expert index. All spans must be
// four-byte aligned and outputs disjoint from each other and the input.
// Requires a nonnull ordered CUDA stream. No allocation or synchronization.
void native_router_top10(const float* logits, int32_t* ids, float* weights, void* stream);
}
include/strata/kernels/ngram.hpp
+194
View File
@@ -0,0 +1,194 @@
// include/strata/kernels/ngram.hpp - P2.S4: the PLE n-gram hash and the per-layer token table.
//
// Two halves that share nothing but the row indices they exchange:
//
// * `ngram_rows` - the HOST hash. Sixteen row indices per token from the last three token ids, computed
// with 64-bit multiply/xor and no tensor op at all. `qwen4exp.cpp` L1092-1124, where the source says
// outright that it is host-side because "ggml has no int64 and no xor".
// * `PleTable` - the row gather. `per_layer_token_embd.weight` is IQ4_NL and is NOT IN THE PACK: it is
// 51.2e9 elements (28.8 GB) in the ORIGINAL second GGUF shard, and it is the only tensor this engine
// reads from the GGUF rather than from the canonical pack.
//
// THE HASH IS THE ONE PLACE IN THIS MODEL WHERE A WRONG BIT IS SILENTLY PLAUSIBLE. Every one of its
// properties has a rival reading that produces a valid index in range:
//
// * XOR, not sum. Two equal tokens with equal multipliers cancel to zero under XOR and double under sum.
// * `%` vocab_size, not `& (vocab_size-1)`. The vocab sizes are 20000003, 20000023, 20000033, ... - not
// powers of two, so a mask would be a different function.
// * the EOS cut propagates FORWARD: an EOS at any position replaces that position AND every older one.
// * a missing predecessor is `LLAMA_TOKEN_NULL = -1`, tested as `t < 0`. Token id 0 is a REAL token.
// * the token's OWN eos does not cut its own context - only the predecessors are examined.
//
// All of those are asserted observable in `ple_parity.cpp` rather than trusted.
#pragma once
#include <cstdint>
#include <string>
namespace strata::kernels {
// ---- geometry, every value confirmed against the artifact (docs/gguf-metadata.txt) ----
inline constexpr int NG_N_EMBD = 2560; // n_embd
inline constexpr int PLE_HEAD_DIM = 160; // qwen4exp.embedding_length_per_layer_input
inline constexpr int NGRAM_SIZE = 3; // qwen4exp.ple.ngram_size
inline constexpr int HEADS_PER_NGRAM = 8; // qwen4exp.ple.heads_per_ngram
inline constexpr int PLE_N_HEADS = (NGRAM_SIZE - 1) * HEADS_PER_NGRAM; // 16
inline constexpr int PLE_CONV_KERNEL = 4; // qwen4exp.ple.conv_kernel
inline constexpr int32_t PLE_EOS_TOKEN_ID = 248044;
inline constexpr int NG_HC = 4; // hyper_connection.count
inline constexpr int NG_HC_DIM = NG_HC * NG_N_EMBD; // 10240
inline constexpr int NG_HIST = (PLE_CONV_KERNEL - 1) * NGRAM_SIZE; // 9
inline constexpr int32_t TOKEN_NULL = -1; // LLAMA_TOKEN_NULL
inline constexpr float NG_RMS_EPS = 1e-6f;
// The table: [160, 320001536] IQ4_NL. ne0 = 160 is the FAST axis, so one row is 160 contiguous elements =
// 5 blocks of 32 at 18 bytes = 90 bytes. The head-slowest flatten then makes 16 rows exactly n_embd = 2560.
inline constexpr uint64_t PLE_TABLE_ROWS = 320001536ull;
inline constexpr int PLE_ROW_BYTES = (PLE_HEAD_DIM / 32) * 18; // 90
/// The artifact's own hash constants, transcribed from `docs/gguf-dump-shard1.txt`:
///
/// layer_multipliers = [23703573157769, 20109073645365, 8052911324071]
/// head_vocab_sizes = [20000003, 20000023, ...] 16 entries, NOT an arithmetic sequence
/// head_offsets = [0, 20000003, 40000026, ...] the running sum of head_vocab_sizes
///
/// `ref/ngram.py::selftest` calls its own copies "synthetic but STRUCTURALLY REAL"; they are in fact
/// `20000003 + 10*i` for the vocab sizes, which the artifact's are not (20000003, 20000023, 20000033 - the
/// differences are 20, 10, ...). A selftest is free to use whatever constants it likes, but a PARITY test
/// against the artifact must use these, and the difference is recorded because it is exactly the kind of
/// thing that gets copied from the reference into an engine.
struct PleConsts {
uint64_t mult[NGRAM_SIZE];
uint64_t vocab[PLE_N_HEADS];
uint64_t offset[PLE_N_HEADS];
};
PleConsts ple_artifact_consts();
/// `mixed_n = (ctx[0]*m[0]) ^ (ctx[1]*m[1]) ^ ... ^ (ctx[n-1]*m[n-1])` in uint64.
///
/// Exposed separately so a test can assert the XOR STRUCTURE directly rather than only its consequences.
uint64_t ngram_mixed(const int64_t* ctx, const uint64_t* mult, int n);
/// The sixteen row indices per token. `qwen4exp.cpp` L1092-1124.
///
/// tokens (n_tokens,) the ubatch's own tokens
/// prev (n_tokens, 2) predecessors OLDEST FIRST, `TOKEN_NULL` where there is none
/// out (n_tokens, 16) row indices into the PLE table
///
/// The window is built once per token and then only its PREFIXES are used, which is what makes the bigram
/// heads (0..7) and the trigram heads (8..15) agree about the recent history.
void ngram_rows(const int32_t* tokens, const int32_t* prev, int n_tokens, const PleConsts& c, uint32_t* out);
/// The IQ4_NL codebook: `{-127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113}`,
/// which is `kvalues_iq4nl` from `ggml-common.h` L1121 and also the manifest's `codebooks.IQ4NL`.
///
/// There is NO `-8` in this version of ggml. Some releases of `dequantize_row_iq4_nl` subtract 8 from the
/// codebook value; at 3cf03257 they do not, and the difference would shift every one of the 51.2e9 elements.
int iq4nl_code(int code);
/// One 90-byte row -> 160 floats. The canonical rule `value = cb[code] * scale + offset` with offset 0,
/// because IQ4_NL has no min term (`has_offset` is false in the manifest).
///
/// THE NIBBLE ORDER IS SPLIT-HALF, NOT INTERLEAVED, and that is the trap here. `dequantize_row_iq4_nl` is
/// y[j] = d * kvalues[qs[j] & 0xf];
/// y[j + QK4_NL/2] = d * kvalues[qs[j] >> 4];
/// so byte j carries elements j and j+16 - NOT 2j and 2j+1, which is the natural reading and which produces
/// a perfectly plausible embedding of the wrong 160 values.
void iq4nl_dequant_row(const uint8_t* row, float* out160);
/// How the table's rows are read (plan v0.3 P2). `Direct` is the default: unbuffered 4 KiB reads from the SSD,
/// so the table never occupies RAM or the OS file cache. `Mmap` is the earlier memory-mapped path, kept as the
/// A/B arm; it returns the same bytes.
///
/// NEVER KEEP THE SHARD MAPPED WHILE READING IT DIRECT: a live section on the same file serializes the unbuffered
/// reads (311 -> 1,575 us per token, bench/results/2026-09-23-p2-ssd-direct). Direct mode drops its own mapping
/// after the header parse; nothing else in the process may hold one.
enum class PleIo { Direct, Mmap };
struct PleIoOptions {
PleIo mode = PleIo::Direct;
uint32_t max_inflight = 64; ///< outstanding SSD reads (decode needs 16; prefill chunks use more)
uint64_t cache_rows = 1u << 20; ///< bounded row cache: 1,048,576 rows x 90 B ~ 95 MB; 0 disables
bool io_thread = true; ///< reads submitted by a worker thread, not the caller
};
/// The PLE table. Held by pointer-to-impl so this header does not drag `<windows.h>` into every
/// translation unit that wants the hash.
class PleTable {
public:
PleTable();
~PleTable();
PleTable(const PleTable&) = delete;
PleTable& operator=(const PleTable&) = delete;
/// Open with an explicit I/O mode. The two-argument `open` below is the default (Direct).
bool open(const std::string& gguf_path, std::string& err, const PleIoOptions& io);
/// The split the plan asks for: `issue` as soon as the token id is known, `collect` just before layer 1
/// needs the rows. `gather` is `issue` followed by `collect`. In Mmap mode `issue` only prefetches.
bool issue(const uint32_t* rows16);
bool collect(float* out2560, std::string& err);
/// Plan v0.3 P5: the rows of `n_tokens` tokens (16 each, `rows` token-major) into `out` (2560 floats per token),
/// as ONE reader request - page dedupe and sort across the whole batch, the reader's full queue depth. Not
/// while a single-token `issue` is pending. The mapped mode gathers row by row.
bool gather_batch(const uint32_t* rows, size_t n_tokens, float* out, std::string& err);
/// Fault injection (Direct mode): every row read completes no earlier than `us` after issue.
void set_injected_delay_us(double us);
/// Maps the ORIGINAL second GGUF shard read-only. The tensor's data does NOT start at file offset 0:
/// the manifest's `shard2_tensor.offset` is relative to the GGUF's DATA SECTION, and the header before it
/// is 192 bytes. Reading at 0 would silently decode the header plus 192 bytes of shifted rows - still
/// plausible IQ4_NL, which is why the reader takes the offset from a real GGUF parse (`GgufFile`, the
/// project's validated reader) and checks that the tensor exactly fills the rest of the file.
bool open(const std::string& gguf_path, std::string& err);
void close();
bool is_open() const;
uint64_t rows() const;
/// 16 row indices -> 2560 floats. The gathered rows are flattened HEAD-SLOWEST: row h's 160 values
/// occupy `out[h*160, (h+1)*160)`, which is what `ggml_get_rows` does and what makes the result a plain
/// n_embd-wide vector for the block that consumes it.
void gather(const uint32_t* rows16, float* out2560) const;
/// One row -> 160 floats, for diagnostics and for the parity test's oracle comparison.
void read_row(uint32_t row, float* out160) const;
/// Bytes actually touched since `open`, for the "how often is it not ready" measurement P2.S4 asks for.
uint64_t bytes_read() const;
/// One line of reader statistics (Direct mode), e.g. for --stats. Empty in Mmap mode.
std::string io_report() const;
PleIo mode() const;
private:
struct Impl;
Impl* impl_ = nullptr;
};
// ================================ THE PLE GATHER'S PREFETCH, AND ITS A/B ARM ================================
//
// `gather` issues its sixteen row reads as sixteen SEPARATE page faults into a 26.8 GB mapping, which measured
// **2.10-2.61 ms per token** and is second only to the layer loop among the token's avoidable terms. The table
// is 90 B per row, so the sixteen rows are 1,440 B - 0.5 MB/s, which is latency and not bandwidth, and it is
// because the sixteen pages are taken one at a time. `PrefetchVirtualMemory` issues them in one call.
//
// **THIS SWITCH EXISTS SO THE CLAIM CAN BE MEASURED RATHER THAN ASSERTED.** Cross-build comparisons in this
// project have repeatedly turned out to be machine drift - `LAYERS` alone moves 48.9 -> 52.3 between runs of
// one binary - so the two arms are alternated inside one session. Off is `--no-ple-prefetch`.
void ple_prefetch_enable(bool on);
bool ple_prefetch_enabled();
// ---- NOTE ON `madvise(MADV_RANDOM)` ------------------------------------------------------------------
// P2.S4 asks for `madvise(MADV_RANDOM)` on the mapping, and on Linux that is exactly right: a token gathers
// 16 rows scattered over 28.8 GB, so read-ahead is pure waste and would evict useful pages.
//
// THERE IS NO LINUX, SO THERE IS NO MADV_RANDOM HERE. Windows has no equivalent of the advice - the closest
// is `FILE_FLAG_RANDOM_ACCESS` at CreateFile time, which suppresses the cache-manager's read-ahead for the
// whole handle, and `PrefetchVirtualMemory` for the positive case. `GgufFile` opens with neither, so this
// build gets the DEFAULT sequential read-ahead on a random access pattern. Saying so is the point: a
// comment claiming MADV_RANDOM here would be a comment about code that does not run. On Linux the advice
// belongs in `GgufFile::open`, which this module does not own.
// ------------------------------------------------------------------------------------------------------
} // namespace strata::kernels
include/strata/kernels/ple.hpp
+119
View File
@@ -0,0 +1,119 @@
// include/strata/kernels/ple.hpp - P2.S4: the PLE block's GPU half (`build_ple`, qwen4exp.cpp L1206-1293).
//
// One token, sixteen gathered rows of 160 -> a residual-stack update:
//
// key = grouped_norm(ple_key @ emb) hc_dim, per-stream RMS then a [n_embd,hc] gamma
// query = grouped_norm(hidden, ple_norm_query)
// value = ple_value @ emb n_embd
// s[c] = sum_d(key[c]*query[c]) / sqrt(n_embd)
// gate[c] = sigmoid(sign(s) * sqrt(max(|s|, 1e-6))) signed square root, then sigmoid
// gated[c] = value * gate[c] value broadcast over the hc streams
// normalized= grouped_norm(gated, ple_norm_conv) <-- THE CONV INPUT
// conv = sum_k kW[k][c] * padded[hist + t - (kern-1-k)*dil][c]
// out = hidden + gated + silu(conv)
//
// **`normalized` IS THE CONV INPUT, AND `ref/ngram.py` SAYS OTHERWISE.** Its `ple_block` docstring calls
// `terms` "the padded gated values", but the source pads `normalized` - the grouped_norm of the gated values,
// not the gated values themselves - and `ple_layer_xcheck.cpp` L107-112 builds it that way from ggml. The
// reference computes `normalized` and then ignores it, passing `terms` in from the caller, so its own code
// does not contradict its docstring and nothing catches it. This header takes the HISTORY OF NORMALIZED
// ROWS, which is what the state actually is.
//
// WHERE THE WEIGHTS LIVE, and the two different layouts in play. `ple_conv1d.weight` has manifest shape
// [4, 10240] with ne0 = 4 fast, so ggml's flat index is `k + 4*c`; `ref/ngram.py`'s `kernel` is a (4, 10240)
// numpy array whose flat index is `k*10240 + c`. They are transposes, and the source's own view
// (`ggml_view_2d(tWc, 1, hc_dim, nb[1], k*nb[0])`) confirms which is meant. `ple_key` and `ple_value` are
// [2560, n] with ne0 = 2560 fast, which IS the reference's row-major orientation.
#pragma once
#include <cstdint>
namespace strata::kernels {
/// Opt in to pinned CUDA BF16/F32 MMVF for the PLE value projection only (default false).
/// Configure before session capture; captured graphs retain their selected projection kernels.
void ple_set_native_bf16(bool enabled);
/// Opt in to pinned CUDA postprojection norms, gate, convolution and residual
/// arithmetic (default false). Set before capture; existing graphs keep their
/// selected kernels. Independent of the key/value projection options.
void ple_set_native_postops(bool enabled);
bool ple_native_postops_enabled();
/// Weights of the single PLE layer (`ple.layers = [1]`). `blk.1.*` in the pack.
///
/// key_codes / key_scales `ple_key.weight` S2, [n_embd, hc_dim], code_bias -1, group 64.
/// The pack stores `scales_fp16: true`, and every S-form kernel in this project
/// (including `s2_gemv_q8`) takes scales as F32 - so the LOADER widens them once,
/// as it already must for the experts. Recorded because "the manifest says fp16"
/// and "the kernel wants float" is a silent 2x-read if nobody says which wins.
/// value_bf16 `ple_value.weight` BF16 as raw bits, [n_embd, n_embd]
/// norm_key/query/conv F32 [hc_dim], NO (1+w) folding - `ref/gdn.py` records the same for `ssm_norm`
/// conv1d_f16 `ple_conv1d.weight` F16, GGML-NATIVE `kW[k + kern*c]`
struct PleWeights {
const uint8_t* key_codes = nullptr; // hc_dim * (n_embd/4) bytes
const float* key_scales = nullptr; // hc_dim * (n_embd/64) floats
const uint16_t* value_bf16 = nullptr; // n_embd * n_embd
const float* norm_key = nullptr; // hc_dim
const float* norm_query = nullptr; // hc_dim
const float* norm_conv = nullptr; // hc_dim
const uint16_t* conv1d_f16 = nullptr; // PLE_CONV_KERNEL * hc_dim, flat index k + 4*c
/// Optional unchanged GGUF Q2_0 [2560,10240] key projection. Nonnull data
/// requires type 42 and caller-owned native_q8_1_bytes(2560) device scratch,
/// disjoint from the PLE workspace, inputs, weights and exported outputs.
/// Use an explicit ordered stream and retain both buffers through graph
/// execution. Null data preserves the canonical key projection.
const void* key_native_data = nullptr;
int key_native_type = -1;
void* key_native_q8_1 = nullptr;
/// Plan v0.3 P6: the IQ model files keep `ple_key` in BF16 ([n_embd, hc_dim], `w[o*n_embd+i]`); when set it
/// replaces both key paths above.
const uint16_t* key_bf16 = nullptr;
};
/// Everything the block produces, in the order `ple_layer_xcheck`'s oracle writes it. Any pointer may be
/// null; the caller that only wants the result passes nothing else. Having the intermediates is what makes
/// the parity test able to say WHICH stage diverged instead of only that the sum did.
/// Every nonnull output region, including result, must be disjoint from the block's entire scratch region.
/// result may alias hidden. The block rejects scratch/output overlap before launching any work.
struct PleOut {
float* key = nullptr; // hc_dim
float* value = nullptr; // n_embd
float* gate = nullptr; // hc
float* gated = nullptr; // hc_dim
float* normalized = nullptr; // hc_dim
float* conv = nullptr; // hc_dim
float* result = nullptr; // hc_dim
};
/// One token. `emb` is the 2560-wide gathered n-gram embedding, `hidden` the hc_dim residual stack, and
/// `hist_rows` the NG_HIST previous NORMALIZED rows, oldest first (the caller owns that cache).
///
/// `hist_rows` IS ROW-FASTEST: `hist_rows[row + NG_HIST*channel]`. That is not a choice - it is
/// `ggml_reshape_3d(state, d_conv-1, conv_channels, n_seqs)`, so ne0 = the history row varies fastest, and
/// `ple_layer_xcheck.cpp` L79-83 records the transposed reading failing ggml's own assert.
///
/// The block does NOT write back to the history: the caller keeps ownership of its state and this stays a
/// pure function of its inputs, which is what makes the chunked-versus-single-shot property testable.
/// Bytes the caller must hand `ple_block` as its workspace: five `hc_dim` buffers (key, query, norm,
/// gated, conv), one `n_embd` (value) and one `hc` (gate), plus the Q8 activation image and the BF16
/// embedding copy. **THE WORKSPACE IS THE CALLER'S BECAUSE THIS FUNCTION IS CAPTURED**: it used to
/// `cudaMalloc` three times and `cudaStreamSynchronize` once per call, and inside a graph both are errors -
/// `ple_block: scratch: operation not permitted when stream is capturing`. Same defect, same fix as
/// `shared_expert`.
uint64_t ple_block_scratch_bytes();
void ple_block(const float* emb, const float* hidden, const float* hist_rows, const PleWeights& w,
PleOut& out, void* scratch, void* stream);
/// Advance the row-fastest normalized history by one token: hist[r,c] = old_hist[r+1,c], then append
/// normalized[c] at row NG_HIST-1. One thread owns each channel, making the in-place shift well-defined.
/// hist and normalized must be nonnull, disjoint device regions of NG_HIST*NG_HC_DIM and NG_HC_DIM floats.
/// Async on stream, allocation-free and capture-safe. Throws on null or overlapping regions.
void ple_history_advance(float* hist, const float* normalized, void* stream);
/// Whether the CUDA path is available. The host hash and the table read work without it; the block does not.
bool ple_block_available();
} // namespace strata::kernels
include/strata/kernels/qsa.hpp
+360
View File
@@ -0,0 +1,360 @@
// include/strata/kernels/qsa.hpp - the QSA (full-attention) layer's cache, indexer and attention, P2.S2.
//
// Spec of record: `ref/qsa.py` (whose own selftest carries eleven PROPERTY checks) and `ref/model.py::_qsa` /
// `_indexer` / `_attend`, transcribed from `build_layer_attn` / `build_attn_qsa` in the source. The
// projections are NOT here: `attn_q`, `attn_k`, `attn_v`, `attn_output`, `indexer.q_proj` and
// `indexer.k_proj` are the existing `s_gemv`/`s2_gemv` kernels, exactly as `gdn_step` leaves them. What lives
// here is everything BETWEEN the projections:
//
// kv_append the new token's K and V into the paged cache
// indexer_key_append the raw indexer tail, and the pooling that fires on block completion
// qsa_index per-CELL indexer scores: Relu per head, sum over heads, + bias, mapped by cell_block
// topk_512 the selection: the `width` largest cells, ties by ascending index, returned ascending
// kv_gather selected cells through the PAGE TABLE into a contiguous scratch
// qsa_attend softmax GQA over the scratch
// qsa_gate_apply `attn * sigmoid(gate)` and the round to the fp16 activation `attn_output` consumes
//
// DECODE ONLY, and that is a deliberate scope statement rather than an omission. `ref/qsa.py::indexer_scores`
// has no mask parameter and says why: "The causal part of the mask is already satisfied: every cached cell
// precedes the query at decode." The plan's `qsa_index` bullet asks for a "block-causal mask", and at decode
// its entire content is the TAIL -> SPARE SLOT mapping (`cell_block`), which is implemented here and tested.
// The per-query form - a query at position i may only see blocks with `block_start + r - 1 <= i`, plus the tail
// tokens - is a PREFILL property: `phases/phase-5-prefill-and-context-reuse.md` L73 builds it as a
// [chunk x blocks] GEMM with per-row top-k. Adding a mask flag here that decode cannot exercise would be a
// parameter with no observable behaviour, and this project has already been bitten by checks that cannot see
// the thing they claim to check.
//
// ============================================================================================
// **THIS WHOLE API IS NOT YET CAPTURABLE, AND IT IS NOT OBVIOUS WHICH PARTS ARE NOT.**
//
// A CUDA graph bakes kernel ARGUMENTS in at capture time. Any argument that CHANGES PER TOKEN is therefore
// frozen at its first-token value on replay - silently, because the graph still runs, still produces finite
// output, and still looks correct. Auditing this header for that property (round 211-212) found FIVE entry
// points, not one:
//
// kv_append int64_t pos every token the new cell's position
// qsa_index int64_t n_bid every token completed blocks
// qsa_index int64_t n_kv every token cached cells
// topk_512 int64_t n_kv every token cached cells
// kv_gather int64_t n_ids every token = qsa_selection_width(n_kv, s)
// qsa_attend int64_t n_ids every token as above
// indexer_key_append int64_t pos FIXED now takes `const int32_t* pos_dev`
//
// WHAT A CAPTURED LAYER WOULD ACTUALLY DO, and it is worse than a crash: `kv_append` would write EVERY token to
// position 0, `qsa_index` and `topk_512` would score exactly one cell, `kv_gather` would gather one cell, and
// attention would be over a single cell for the whole sequence. Finite, plausible, and indistinguishable from
// "the model has no context" - a symptom that would be blamed on the KV cache, the selection, or the weights.
//
// THE FIX IS ONE DEVICE BUFFER, NOT SIX SIGNATURES. Every one of these values is a pure function of the
// current position and the geometry, so they belong together in a small device-resident step state that is
// written once per token:
//
// step[0] = pos step[1] = n_kv = pos+1
// step[2] = n_bid step[3] = width = qsa_selection_width(n_kv, s)
//
// Each kernel reads the entries it needs, every argument list becomes constant across tokens, and the whole
// QSA path captures. Writing it as six separate pointer parameters would leave six places to forget and no
// single thing to check.
// ============================================================================================
//
// ---- LAYOUTS. Every one of these is a runtime buffer and not a weight, so the layout is a free choice - but
// it must be WRITTEN DOWN, because the alternatives index the same numbers and a mix-up is silent.
//
// K/V CACHE: FP16, paged, `[page][kv_head][page_size][head_dim]` with head_dim fastest. The plan says "FP16
// first" and fp16 is also what makes 32K of context fit: 12 layers x 2 heads x 256 dims x 2 (K and V) x 2 B x
// 32768 = 805 MiB, against 1.6 GiB in f32. The page granule is coarse (the shape's `page_size`) for the
// reason `Memory/Qwen-5070-engine-campaign/src/plans/paged-kv-cache.md` L146-170 gives: this engine serves a
// handful of sequences and not hundreds, so fine granularity buys nothing and costs per-token indirection. The
// physical row for logical cell `t` and head `h` is
//
// page = table[t / page_size]; (page * n_head_kv + h) * page_size + (t % page_size)
//
// i.e. ONE HEAD's worth of consecutive cells is contiguous within a page. That is the property attention
// wants; the alternative (cell-major inside the page) makes a cell's two heads contiguous instead, which is
// the property `kv_append` wants and nowhere near as valuable. **Nothing outside this header may assume the
// mapping**: `qsa_attend` reads a gathered scratch and never a pool, and the test runs the same fixture at
// page_size 1, 4 and 512 with a permuted page table and requires the attention output to be IDENTICAL, so the
// granule cannot leak into the result.
//
// INDEXER KEYS: FP32, not fp16 and not the plan's INT8. These carry no tolerance anywhere in the contract:
// `docs/capture-format.md` L127 makes `indexer_ids` an EXACT key, because a different selected cell is a
// different attention input. The key is a pooled MEAN of four raw keys followed by a norm and a rotation, so
// an fp16 store perturbs every score by ~2^-11 relative, and the selection boundary is exactly where that
// matters. The cost of fp32 is 512 B per block per layer = 128 B/token/layer = 1.5 KiB/token over the 12 QSA
// layers, i.e. 50 MiB at 32K: cheap next to a correctness contract. The plan's INT8 indexer key is a Phase 3
// optimisation and it has to be justified by a measured selection-flip rate against this baseline.
//
// The RAW tail is fp32 for the same reason and because it is only `idx_block - 1` rows.
#pragma once
#include <cstdint>
namespace strata::kernels {
/// Geometry of one QSA layer. The real artifact: n_head 24, n_head_kv 2, head_dim 256, n_rot 64,
/// idx_n_head 4, idx_dim 128, idx_block 4, idx_top_k 2048, freq_base 1e7, rms_eps 1e-6.
struct QsaShapes {
int64_t n_head = 0; ///< query heads (24). n_head must be a multiple of n_head_kv
int64_t n_head_kv = 0; ///< key/value heads (2). The q -> kv map is `q / (n_head/n_head_kv)`, NOT `q % kv`
int64_t head_dim = 0; ///< 256, and must be a multiple of 4
int64_t n_rot = 0; ///< 64 rotated dims of head_dim (and of idx_dim), NEOX pairs
int64_t idx_n_head = 0; ///< indexer query heads (4)
int64_t idx_dim = 0; ///< indexer key width (128)
int64_t idx_block = 0; ///< r = 4: the indexer pools one key per 4 cells
int64_t idx_top_k = 0; ///< 2048 cells of selection budget
int64_t page_size = 0; ///< KV page granule in CELLS; 1 is a legal unpaged configuration
};
/// The real artifact's geometry, so no caller transcribes it by hand.
inline QsaShapes qsa_real_shapes() {
QsaShapes s;
s.n_head = 24; s.n_head_kv = 2; s.head_dim = 256; s.n_rot = 64;
s.idx_n_head = 4; s.idx_dim = 128; s.idx_block = 4; s.idx_top_k = 2048;
s.page_size = 512;
return s;
}
/// The one legal RMSNorm epsilon for this artifact (`attention.layer_norm_rms_epsilon`).
inline float qsa_rms_eps() { return 1e-6f; }
/// `rope.freq_base`, no rope.scaling keys, so freq_scale = 1 and no YaRN.
inline double qsa_freq_base() { return 1e7; }
/// The selection width: `min(n_kv, idx_top_k + idx_block - 1)`.
///
/// The `r - 1 = 3` is not slack, it is the incomplete tail: the budget is whole blocks PLUS whatever cells do
/// not fill one. So the plan's "selection skipped when context <= 2,048" is conservative by two - the source's
/// own bound is 2,051, and `ref/qsa.py` L292-294 records the disagreement. Below the bound the selection is
/// the IDENTITY, which `topk_512` is required to reproduce exactly rather than approximately.
inline int64_t qsa_selection_width(int64_t n_kv, const QsaShapes& s) {
const int64_t w = s.idx_top_k + s.idx_block - 1;
return n_kv < w ? n_kv : w;
}
/// THE PER-TOKEN VALUES EVERY QSA KERNEL NEEDS, IN DEVICE MEMORY.
///
/// This exists so the QSA layer can be a CUDA graph. A graph bakes kernel ARGUMENTS in at capture time, so any
/// per-token argument is frozen at its first-token value on replay - silently, because the graph still runs and
/// still produces finite output. Six entry points used to take one of these as a host scalar; every one of
/// them now reads it from here instead, which makes every argument list constant across tokens.
///
/// WHAT IT PREVENTS, concretely: with `pos` frozen, `kv_append` writes every token to cell 0, `qsa_index` and
/// `topk_512` score one cell, `kv_gather` gathers one cell, and attention runs over a single cell for the whole
/// sequence. Finite, plausible, and indistinguishable from "the model does not use its context".
///
/// The four entries are NOT independent - `n_kv = pos+1` and `width = qsa_selection_width(n_kv, s)` - which is
/// exactly why they live in one buffer with one writer: `qsa_step_write` computes all four from `pos`, so a
/// caller cannot update one and forget another.
///
/// It is written by a plain H2D copy of four int32, which IS capturable and re-reads its source at replay
/// (`bench/micro/pinned_capture.cu` case B) - so the same buffer works in a captured graph and in a direct
/// launch.
enum QsaStep : int {
kStepPos = 0, ///< the new cell's sequence position
kStepNKv = 1, ///< pos + 1: how many cells are in the cache
kStepNBid = 2, ///< n_kv / idx_block: how many blocks have COMPLETED
kStepWidth = 3, ///< qsa_selection_width(n_kv, s): how many cells the selection returns
kStepCount = 4,
};
/// `step` is (kStepCount,) int32 DEVICE memory. Pointers to it are passed to the entry points below; it is
/// allocated by the caller (it is per-sequence state, not scratch).
///
/// Returns the bytes one step buffer needs, so a caller sizing an arena does not transcribe `kStepCount`.
inline uint64_t qsa_step_bytes() { return sizeof(int32_t) * (uint64_t) kStepCount; }
/// Fills a HOST array of `kStepCount` int32 from `pos`, so a caller can upload it in one copy. Keeping the
/// arithmetic here rather than at the call site is the point of the buffer: `n_kv` and `width` are derived,
/// and a caller computing them by hand is a caller who can compute them differently.
void qsa_step_fill(int32_t* host_step, int64_t pos, const QsaShapes& s);
/// One token's K and V into the paged FP16 cache, at logical cell `pos`.
///
/// k_pool, v_pool `[page][kv_head][page_size][head_dim]` fp16, page_size*n_pages cells each
/// page_table (n_pages,) i32: logical page -> physical page. NOT assumed to be the identity
/// kcur, vcur (n_head_kv, head_dim) f32. `kcur` must ALREADY be normed and rotated - V is never rotated
///
/// The fp32 -> fp16 conversion is `f16_from_f32` (see f16_bits.hpp) and not `__float2half`.
void kv_append(uint16_t* k_pool, uint16_t* v_pool, const int32_t* page_table, int64_t pos,
const float* kcur, const float* vcur, const QsaShapes& s, void* stream);
/// The selected cells through the page table into a contiguous scratch.
///
/// ids (n_ids,) i32, logical cells, ASCENDING and unique as `topk_512` returns them
/// k_scratch etc (n_ids, n_head_kv, head_dim) fp16, written in full
///
/// A cell that is not gathered does not exist for `qsa_attend`, which is how the top-k mask is realised: the
/// source fills the kq mask with -inf and leaves the selected cells at zero, so an unselected cell has weight
/// exactly zero even when it would score highest (`ref/qsa.py` PROPERTY 8 tests precisely that).
void kv_gather(const uint16_t* k_pool, const uint16_t* v_pool, const int32_t* page_table, const int32_t* ids,
int64_t n_ids, const QsaShapes& s, uint16_t* k_scratch, uint16_t* v_scratch, void* stream);
/// WHERE THE INDEXER'S PER-SEQUENCE STATE LIVES. One set per QSA layer per sequence.
///
/// tail (idx_block - 1, idx_dim) f32 the RAW tail: cells of the block being filled, slot = pos % idx_block
/// dead (idx_dim,) f32 the spare slot's key, = rms_norm(raw key of cell 0), CONSTANT
/// pooled (max_cells/idx_block + 1, idx_dim) f32 block b at row b, the spare slot at row n_bid
/// block_pos (1,) i32 scratch: the kernel writes the completed block's first-cell
/// position here so the rotation can read it ON THE DEVICE
///
/// **Every pointer in this struct is a DEVICE pointer**, and so are every other pointer this header's entry
/// points take: `rope_neox_apply` indexes its cos/sin table and its position array on the device, so a host
/// scalar passed as `pos` is an illegal access rather than a convenience. `block_pos` exists for exactly that
/// reason - the rotation position is derived from the host's `pos`, so it has to be handed to the device
/// through memory, which is also the only form a CUDA graph can replay.
///
/// The spare slot MOVES: it is at row `n_bid`, and `n_bid` grows by one on every block completion. It is also
/// constant in value, because it is derived from cell 0's raw key and `rope_neox(., 0)` is the identity - which
/// is what makes a tail of three keys sufficient. A design that pooled only from the tail could not produce
/// this key at all after the first block.
struct QsaIndexerBuffers {
float* tail = nullptr;
float* dead = nullptr;
float* pooled = nullptr;
int32_t* block_pos = nullptr;
};
/// Append one raw indexer key (IDX_DIM wide, fp32, NEVER normed and NEVER rotated - `ref/qsa.py` L201-210) and
/// run the pooling if this cell COMPLETES a block.
///
/// `pooled[b] = rope_neox(rms_norm(mean(raw[b*r .. b*r+r-1]), w_k_norm), pos of cell b*r)` for b = pos/r on a
/// completion, and the spare slot `pooled[n_bid] = rope_neox(rms_norm(raw[0], w_k_norm), 0)`. The rotation
/// position is the block's FIRST cell, which `ref/qsa.py` PROPERTY 4 pins and the test re-pins: rotating at the
/// block's LAST cell keeps every shape and every magnitude.
///
/// **`pos_dev` IS A DEVICE POINTER AND THIS FUNCTION IS CAPTURABLE.** It used to take `int64_t pos` as a host
/// scalar, and a CUDA graph bakes kernel ARGUMENTS in at capture time - so a captured call would have replayed
/// the FIRST token's position forever: `slot = pos % idx_block` would never advance, no block would ever
/// complete, the indexer would keep pooling cells 0-3, and the selection would stay pinned to the first block.
/// Not a crash and not a small error - a plausible attention output over the wrong cells. Worse than a stale
/// POINTER, which this project's pitfalls list already warns about, because a stale pointer tends to fault while
/// a stale VALUE looks like a working graph.
///
/// It could not be fixed by swapping the scalar for a pointer alone: `pos` was ALSO read on the host, to decide
/// whether to run the completion rotation at all and which pooled row to rotate. So the rotation moved INTO
/// the pooling kernel, which is why this now takes `cos_tab`/`sin_tab` and `n_rot`, and why `block_pos` is kept
/// only as a readable record of the rotation position rather than as the input to a second launch.
///
/// `pos_base` IS THE SEQUENCE'S FIRST CELL'S POSITION, so cell j sits at position `pos_base + j`. The kernel
/// cannot derive it: `QsaCache.pos` is what the reference rotates by, and equating a cell's position with its
/// index is only right for a sequence that starts at 0. The test runs the fixture at `pos_base = 100`, where
/// the two readings are 514% apart. `pos_base` is CONSTANT for a sequence, so unlike `pos_dev` it is safe to
/// bake into a graph.
void indexer_key_append(const float* raw, const int32_t* pos_dev, int32_t pos_base, const float* w_k_norm,
float eps, const QsaIndexerBuffers& b, const QsaShapes& s, const float* cos_tab,
const float* sin_tab, void* stream);
/// Per-CELL indexer scores, `n_kv = pos + 1` of them, fp32.
///
/// pooled (n_bid + 1, idx_dim) block keys, the spare slot last, as `indexer_key_append` maintains them
/// q_idx (idx_n_head, idx_dim) ALREADY normed and rotated by the caller
/// bias (n_bid + 1,) or null optional additional per-block bias
///
/// score[j] = bias[cb] + sum_h Relu(dot(pooled[cb], q_idx[h])), cb = min(j / r, n_bid)
/// score[j] += 1e9f for incomplete-tail cells (j >= (n_kv / r) * r).
/// The tail bias matches llama.cpp's causal decode input and applies even when bias is null.
/// Completed blocks receive no tail bias; the step API derives this from each replay's n_kv.
///
/// The Relu is PER HEAD and the sum is over heads AFTERWARDS: `relu(sum_h dot_h)` is the rival reading, and it
/// is a different number whenever any head's dot is negative. Flipping one query head's sign can never lower
/// the score, which is the property the test asserts. There is NO softmax and NO 1/sqrt(d) here.
void qsa_index(const float* pooled, int64_t n_bid, const float* q_idx, const float* bias, const QsaShapes& s,
int64_t n_kv, float* cell_scores, void* stream);
/// The largest `qsa_selection_width(n_kv)` cells of `cell_scores`, ascending, into `ids`.
///
/// Ties are broken by ASCENDING INDEX. ggml does not specify that, and `ref/qsa.py` is explicit about it so a
/// mismatch is at least diagnosable; the host test compares against `np.lexsort((arange, -scores))`, which is
/// the same rule, implemented independently of the kernel's threshold search.
///
/// `cap` is the caller's buffer size in ids and must be >= the width; the kernel writes EXACTLY `width` and
/// refuses rather than truncating. The `-1` padding that `docs/capture-format.md` L66-73 describes is the
/// capture writer's job, not this kernel's: the padded width is a property of the run, not of one token.
///
/// `n_kv` must be <= 32768 (`kTopkMaxCells`): the selection runs in ONE block, which is what makes the
/// ascending tie rule deterministic without a second pass over global memory. Phase 3 replaces this with a
/// radix select and Phase 7 raises the cap; until then a longer context is refused loudly.
inline constexpr int64_t kTopkMaxCells = 32768;
void topk_512(const float* cell_scores, int64_t n_kv, const QsaShapes& s, int64_t cap, int32_t* ids,
void* stream);
/// Softmax GQA over ONLY the gathered cells. `q` is (n_head, head_dim) f32, already normed and rotated.
///
/// scores[j] = dot(k[j][q / (n_head/n_head_kv)], q[h]) / sqrt(head_dim)
/// w = softmax_j(scores)
/// attn[h] = sum_j w[j] * v[j][q / (n_head/n_head_kv)]
///
/// The q -> kv head map is INTEGER DIVISION (`ggml-cpu/ops.cpp` L8729: `iv2 = iq2 / rv2`), not `q % n_head_kv`
/// - with 24 and 2 the two give [0 x12, 1 x12] and [0,1,0,1,...], and both produce a well-formed result.
///
/// `attn` is (n_head, head_dim) f32 and NOT yet gated; `weights`, when non-null, is (n_head, n_ids) f32, which
/// is what `docs/capture-format.md`'s `L{n}.qsa_scores` records. Produced here rather than recomputed by the
/// caller so the capture cannot drift from the attention it claims to describe. An EMPTY `ids` is legal and
/// gives a zero `attn` (no selected cell has no softmax to take).
void qsa_attend(const float* q, const uint16_t* k_scratch, const uint16_t* v_scratch, int64_t n_ids,
const QsaShapes& s, float* attn, float* weights, void* stream);
/// `attn *= sigmoid(gate)` in f32, then ONE round to fp16 for the `attn_output` GEMV.
///
/// q_full (n_head, 2*head_dim) f32 the raw `attn_q` projection, BEFORE any norm
/// gate = q_full[h][head_dim + d] - the SECOND half of each head's 2*head_dim block
/// out (n_head, head_dim) fp16
///
/// The split is NOT element-interleaved (`ref/qsa.py` PROPERTY 10 pins that: `q = per_head[:, :head_dim]`,
/// `gate = per_head[:, head_dim:]`), and the rounding happens AFTER the multiply so there is one rounding
/// rather than two. SIGMOID, not SiLU: this artifact's QSA gate is a sigmoid.
void qsa_gate_apply(const float* attn, const float* q_full, const QsaShapes& s, uint16_t* out, void* stream);
/// THE SAME MULTIPLY WITHOUT THE fp16 ROUND, for the caller whose weight wants Q8_K.
///
/// `attn_output` is a K-quant (Q4_K/Q5_K/Q6_K in every QSA layer of this artifact) and its `vec_dot_type` is
/// **Q8_K, not fp16** - `docs/activation-contract.md` is explicit, and the measured cost of getting it wrong is
/// 0.66-1.41% per GEMV against a 1e-3 tolerance, which is above what Gate C1 allows. The fp16 form exists
/// because that is what the engine did first; this one exists because that was a gap, and the caller quantizes
/// to Q8_K between the two calls.
///
/// The arithmetic is IDENTICAL to `qsa_gate_apply`'s up to the last step, and the multiply is in double in both
/// - so switching between them changes exactly one rounding, which is the point.
void qsa_gate_apply_f32(const float* attn, const float* q_full, const QsaShapes& s, float* out, void* stream);
// ================= THE CAPTURABLE ENTRY POINTS =================
//
// These take the per-token counts from `step` AND size every launch from a CONSTANT capacity, which is the two
// things a CUDA graph requires. The host-scalar entry points above are thin wrappers over these: they fill a
// module-level step buffer and pass this token's counts as the capacities, so they behave exactly as before and
// stay convenient for a direct launch or a test.
//
// **THE LAYER MUST USE THESE.** The wrappers are correct for a direct launch and WRONG inside a graph: their
// grid and their shared size change per token, and a graph bakes both in. A captured `qsa_index` launched for
// one block leaves every other pooled row unscored; a captured `qsa_attend` sized for one cell reads past its
// shared buffer once the sequence grows. Both failures are silent.
//
// The kernels already guard for the surplus: `qsa_index_kernel` returns for `b > n_bid`, `kv_gather_kernel`
// returns past its device `n_ids`, and `qsa_attend_kernel` handles an empty selection itself. So passing a
// capacity LARGER than this token needs is correct, not merely safe.
/// `max_blocks` must be at least `(max_cells / idx_block) + 2`.
void qsa_index_step(const float* pooled, const float* q_idx, const float* bias, const QsaShapes& s,
const int32_t* step, int64_t max_blocks, float* cell_scores, void* stream);
/// `cap` must be `qsa_selection_width(kTopkMaxCells, s)` - the largest width the kernel can ever write.
void topk_512_step(const float* cell_scores, const QsaShapes& s, int64_t cap, const int32_t* step,
int32_t* ids, void* stream);
/// `max_ids` must be the CAPACITY of `k_scratch`/`v_scratch` in ids.
void kv_gather_step(const uint16_t* k_pool, const uint16_t* v_pool, const int32_t* page_table,
const int32_t* ids, const int32_t* step, int64_t max_ids, const QsaShapes& s,
uint16_t* k_scratch, uint16_t* v_scratch, void* stream);
/// `max_ids` must be at least the largest selection the geometry can produce, because it sizes the SHARED
/// memory the softmax works in.
void qsa_attend_step(const float* q, const uint16_t* k_scratch, const uint16_t* v_scratch,
const int32_t* step, int64_t max_ids, const QsaShapes& s, float* attn, float* weights,
void* stream);
/// The cell's position comes from `step`, so this is the form a graph may contain.
void kv_append_step(uint16_t* k_pool, uint16_t* v_pool, const int32_t* page_table, const int32_t* step,
const float* kcur, const float* vcur, const QsaShapes& s, void* stream);
} // namespace strata::kernels
include/strata/kernels/qsa_decode_attn.hpp
+44
View File
@@ -0,0 +1,44 @@
// include/strata/kernels/qsa_decode_attn.hpp - plan v0.3 P3/P7: split-K decode attention over the selected cells.
//
// The first attention kernel runs one block per query head (24 blocks on a 48-SM part), reads each key row with one
// thread (uncoalesced) and walks every selected cell serially for the values; at a 4K context (2,051 selected
// cells) that costs ~0.3-0.5 ms per QSA layer. This one reads the KV POOLS directly through the page table and the
// selection ids (no gather copy), in chunks of CHUNK cells per block, and serves all `n_head / n_head_kv` query
// heads that share a KV head from one read of the chunk:
//
// grid (ceil(cap / CHUNK), n_head_kv): s = q.k * scale for 12 heads x CHUNK cells, chunk max m and sum l,
// acc = p . V -> partials
// grid (n_head): merge the chunks with the usual log-sum-exp rescale
//
// FP16 pools or INT8 pools (codes + fp16 scale per 64 values, `kv_q8.hpp`). The grid is sized by the capacity;
// the real count comes from `step[kStepWidth]` as for every other capturable QSA kernel.
#pragma once
#include "strata/kernels/qsa.hpp"
#include <cstdint>
namespace strata::kernels {
struct QsaAttnPools {
const uint16_t* k_pool = nullptr; ///< fp16 [page][kv_head][page_size][head_dim], or null when int8
const uint16_t* v_pool = nullptr;
const int8_t* k_q = nullptr; ///< int8 codes, same layout
const int8_t* v_q = nullptr;
const uint16_t* k_scale = nullptr; ///< fp16 [page][kv_head][page_size][head_dim / 64]
const uint16_t* v_scale = nullptr;
const int32_t* page_table = nullptr;
};
/// Scratch floats for `cap` selected cells: partial accumulators, maxima and sums.
uint64_t qsa_decode_attn_scratch_floats(int64_t cap, const QsaShapes& s);
void qsa_decode_attn_step(const float* q, const QsaAttnPools& pools, const int32_t* ids, const int32_t* step,
int64_t cap, const QsaShapes& s, float* scratch, float* attn, void* stream);
/// Plan v0.3 P5: `n_q` queries at once, each with its own selection: q [n_q, n_head, 256], ids [n_q, cap], steps
/// [n_q, kStepCount], attn [n_q, n_head, 256]; scratch is `n_q` times the single-query size.
void qsa_decode_attn_batch(const float* q, const QsaAttnPools& pools, const int32_t* ids, const int32_t* steps,
int64_t cap, const QsaShapes& s, float* scratch, float* attn, int64_t n_q, void* stream);
} // namespace strata::kernels
include/strata/kernels/qsa_select.hpp
+33
View File
@@ -0,0 +1,33 @@
// include/strata/kernels/qsa_select.hpp - plan v0.3 P5/P7: the QSA indexer's block scores and top-k selection,
// for many queries at once and at long context.
//
// The per-token path scores every pooled block with one FP64 block per row (qsa_index_kernel) and selects with a
// single-block kernel that makes 32 bit-serial passes over every CELL (topk_kernel): ~0.1-0.2 ms per query at 32K,
// once per QSA layer and token - 1-2 ms of a decode token and most of a 32K prompt's time. Here:
//
// qsa_block_scores : one warp per (query, block), FP32; relu per indexer head, summed; block n_bid (the
// incomplete tail) scores the `dead` key and gets +1e9 when it has cells - which is exactly
// what the per-token path reads there, because pooled[n_bid] holds `dead` at that time (in
// a chunk it may already hold a block completed later, so it is not read).
// qsa_block_topk : one block per query; a 4-pass radix select over the query's n_bid + 1 blocks, each
// weighted by its cell count, then the cells emitted in ascending order with ties to the
// lowest index - the same selection as topk_kernel, over a quarter of the elements.
//
// Queries carry their own step record (pos, n_kv, n_bid, width) as everywhere else in QSA.
#pragma once
#include "strata/kernels/qsa.hpp"
#include <cstdint>
namespace strata::kernels {
/// scores [nq, max_blocks]; q_idx [nq, idx_n_head, idx_dim] (normed and rotated); steps [nq, kStepCount].
void qsa_block_scores(const float* pooled, const float* dead, const float* q_idx, const int32_t* steps, int64_t nq,
int64_t max_blocks, const QsaShapes& s, float* scores, void* stream);
/// ids [nq, cap] (cells, ascending); `cap` >= the largest selection width.
void qsa_block_topk(const float* scores, const int32_t* steps, int64_t nq, int64_t max_blocks, int64_t cap,
const QsaShapes& s, int32_t* ids, void* stream);
} // namespace strata::kernels
include/strata/kernels/quantize_act.hpp
+58
View File
@@ -0,0 +1,58 @@
// include/strata/kernels/quantize_act.hpp - activation quantization, host-callable (P2.S2).
//
// `ggml_mul_mat` converts src1 (the ACTIVATION) to the weight's `vec_dot_type` before the dot product. The
// CPU expert path already does this (`bench/micro/cpu_s2.cpp`, int8 + vpdpbusd, parity 1.461e-06); the GPU
// kernels take FP16 activations. Leaving the two inconsistent means the same token computes different
// numbers for different experts, so this is the GPU half of the same rule.
//
// Q2_0 -> Q8_0 is the case that matters most: the routed experts are all Q2_0, which is 31.64 GiB of the 38
// GiB pack. Q8_K is the OTHER HALF of `VEC_DOT_TYPE` (Q3_K/Q4_K/Q5_K/Q6_K/IQ4_XS) and covers 2.89 GiB of
// dense weights - the attention projections and `ssm_out`, i.e. the numerically sensitive ones.
//
// `docs/activation-contract.md` records the decision that BOTH must exist and be used: an FP16 activation
// against a quantized weight disagrees with the oracle by 0.6-0.9%, which is 6-9x P2.S2's 1e-3 tolerance.
#pragma once
#include <cstdint>
namespace strata::kernels {
// x (n floats) -> ggml's block_q8_0 layout: 34 bytes per 32 elements, { fp16 d ; int8 qs[32] }.
// n must be a multiple of 32. `stream` may be null, in which case the call synchronises.
void quantize_q8_0(const float* x, uint8_t* blocks, int64_t n, void* stream);
// **R4.2h: THE SAME BLOCKS, PLUS THE CPU's fp32 SCALE - and the CPU's rounding rule with it.** The engine
// computes each routed expert twice (misses on the CPU through `act_quant_q8_1`, hits on the GPU through
// `quantize_q8_0`) and the two disagreed: `ActQ::scale` is fp32 while a `block_q8_0` stores fp16 `d`, which is
// **4.761e-04 relative on 80 of 80 chunks** (`bench/micro/act_quant_parity.cu`). This variant exists so the
// hit path can use the CPU's multiplier. Writes `blocks[n/32 * 34]` and `scales[n/32]`; `scales` must not be
// null. `quantize_q8_0` is unchanged and still matches ggml's bytes, which is what `moe_hit_parity` checks.
void quantize_q8_0_scaled(const float* x, uint8_t* blocks, float* scales, int64_t n, void* stream);
// The inverse, for round-trip checks: each element becomes `q * d16`.
void dequant_q8_0(const uint8_t* blocks, float* x, int64_t n, void* stream);
// x (n floats) -> ggml's block_q8_K layout: 292 bytes per 256 elements
// { float d ; int8_t qs[256] ; int16_t bsums[16] }
// n must be a multiple of 256. `stream` may be null, in which case the call synchronises.
//
// This is `quantize_row_q8_K_ref` (ggml-quants.c L2768) transcribed, including the two things that look like
// mistakes and are not:
//
// * the scale is `-127/max`, NOT `-128/max`. The source carries the `-128` version COMMENTED OUT with the
// note that IQ2_XXS needs the change for an awkward AVX implementation. Getting this wrong is a 0.79%
// scaling error - the same order as the quantisation step, so it yields a working-looking activation that
// ggml never sees. It cost round 120.
// * `max` is the SIGNED value at the position of the largest magnitude, so a POSITIVE max gives a NEGATIVE
// iscale and the `+max` element maps to -127 while the `-max` element maps to +127.
//
// The rounding is `nearest_int`, which is round-half-to-EVEN via the 12582912.0f magic number - not
// `roundf`'s half-away-from-zero. They differ only on exact ties, i.e. almost never on real data, which is
// precisely why the wrong one would survive any amount of end-to-end testing; the parity test constructs
// exact ties on purpose.
void quantize_q8_K(const float* x, uint8_t* blocks, int64_t n, void* stream);
// The inverse: each element becomes `d * qs[j]`, with `d = 1/iscale` stored as f32.
void dequant_q8_K(const uint8_t* blocks, float* x, int64_t n, void* stream);
} // namespace strata::kernels
include/strata/kernels/rope.hpp
+50
View File
@@ -0,0 +1,50 @@
// include/strata/kernels/rope.hpp - NEOX partial RoPE, host-callable (P2.S2).
//
// Pairing is NEOX: (i, i + n_rot/2), NOT the adjacent pair (2i, 2i+1). The rotation is PARTIAL - only the
// first n_rot of head_dim are touched, 64 of 256 in this artifact - and the tail passes through.
//
// The cos/sin table is built on the HOST in float64 (see src/kernels/cuda/rope.cu for why) and the kernel is
// a pure rotation over it. The table depends on (n_rot, theta, position) and not on the token, so it is worth
// caching; `build_rope_table` fills `max_pos` positions of `n_rot/2` pairs each.
#pragma once
#include <cstdint>
#if defined(__CUDACC__)
#define STRATA_ROPE_HD __host__ __device__
#else
#define STRATA_ROPE_HD
#endif
namespace strata::kernels {
/// ONE NEOX PAIR: `(a, b) -> (a*c - b*s, a*s + b*c)`.
///
/// This exists so the pairing and the sign convention live in ONE place. `rope_neox_apply` walks a whole row
/// with it, and `indexer_key_append`'s pooling kernel uses it to rotate the single pooled row it has just
/// written - two call sites that MUST agree, because the indexer's key is compared against query scores
/// downstream and a mismatched rotation produces a well-formed key pointing the wrong way.
///
/// It is `__host__ __device__` (the `f16_bits.hpp` pattern) rather than device-only so a host-side reference in
/// a test can call THE SAME function instead of transcribing it. A transcription is exactly how the pairing
/// convention gets written down wrong.
///
/// In-place is safe: the caller writes both outputs from the two inputs it read, and under NEOX no two pairs
/// share an element (`i` pairs with `i + n_rot/2`, so the halves are disjoint).
STRATA_ROPE_HD inline void rope_neox_pair(float a, float b, float c, float s, float& oa, float& ob) {
oa = a * c - b * s;
ob = a * s + b * c;
}
void build_rope_table(int n_rot, double theta, int max_pos, float* cos_tab, float* sin_tab);
// `x` and `out` are (rows, head_dim) and `pos` is (rows,) - one position per row, so a batch of heads at
// different sequence positions is one call. `cos_tab`/`sin_tab` are (max_pos, n_rot/2). `out` may equal `x`.
//
// **`pos` IS A DEVICE ARRAY, NOT A HOST SCALAR**, and that is a graph-capture requirement rather than a style
// choice: a CUDA graph bakes kernel arguments in at capture time, so a host scalar position would replay the
// first token's value forever. See the same note on `indexer_key_append` in `qsa.hpp`.
void rope_neox_apply(const float* x, float* out, int64_t rows, int head_dim, int n_rot, const float* cos_tab,
const float* sin_tab, const int* pos, void* stream);
} // namespace strata::kernels
include/strata/kernels/router_top10.hpp
+20
View File
@@ -0,0 +1,20 @@
// include/strata/kernels/router_top10.hpp - the MoE router, host-callable (P2.S2).
//
// logits (n_tokens, n_expert) -> the top-k experts per token with renormalised weights. Semantics are
// `ref/moe.py::router`'s: softmax over ALL experts, stable descending argsort (ties break by index), gather,
// then renormalise with ggml's 2**-14 lower clamp.
//
// The BF16 GEMV that PRODUCES the logits is a separate kernel and is not here; this handles the routing only,
// which is why `logits` is taken as an f32 buffer the caller has already filled.
#pragma once
#include <cstdint>
namespace strata::kernels {
// `stream` may be null, in which case the call synchronises before returning. `ids` is (n_tokens, k) int32 and
// `weights` is (n_tokens, k) f32, both device pointers. k must be <= 64.
void router_top10(const float* logits, int n_tokens, int n_expert, int k, int* ids, float* weights,
void* stream);
} // namespace strata::kernels
include/strata/kernels/s2_expert_grouped.hpp
+118
View File
@@ -0,0 +1,118 @@
// include/strata/kernels/s2_expert_grouped.hpp - R4's GPU expert tier.
//
// **WHAT THIS IS FOR.** The CPU expert pool is the engine's largest single cost - 663.6 MB of expert bytes per
// token at ~40 GB/s = 16.2 ms on a ~53 ms token - and `test_expert_pool` measures the pipeline hiding
// **1.055 ms of 19.035**. The CPU work is 96% exposed because the residual chain is strictly serial, so the
// only way past it is to stop reading those bytes on the CPU: put experts in VRAM and compute them on the GPU.
//
// **WHY GROUPED, AND WHY THAT WORD IS THE WHOLE DESIGN.** A layer routes ten experts. The obvious
// implementation is a loop over hits calling the per-matrix kernel, and it is wrong by arithmetic: five
// launches per expert at ~6.45 hits and 48 layers is **~1,550 launches per token**, which at a real-stream
// launch cost of ~3.6 us is 5.6 ms - most of the 10.4 ms the tier is supposed to save. So the hits are grouped
// into ONE launch per projection:
//
// 1. `gu_kernel` n_hits x 2 x FF rows - gate and up, together
// 2. `swiglu_kernel` n_hits x FF - silu(gate) * up
// 3. `quantize_q8_0` (the existing one) - the down weight's own activation contract
// 4. `down_kernel` n_hits x H rows
//
// **four launches per layer, not five per expert.**
//
// `h_layer` held-out is **0.0456** - only 4.6% of (layer, token) pairs have all ten experts resident - so this
// cannot be a per-layer grouped kernel with a fixed grid, and the hit list is passed **as data** so the caller
// can keep one captured graph per layer with the maximum grid and a device-side count.
#pragma once
#include <cstdint>
namespace strata::kernels {
/// Bytes of caller-owned scratch `moe_hit_grouped_s2` needs for `n_hits` experts.
///
/// Two regions: the gate/up activations (`n_hits x 2 x FF` floats) and the quantized intermediate
/// (`n_hits x (FF/32) x 34` bytes). Carved from the caller's arena rather than allocated - a `cudaMalloc`
/// on the token path is forbidden by P2.T10 and illegal during stream capture, and this is called from both.
uint64_t moe_hit_grouped_scratch_bytes(int64_t n_hits, int64_t n_embd, int64_t n_ff);
/// **ONE GROUPED EXPERT EVALUATION.** For each hit `h` in `[0, n_hits)`, evaluates the expert held in
/// `blob_base + slot_index[h] * blob_bytes` against `x_q8_0` and writes `n_embd` floats to
/// `out + dst_index[h] * n_embd`.
///
/// * `blob_base` the device slot arena; slot `s` starts at `blob_base + s * blob_bytes`
/// * `slot_index` `n_hits` `int32_t` slot indices on the DEVICE, so the caller can change which experts are
/// resident without recapturing anything
/// * **`dst_index` where each hit's answer goes, also on the device and also `n_hits` long.**
///
/// **THIS IS NOT A CONVENIENCE - IT IS WHAT LETS THE TWO HALVES MEET.** A layer routes ten
/// experts and only some are resident, so the CPU produces the misses at their ROUTED indices
/// (the router's order is what `moe_combine` weights against) and the GPU produces the hits in
/// whatever order the cache happens to hold them. Writing the hits sequentially would land
/// them on the misses' rows; `dst_index` puts hit `h`'s answer on the row its expert was routed
/// to, so the CPU can zero the hit rows, the GPU can fill them, and `moe_combine` needs no
/// change at all.
/// * `x_q8_0` `(n_embd/32)` `block_q8_0`, 34 bytes each, from `quantize_q8_0`
/// * `scratch` `moe_hit_grouped_scratch_bytes(n_hits, n_embd, n_ff)`, caller-owned
/// * `out` `(k, n_embd)` floats - the SAME buffer the misses were written into
/// * **`x_scales` R4.2h, OPTIONAL, and null keeps the previous behaviour exactly.** When non-null it is
/// `n_embd/32` **fp32** activation scales from `quantize_q8_0_scaled`, and they REPLACE the
/// block's fp16 `d` as the activation multiplier - for the gate/up projection and for the
/// intermediate's own quantization in step 3.
///
/// **THIS EXISTS BECAUSE A HIT AND A MISS MUST PRODUCE THE SAME NUMBER.** The CPU pool
/// multiplies by the fp32 `ActQ::scale` (`cpu/expert.cpp:92`); this kernel read fp16 `d` out
/// of the `block_q8_0`. `bench/micro/act_quant_parity.cu` measured the gap with both real
/// implementations linked: **80 of 80 chunks differ, max relative 4.761e-04** - 2^-11, pure
/// fp16 rounding - which is why enabling the cache changed the generated tokens. Pass null
/// and this is bit-for-bit the kernel `moe_hit_parity` has always checked at 6.279e-08.
///
/// **THE BLOB LAYOUT IS `expert.hpp`'s AND IS NOT RESTATED HERE.** Gate/up codes are interleaved by row
/// (`2r` = gate, `2r+1` = up), the down codes follow, and each plane's scales are fp16 - which is the ONE
/// difference from `s_gemv_q8`'s S-forms, whose scales are fp32. That difference is why this is a separate
/// kernel rather than a call into the dense path.
void moe_hit_grouped_s2(const uint8_t* blob_base, const int32_t* slot_index, const int32_t* dst_index,
int64_t n_hits, int64_t blob_bytes, const uint8_t* x_q8_0, void* scratch, float* out,
void* stream, const float* x_scales = nullptr);
/// Experimental CPU-order implementation. Eight CUDA lanes reproduce the CPU VNNI accumulator lanes,
/// fused multiply-add order, separate correction, and final horizontal reduction. The input must use
/// the CPU quantization contract and fp32 scales. SiLU uses accurate CUDA expf, which can differ from
/// the CPU math library by last bits. End-to-end correctness must be checked against the C1 oracle.
/// Validate both intermediate quantization and final outputs.
/// Optional gate_up_trace receives [all gates | all ups] before SiLU, for parity diagnostics only.
/// Plan v0.3 P4 token graph: the hit path with NO host step. `moe_hit_select` builds this layer's hit list
/// from the routed `ids` and the static residency row (`res_row[e]` = slot or -1) into `slot`/`dst`/`count`
/// (device); `moe_hit_grouped_s2_dev` is `moe_hit_grouped_s2` sized for `cap` hits that reads the real count
/// from `d_count`; `moe_hit_add` adds each hit's row of `hit_out` into `parts` (rows the CPU left at zero).
void moe_hit_select(const int32_t* ids, const int32_t* res_row, int k, int n_expert, int32_t* slot, int32_t* dst,
int32_t* count, void* stream);
void moe_hit_grouped_s2_dev(const uint8_t* blob_base, const int32_t* slot_index, const int32_t* dst_index,
const int32_t* d_count, int64_t cap, int64_t blob_bytes, const uint8_t* x_q8_0,
void* scratch, float* out, void* stream, const float* x_scales);
void moe_hit_add(float* parts, const float* hit_out, const int32_t* dst, const int32_t* count, int64_t cap,
int64_t n_embd, void* stream);
/// Plan v0.3 P6 verify window: `moe_hit_select` over `n` <= 128 routed entries (T tokens x k, flattened), and the
/// hit kernel with one activation PER TOKEN - entry `dst` reads token `dst / k_per_token`'s rows of `x_q8_0`
/// ((n_embd/32)*34 bytes each) and `x_scales` (n_embd/32 floats each). Per hit, bitwise `moe_hit_grouped_s2_dev`.
/// Plan v0.3 P6: routed experts GROUPED. Group g's blob is at device address grp_ptr[g] (a VRAM slot, or a mapped
/// host blob the kernel reads over PCIe); its entries are [grp_start[g], grp_start[g+1]) with ent_tok (the token
/// whose activation row of x_q8_0/x_scales it reads) and ent_dst (its row of `out`). Each row of a blob is read
/// once for all of its entries. Counts are read from device memory; every entry is bitwise the per-entry kernel.
void moe_grouped_s2(const unsigned long long* grp_ptr, const int32_t* grp_start, const int32_t* n_groups,
const int32_t* ent_dst, const int32_t* ent_tok, int64_t cap_groups, int64_t cap_entries,
const uint8_t* x_q8_0, const float* x_scales, void* scratch, float* out, void* stream);
/// Plan v0.3 P6: the groups for `moe_grouped_s2` when every expert id is resident at base + id * blob (the MTP
/// layer's 512 experts): counts[0] groups, counts[1] entries; entry dst = routing index, tok = index / k_per_tok.
void moe_group_resident(const int32_t* ids, int n, int k_per_tok, const uint8_t* base, int64_t blob,
unsigned long long* grp_ptr, int32_t* grp_start, int32_t* counts, int32_t* ent_dst,
int32_t* ent_tok, void* stream);
void moe_hit_select_multi(const int32_t* ids, const int32_t* res_row, int n, int n_expert, int32_t* slot, int32_t* dst,
int32_t* count, void* stream);
void moe_hit_grouped_s2_multi(const uint8_t* blob_base, const int32_t* slot_index, const int32_t* dst_index,
const int32_t* d_count, int64_t cap, int64_t blob_bytes, const uint8_t* x_q8_0,
const float* x_scales, int k_per_token, void* scratch, float* out, void* stream);
void moe_hit_grouped_s2_cpu_order(const uint8_t* blob_base, const int32_t* slot_index,
const int32_t* dst_index, int64_t n_hits, int64_t blob_bytes,
const uint8_t* x_q8_0, void* scratch, float* out, void* stream,
const float* x_scales, float* gate_up_trace = nullptr);
} // namespace strata::kernels
include/strata/kernels/s2_gemv.hpp
+25
View File
@@ -0,0 +1,25 @@
// include/strata/kernels/s2_gemv.hpp - the S2 GEMV, host-callable (P2.S2).
//
// The expert matvec is the hottest thing in the engine: every token runs 48 layers x 10 experts x 3 roles of
// it, against weights that live in DRAM at 31.64 GiB. Layout is `docs/pack-format.md`'s S2: codes packed 4
// per byte LSB-first, one FP32 scale per 64 elements. GGUF's ne[0] varies fastest, so a tensor is [n_in,
// n_out] with n_in CONTIGUOUS - which is why one thread per output row reads a contiguous run of blocks.
//
// Activations are FP16 per the phase spec ("FP16 activations in"). The OUTPUT here is FP32 rather than FP16:
// the intended FP16 out needs an f32->fp16 conversion, and rather than write one under a parity test that
// would then be testing two new things at once, this step keeps the output width the one that can be checked
// exactly. The FP16 output and its 1e-3 tolerance is the next increment, and this file says so rather than
// leaving a reader to wonder.
#pragma once
#include <cstdint>
namespace strata::kernels {
// y[o] = sum_i fp16(x[i]) * W[i][o], with W in S2 canonical planes.
// `codes` is n_out * (n_in/64) * 16 bytes, `scales` n_out * (n_in/64) floats, `x` n_in uint16 fp16 patterns,
// `y` n_out floats. All device pointers. n_in must be a multiple of 64.
void s2_gemv(const uint16_t* x, const uint8_t* codes, const float* scales, float* y, int64_t n_in,
int64_t n_out);
} // namespace strata::kernels
include/strata/kernels/s2_gemv_q8.hpp
+22
View File
@@ -0,0 +1,22 @@
// include/strata/kernels/s2_gemv_q8.hpp - the S2 GEMV over Q8_0-QUANTIZED activations.
//
// `ggml_mul_mat` converts src1 to the weight's `vec_dot_type`, which for Q2_0 is Q8_0. The other CUDA kernels
// here take FP16 activations, and the CPU expert path already quantizes to int8 - so this is the missing GPU
// half of the same rule, and without it the two paths disagree on the same token.
//
// CORRECTNESS VERSION: the activation is dequantized inside the loop. The speed win here is `__dp4a` (four
// int8 MACs per instruction) and belongs to Phase 3, once the numerics are settled.
#pragma once
#include <cstdint>
namespace strata::kernels {
// `act` (n_in/32) ggml Q8_0 blocks, 34 bytes each, as `quantize_q8_0` produces
// `codes` n_out * (n_in/4) bytes, one 2-bit code byte per quad
// `scales` n_out * (n_in/64)
// `y` n_out floats
void s2_gemv_q8(const uint8_t* act, const uint8_t* codes, const float* scales, float* y, int64_t n_in,
int64_t n_out, int threads_per_row, void* stream);
} // namespace strata::kernels
include/strata/kernels/s_gemv.hpp
+133
View File
@@ -0,0 +1,133 @@
// include/strata/kernels/s_gemv.hpp - the S-family GEMV, host-callable (P2.S2).
//
// ONE kernel for S2, S4 and S8, because `docs/pack-format.md` gives them ONE decode:
//
// value = cb[code] * scale + offset
//
// with a per-TYPE codebook, code bias, group size and offset presence. That uniformity is the point of the
// canonical form - 13 source types collapse to three code widths and a handful of attributes - so the kernel
// takes those attributes as ARGUMENTS rather than being specialised per source type. The manifest already
// carries them per tensor (round 163 added the tests that keep it honest).
#pragma once
#include <cstdint>
namespace strata::kernels {
// Which table turns a code into a number. `Affine` is `code + bias`; `Iq4Nl` is the non-linear
// `kvalues_iq4nl`, which no bias can express - the reason the canonical form carries a codebook at all.
enum class Codebook : int { Affine = 0, Iq4Nl = 1 };
struct SForm {
int code_bits = 2; // 2, 4 or 8
int code_bias = -1; // subtracted from the code IN THE INTEGER DOMAIN before the multiply
int group_elems = 64; // elements per scale (and per offset, when present)
Codebook codebook = Codebook::Affine;
bool has_offset = false; // only the K-quants with a min (Q4_K, Q5_K) have one
/// WHICH ACTIVATION THIS WEIGHT'S `vec_dot_type` IS: 0 = Q8_0, 1 = Q8_K.
///
/// **IT IS AN ATTRIBUTE OF THE CANONICAL FORM AND IT CANNOT BE DERIVED FROM THE OTHERS.** Q5_0 and Q5_K
/// are both 8-bit with bias -16 (only `has_offset` differs); IQ4_NL and IQ4_XS are both 4-bit with the
/// IQ4NL codebook and no offset, and differ in NOTHING this struct otherwise holds. So it is CARRIED from
/// the manifest's `source_type` rather than reconstructed - see `docs/activation-contract.md` and the
/// `act_kind` column `tools/pack_index.py` writes.
int act_kind = 0;
};
// y[o] = sum_i fp16(x[i]) * W[i][o]. GGUF's ne[0] varies fastest, so a tensor is [n_in, n_out] with n_in
// CONTIGUOUS: row o is a contiguous run of n_in/group_elems groups.
//
// `codes` n_out * (n_in / (8/code_bits)) bytes
// `scales` n_out * (n_in / group_elems) floats
// `offset` the same count as scales, or nullptr when `has_offset` is false
// `y` n_out floats
void s_gemv(const uint16_t* x, const uint8_t* codes, const float* scales, const float* offset, float* y,
int64_t n_in, int64_t n_out, const SForm& form);
// The same computation with the ROW SPLIT ACROSS THREADS: one block of `threads_per_row` threads per output
// row, each accumulating a strided subset, reduced through shared memory.
//
// WHY BOTH EXIST. `s_gemv` gives one thread per output row, so its parallelism is the OUTPUT WIDTH - and the
// expert gate/up shape is [2560 x 640], which is 640 threads over 48 SMs, 13 per SM. Measured, that
// orientation ran 4.6x slower than [640 x 2560] with identical weights (Memory/LEDGER.md L1). This variant
// makes the parallelism independent of the output width.
//
// It is NOT bit-identical to `s_gemv`: the partial sums are added in a different order, so the parity test
// carries a tolerance for exactly this reason. Both are kept and both are tested, because the naive one is
// the reference the split one is checked against.
void s_gemv_split(const uint16_t* x, const uint8_t* codes, const float* scales, const float* offset,
float* y, int64_t n_in, int64_t n_out, const SForm& form, int threads_per_row);
// S2 ONLY, with the code load amortised over four consecutive elements (they share one byte) and the
// activation load over four halves (one 64-bit word). S2 is 31.64 GiB of the 38 GiB pack, so it is the kernel
// whose instruction count matters. No `offset` - S2 has none, and a parameter that is always null is a
// parameter that will eventually be passed by mistake.
void s2_gemv_quads(const uint16_t* x, const uint8_t* codes, const float* scales, float* y, int64_t n_in,
int64_t n_out, int threads_per_row);
// S2 with two further instruction reductions: `x` staged in SHARED memory once per block instead of re-read
// from global by every output row, and the code unpack replaced by a 256-entry constant-memory table of
// float4 (one broadcast load per four elements instead of shifts, masks and an int-to-float convert).
// `stage_x` selects the staging so the two can be measured against each other rather than assumed.
void s2_gemv_fast(const uint16_t* x, const uint8_t* codes, const float* scales, float* y, int64_t n_in,
int64_t n_out, int threads_per_row, bool stage_x);
// The same computation as `s_gemv_split` but launched on `stream` (a cudaStream_t passed as void*, so this
// header stays free of CUDA types) and WITHOUT synchronising. It exists because the architecture's premise -
// that the expert stream overlaps with the compute that consumes it - cannot be measured through an entry
// point that waits for every stream. The synchronous ones remain the default; this is for the loop that has
// to overlap, and for the test that establishes whether overlapping is worth building.
void s_gemv_split_async(const uint16_t* x, const uint8_t* codes, const float* scales, const float* offset,
float* y, int64_t n_in, int64_t n_out, const SForm& form, int threads_per_row,
void* stream);
/// THE SAME GEMV OVER **Q8_K** ACTIVATIONS, which is the contract for every K-quant and IQ-quant weight.
///
/// `docs/activation-contract.md` settles this: `ggml_mul_mat` converts the activation to the weight's
/// `vec_dot_type`, and for Q3_K/Q4_K/Q5_K/Q6_K/IQ4_XS that type is Q8_K - NOT fp16. The two differ by
/// 0.6-0.9% per GEMV against a 1e-3 tolerance, and the K-quants are 2.89 GiB of the dense weights: the
/// attention projections and `ssm_out`, i.e. the numerically sensitive ones.
///
/// `x` is the block_q8_K layout, 292 bytes per 256 elements: `{ f32 d ; int8_t qs[256] ; int16_t bsums[16] }`.
/// `bsums` is NOT read here - it exists for ggml's AVX2 path - but the STRIDE includes it, so a buffer
/// produced by `quantize_q8_K` can be passed straight in.
///
/// The activation is dequantized on the fly (`d * qs[i]`), one multiply per element, exactly as the fp16 path
/// does one `__half2float` per element. The weight decode is the SAME code as `s_gemv` - one implementation,
/// because a second decode of the canonical form is a second thing to get wrong.
void s_gemv_q8k(const uint8_t* x_q8k, const uint8_t* codes, const float* scales, const float* offset,
float* y, int64_t n_in, int64_t n_out, const SForm& form, void* stream);
/// The row-split form, for output widths too small to fill the machine.
///
/// **THERE IS NO `threads_per_row` HERE, AND THAT IS A FIX RATHER THAN AN OVERSIGHT.** The parameter used to be
/// in this signature and the implementation threw it away - `(void) threads_per_row; // informational` - so a
/// caller reading the declaration reasonably believed it was tunable, and round 224 spent a whole sweep
/// measuring it and concluding from the flat result that the kernel was instruction-bound. The sweep could not
/// have shown anything else: **it was measuring a knob that is not connected.**
///
/// This form is WARP per output row, which is what the shape wants, and the parameter is gone from the
/// interface so that no one can believe otherwise. `s_gemv_split`/`s_gemv_split_async` DO take threads_per_row
/// and DO use it, if a caller wants to tune that shape.
void s_gemv_q8k_split(const uint8_t* x_q8k, const uint8_t* codes, const float* scales, const float* offset,
float* y, int64_t n_in, int64_t n_out, const SForm& form, void* stream);
/// THE SAME, OVER A **Q8_0** ACTIVATION - `block_q8_0`, 34 bytes per 32 elements, from `quantize_q8_0`.
///
/// **THIS IS THE KERNEL THAT WAS MISSING, AND ITS ABSENCE WAS DESCRIBED INCORRECTLY FOR SEVERAL ROUNDS.** The
/// legacy block formats (Q2_0, Q4_0, Q5_0, Q8_0, IQ4_NL) have `vec_dot_type` Q8_0, not Q8_K. Every one of them
/// in this pack except `ffn_down_shexp` is `Q2_0`, whose attributes happen to BE S2's (2 bits, group 64,
/// bias -1) - so `s2_gemv_q8`, which hardcodes those, covered them by coincidence. `ffn_down_shexp` is
/// IQ4_NL/Q4_0/Q5_0/Q8_0 in every layer with **group 32 and bias -16/0/-8**, and its `n_in` is **640**.
///
/// 640 is a multiple of 32 and NOT of 256, so **Q8_K is structurally impossible for it rather than merely
/// unimplemented** - which is what the "Q8_K is not implemented" note in `shared_expert.hpp` was really
/// pointing at.
///
/// `n_in` must be a multiple of 32. `group_elems` must be a power of two and a multiple of four, the same two
/// requirements `s_gemv_q8k_split` has and for the same reasons.
void s_gemv_q8_0_split(const uint8_t* x_q8_0, const uint8_t* codes, const float* scales, const float* offset,
float* y, int64_t n_in, int64_t n_out, const SForm& form, void* stream);
} // namespace strata::kernels
include/strata/kernels/sampler.hpp
+43
View File
@@ -0,0 +1,43 @@
// include/strata/kernels/sampler.hpp - the sampler chain, host-callable (P2.S2).
//
// penalties -> top_k -> top_p -> temperature -> pick
//
// The ORDER is specified in docs/sampling.md, transcribed from llama.cpp's own chain. Two facts there are
// easy to get backwards: temperature comes AFTER the truncation filters, and penalties come AFTER temperature.
//
// PENALTIES ARE NOT IMPLEMENTED YET: they need the token history, which is not in this signature, so a caller
// passing a non-zero `penalty_last_n` gets a loud failure rather than a silently unpenalised sample.
#pragma once
#include <cstdint>
namespace strata::kernels {
struct SamplerParams {
int top_k = 20; // 0 disables the filter
float top_p = 0.95f; // 1.0 disables the filter
float temperature = 1.0f; // <= 0 means greedy
int min_keep = 1; // top_p keeps at least this many
int penalty_last_n = 0; // 0 disables; the window over `history` to count occurrences in
float penalty_repeat = 1.0f;
float penalty_freq = 0.0f;
float penalty_present = 0.0f;
uint64_t seed = 0; // drives Philox, which is counter-based on (seed, token index)
uint64_t counter = 0; // absolute draw index of row 0; advance across decode calls
bool greedy = false;
};
// logits (n_tokens, n_vocab) -> one sampled token id per row in `out`.
//
// **`logits` AND `out` ARE DEVICE POINTERS.** This is a CUDA kernel launch, not a host function, and nothing
// in the parameter names or the types says so - `const float* logits` reads exactly like a host buffer.
// Measured: passing host memory for `logits` faults inside the kernel with an ILLEGAL MEMORY ACCESS that is
// reported by whatever synchronising call happens NEXT, which will be somewhere else entirely and will name a
// buffer that has nothing to do with it. A sticky async fault does not know where it came from.
//
// `history` is (n_tokens, history_len) int32, the most recent tokens for each row with any unused slots set to
// -1; only the last `p.penalty_last_n` of each row are counted. Pass nullptr and 0 when no penalties apply.
void sample_tokens(const float* logits, int n_tokens, int n_vocab, const int* history, int history_len,
const SamplerParams& p, int* out, void* stream);
} // namespace strata::kernels
include/strata/kernels/shared_expert.hpp
+101
View File
@@ -0,0 +1,101 @@
// include/strata/kernels/shared_expert.hpp - the shared expert, host-callable (P2.S2).
//
// h = silu(x @ gate_shexp.T) * (x @ up_shexp.T) <- SILU ON GATE
// h = h @ down_shexp.T
// g = sigmoid(x @ gate_inp_shexp) one SCALAR per token
// return h * g[:, None]
//
// The result is ADDED to the routed output, not weighted against it. Shape note: `n_ff` for the shared expert
// matches the routed experts' width (640) in this artifact.
//
// The SForm arguments describe canonical planes for the three expert projections and their historical
// CPU-derived activation images (Q8_0 or Q8_K). The default scalar gate separately consumes BF16 activations.
// The opt-in native BF16 gate instead consumes the original F32 activation and uses the pinned CUDA MMVF
// contract. Optional native projection weights select the pinned CUDA Q8_1 MMVQ contract independently.
#pragma once
#include <cstdint>
#include "strata/kernels/s_gemv.hpp"
namespace strata::kernels {
/// Opt in to pinned CUDA BF16/F32 MMVF and FP32 sigmoid for the shared scalar gate (default false).
/// Configure before session capture; captured graphs retain their selected kernels.
/// In this mode shared_expert requires its optional unrounded x_f32 input.
void shared_expert_set_native_bf16(bool enabled);
/// Optional native GGUF projections. Each supported type with nonnull data
/// replaces only that canonical projection; absent or unsupported entries fall
/// back independently. Any active entry selects pinned FP32 SwiGLU arithmetic.
/// Caller owns device weights and Q8_1 scratch for the largest active input
/// (n_embd for gate/up, n_ff for down), sized by native_q8_1_bytes. All storage
/// must survive graph execution; use one ordered nonnull stream with no overlap.
struct NativeSharedWeights {
int gate_type = -1, up_type = -1, down_type = -1;
const void* gate_data = nullptr;
const void* up_data = nullptr;
const void* down_data = nullptr;
void* q8_1 = nullptr;
};
/// Bytes of caller-owned scratch `shared_expert` needs. **THE KERNEL USED TO `cudaMalloc` FOUR BUFFERS ON
/// EVERY CALL**, which is two separate violations in one line: `cudaMalloc` is ILLEGAL INSIDE A CAPTURE
/// (`cudaErrorStreamCaptureUnsupported`), so the MoE - and therefore the whole per-layer graph - could not be
/// captured at all; and P2.T10 requires ZERO token-path allocations, which this was. Neither is visible in a
/// test that only checks the output, and the capture test found it in one run.
uint64_t shared_expert_scratch_bytes(int64_t n_ff);
/// `x_q8_0` and `x_q8k` are the TWO QUANTIZED IMAGES of the same activation, and which one a projection uses is
/// decided per weight by its `SForm::act_kind`. There is no fp16 activation here any more.
///
/// **THE fp16 PATH WAS A REAL ERROR AND IT WAS DESCRIBED WRONG FOR SEVERAL ROUNDS.** This header used to say
/// the weights "are Q3_K / IQ4_XS / Q5_0 whose `vec_dot_type` is Q8_K, which is not implemented". In fact
/// `ffn_down_shexp` is IQ4_NL/Q4_0/Q5_0/Q8_0/Q2_0 in every layer - **no K-quant at all** - so its contract is
/// Q8_0; and its `n_in` is 640, which is not a multiple of 256, so **Q8_K is structurally impossible for it**.
/// The gate/up weights are Q2_0 on 21 and 13 layers (Q8_0 again) and K-quants on the rest.
///
/// x_f32 is required for the native BF16 gate or native gate/up projections.
/// Native down reads the unrounded F32 SwiGLU intermediate directly. No allocation occurs.
///
/// Measured cost of the wrong choice: 0.66-1.41% per GEMV. `s_gemv_q8_0_split` was written for the Q8_0 half
/// (LEDGER L55) and `s2_gemv_q8` already served the Q2_0 half.
void shared_expert(const uint8_t* x_q8_0, const uint8_t* x_q8k, const uint16_t* x_bf16, const SForm& gate_form,
const uint8_t* gate_codes, const float* gate_scales, const float* gate_off,
const SForm& up_form, const uint8_t* up_codes, const float* up_scales, const float* up_off,
const SForm& down_form, const uint8_t* down_codes, const float* down_scales,
const float* down_off, const uint16_t* gate_inp_bf16, float* scratch, float* out,
int64_t n_embd, int64_t n_ff, int tpr, void* stream, const float* x_f32 = nullptr,
const NativeSharedWeights* native = nullptr);
/// Plan v0.3 P6: the shared expert for `n_tok` <= 8 tokens (a verify window) with all three projections native:
/// multi-column MMVQ, so the weights are read once; every token is bitwise `shared_expert` on that token.
/// `x` (n_tok, n_embd) f32, `x_bf16` the same rounded (only read when the native BF16 gate is off), `gate`/`up`
/// (n_tok, n_ff) scratch, `g` n_tok floats, `out` (n_tok, n_embd). `nw.q8_1` must hold n_tok columns of n_embd.
void shared_expert_multi(int n_tok, const float* x, const uint16_t* x_bf16, const NativeSharedWeights& nw,
const uint16_t* gate_inp_bf16, float* gate, float* up, float* g, float* out, int64_t n_embd,
int64_t n_ff, void* stream);
/// The MoE block's final combination, `ref/moe.py::moe` L156:
///
/// y = sum over the k routed experts of w[k] * parts[k] + shared
///
/// `parts` is (k, n_embd) contiguous - one row per SELECTED expert, in the router's own order - and `shared`
/// may be null for a layer with no shared expert.
///
/// TWO THINGS A READER GETS WRONG HERE, and they are why this is not inlined at the call site:
///
/// * The ROUTED outputs are router-WEIGHTED; the SHARED output is ADDED PLAIN. `ref/moe.py` L116 says so
/// outright - "router-weighted, not renormalised against it" - and weighting the shared term too is a
/// one-line change that produces a plausible number.
/// * The weights come from `router_top10` ALREADY renormalised (they sum to 1 over the selected experts).
/// Normalising again here would be a second, different normalisation, and it is invisible whenever the
/// router's weights happen to sum near 1 anyway.
///
/// One thread per output element, k accumulations each; k is 10 at the real geometry. This is where the CPU
/// expert pool's outputs are summed: `ExpertPool` carries a `weight` per job and deliberately does NOT apply
/// it, because the sum has to happen once over all k and not once per worker.
void moe_combine(const float* parts, const float* weights, const float* shared, float* y, int64_t n_embd,
int64_t k, void* stream);
} // namespace strata::kernels
include/strata/kernels/verify_kernels.hpp
+80
View File
@@ -0,0 +1,80 @@
// include/strata/kernels/verify_kernels.hpp - plan v0.3 P6: kernels for the speculative VERIFY window, where T
// tokens (the last accepted token and T-1 drafts) go through a layer at once.
//
// The GDN recurrence is the one piece of state that cannot simply be overwritten by the next window, so its
// kernels come in two halves:
//
// * the verify half reads the conv history and the recurrent state, runs all T tokens in order (bitwise the
// single-token kernels' arithmetic per token) and writes ONLY the per-token outputs - the state is left
// untouched;
// * the commit half, launched once acceptance is known, replays the first `n_keep` tokens from the inputs the
// verify half stored and writes the state. Rejected tokens therefore never touch the state and there is no
// snapshot of the 113 MB recurrent state to keep.
//
// Everything that varies per window (token ids, `n_keep`) is read from DEVICE memory so the kernels can be captured.
#pragma once
#include <cstdint>
namespace strata::kernels {
inline constexpr int kVerifyMaxT = 8;
/// For token t of T: conv over [history(3) | qkv_0 .. qkv_t] -> SiLU -> L2 norm of the q/k heads -> h[t].
/// `history` is NOT written. Bitwise `fused_gdn_conv_l2` per token.
void gdn_conv_l2_multi(const float* history, const float* qkv, const float* conv_w, float* h, int channels,
int qk_heads, float eps, int n_tok, void* stream, int t_begin = 0);
/// history <- the last 3 entries of [history | qkv_0 .. qkv_{n-1}], n = *n_keep (0 leaves it as it was).
void gdn_conv_commit(float* history, const float* qkv, int channels, const int32_t* n_keep, void* stream);
/// alpha/beta for T columns of x (T, n_embd): gate (T, h_v), beta (T, h_v). Bitwise `fused_gdn_ab` per column.
void gdn_ab_multi(const float* x, const uint16_t* w_alpha, const uint16_t* w_beta, const float* dt, const float* ssm_a,
float* gate, float* beta, int n_embd, int h_v, int n_tok, void* stream);
/// The recurrence + output norm for T tokens (h = (T, conv_channels) as q|k|v, gate/beta (T, h_v), z/y
/// (T, value_dim)). With `n_keep == nullptr` the state is read and NOT written (verify); otherwise the first
/// *n_keep tokens are run and the state is written (commit; `y` may be scratch). Bitwise `fused_gdn_step_norm`.
void gdn_step_norm_multi(float* state, const float* h, int conv_channels, const float* gate, const float* beta,
const float* z, const float* gamma, float eps, float* y, int h_k, int h_v, int n_tok,
const int32_t* n_keep, void* stream, int t_out_begin = 0);
/// Spin until *flag >= value (a mapped host flag). The value is fixed at capture, so several rings can be
/// outstanding at once (the split verify window keeps two).
void wait_flag_ge(const uint32_t* flag, uint32_t value, void* stream);
/// Rows of the S2/S4/S8 embedding for T token ids read from DEVICE memory; out (T, n). Bitwise `embedding_gather`.
void embedding_gather_dev(const uint8_t* codes, const float* scales, const float* offsets, const int32_t* tokens,
int n_tok, int64_t n, int code_bits, int code_bias, int group_elems, uint64_t row_codes,
uint64_t row_groups, float* out, void* stream);
/// R[t][c][:] = x[t][:] for the hc streams.
void broadcast_streams(const float* x, float* R, int64_t n_embd, int hc, int n_tok, void* stream);
/// dst[:n] = src[(*index) * stride + :n] (index read from device memory; a negative index copies nothing).
void copy_indexed(float* dst, const float* src, int64_t stride, const int32_t* index, int64_t n, void* stream);
/// Plan v0.3 P6: copy *n (device memory) blobs of `blob_bytes` from mapped host memory (src[k], device aliases)
/// into dst + k * blob_bytes with coalesced 16-byte loads - the PCIe share of a layer's missed experts, staged
/// into VRAM before the grouped expert kernel reads them. Launched for a capacity of `cap` blobs.
void fetch_blobs(const unsigned long long* src, const int32_t* n, uint8_t* dst, int64_t blob_bytes, int cap, void* stream);
/// ptr[k] = base + k * blob_bytes for k < *n (the staged copies `fetch_blobs` made).
void rebase_ptrs(unsigned long long* ptr, const int32_t* n, uint8_t* base, int64_t blob_bytes, void* stream);
// ---- the MTP draft layer (src/core/mtp.cpp)
/// R[t][c][:] = h[t][c][:] + e[t][:] (the embedding branch added to every stream).
void add_streams_broadcast(const float* h, const float* e, float* R, int64_t n_embd, int hc, int n_tok, void* stream);
/// Every routed expert is resident (slot = expert id): slot[i] = ids[i], dst[i] = i, *count = n.
void ident_hits(const int32_t* ids, int n, int32_t* slot, int32_t* dst, int32_t* count, void* stream);
/// The draft chain's next input: R_dst[:] = R_src[row], tok_dst[0] = ids[row], out[j] = ids[row], with
/// row = *row_dev (device memory). `out` may be mapped host memory.
void mtp_select(const float* R_src, int64_t R_stride, const int32_t* ids, const int32_t* row_dev, float* R_dst,
int32_t* tok_dst, int32_t* out, int j, void* stream, const float* probs = nullptr,
float* out_p = nullptr);
/// dst row i = src row ids[i] (row_bytes each, multiple of 16), for n rows.
void gather_rows(const uint8_t* src, int64_t row_bytes, const int32_t* ids, int64_t n, uint8_t* dst, void* stream);
/// ids[t] = table[ids[t]] for n entries (a subset index back to a token id).
void map_ids(int32_t* ids, const int32_t* table, int n, void* stream);
/// probs[t] = softmax(logits[t])[ids[t]] for n_rows rows of n_vocab (the probability of each row's argmax).
void row_top_prob(const float* logits, int n_rows, int n_vocab, const int32_t* ids, float* probs, void* stream);
/// Dense-attention step records for `n` cells: [cell, cell+1, (cell+1)/4, cell+1] from cells[i] (device memory).
void dense_steps(const int32_t* cells, int n, int32_t* steps, void* stream);
/// A sliding attention window: for `n` step records (kStepCount ints each, n_kv at [1]) the selection becomes the
/// last `window` cells: ids[q * ids_stride + j] = max(0, n_kv - window) + j and the record's width = the count.
void window_ids(int32_t* steps, int n, int window, int32_t* ids, int64_t ids_stride, void* stream);
} // namespace strata::kernels
include/strata/ngram/ple_reader.hpp
+86
View File
@@ -0,0 +1,86 @@
// include/strata/ngram/ple_reader.hpp - plan v0.3 P2: the n-gram table read straight from the SSD.
//
// The table is 320,001,536 rows of 90 bytes (26.8 GiB) and is never held in RAM: every row comes from an
// unbuffered 4 KiB read (platform::DirectFile). A token needs 16 rows on 16 different pages, and all 16 depend
// on the token itself, so the only time to hide them is the embedding plus layer 0. Hence the split API:
//
// Ticket t = reader.issue(rows, n, out_raw); // as soon as the token id is known
// ... // embedding, layer 0
// reader.collect(t); // before layer 1 reads the rows
//
// `issue` also serves prefill: pass all 16 x N rows of a chunk; pages are deduplicated, sorted by offset and
// kept at most `max_inflight` deep, so a chunk's reads can run while the previous chunk computes.
//
// THE ROW CACHE IS NOT THE TABLE. It keeps rows this process has already fetched (90 bytes each, bounded,
// clock eviction). Measured on the frozen corpus (bench/results/2026-09-23-ngram-io): 1M rows (~95 MB)
// would serve up to ~82% of reads; within one long prompt 20-34% of rows recur. Capacity 0 disables it.
#pragma once
#include "strata/platform/direct_file.hpp"
#include <cstdint>
#include <string>
#include <vector>
namespace strata::ngram {
inline constexpr uint32_t ROW_BYTES = 90;
inline constexpr uint32_t PAGE = 4096;
struct ReaderStats {
uint64_t requests = 0; ///< rows asked for
uint64_t cache_hits = 0; ///< rows served from the row cache
uint64_t dedup_rows = 0; ///< rows that shared a page already being read in the same ticket
uint64_t reads = 0; ///< SSD read requests issued
uint64_t bytes = 0; ///< bytes read from the SSD
double wait_us = 0; ///< time `collect` spent blocked
double submit_us = 0; ///< time spent inside the read submission call (non-zero = it blocks)
double read_us_sum = 0; ///< sum of per-read latencies (issue to completion)
uint64_t late_injected = 0; ///< reads delayed by fault injection
std::vector<float> read_us; ///< last <= 65,536 read latencies, for percentiles
double percentile(double q) const;
};
class PleReader {
public:
struct Ticket {
uint32_t id = 0;
};
PleReader();
~PleReader();
PleReader(const PleReader&) = delete;
PleReader& operator=(const PleReader&) = delete;
/// `table_offset` is the byte offset of row 0 in the file and `n_rows` the row count; both come from a
/// validated GGUF parse (PleTable::open checks the table exactly fills the file from there).
/// `io_thread` (default): a worker thread submits and reaps reads, so `issue` costs the caller no ReadFile
/// calls. false: the caller's thread does it (A/B arm).
bool open(const std::string& path, uint64_t table_offset, uint64_t n_rows, uint32_t max_inflight,
uint64_t cache_rows, std::string& err, bool io_thread = true);
void close();
bool is_open() const;
/// Start fetching `n` rows; row i's 90 raw bytes land at `out_raw + 90 * i`. `out_raw` must stay valid
/// until `collect` returns. Out-of-range rows produce 90 zero bytes (the mmap path's behaviour).
Ticket issue(const uint32_t* rows, size_t n, uint8_t* out_raw);
/// Block until every row of the ticket is in `out_raw`. Returns false on an I/O error (message in `err`).
bool collect(Ticket t, std::string& err);
/// Fault injection for tests and for the plan's P2 exit check: every read completes no earlier than
/// `delay_us` after it was issued. 0 disables.
void set_injected_delay_us(double delay_us);
/// Read with no ticket in flight (the worker updates these while reads are outstanding).
const ReaderStats& stats() const;
void reset_stats();
uint64_t cache_capacity() const;
uint64_t cache_size() const;
private:
struct Impl;
Impl* impl_ = nullptr;
};
} // namespace strata::ngram
include/strata/plan/plan.hpp
+168
View File
@@ -0,0 +1,168 @@
// include/strata/plan/plan.hpp - P1.S9 census + memory planner.
//
// This is the component that makes one binary adapt to the GPU it is running on. The engine has 33.97 GB of
// expert weights (has to live in DRAM and be streamed) and a fixed VRAM budget that has to hold the dense
// weights, the KV cache, the indexer keys, the recurrent state and as many expert-cache slots as remain.
// Every one of those competes, so the plan is an arithmetic result rather than a configuration choice, and
// it must REFUSE rather than silently overcommit.
//
// DERIVATION POLICY. `phases/00-README.md` §2 gives three rows of expected values but no formula, so the
// arithmetic here is derived from the model's geometry and the table is used as a CHECK. Where the two
// disagree the geometry wins and the disagreement is reported - a table transcribed by hand is not a
// specification, and fitting formulas to it would hide exactly the errors this is supposed to catch.
#pragma once
#include <cstdint>
#include <stdexcept>
#include <string>
#include <vector>
namespace strata::plan {
// ---- model geometry (the real Qwen3.8-Flash-Next artifact) ------------------
struct Geometry {
int n_layers = 48;
int n_qsa_layers = 12; // layers 3, 7, ... 47 - the full-attention ones
int n_kv_heads = 2;
int head_dim = 256;
int kv_group = 64; // INT8 KV is grouped-64: one fp16 scale per 64 elements
int qsa_block = 4; // the indexer pools one key per 4-token block
/// THE INDEXER'S KEY STORE, and the trap that has now been walked into twice in opposite directions.
///
/// `indexer.q_proj.weight` is [2560, 512] = **4 QUERY heads** of 128. `indexer.k_proj.weight` is
/// [2560, 128] = **ONE shared KEY head**. The queries are not cached; only the key is. So the cache
/// holds one 128-wide key per TOKEN, pooled to one per `qsa_block`-token block - which
/// `include/strata/kernels/qsa.hpp` states directly: `pooled (max_cells/idx_block + 1, idx_dim)`, and
/// its own note that the fp32 store is "50 MiB at 32K" (12 layers x 8192 blocks x 128 x 4 B = 48 MiB).
///
/// original indexer_heads(1) * head_dim(256) / 4 = 64 B/token/layer 2x too big
/// round 194 indexer_heads(4) * key_dim(128) / 4 = 128 B/token/layer 4x too big
/// correct key_heads(1) * key_dim(128) / 4 = 32 B/token/layer
///
/// Round 194 "fixed" this by reading `indexer.head_count = 4` out of the metadata and multiplying by it -
/// right about the metadata, wrong about what was being counted. **A head count in the metadata does not
/// say WHICH tensor it counts.**
int indexer_q_heads = 4; // query heads: they multiply the SCORES, and cost no cache at all
int indexer_key_heads = 1; // k_proj is [n_embd, 128]: ONE shared key head
int indexer_key_dim = 128; // the key width, and the width of the thing that IS cached
/// GDN geometry. `docs/semantics.md`: `ssm.state_size = 128`, `ssm.time_step_rank = 48` (the v heads),
/// `ssm.conv_kernel = 4`, and `n_embd_r() = (d_conv-1) * (ssm_d_inner + 2*n_group*d_state) = 3 * 10240`.
int ssm_state_size = 128;
int ssm_v_heads = 48;
int ssm_d_conv = 4;
int ssm_conv_channels = 10240;
/// 36 of the 48 layers are GDN; the rest are QSA. Declared as a derivation rather than a third constant,
/// because a `n_gdn_layers = 36` alongside `n_layers` and `n_qsa_layers` is a third thing to keep in step.
int n_gdn_layers() const { return n_layers - n_qsa_layers; }
};
// ---- byte costs ------------------------------------------------------------
struct Costs {
uint64_t expert_blob = 1382400; // one expert's gate_up+down, codes and scales (architecture §3.2)
uint64_t dense_bytes = 0; // from the manifest: every non-expert tensor, canonical form
uint64_t embd_bytes = 0; // token embedding
uint64_t mtp_bytes = 0; // multi-token-prediction head, when present
uint64_t workspace_bytes = 0;
uint64_t state_bytes = 0; // GDN recurrence + conv history + replay log, per sequence
};
// GDN recurrence + conv history, bytes per sequence. DERIVED, for the same reason everything else here is.
//
// This term was MISSING from the plan until round 198: `Costs::state_bytes` exists and `--state` exists, but
// nothing passed it, so the planner was fed zero and quietly planned without it. 36 GDN layers of fp32 state
// is **117.7 MB, or 85 expert cache slots** - not a rounding error.
//
// recurrence (S, h_v, S) fp32 = 128 * 48 * 128 * 4 = 3,145,728 B per layer
// conv history (d_conv-1, C) fp32 = 3 * 10240 * 4 = 122,880 B per layer
//
// It is a FIXED cost and not a cache: it is state, it is resident for the life of a sequence, and it cannot be
// evicted to make room for an expert. A planner that omits it returns a plan it cannot honour, which is the
// failure `DoesNotClose` exists to prevent.
inline uint64_t state_bytes(const Geometry& g) {
const uint64_t per_layer = (uint64_t) g.ssm_state_size * (uint64_t) g.ssm_v_heads *
(uint64_t) g.ssm_state_size * sizeof(float) +
(uint64_t) (g.ssm_d_conv - 1) * (uint64_t) g.ssm_conv_channels * sizeof(float);
return (uint64_t) g.n_gdn_layers() * per_layer;
}
// INT8 KV + indexer keys, bytes per token. Both terms come from the geometry above and nothing else.
inline uint64_t kv_bytes_per_token(const Geometry& g) {
const uint64_t kv_elems = 2ull * g.n_kv_heads * g.head_dim; // K and V
const uint64_t kv_scales = (kv_elems / (uint64_t)g.kv_group) * 2; // fp16 scale per group
const uint64_t per_layer = kv_elems * 1 + kv_scales; // INT8: one byte per element
// indexer: ONE cached key per block, `idx_dim` wide, SHARED across the query heads. See the note on
// `indexer_key_heads` above - this term has been wrong in both directions and the derivation is why.
const uint64_t idx_layer =
(uint64_t) g.indexer_key_heads * (uint64_t) g.indexer_key_dim / (uint64_t) g.qsa_block;
return (uint64_t)g.n_qsa_layers * (per_layer + idx_layer);
}
// ---- the plan --------------------------------------------------------------
struct Plan {
uint64_t max_context = 0;
uint64_t kv_bytes = 0; // KV + indexer keys at max_context
uint64_t state_bytes = 0; // GDN recurrence + conv history - FIXED, not evictable
uint64_t cache_bytes = 0; // what is left for expert cache slots
uint64_t cache_slots = 0;
uint64_t vram_budget = 0;
uint64_t vram_used = 0;
uint64_t dram_experts = 0; // the whole expert arena, streamed
std::vector<std::string> notes;
};
// Thrown when the requested context and the fixed costs cannot both fit. A planner that returns a plan it
// cannot honour is worse than one that fails, because the failure then happens at token 4000 instead of at
// startup.
struct DoesNotClose : std::runtime_error {
explicit DoesNotClose(const std::string& w) : std::runtime_error(w) {}
};
// The VRAM budget that the KV cache and the expert cache share. This is the one number taken from
// `00-README.md` §2 rather than derived: the table's three rows all leave `KV + cache` constant, and that
// constant is what the planner is dividing. Expressed in bytes, decimal, matching how the table rounds.
inline uint64_t vram_pool_bytes() {
return 5943000000ull;
}
inline Plan make_plan(uint64_t max_context, const Geometry& g, const Costs& c,
uint64_t pool = vram_pool_bytes()) {
Plan p;
p.max_context = max_context;
p.vram_budget = pool;
p.kv_bytes = kv_bytes_per_token(g) * max_context;
p.state_bytes = c.state_bytes;
p.dram_experts = c.dense_bytes + c.embd_bytes; // costs that never enter the pool
const uint64_t fixed = c.dense_bytes + c.embd_bytes + c.mtp_bytes + c.workspace_bytes + c.state_bytes;
if (p.kv_bytes + fixed > pool) {
throw DoesNotClose("max-context " + std::to_string(max_context) + " needs " +
std::to_string(p.kv_bytes + fixed) + " B of the " + std::to_string(pool) +
" B VRAM pool; reduce --max-context");
}
p.cache_bytes = pool - p.kv_bytes - fixed;
p.cache_slots = p.cache_bytes / c.expert_blob;
p.vram_used = pool - p.cache_bytes % c.expert_blob; // slot remainder is unusable, not free
if (p.cache_slots == 0) {
throw DoesNotClose("no room for even one expert cache slot at max-context " +
std::to_string(max_context));
}
return p;
}
inline std::string to_string(const Plan& p) {
auto gb = [](uint64_t b) { return std::to_string(b / 1e9); };
std::string s;
s += "memory plan --max-context " + std::to_string(p.max_context) + "\n";
s += " KV + indexer keys " + gb(p.kv_bytes) + " GB\n";
s += " recurrent state " + gb(p.state_bytes) + " GB (GDN recurrence + conv history, NOT evictable)\n";
s += " expert cache " + gb(p.cache_bytes) + " GB (" + std::to_string(p.cache_slots) +
" slots of 1,382,400 B)\n";
s += " VRAM pool " + gb(p.vram_budget) + " GB\n";
s += " expert arena (DRAM) " + gb(p.dram_experts) + " GB\n";
for (const std::string& n : p.notes) s += " note: " + n + "\n";
return s;
}
} // namespace strata::plan
include/strata/platform/direct_file.hpp
+59
View File
@@ -0,0 +1,59 @@
// include/strata/platform/direct_file.hpp - plan v0.3 P1/P2: unbuffered asynchronous file reads.
//
// The n-gram table (26.8 GiB) stays on the SSD and must never occupy RAM, including the OS file cache. A memory
// map cannot promise that; an unbuffered read can. On Windows this is FILE_FLAG_NO_BUFFERING | OVERLAPPED with
// an I/O completion port; on Linux, O_DIRECT (synchronous pread for now; io_uring is phase L).
//
// Contract of every read: offset, length and buffer address are multiples of `alignment()` (4096 here).
// Reads past end of file return the bytes that exist; `Completion::bytes` says how many.
#pragma once
#include <cstdint>
#include <string>
namespace strata::platform {
struct Completion {
uint64_t tag = 0; ///< the caller's tag from `submit`
uint32_t bytes = 0; ///< bytes transferred (short only at end of file)
bool ok = false;
};
class DirectFile {
public:
DirectFile();
~DirectFile();
DirectFile(const DirectFile&) = delete;
DirectFile& operator=(const DirectFile&) = delete;
bool open(const std::string& path, std::string& err);
void close();
bool is_open() const;
uint64_t size() const;
static constexpr uint32_t alignment() { return 4096; }
/// Queue one read. Returns false (and sets `err`) if the request could not be queued; a queued request
/// always produces exactly one completion.
bool submit(uint64_t offset, void* buffer, uint32_t length, uint64_t tag, std::string& err);
/// Wait for up to `max` completions; returns how many were written to `out`. `timeout_ms` < 0 waits
/// forever; 0 polls. A `wake()` shows up as one completion with tag `WAKE_TAG`.
int wait(Completion* out, int max, int timeout_ms);
/// Make a thread blocked in `wait` return (thread-safe). Used to hand it newly queued work.
static constexpr uint64_t WAKE_TAG = ~0ull;
void wake();
/// Allocation helpers for aligned read buffers (page-aligned, never touched by the cache manager).
static void* alloc_aligned(size_t bytes);
static void free_aligned(void* p);
private:
struct Impl;
Impl* impl_ = nullptr;
};
/// Monotonic time in microseconds, for latency accounting.
double now_us();
} // namespace strata::platform
include/strata/platform/memory.hpp
+26
View File
@@ -0,0 +1,26 @@
// include/strata/platform/memory.hpp - plan v0.3 P0.1/P1: keep a large host region resident.
//
// `cudaHostRegister` refuses the 31.6 GiB expert arena on Windows, and an unlocked arena is trimmed under memory
// pressure (the CPU pool then swung 2x between runs). With the n-gram table out of RAM there is headroom to lock
// it instead: raise the process's minimum working set by the region's size, then VirtualLock it (Windows needs
// only SeIncreaseWorkingSetPrivilege, which ordinary accounts hold). Linux: mlock.
#pragma once
#include <cstdint>
#include <string>
namespace strata::platform {
struct LockResult {
bool ok = false;
uint64_t locked_bytes = 0; ///< may be less than requested; the rest stays pageable
std::string note; ///< what was done or why it failed, for the startup print
};
/// Lock [p, p + bytes) into physical memory. Partial success is reported, not hidden.
LockResult lock_resident(void* p, uint64_t bytes);
/// Undo lock_resident for the same region (best effort).
void unlock_resident(void* p, uint64_t bytes);
} // namespace strata::platform
include/strata/prefill/gemm.hpp
+54
View File
@@ -0,0 +1,54 @@
// include/strata/prefill/gemm.hpp - plan v0.3 P5: the batched projections of prompt processing.
//
// Every projection of a chunk of T tokens is Y[T, N] = X[T, K] . W[N, K]^T with W row-major (the GGUF / pack layout)
// and FP32 outputs. Weights are BF16 on the device - either already (the pack's BF16 tensors) or dequantized from
// their native GGUF blocks into a reusable scratch (`dequant_bf16`) right before the product - and activations are
// rounded to BF16, which is also what llama.cpp's batched CUDA path does. Tensor-core GEMM through cuBLAS.
#pragma once
#include <cstdint>
#include <string>
namespace strata::prefill {
class Gemm {
public:
Gemm() = default;
~Gemm();
Gemm(const Gemm&) = delete;
Gemm& operator=(const Gemm&) = delete;
/// `scratch_elems`: BF16 elements of the dequantization scratch (the largest weight dequantized at once).
bool init(void* stream, int64_t scratch_elems, std::string& err);
/// The same with caller-owned device buffers (the prompt path borrowing expert-cache slots).
bool init_external(void* stream, uint16_t* scratch, int64_t scratch_elems, void* workspace, size_t ws_bytes,
std::string& err);
/// Y[T, N] (fp32, row stride ldy) = X[T, K] (bf16, row-major) . W[N, K]^T (bf16, row-major). `beta` = 1 adds.
void bf16(const uint16_t* X, const uint16_t* W, float* Y, int64_t T, int64_t N, int64_t K, int64_t ldy = 0,
float beta = 0.0f);
/// Y = X . W^T with both in FP16 (bits).
void f16(const uint16_t* X, const uint16_t* W, float* Y, int64_t T, int64_t N, int64_t K, int64_t ldy = 0,
float beta = 0.0f);
/// W given as native GGUF blocks of `ggml_type`, dequantized to FP16 in the scratch, X in FP16.
void native(const uint16_t* X, int ggml_type, const void* W_blocks, float* Y, int64_t T, int64_t N, int64_t K,
int64_t ldy = 0, float beta = 0.0f);
uint16_t* scratch() const { return scratch_; }
int64_t scratch_elems() const { return scratch_elems_; }
void* stream() const { return stream_; }
private:
void* handle_ = nullptr;
void* stream_ = nullptr;
uint16_t* scratch_ = nullptr;
int64_t scratch_elems_ = 0;
void* workspace_ = nullptr;
bool external_ = false;
};
} // namespace strata::prefill
include/strata/prefill/kernels.hpp
+74
View File
@@ -0,0 +1,74 @@
// include/strata/prefill/kernels.hpp - plan v0.3 P5: the element-wise and stateful kernels of a prompt chunk.
//
// T tokens at a time, all row-major with the token as the slow index. The arithmetic mirrors the per-token native
// kernels (the decode path), batched; where a step is a recurrence in time (the GDN conv and state, the PLE conv)
// the kernel walks the chunk's tokens in order inside one launch.
#pragma once
#include <cstdint>
namespace strata::prefill {
// ---- hyper-connection (n_embd 2560, hc 4, hc_lr 320)
/// xn[t, c*2560 + d] = R[t,c,d] * rsqrt(mean_d R[t,c,:]^2 + eps) * w_norm[c*2560 + d]; also its BF16 image.
void gr_norm(const float* R, const float* w_norm, float eps, float* xn, uint16_t* xn16, int64_t T, void* stream);
/// lo16[t, k] = bf16(silu(lo[t, k] / hc))
void gr_silu(const float* lo, uint16_t* lo16, int64_t T, void* stream);
/// mixed[t, d] = mean_c xn[t, c, d] * sigmoid(gated[t, c, d]); FP32, BF16 and FP16 (either image may be null).
void gr_mix(const float* xn, const float* gated, float* mixed, uint16_t* mixed16, int64_t T, void* stream,
uint16_t* mixed_h = nullptr);
/// R[t, c, d] += bo[t, d] * 2 sigmoid(inj[t, c] / hc) (inj has row stride inj_ld)
void gr_write(float* R, const float* bo, const float* inj, int64_t inj_ld, int64_t T, void* stream);
/// R[t, c, :] = e[t, :] for all four streams (the embedding broadcast).
void gr_broadcast(const float* e, float* R, int64_t T, void* stream);
// ---- GDN (state 128, 16 k heads, 48 v heads, 10240 conv channels, 4 taps)
/// gate[t,h] = softplus(ab[t,h] + dt[h]) * ssm_a[h]; beta[t,h] = sigmoid(ab[t, 48 + h]) (ab: [T, 96])
void gdn_gates(const float* ab, const float* dt, const float* ssm_a, float* gate, float* beta, int64_t T, void* stream);
/// The 4-tap causal conv + SiLU over the chunk (history [C][3] in, updated to the chunk's last three inputs), then
/// the L2 norm of the q and k heads of every token. h: [T, C].
void gdn_conv(float* history, const float* qkv, const float* conv_w, float* h, int64_t T, float eps, void* stream);
/// The recurrence over the chunk, block per value head, state in registers; y[t] = rmsnorm(o) * gamma * sigmoid(z)
/// (FP32 and FP16 bits: the out projection is quantized).
void gdn_recurrence(float* state, const float* h, const float* gate, const float* beta, const float* z,
const float* gamma, float eps, float* y, uint16_t* y16, int64_t T, void* stream);
// ---- MoE
/// softmax over 512, top-10 (ties to the lower id), weights renormalised over the ten (the native router).
void route(const float* logits, int32_t* ids, float* weights, int64_t T, void* stream);
/// Expert blob (Strata pack layout, Q2_0) -> BF16 matrices: gate/up interleaved [1280, 2560], down [2560, 640].
void blob_dequant(const uint8_t* blob, uint16_t* gu16, uint16_t* down16, void* stream);
/// h16[n, r] = fp16(silu(gu[n, 2r]) * gu[n, 2r + 1]) (the interleaved expert gate/up)
void swiglu_interleaved(const float* gu, uint16_t* h16, int64_t n, void* stream);
/// h16[n, r] = fp16(silu(g[n, r]) * u[n, r]) (the shared expert, gate and up separate, width 640)
void swiglu_pair(const float* g, const float* u, uint16_t* h16, int64_t n, void* stream);
/// Gather rows: dst16[i, :] = x16[src[i], :] (n rows of `width` BF16).
void gather_rows16(const uint16_t* x16, const int32_t* src, uint16_t* dst16, int64_t n, int64_t width, void* stream);
/// bo[t, :] = shared[t, :] * sigmoid(sg[t]) + sum_k w[t, k] * D[slot[t, k], :]
void moe_combine(const float* D, const int32_t* slot, const float* w, const float* shared, const float* sg, float* bo,
int64_t T, void* stream);
// ---- QSA helpers
/// In place: x[r, :] = x[r, :] * rsqrt(mean x^2 + eps) * w over rows of `cols` (row stride `ld`).
void rms_rows(float* x, const float* w, int64_t rows, int64_t cols, int64_t ld, float eps, void* stream);
/// NEOX rotary on the first 64 dims of each head: x [T, heads, dim] at positions pos0 + t.
void rope(float* x, int64_t T, int64_t heads, int64_t dim, int64_t ld, int64_t pos0, float freq_base, void* stream);
/// q_full [T, 24, 512] (q | gate per head) -> q [T, 24, 256]
void split_q(const float* q_full, float* q, int64_t T, void* stream);
/// attn[t, h, d] *= sigmoid(q_full[t, h, 256 + d]) -> out16 (fp16 bits: the o-projection is quantized)
void gate_attn(const float* attn, const float* q_full, uint16_t* out16, int64_t T, void* stream);
/// K and V of T consecutive cells (positions pos0..pos0+T-1; K normed and rotated) into the paged pools: FP16
/// (`k_pool`/`v_pool`) or INT8 codes + FP16 scale per 64 (`k_q`...), the decode append's arithmetic.
/// K, V: [T, 2, 256].
void kv_append(const float* K, const float* V, int64_t T, int64_t pos0, const int32_t* page_table, int64_t page_size,
uint16_t* k_pool, uint16_t* v_pool, int8_t* k_q, int8_t* v_q, uint16_t* k_scale, uint16_t* v_scale,
void* stream);
/// fp32 -> fp16 bits and fp32 -> bf16, n elements (the two activation images of the prompt GEMMs).
void to_f16(const float* x, uint16_t* y, int64_t n, void* stream);
void to_bf16(const float* x, uint16_t* y, int64_t n, void* stream);
/// Expert blob -> FP16 (Q2_0 values are exact in FP16).
void blob_dequant_f16(const uint8_t* blob, uint16_t* gu16, uint16_t* down16, void* stream);
} // namespace strata::prefill
include/strata/prefill/prefill.hpp
+75
View File
@@ -0,0 +1,75 @@
// include/strata/prefill/prefill.hpp - plan v0.3 P5: batched prompt processing.
//
// The prompt's positions [pos0, pos0 + n) are processed in chunks of `chunk` tokens through all 48 layers, leaving
// the session state (GDN recurrence and conv state, QSA KV pools and indexer, PLE history) where the token path
// would have left it; the decode loop then continues with the next token. Per layer: the projections are
// tensor-core GEMMs (quantized weights dequantized to FP16 on the fly, BF16 weights as they are), the recurrences
// walk the chunk inside one kernel, and the routed experts are grouped by expert: resident ones are read from the
// VRAM tier, the others streamed from the host arena through a pinned ring on a copy stream.
//
// Requires the native weights (`--native`): every quantized projection must carry its GGUF blocks.
#pragma once
#include "strata/core/expert_cache.hpp"
#include "strata/core/expert_source.hpp"
#include "strata/core/layer.hpp"
#include "strata/core/session.hpp"
#include "strata/core/weights.hpp"
#include <cstdint>
#include <functional>
#include <memory>
#include <string>
namespace strata::prefill {
struct PrefillStats {
int64_t tokens = 0;
int64_t chunks = 0;
double ms_total = 0;
double ms_experts_host = 0; ///< host time staging non-resident experts
int64_t experts_streamed = 0; ///< expert blobs copied host -> device
int64_t experts_dma = 0; ///< ...of which straight from the pinned arena (no CPU copy)
int64_t experts_resident = 0; ///< expert-layer groups served from the VRAM tier
double ms_ple = 0;
};
class Prefill {
public:
Prefill();
~Prefill();
Prefill(const Prefill&) = delete;
Prefill& operator=(const Prefill&) = delete;
/// `host_res`: the static residency table (n_layers x n_expert, slot or -1) or null; `cache` its slots.
/// `borrow`/`borrow_bytes`: device memory to carve every buffer from (the top slots of the expert cache,
/// lent for the prompt and refilled after it); null = allocate normally.
bool init(const core::WeightTable& wt, const core::ModelGeometry& g, core::SessionState& ss,
core::ExpertSource* src, const core::ExpertCache* cache, const int32_t* host_res, int64_t chunk,
void* stream, std::string& err, void* borrow = nullptr, uint64_t borrow_bytes = 0);
/// Device bytes `init` needs for a chunk of `chunk` tokens (what a borrowed region must hold).
static uint64_t bytes_needed(const core::ModelGeometry& g, const core::SessionState& ss, int64_t chunk);
/// Positions [pos0, pos0 + n) holding `tokens`; `ss.ple_prev` must be the two tokens before pos0 (oldest
/// first, -1 for none) and is advanced to the last two of these.
bool run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& err);
const PrefillStats& stats() const { return stats_; }
/// Plan v0.3 P6: called after every chunk with the chunk's final multi-stream residual rows (device,
/// T x hc*n_embd, valid until the next chunk) and the chunk's first position; the MTP draft layer builds its
/// K/V from them. The prefill stream is synchronized before the call.
std::function<bool(const float* R_rows, int64_t T, int64_t pos0, std::string& err)> on_chunk;
/// The vision path: HOST rows (n_embd floats) indexed by absolute position, read in place of the token
/// embedding where non-null (an image's <|image_pad|> cells). Null (default): every position embeds its token.
const float* const* embd_rows = nullptr;
private:
struct Impl;
std::unique_ptr<Impl> impl_;
PrefillStats stats_;
};
} // namespace strata::prefill
include/strata/program/logits_selection.hpp
+30
View File
@@ -0,0 +1,30 @@
#pragma once
#include <charconv>
#include <cstdint>
#include <string_view>
namespace strata::program::logits_selection {
// Storage selection only: callers must still condition on every input token.
// Positive decimal int64, with no sign, whitespace, suffix or overflow.
inline bool parse_stride(std::string_view text, int64_t& result) {
if (text.empty() || text.front() < '0' || text.front() > '9') return false;
int64_t value = 0;
const auto parsed = std::from_chars(text.data(), text.data() + text.size(), value);
if (parsed.ec != std::errc{} || parsed.ptr != text.data() + text.size() || value <= 0) return false;
result = value;
return true;
}
// Zero-based multiples of stride, followed by the final input position once.
inline constexpr bool selected(int64_t position, int64_t total, int64_t stride) {
return total > 0 && stride > 0 && position >= 0 && position < total &&
(position % stride == 0 || position == total - 1);
}
inline constexpr int64_t row_count(int64_t total, int64_t stride) {
return total > 0 && stride > 0 ? (total - 1) / stride + 1 + ((total - 1) % stride != 0) : 0;
}
} // namespace strata::program::logits_selection
include/strata/spec/controller.hpp
+73
View File
@@ -0,0 +1,73 @@
// include/strata/spec/controller.hpp - plan v0.3 P6: which drafter, and how many draft tokens, this step.
//
// Each step the controller maximizes E[tokens committed] / T(step) over the choices
// none (k = 0) | suffix lookup with k <= its proposal | MTP with k <= K_MAX
// using running estimates of per-position acceptance and a cost model of a (k+1)-token verify pass. It keeps
// k = 0 unless the best choice beats plain decoding by `min_gain` (plan: 5%).
//
// Acceptance: for MTP, position i's conditional acceptance p_i (EMA). For lookup, one conditional acceptance per
// match-length bucket (a long match is far more reliable than a trigram), also an EMA. `observe` updates them
// from how many drafted tokens the verify pass accepted.
//
// Cost: T(n tokens) = dense(n) + experts(n) + sync + draft. The defaults are the 23 Sep measurements:
// dense ratio per token count from native MMVQ (bench/results/2026-09-23-mmvq-multi), distinct experts per window
// from the routing trace (2026-09-23-spec-economics), extra-token CPU cost 0.2-0.35 (2026-09-23-cpu-expert-multi).
// Replace them with in-engine measurements through `CostModel`.
#pragma once
#include <array>
#include <cstdint>
namespace strata::spec {
inline constexpr int K_MAX = 8;
struct CostModel {
double dense_ms = 11.0; // one-token dense pass (plan P3 target)
std::array<double, K_MAX + 1> dense_ratio{1.0, 1.05, 1.3, 1.45, 1.85, 2.2, 2.6, 3.0, 3.3}; // by n = k+1 (index n-1)
double cpu_all_miss_ms = 15.8; // 480 experts on the CPU
double hit_rate = 0.55; // share of distinct experts served from VRAM
std::array<double, K_MAX + 1> distinct_ratio{1.0, 1.70, 2.31, 2.88, 3.40, 3.89, 4.35, 4.80, 5.2}; // U(n)/U(1)
double extra_use_cost = 0.25; // each extra token on an expert, fraction of a read
double sync_ms = 2.4;
double mtp_draft_ms = 1.2; // per MTP draft token
double lookup_draft_ms = 0.01;
/// Step time for verifying n = k + 1 tokens, plus drafting k tokens with the given source.
double step_ms(int k, bool mtp) const;
};
enum class Source { None, Lookup, Mtp };
struct Choice {
Source source = Source::None;
int k = 0;
double expected_tokens = 1.0;
double tokens_per_ms = 0.0;
};
class Controller {
public:
explicit Controller(CostModel cost = {}, double min_gain = 0.05, double ema = 0.05);
/// `lookup_available` tokens the suffix drafter can propose now, from a match of `lookup_match` tokens.
Choice choose(int lookup_available, int lookup_match, bool mtp_ready) const;
/// After verification: `accepted` of the `k` drafted tokens matched (a prefix).
void observe(const Choice& c, int accepted, int lookup_match);
double mtp_accept(int position) const { return mtp_p_[position]; } // position 0 = first draft token
double lookup_accept(int match) const { return lookup_q_[bucket(match)]; }
const CostModel& cost() const { return cost_; }
private:
static int bucket(int match) { return match >= 16 ? 3 : match >= 8 ? 2 : match >= 5 ? 1 : 0; }
static double expected(const double* p, int k, bool conditional_same);
CostModel cost_;
double min_gain_, ema_;
std::array<double, K_MAX> mtp_p_{};
std::array<double, 4> lookup_q_{};
};
} // namespace strata::spec
include/strata/spec/suffix_drafter.hpp
+53
View File
@@ -0,0 +1,53 @@
// include/strata/spec/suffix_drafter.hpp - plan v0.3 P6: the suffix-lookup drafter (no weights, no GPU).
//
// Proposes the tokens that followed the longest earlier occurrence of the sequence's current suffix: when the
// output quotes its input (edits, refactors, repeated code) the continuation is usually exact, and a verify pass
// accepts many tokens at once. It needs no model and costs microseconds.
//
// Index: every trigram of the history (prompt + accepted output) maps to its WAYS most recent end positions in a
// fixed-size open-addressing table, so memory is bounded (~20 bytes per history token) and appends are O(1).
// A proposal checks those candidates, extends each match backwards up to `max_match`, and takes the longest
// (most recent on ties). Matches shorter than `min_match` propose nothing.
#pragma once
#include <cstddef>
#include <cstdint>
#include <vector>
namespace strata::spec {
class SuffixDrafter {
public:
static constexpr int WAYS = 4;
explicit SuffixDrafter(int min_match = 3, int max_match = 32, size_t capacity_tokens = 1u << 19);
void reset();
/// Add tokens to the history (the prompt, then every accepted token).
void append(const int32_t* tokens, size_t n);
void append(int32_t token) { append(&token, 1); }
/// Write up to `max_k` proposed next tokens to `out`; returns how many (0 = no match of at least min_match).
int propose(int max_k, int32_t* out);
/// Length of the match behind the last proposal (0 if none).
int last_match() const { return last_match_; }
size_t size() const { return hist_.size(); }
private:
struct Slot {
uint64_t key = 0; // trigram hash + 1 (0 = empty)
uint32_t pos[WAYS] = {}; // end positions, most recent first
uint8_t n = 0;
};
Slot* find_slot(uint64_t key, bool insert);
uint64_t key_at(size_t end) const;
int min_match_, max_match_;
std::vector<int32_t> hist_;
std::vector<Slot> table_;
size_t mask_ = 0;
int last_match_ = 0;
};
} // namespace strata::spec
ref/load.py
+212
View File
@@ -0,0 +1,212 @@
"""ref/load.py - P1.S4: GGUF -> FP32 tensors on demand, with LAZY per-expert dequantization.
The plan is explicit: "experts are dequantized lazily per (layer, expert) - do not materialize 120B
parameters". That is the whole design constraint here, and it is what `expert()` below exists for: the
routed-expert tensors are fused `[in, out, experts]`, and touching one expert must not touch the other
511.
Dequantization itself is NOT reimplemented. `gguf-py` handles the twelve standard encodings and is
already the reference this project validates its C++ against; `Q2_0` (type 42) is the one gguf-py
cannot represent at all, so it is served by this project's own decoder. That split is deliberate: it
means a bug in `tools/gguf_writer.dequantize_q2_0` cannot hide behind gguf-py and vice versa.
from ref.load import Gguf
g = Gguf("path.gguf")
w = g.tensor("blk.0.attn_gate.weight") # (out, in) fp32, whole tensor
e = g.expert("blk.0.ffn_gate_exps.weight", 7) # (out, in) fp32, ONE expert only
"""
from __future__ import annotations
import pathlib
import re
import sys
import numpy as np
sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent / "tools"))
from gguf_reader import GGUFFile # noqa: E402
from gguf_writer import dequantize_q2_0 # noqa: E402
# gguf-py names differ from our reader's names in a few places; map only what differs.
_GGUF_PY_NAME = {"Q4_0": "Q4_0", "Q8_0": "Q8_0", "Q4_K": "Q4_K", "Q5_K": "Q5_K", "Q6_K": "Q6_K",
"Q3_K": "Q3_K", "IQ4_NL": "IQ4_NL", "IQ4_XS": "IQ4_XS", "Q5_0": "Q5_0"}
# (elements per block, bytes per block) - mirrors block_geometry() in src/artifact/gguf_reader.cpp
GEOM = {"F32": (1, 4), "F16": (1, 2), "BF16": (1, 2),
"Q4_0": (32, 18), "Q5_0": (32, 22), "Q8_0": (32, 34), "IQ4_NL": (32, 18),
"Q3_K": (256, 110), "Q4_K": (256, 144), "Q5_K": (256, 176), "Q6_K": (256, 210),
"IQ4_XS": (256, 136), "Q2_0": (64, 18)}
def _dequant_flat(type_name: str, raw: np.ndarray) -> np.ndarray:
"""One dequantization dispatch, shared by tensor()/expert()/rows().
Q2_0 is this project's own decoder because gguf-py's enum cannot represent type 42 at all;
everything else goes through gguf-py, which is already the reference this project validates its
C++ against. Keeping one dispatch means a type added in one path cannot be forgotten in another.
"""
if type_name == "Q2_0":
return dequantize_q2_0(raw.tobytes())
if type_name == "BF16":
return _bf16_to_f32(raw)
if type_name == "F32":
return raw.view(np.float32)
if type_name == "F16":
return raw.view(np.float16).astype(np.float32)
from gguf import quants, GGMLQuantizationType as Q
return quants.dequantize(raw, Q[_GGUF_PY_NAME[type_name]]).astype(np.float32)
def _bf16_to_f32(raw: np.ndarray) -> np.ndarray:
return (raw.view(np.uint16).astype(np.uint32) << 16).view(np.float32)
def _discover_shards(path) -> list[pathlib.Path]:
"""Resolve a GGUF argument to its shard list.
The real artifact is TWO files and the split is not incidental: shard 2 holds exactly one tensor,
`per_layer_token_embd.weight` (26.82 GiB). A reader that opens one file can therefore load the
whole model EXCEPT the PLE table and will not fail until a forward pass reaches layer 1 - so
shard discovery is not a convenience, it is a correctness requirement for the real model.
Accepts a directory (all `*.gguf` in it), any single shard (siblings found by the
`-NNNNN-of-MMMMM` suffix), or an explicit list.
"""
if isinstance(path, (list, tuple)):
return [pathlib.Path(p) for p in path]
p = pathlib.Path(path)
if p.is_dir():
# A directory is ambiguous in practice: the project root holds BOTH the model shards (in a
# subdirectory) and an unrelated `mmproj-*.gguf`. Globbing `*.gguf` there silently loaded
# the multimodal projector instead of the model. So prefer an explicit shard set, accept a
# single unambiguous file, and REFUSE to guess otherwise.
shards = sorted(p.glob("*-[0-9][0-9][0-9][0-9][0-9]-of-[0-9][0-9][0-9][0-9][0-9].gguf"))
if shards:
totals = {re.match(r"^.*-(\d{5})-of-(\d{5})\.gguf$", s.name).group(2) for s in shards}
if len(totals) != 1:
raise ValueError(f"{p}: shard files from more than one split: {sorted(totals)}")
return shards
found = sorted(p.glob("*.gguf"))
if len(found) == 1:
return found
if not found:
raise FileNotFoundError(f"no .gguf files in {p}")
raise ValueError(
f"{p}: {len(found)} .gguf files and none is a shard set, so which is the model is "
f"ambiguous: {[f.name for f in found]}. Pass one file explicitly.")
m = re.match(r"^(.*)-(\d{5})-of-(\d{5})\.gguf$", p.name)
if m:
stem, _, total = m.groups()
sibs = sorted(p.parent.glob(f"{stem}-*-of-{total}.gguf"))
if len(sibs) != int(total):
raise FileNotFoundError(
f"{p.name}: expected {total} shards, found {len(sibs)} in {p.parent}")
return sibs
return [p]
class Gguf:
"""Read-only, demand-paged view of a (possibly multi-shard) GGUF.
Nothing is dequantized until asked for. `g` is the FIRST shard's `GGUFFile`, so metadata and the
tensor directory are reachable as before; reads route to whichever shard owns the tensor.
"""
def __init__(self, path):
self.paths = _discover_shards(path)
self.path = self.paths[0]
self.shards: list[tuple[GGUFFile, object]] = []
self._by_name: dict[str, tuple[int, object]] = {}
for i, sp in enumerate(self.paths):
gf = GGUFFile(sp)
self.shards.append((gf, sp.open("rb")))
for t in gf.tensors:
if t.name in self._by_name:
raise ValueError(f"{sp.name}: duplicate tensor {t.name} across shards")
self._by_name[t.name] = (i, t)
self.g = self.shards[0][0] # metadata / directory / alignment live in shard 1
# ---- raw bytes
def _raw(self, name: str, nbytes: int, byte_offset: int = 0) -> bytes:
i, t = self._by_name[name]
gf, fh = self.shards[i]
fh.seek(gf.data_start + t.offset + byte_offset)
return fh.read(nbytes)
def info(self, name: str):
return self._by_name[name][1]
def row_bytes(self, t) -> int:
el, by = GEOM[t.type_name]
return t.shape[0] // el * by
# ---- full tensors
def tensor(self, name: str) -> np.ndarray:
"""Dequantize a whole tensor to FP32, shape as GGUF order reversed to (out, in[, experts])."""
t = self.info(name)
el, by = GEOM[t.type_name]
if t.elements % el:
raise ValueError(f"{name}: {t.elements} not a multiple of block {el}")
raw = np.frombuffer(self._raw(name, t.elements // el * by), dtype=np.uint8)
flat = _dequant_flat(t.type_name, raw)
# GGUF dim 0 varies fastest; numpy's last axis does, so reverse and transpose to (out, in).
return flat.reshape(tuple(reversed(t.shape)))
# ---- the lazy row path
def rows(self, name: str, idx) -> np.ndarray:
"""Gather specific ROWS of a 2-D tensor without decoding the rest of it.
This exists for `per_layer_token_embd.weight`, which is `[160, 320001536]` IQ4_NL = **26.82
GiB** - 43% of the artifact, the entire second shard - and is read 16 rows per token by one
layer (see round 98's finding). Calling `tensor()` on it would dequantize 320 million rows to
FP32, about 205 GB, to use 2.5 KB of them.
A row is contiguous in GGUF (ne0 varies fastest), so one row is `row_bytes` at
`row * row_bytes`. The type's block size must divide the row width; for this table it does
(160 = 5 x 32 for IQ4_NL).
"""
t = self.info(name)
if len(t.shape) != 2:
raise ValueError(f"{name}: rows() needs a 2-D tensor, got shape {t.shape}")
el, by = GEOM[t.type_name]
width = t.shape[0]
if width % el:
raise ValueError(f"{name}: row width {width} is not a multiple of block {el}")
rb = self.row_bytes(t)
idx = np.asarray(idx, dtype=np.int64).ravel()
if idx.size and (idx.min() < 0 or idx.max() >= t.shape[1]):
raise ValueError(f"{name}: row index out of range (0..{t.shape[1] - 1})")
out = np.empty((idx.size, width), dtype=np.float32)
for k, r in enumerate(idx):
raw = np.frombuffer(self._raw(name, rb, int(r) * rb), dtype=np.uint8)
out[k] = _dequant_flat(t.type_name, raw)
return out
# ---- the lazy path
def expert(self, name: str, expert_id: int) -> np.ndarray:
"""One expert out of a fused `[in, out, experts]` tensor, WITHOUT touching the others.
In GGUF order the expert index is the LAST axis, so expert e's matrix occupies a contiguous
run of `in*out` elements at offset e*in*out. Only that run is read and decoded - which is the
plan's "do not materialize 120B parameters" requirement, and the reason the engine can run at
all on a machine where the expert set is 34 GB.
"""
t = self.info(name)
if len(t.shape) != 3:
raise ValueError(f"{name}: not a fused expert tensor (shape {t.shape})")
in_f, out_f, n_exp = t.shape
if not 0 <= expert_id < n_exp:
raise ValueError(f"{name}: expert {expert_id} out of range (0..{n_exp - 1})")
el, by = GEOM[t.type_name]
per_expert = in_f * out_f
if per_expert % el:
raise ValueError(f"{name}: per-expert element count {per_expert} not a multiple of {el}")
# byte offset of expert e: whole experts before it, in blocks
start = expert_id * (per_expert // el) * by
raw = np.frombuffer(self._raw(name, per_expert // el * by, start), dtype=np.uint8)
flat = _dequant_flat(t.type_name, raw)
return flat.reshape((out_f, in_f))
def close(self) -> None:
for _, fh in self.shards:
fh.close()
serve/__init__.py
+1
View File
@@ -0,0 +1 @@
"""serve - plan v0.3 P8: the API layer (text in, token ids to the engine, text deltas out)."""
serve/chat_golden.json
+197
View File
@@ -0,0 +1,197 @@
[
{
"name": "single user",
"messages": [
{
"role": "user",
"content": "Hello"
}
],
"rendered": "<|im_start|>system\nReasoning effort is set to xhigh. Please think carefully through the task, validate key assumptions, consider plausible alternatives, and prioritize correctness, consistency, and clarity in the final answer.<|im_end|>\n<|im_start|>user\nHello<|im_end|>\n<|im_start|>assistant\n<think>\n",
"error": null,
"hf": "<|im_start|>system\nReasoning effort is set to xhigh. Please think carefully through the task, validate key assumptions, consider plausible alternatives, and prioritize correctness, consistency, and clarity in the final answer.<|im_end|>\n<|im_start|>user\nHello<|im_end|>\n<|im_start|>assistant\n<think>\n"
},
{
"name": "system + user",
"messages": [
{
"role": "system",
"content": "You are a helpful assistant."
},
{
"role": "user",
"content": "What is 2+2?"
}
],
"rendered": "<|im_start|>system\nReasoning effort is set to xhigh. Please think carefully through the task, validate key assumptions, consider plausible alternatives, and prioritize correctness, consistency, and clarity in the final answer.\n\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\nWhat is 2+2?<|im_end|>\n<|im_start|>assistant\n<think>\n",
"error": null,
"hf": "<|im_start|>system\nReasoning effort is set to xhigh. Please think carefully through the task, validate key assumptions, consider plausible alternatives, and prioritize correctness, consistency, and clarity in the final answer.\n\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\nWhat is 2+2?<|im_end|>\n<|im_start|>assistant\n<think>\n"
},
{
"name": "multi-turn",
"messages": [
{
"role": "system",
"content": "Be brief."
},
{
"role": "user",
"content": "Hi"
},
{
"role": "assistant",
"content": "Hello! How can I help?"
},
{
"role": "user",
"content": "Explain BPE in one sentence."
}
],
"rendered": "<|im_start|>system\nReasoning effort is set to xhigh. Please think carefully through the task, validate key assumptions, consider plausible alternatives, and prioritize correctness, consistency, and clarity in the final answer.\n\nBe brief.<|im_end|>\n<|im_start|>user\nHi<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\nHello! How can I help?<|im_end|>\n<|im_start|>user\nExplain BPE in one sentence.<|im_end|>\n<|im_start|>assistant\n<think>\n",
"error": null,
"hf": "<|im_start|>system\nReasoning effort is set to xhigh. Please think carefully through the task, validate key assumptions, consider plausible alternatives, and prioritize correctness, consistency, and clarity in the final answer.\n\nBe brief.<|im_end|>\n<|im_start|>user\nHi<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\nHello! How can I help?<|im_end|>\n<|im_start|>user\nExplain BPE in one sentence.<|im_end|>\n<|im_start|>assistant\n<think>\n"
},
{
"name": "assistant with thinking",
"messages": [
{
"role": "user",
"content": "What is 17*23?"
},
{
"role": "assistant",
"content": "Let me compute.",
"reasoning_content": "17*23 = 17*20 + 17*3 = 340+51 = 391"
},
{
"role": "user",
"content": "Thanks"
}
],
"rendered": "<|im_start|>system\nReasoning effort is set to xhigh. Please think carefully through the task, validate key assumptions, consider plausible alternatives, and prioritize correctness, consistency, and clarity in the final answer.<|im_end|>\n<|im_start|>user\nWhat is 17*23?<|im_end|>\n<|im_start|>assistant\n<think>\n17*23 = 17*20 + 17*3 = 340+51 = 391\n</think>\n\nLet me compute.<|im_end|>\n<|im_start|>user\nThanks<|im_end|>\n<|im_start|>assistant\n<think>\n",
"error": null,
"hf": "<|im_start|>system\nReasoning effort is set to xhigh. Please think carefully through the task, validate key assumptions, consider plausible alternatives, and prioritize correctness, consistency, and clarity in the final answer.<|im_end|>\n<|im_start|>user\nWhat is 17*23?<|im_end|>\n<|im_start|>assistant\n<think>\n17*23 = 17*20 + 17*3 = 340+51 = 391\n</think>\n\nLet me compute.<|im_end|>\n<|im_start|>user\nThanks<|im_end|>\n<|im_start|>assistant\n<think>\n"
},
{
"name": "no generation prompt",
"messages": [
{
"role": "user",
"content": "Hello"
},
{
"role": "assistant",
"content": "Hi there"
}
],
"add_generation_prompt": false,
"rendered": "<|im_start|>system\nReasoning effort is set to xhigh. Please think carefully through the task, validate key assumptions, consider plausible alternatives, and prioritize correctness, consistency, and clarity in the final answer.<|im_end|>\n<|im_start|>user\nHello<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\nHi there<|im_end|>\n",
"error": null,
"hf": "<|im_start|>system\nReasoning effort is set to xhigh. Please think carefully through the task, validate key assumptions, consider plausible alternatives, and prioritize correctness, consistency, and clarity in the final answer.<|im_end|>\n<|im_start|>user\nHello<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\nHi there<|im_end|>\n"
},
{
"name": "thinking disabled",
"messages": [
{
"role": "user",
"content": "Hello"
}
],
"enable_thinking": false,
"rendered": "<|im_start|>user\nHello<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n",
"error": null,
"hf": "<|im_start|>user\nHello<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"
},
{
"name": "unicode content",
"messages": [
{
"role": "user",
"content": "你好，世界 😀"
}
],
"rendered": "<|im_start|>system\nReasoning effort is set to xhigh. Please think carefully through the task, validate key assumptions, consider plausible alternatives, and prioritize correctness, consistency, and clarity in the final answer.<|im_end|>\n<|im_start|>user\n你好，世界 😀<|im_end|>\n<|im_start|>assistant\n<think>\n",
"error": null,
"hf": "<|im_start|>system\nReasoning effort is set to xhigh. Please think carefully through the task, validate key assumptions, consider plausible alternatives, and prioritize correctness, consistency, and clarity in the final answer.<|im_end|>\n<|im_start|>user\n你好，世界 😀<|im_end|>\n<|im_start|>assistant\n<think>\n"
},
{
"name": "tool call and response",
"messages": [
{
"role": "user",
"content": "Weather in Oslo?"
},
{
"role": "assistant",
"content": "",
"tool_calls": [
{
"type": "function",
"function": {
"name": "get_weather",
"arguments": {
"city": "Oslo"
}
}
}
]
},
{
"role": "tool",
"content": "{\"temp\": 3}",
"name": "get_weather"
},
{
"role": "user",
"content": "And in Bergen?"
}
],
"tools": [
{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get the weather",
"parameters": {
"type": "object",
"properties": {
"city": {
"type": "string"
}
},
"required": [
"city"
]
}
}
}
],
"rendered": "<|im_start|>system\nReasoning effort is set to xhigh. Please think carefully through the task, validate key assumptions, consider plausible alternatives, and prioritize correctness, consistency, and clarity in the final answer.\n\n# Tools\n\nYou have access to the following functions:\n\n<tools>\n{\"type\": \"function\", \"function\": {\"name\": \"get_weather\", \"description\": \"Get the weather\", \"parameters\": {\"type\": \"object\", \"properties\": {\"city\": {\"type\": \"string\"}}, \"required\": [\"city\"]}}}\n</tools>\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\nvalue_1\n</parameter>\n<parameter=example_parameter_2>\nThis is the value for the second parameter\nthat can span\nmultiple lines\n</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\nReminder:\n- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\n- Required parameters MUST be specified\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n</IMPORTANT><|im_end|>\n<|im_start|>user\nWeather in Oslo?<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n<tool_call>\n<function=get_weather>\n<parameter=city>\nOslo\n</parameter>\n</function>\n</tool_call><|im_end|>\n<|im_start|>user\n<tool_response>\n{\"temp\": 3}\n</tool_response><|im_end|>\n<|im_start|>user\nAnd in Bergen?<|im_end|>\n<|im_start|>assistant\n<think>\n",
"error": null,
"hf": "<|im_start|>system\nReasoning effort is set to xhigh. Please think carefully through the task, validate key assumptions, consider plausible alternatives, and prioritize correctness, consistency, and clarity in the final answer.\n\n# Tools\n\nYou have access to the following functions:\n\n<tools>\n{\"type\": \"function\", \"function\": {\"name\": \"get_weather\", \"description\": \"Get the weather\", \"parameters\": {\"type\": \"object\", \"properties\": {\"city\": {\"type\": \"string\"}}, \"required\": [\"city\"]}}}\n</tools>\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\nvalue_1\n</parameter>\n<parameter=example_parameter_2>\nThis is the value for the second parameter\nthat can span\nmultiple lines\n</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\nReminder:\n- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\n- Required parameters MUST be specified\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n</IMPORTANT><|im_end|>\n<|im_start|>user\nWeather in Oslo?<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n<tool_call>\n<function=get_weather>\n<parameter=city>\nOslo\n</parameter>\n</function>\n</tool_call><|im_end|>\n<|im_start|>user\n<tool_response>\n{\"temp\": 3}\n</tool_response><|im_end|>\n<|im_start|>user\nAnd in Bergen?<|im_end|>\n<|im_start|>assistant\n<think>\n"
},
{
"name": "empty user",
"messages": [
{
"role": "user",
"content": ""
}
],
"rendered": "<|im_start|>system\nReasoning effort is set to xhigh. Please think carefully through the task, validate key assumptions, consider plausible alternatives, and prioritize correctness, consistency, and clarity in the final answer.<|im_end|>\n<|im_start|>user\n<|im_end|>\n<|im_start|>assistant\n<think>\n",
"error": null,
"hf": "<|im_start|>system\nReasoning effort is set to xhigh. Please think carefully through the task, validate key assumptions, consider plausible alternatives, and prioritize correctness, consistency, and clarity in the final answer.<|im_end|>\n<|im_start|>user\n<|im_end|>\n<|im_start|>assistant\n<think>\n"
},
{
"name": "whitespace and newlines",
"messages": [
{
"role": "user",
"content": " line1\n line2\n\n\nend "
}
],
"rendered": "<|im_start|>system\nReasoning effort is set to xhigh. Please think carefully through the task, validate key assumptions, consider plausible alternatives, and prioritize correctness, consistency, and clarity in the final answer.<|im_end|>\n<|im_start|>user\nline1\n line2\n\n\nend<|im_end|>\n<|im_start|>assistant\n<think>\n",
"error": null,
"hf": "<|im_start|>system\nReasoning effort is set to xhigh. Please think carefully through the task, validate key assumptions, consider plausible alternatives, and prioritize correctness, consistency, and clarity in the final answer.<|im_end|>\n<|im_start|>user\nline1\n line2\n\n\nend<|im_end|>\n<|im_start|>assistant\n<think>\n"
}
]
serve/chat_template.jinja
+184
View File
@@ -0,0 +1,184 @@
{%- set image_count = namespace(value=0) %}
{%- set video_count = namespace(value=0) %}
{%- macro render_content(content, do_vision_count, is_system_content=false) %}
{%- if content is string %}
{{- content }}
{%- elif content is iterable and content is not mapping %}
{%- for item in content %}
{%- if 'image' in item or 'image_url' in item or item.type == 'image' %}
{%- if is_system_content %}
{{- raise_exception('System message cannot contain images.') }}
{%- endif %}
{%- if do_vision_count %}
{%- set image_count.value = image_count.value + 1 %}
{%- endif %}
{%- if add_vision_id %}
{{- 'Picture ' ~ image_count.value ~ ': ' }}
{%- endif %}
{{- '<|vision_start|><|image_pad|><|vision_end|>' }}
{%- elif 'video' in item or item.type == 'video' %}
{%- if is_system_content %}
{{- raise_exception('System message cannot contain videos.') }}
{%- endif %}
{%- if do_vision_count %}
{%- set video_count.value = video_count.value + 1 %}
{%- endif %}
{%- if add_vision_id %}
{{- 'Video ' ~ video_count.value ~ ': ' }}
{%- endif %}
{{- '<|vision_start|><|video_pad|><|vision_end|>' }}
{%- elif 'text' in item %}
{{- item.text }}
{%- else %}
{{- raise_exception('Unexpected item type in content.') }}
{%- endif %}
{%- endfor %}
{%- elif content is none or content is undefined %}
{{- '' }}
{%- else %}
{{- raise_exception('Unexpected content type.') }}
{%- endif %}
{%- endmacro %}
{%- if not messages %}
{{- raise_exception('No messages provided.') }}
{%- endif %}
{%- set sysns = namespace(count=0, text='') %}
{%- for message in messages %}
{%- if sysns.count == loop.index0 and (message.role == 'system' or message.role == 'developer') %}
{%- set sys_content = render_content(message.content, false, true)|trim %}
{%- if sys_content %}
{%- set sysns.text = sysns.text + ('\n' if sysns.text else '') + sys_content %}
{%- endif %}
{%- set sysns.count = sysns.count + 1 %}
{%- endif %}
{%- endfor %}
{%- set num_sys = sysns.count %}
{%- set merged_system = sysns.text %}
{%- set reasoning_instructions = '' %}
{%- if enable_thinking is undefined or enable_thinking is true %}
{%- set resolved_reasoning_effort = reasoning_effort|default('xhigh') %}
{%- if resolved_reasoning_effort == 'high' %}
{%- set resolved_reasoning_effort = 'xhigh' %}
{%- endif %}
{%- if resolved_reasoning_effort not in ('xhigh', 'medium', 'low') %}
{{- raise_exception('Unexpected reasoning effort ' ~ reasoning_effort ~ '. Supported types are xhigh (default), medium, and low.') }}
{%- endif %}
{%- if resolved_reasoning_effort == 'xhigh' %}
{%- set reasoning_instructions = 'Reasoning effort is set to xhigh. Please think carefully through the task, validate key assumptions, consider plausible alternatives, and prioritize correctness, consistency, and clarity in the final answer.' %}
{%- elif resolved_reasoning_effort == 'low' %}
{%- set reasoning_instructions = 'Reasoning effort is set to low. Keep your thinking brief and focused, moving directly to the conclusion without unnecessary elaboration.' %}
{%- endif %}
{%- endif %}
{%- if tools and tools is iterable and tools is not mapping %}
{{- '<|im_start|>system\n' }}
{%- if reasoning_instructions %}
{{- reasoning_instructions + '\n\n' }}
{%- endif %}
{{- "# Tools\n\nYou have access to the following functions:\n\n<tools>" }}
{%- for tool in tools %}
{{- "\n" }}
{{- tool | tojson }}
{%- endfor %}
{{- "\n</tools>" }}
{{- '\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\nvalue_1\n</parameter>\n<parameter=example_parameter_2>\nThis is the value for the second parameter\nthat can span\nmultiple lines\n</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\nReminder:\n- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\n- Required parameters MUST be specified\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n</IMPORTANT>' }}
{%- if merged_system %}
{{- '\n\n' + merged_system }}
{%- endif %}
{{- '<|im_end|>\n' }}
{%- else %}
{%- if merged_system %}
{{- '<|im_start|>system\n' + (reasoning_instructions + '\n\n' if reasoning_instructions else '') + merged_system + '<|im_end|>\n' }}
{%- elif reasoning_instructions %}
{{- '<|im_start|>system\n' + reasoning_instructions + '<|im_end|>\n' }}
{%- endif %}
{%- endif %}
{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
{%- for message in messages[::-1] %}
{%- set index = (messages|length - 1) - loop.index0 %}
{%- if ns.multi_step_tool and message.role == "user" %}
{%- set content = render_content(message.content, false)|trim %}
{%- if not(content.startswith('<tool_response>') and content.endswith('</tool_response>')) %}
{%- set ns.multi_step_tool = false %}
{%- set ns.last_query_index = index %}
{%- endif %}
{%- endif %}
{%- endfor %}
{%- for message in messages %}
{%- if loop.index0 >= num_sys %}
{%- set content = render_content(message.content, true)|trim %}
{%- if message.role == "system" or message.role == "developer" %}
{{- raise_exception('System message must be at the beginning.') }}
{%- elif message.role == "user" %}
{{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
{%- elif message.role == "assistant" %}
{%- set reasoning_content = '' %}
{%- if message.reasoning_content is string %}
{%- set reasoning_content = message.reasoning_content %}
{%- endif %}
{%- set reasoning_content = reasoning_content|trim %}
{%- if preserve_thinking is undefined or preserve_thinking is true or loop.index0 > ns.last_query_index %}
{{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content + '\n</think>\n\n' + content }}
{%- else %}
{{- '<|im_start|>' + message.role + '\n' + content }}
{%- endif %}
{%- if message.tool_calls and message.tool_calls is iterable and message.tool_calls is not mapping %}
{%- for tool_call in message.tool_calls %}
{%- if tool_call.function is defined %}
{%- set tool_call = tool_call.function %}
{%- endif %}
{%- if tool_call.name is not defined or tool_call.name is none %}
{{- raise_exception('Tool call is missing a function name.') }}
{%- endif %}
{%- if loop.first %}
{%- if content|trim %}
{{- '\n\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
{%- else %}
{{- '<tool_call>\n<function=' + tool_call.name + '>\n' }}
{%- endif %}
{%- else %}
{{- '\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
{%- endif %}
{%- if tool_call.arguments is mapping %}
{%- for args_name, args_value in tool_call.arguments|items %}
{{- '<parameter=' + args_name + '>\n' }}
{%- set args_value = args_value | string if args_value is string else args_value | tojson | safe %}
{{- args_value }}
{{- '\n</parameter>\n' }}
{%- endfor %}
{%- elif tool_call.arguments is string %}
{%- if tool_call.arguments|trim %}
{{- raise_exception('Tool call arguments for function "' + (tool_call.name | string) + '" were passed as a JSON string. Parse them into an object before calling apply_chat_template.') }}
{%- endif %}
{%- elif tool_call.arguments is defined and tool_call.arguments is not none %}
{{- raise_exception('Tool call arguments for function "' + (tool_call.name | string) + '" must be an object/mapping or a JSON string.') }}
{%- endif %}
{{- '</function>\n</tool_call>' }}
{%- endfor %}
{%- endif %}
{{- '<|im_end|>\n' }}
{%- elif message.role == "tool" %}
{%- if loop.previtem and loop.previtem.role != "tool" %}
{{- '<|im_start|>user' }}
{%- endif %}
{{- '\n<tool_response>\n' }}
{{- content }}
{{- '\n</tool_response>' }}
{%- if not loop.last and loop.nextitem.role != "tool" %}
{{- '<|im_end|>\n' }}
{%- elif loop.last %}
{{- '<|im_end|>\n' }}
{%- endif %}
{%- else %}
{{- raise_exception('Unexpected message role.') }}
{%- endif %}
{%- endif %}
{%- endfor %}
{%- if add_generation_prompt %}
{{- '<|im_start|>assistant\n' }}
{%- if enable_thinking is defined and enable_thinking is false %}
{{- '<think>\n\n</think>\n\n' }}
{%- else %}
{{- '<think>\n' }}
{%- endif %}
{%- endif %}
{#- Unsloth fixes - developer role, merged system messages, tool calling #}
serve/frontend.py
+344
View File
@@ -0,0 +1,344 @@
"""serve/frontend.py - plan v0.3 P8: the request/response layer between API clients and the engine.
Everything here is text: no model, no GPU. The engine consumes token ids and produces text deltas; this module
* normalizes OpenAI Chat Completions and Anthropic Messages requests into the chat template's message form,
* renders the model's own Jinja chat template (pack `tokenizer/chat_template.jinja`) exactly as Hugging Face
does (checked against the pack's `chat_golden.json`, 10 cases incl. thinking and tools),
* parses the streamed output incrementally into reasoning (`<think>...</think>`), content and tool calls in
the template's XML form (`<tool_call><function=NAME><parameter=P>VALUE</parameter>...</function></tool_call>`),
never emitting a partial tag to the client.
Design note: plan v0.3 describes a single C++ process. The frontend is Python first because the template is Jinja
and the formats change often; the engine boundary is token ids in, text deltas out, which keeps a later C++ port
(or a pybind wrapper) a local change.
"""
from __future__ import annotations
import json
import uuid
from dataclasses import dataclass, field
from pathlib import Path
import jinja2
from jinja2.sandbox import ImmutableSandboxedEnvironment
# ------------------------------------------------------------------------------------------------ template
class ChatTemplate:
"""The model's chat template, rendered with the same Jinja settings as transformers' apply_chat_template."""
def __init__(self, path: str | Path):
def raise_exception(message):
raise jinja2.exceptions.TemplateError(message)
def tojson(x, ensure_ascii=False, indent=None, separators=None, sort_keys=False):
return json.dumps(x, ensure_ascii=ensure_ascii, indent=indent, separators=separators, sort_keys=sort_keys)
env = ImmutableSandboxedEnvironment(trim_blocks=True, lstrip_blocks=True, extensions=["jinja2.ext.loopcontrols"])
env.filters["tojson"] = tojson
env.globals["raise_exception"] = raise_exception
self.template = env.from_string(Path(path).read_text(encoding="utf-8"))
def render(self, messages: list[dict], tools: list[dict] | None = None, add_generation_prompt: bool = True,
**kwargs) -> str:
return self.template.render(messages=messages, tools=tools, add_generation_prompt=add_generation_prompt,
**kwargs)
# ------------------------------------------------------------------------------------------------ requests
def _text_of(content) -> str:
if content is None:
return ""
if isinstance(content, str):
return content
return "".join(part.get("text", "") for part in content if isinstance(part, dict) and part.get("type") in
("text", "input_text", None))
IMAGE_PARTS = ("image_url", "input_image", "image")
# Thinking levels. The model's template knows low, medium and xhigh (its default; "high" means xhigh), and
# enable_thinking=false for none. Clients spell these many ways; everything maps onto those four.
EFFORT = {"none": None, "off": None, "minimal": None, "disabled": None, "false": None,
"low": "low", "medium": "medium", "high": "xhigh", "xhigh": "xhigh", "max": "xhigh", "maximum": "xhigh"}
def effort_kwargs(value) -> dict:
"""A reasoning effort as given by a client -> the template's kwargs. Unknown values are a 400, not a crash."""
if value is None or value == "":
return {}
if value is False:
return {"enable_thinking": False}
key = str(value).strip().lower()
if key not in EFFORT:
raise ValueError(f"unknown reasoning effort {value!r}: use none, low, medium or high")
level = EFFORT[key]
return {"enable_thinking": False} if level is None else {"reasoning_effort": level}
def budget_effort(tokens) -> dict:
"""Anthropic's thinking budget (budget_tokens) -> a level: under 2K low, under 8K medium, else high."""
try:
n = int(tokens)
except (TypeError, ValueError):
return {}
return {"reasoning_effort": "low" if n < 2048 else "medium" if n < 8192 else "xhigh"}
def _has_image(content) -> bool:
return isinstance(content, list) and any(isinstance(p, dict) and p.get("type") in IMAGE_PARTS for p in content)
def _image_source(part: dict) -> str:
"""An image part's source as one string: a data: URL, an http(s) URL or a local file path.
OpenAI: {"type": "image_url", "image_url": {"url": ...}} (or "image_url": "..."), Responses-style
{"type": "input_image", "image_url": ...}; Anthropic: {"type": "image", "source": {"type": "base64",
"media_type": ..., "data": ...}} or {"source": {"type": "url", "url": ...}}."""
if part.get("type") == "image":
src = part.get("source") or {}
if src.get("type") == "base64":
return f"data:{src.get('media_type', 'image/png')};base64,{src.get('data', '')}"
return src.get("url") or src.get("path") or ""
url = part.get("image_url")
if isinstance(url, dict):
url = url.get("url")
return url or ""
def _parts_of(content):
"""Message content for the template: a string when there is no image (unchanged behaviour), otherwise the
template's list form - text items and image items, in order - whose image items carry their source."""
if not _has_image(content):
return _text_of(content)
items = []
for part in content:
if not isinstance(part, dict):
continue
if part.get("type") in IMAGE_PARTS:
items.append({"type": "image", "source": _image_source(part)})
elif part.get("type") in ("text", "input_text", None) and "text" in part:
items.append({"type": "text", "text": part.get("text", "")})
return items
def images_of(messages: list[dict]) -> list[str]:
"""The image sources of the rendered conversation, in prompt order (the template renders one
<|vision_start|><|image_pad|><|vision_end|> per image item, message by message)."""
return [item["source"] for m in messages if isinstance(m.get("content"), list)
for item in m["content"] if item.get("type") == "image"]
def openai_to_messages(req: dict) -> tuple[list[dict], list[dict] | None, dict]:
"""OpenAI Chat Completions -> (template messages, template tools, template kwargs)."""
messages = []
for m in req.get("messages", []):
role = m.get("role")
if role == "developer":
role = "system"
out = {"role": role, "content": _parts_of(m.get("content")) if role == "user" else _text_of(m.get("content"))}
if m.get("reasoning_content"):
out["reasoning_content"] = m["reasoning_content"]
if m.get("tool_calls"):
calls = []
for c in m["tool_calls"]:
fn = c.get("function", c)
args = fn.get("arguments")
if isinstance(args, str): # the template requires a mapping, not a JSON string
args = json.loads(args) if args.strip() else {}
calls.append({"function": {"name": fn.get("name"), "arguments": args or {}}})
out["tool_calls"] = calls
messages.append(out)
tools = [t.get("function", t) if isinstance(t, dict) and t.get("type") == "function" else t
for t in req.get("tools") or []] or None
kwargs = {}
# OpenAI Chat Completions: "reasoning_effort"; Responses style: "reasoning": {"effort": ...}
reasoning = req.get("reasoning") if isinstance(req.get("reasoning"), dict) else {}
kwargs.update(effort_kwargs(req.get("reasoning_effort") or reasoning.get("effort")))
# the vLLM / llama.cpp convention: {"chat_template_kwargs": {"enable_thinking": false, "reasoning_effort": "low"}}
for k, v in (req.get("chat_template_kwargs") or {}).items():
if k == "enable_thinking" and not v:
kwargs = {"enable_thinking": False}
elif k == "reasoning_effort" and "enable_thinking" not in kwargs:
kwargs.update(effort_kwargs(v))
return messages, tools, kwargs
def anthropic_to_messages(req: dict) -> tuple[list[dict], list[dict] | None, dict]:
"""Anthropic Messages -> (template messages, template tools, template kwargs)."""
messages = []
system = req.get("system")
if system:
messages.append({"role": "system", "content": _text_of(system)})
for m in req.get("messages", []):
content = m.get("content")
if isinstance(content, str):
messages.append({"role": m["role"], "content": content})
continue
if m.get("role") == "user" and _has_image(content) and not any(
isinstance(b, dict) and b.get("type") == "tool_result" for b in content):
messages.append({"role": "user", "content": _parts_of(content)})
continue
text, reasoning, calls = [], [], []
for block in content or []:
kind = block.get("type")
if kind == "text":
text.append(block.get("text", ""))
elif kind == "thinking":
reasoning.append(block.get("thinking", ""))
elif kind == "tool_use":
calls.append({"function": {"name": block.get("name"), "arguments": block.get("input") or {}}})
elif kind == "tool_result":
messages.append({"role": "tool", "content": _text_of(block.get("content"))})
if text or calls or reasoning:
out = {"role": m["role"], "content": "".join(text)}
if reasoning:
out["reasoning_content"] = "".join(reasoning)
if calls:
out["tool_calls"] = calls
messages.append(out)
tools = [{"name": t["name"], "description": t.get("description", ""), "parameters": t.get("input_schema", {})}
for t in req.get("tools") or []] or None
kwargs = {}
# Anthropic: "thinking": {"type": "disabled"} or {"type": "enabled", "budget_tokens": N};
# "output_config": {"effort": "low" | "medium" | "high"}
thinking = req.get("thinking")
effort = (req.get("output_config") or {}).get("effort") if isinstance(req.get("output_config"), dict) else None
if isinstance(thinking, dict) and thinking.get("type") == "disabled":
kwargs["enable_thinking"] = False
elif effort:
kwargs.update(effort_kwargs(effort))
elif isinstance(thinking, dict) and thinking.get("budget_tokens"):
kwargs.update(budget_effort(thinking["budget_tokens"]))
return messages, tools, kwargs
# ------------------------------------------------------------------------------------------------ output parser
@dataclass
class ToolCall:
name: str
arguments: dict
id: str = field(default_factory=lambda: "call_" + uuid.uuid4().hex[:24])
@dataclass
class Event:
kind: str # "reasoning" | "content" | "tool_call"
text: str = ""
call: ToolCall | None = None
THINK_END = "</think>"
CALL_START = "<tool_call>"
CALL_END = "</tool_call>"
def parse_tool_call(body: str, schema: dict | None = None) -> ToolCall:
"""`<function=NAME>\\n<parameter=P>\\nVALUE\\n</parameter>...</function>` -> ToolCall. Values are JSON-decoded
when the tool's schema says the parameter is not a string (or, without a schema, when they parse as JSON
objects/arrays/numbers/booleans)."""
body = body.strip()
if not body.startswith("<function=") or ">" not in body:
raise ValueError("malformed tool call: " + body[:80])
name = body[len("<function="):body.index(">")]
rest = body[body.index(">") + 1:]
props = ((schema or {}).get("parameters") or {}).get("properties") or {}
args = {}
while "<parameter=" in rest:
rest = rest[rest.index("<parameter=") + len("<parameter="):]
pname = rest[:rest.index(">")]
rest = rest[rest.index(">") + 1:]
end = rest.find("</parameter>")
value = rest[:end] if end >= 0 else rest
rest = rest[end + len("</parameter>"):] if end >= 0 else ""
if value.startswith("\n"):
value = value[1:]
if value.endswith("\n"):
value = value[:-1]
declared = (props.get(pname) or {}).get("type")
if declared == "string":
args[pname] = value
else:
try:
args[pname] = json.loads(value)
except ValueError:
args[pname] = value
return ToolCall(name=name, arguments=args)
class OutputParser:
"""Incremental parser of the model's text. Feed deltas; get events. A tag split across deltas is held back
until it is complete, so clients never see `<tool_` or `</thi`."""
def __init__(self, thinking: bool = True, tools: list[dict] | None = None):
self.state = "reasoning" if thinking else "content"
self.buf = ""
self.lead = False
self.schemas = {t.get("name"): t for t in tools or []}
def _hold(self, text: str, tags: tuple[str, ...]) -> int:
"""Length of the longest suffix of `text` that is a proper prefix of one of `tags`."""
best = 0
for tag in tags:
for n in range(1, len(tag)):
if text.endswith(tag[:n]):
best = max(best, n)
return best
def feed(self, delta: str) -> list[Event]:
self.buf += delta
out: list[Event] = []
while True:
if self.state == "reasoning":
i = self.buf.find(THINK_END)
if i < 0:
keep = self._hold(self.buf, (THINK_END,))
if len(self.buf) > keep:
out.append(Event("reasoning", self.buf[:len(self.buf) - keep]))
self.buf = self.buf[len(self.buf) - keep:]
return out
if i:
out.append(Event("reasoning", self.buf[:i]))
self.buf = self.buf[i + len(THINK_END):]
self.state, self.lead = "content", True
elif self.state == "content":
if self.lead: # newlines right after </think> or a call
stripped = self.buf.lstrip("\n")
if not stripped:
self.buf = ""
return out
self.buf, self.lead = stripped, False
i = self.buf.find(CALL_START)
if i < 0:
# Hold a partial tag AND the newlines before it: if a tool call follows, they are dropped,
# so emitting them early would make streamed and whole outputs differ.
j = len(self.buf) - self._hold(self.buf, (CALL_START,))
while j > 0 and self.buf[j - 1] == "\n":
j -= 1
if j > 0:
out.append(Event("content", self.buf[:j]))
self.buf = self.buf[j:]
return out
if i and self.buf[:i].strip():
out.append(Event("content", self.buf[:i].rstrip("\n")))
self.buf = self.buf[i + len(CALL_START):]
self.state = "call"
else:
i = self.buf.find(CALL_END)
if i < 0:
return out
body = self.buf[:i]
self.buf = self.buf[i + len(CALL_END):]
name = body.strip()[len("<function="):].split(">", 1)[0]
out.append(Event("tool_call", call=parse_tool_call(body, self.schemas.get(name))))
self.state, self.lead = "content", True
def finish(self) -> list[Event]:
"""End of generation: flush whatever is held (an unterminated tool call is returned as content)."""
out = []
if self.buf:
kind = {"reasoning": "reasoning", "content": "content"}.get(self.state, "content")
text = self.buf if self.state != "call" else CALL_START + self.buf
out.append(Event(kind, text))
self.buf = ""
return out
serve/index.html
+152
View File
@@ -0,0 +1,152 @@
<!doctype html>
<html lang="en">
<head>
<meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<title>Strata</title>
<style>
:root { --bg:#f6f7f9; --card:#fff; --fg:#1b1f24; --mute:#667085; --line:#e3e6ea; --acc:#2f6fed; --me:#e9f0ff; --ok:#1a7f37; }
@media (prefers-color-scheme: dark) {
:root { --bg:#111418; --card:#1a1e24; --fg:#e7e9ec; --mute:#9aa3ae; --line:#2b313a; --acc:#6b9bff; --me:#1f2b44; --ok:#3fb950; }
}
* { box-sizing: border-box; }
body { margin:0; background:var(--bg); color:var(--fg); font:15px/1.5 system-ui, -apple-system, "Segoe UI", sans-serif; }
main { max-width:860px; margin:0 auto; padding:16px; display:flex; flex-direction:column; gap:12px; min-height:100vh; }
.card { background:var(--card); border:1px solid var(--line); border-radius:12px; padding:14px 16px; }
h1 { font-size:20px; margin:0; display:flex; align-items:center; gap:8px; }
.dot { width:10px; height:10px; border-radius:50%; background:var(--mute); display:inline-block; }
.dot.on { background:var(--ok); }
.mute { color:var(--mute); font-size:13px; }
code { background:var(--bg); border:1px solid var(--line); border-radius:6px; padding:1px 6px; font-size:13px; }
.grid { display:grid; grid-template-columns: max-content 1fr; gap:4px 12px; margin-top:8px; font-size:14px; }
#log { flex:1; display:flex; flex-direction:column; gap:10px; overflow-y:auto; min-height:240px; }
.msg { padding:10px 12px; border-radius:10px; white-space:pre-wrap; word-wrap:break-word; }
.msg.user { background:var(--me); align-self:flex-end; max-width:85%; }
.msg.bot { background:var(--card); border:1px solid var(--line); }
.msg img { max-width:220px; max-height:160px; display:block; border-radius:6px; margin-bottom:6px; }
details { color:var(--mute); font-size:13px; margin-bottom:6px; }
form { display:flex; gap:8px; align-items:flex-end; flex-wrap:wrap; }
textarea { flex:1; min-width:200px; min-height:52px; max-height:200px; resize:vertical; padding:10px; border-radius:10px;
border:1px solid var(--line); background:var(--card); color:var(--fg); font:inherit; }
button, label.btn { border:1px solid var(--line); background:var(--card); color:var(--fg); border-radius:10px; padding:10px 14px;
font:inherit; cursor:pointer; }
button.primary { background:var(--acc); border-color:var(--acc); color:#fff; }
button:disabled { opacity:.5; cursor:default; }
.row { display:flex; gap:14px; align-items:center; flex-wrap:wrap; font-size:13px; color:var(--mute); }
#att { font-size:13px; color:var(--mute); }
</style>
</head>
<body>
<main>
<div class="card">
<h1><span id="dot" class="dot"></span> Strata <span id="state" class="mute">connecting ...</span></h1>
<div class="grid">
<span class="mute">Model</span><span id="model">-</span>
<span class="mute">Context</span><span id="ctx">-</span>
<span class="mute">Images</span><span id="img">-</span>
<span class="mute">OpenAI API</span><span><code id="oai"></code> (any API key<span id="keynote"></span>)</span>
<span class="mute">Anthropic API</span><span><code id="ant"></code></span>
</div>
<p class="mute" style="margin:10px 0 0">Point any app with an "OpenAI-compatible" or "Anthropic" provider at these
addresses. Close the Strata window to stop the model.</p>
</div>
<div id="log" class="card"></div>
<div class="card">
<form id="f">
<textarea id="q" placeholder="Ask something ... (Enter sends, Shift+Enter = new line)"></textarea>
<label class="btn" id="attbtn" hidden>Picture<input id="file" type="file" accept="image/png,image/jpeg,image/gif,image/bmp" hidden></label>
<button class="primary" id="send">Send</button>
<button type="button" id="stop" disabled>Stop</button>
</form>
<div class="row" style="margin-top:8px">
<label>Thinking <select id="think" style="padding:4px;border-radius:8px;border:1px solid var(--line);background:var(--card);color:var(--fg)">
<option value="none">off - answer at once</option>
<option value="low">low</option>
<option value="medium" selected>medium</option>
<option value="high">high - best on hard questions, slowest</option>
</select></label>
<span id="att"></span>
<button type="button" id="clear" style="padding:4px 10px">New chat</button>
<input id="key" type="password" placeholder="API key" hidden style="padding:6px;border-radius:8px;border:1px solid var(--line)">
</div>
</div>
</main>
<script>
const $ = id => document.getElementById(id);
const base = location.origin;
$("oai").textContent = base + "/v1";
$("ant").textContent = base + "/v1/messages";
let history = [], image = null, ctrl = null;
async function status() {
try {
const h = await (await fetch("/health")).json();
$("dot").classList.add("on"); $("state").textContent = "running";
$("model").textContent = h.model; $("ctx").textContent = h.max_context.toLocaleString() + " tokens";
$("img").textContent = h.images ? "on" : "off (run setup again and answer yes to images)";
$("attbtn").hidden = !h.images;
if (h.api_key) { $("key").hidden = false; $("keynote").textContent = " - this server needs its key"; }
} catch (e) { $("state").textContent = "not reachable - is the Strata window still open?"; }
}
status();
function bubble(cls, text) {
const d = document.createElement("div"); d.className = "msg " + cls; d.textContent = text || "";
$("log").appendChild(d); $("log").scrollTop = $("log").scrollHeight; return d;
}
$("file").onchange = () => {
const f = $("file").files[0]; if (!f) return;
const r = new FileReader();
r.onload = () => { image = r.result; $("att").textContent = "picture: " + f.name; };
r.readAsDataURL(f);
};
$("clear").onclick = () => { history = []; $("log").innerHTML = ""; };
$("stop").onclick = () => ctrl && ctrl.abort();
$("q").addEventListener("keydown", e => { if (e.key === "Enter" && !e.shiftKey) { e.preventDefault(); $("f").requestSubmit(); } });
$("f").onsubmit = async e => {
e.preventDefault();
const text = $("q").value.trim(); if (!text) return;
$("q").value = "";
const content = image ? [{type: "image_url", image_url: {url: image}}, {type: "text", text}] : text;
const u = bubble("user", "");
if (image) { const im = document.createElement("img"); im.src = image; u.appendChild(im); }
u.appendChild(document.createTextNode(text));
history.push({role: "user", content});
image = null; $("att").textContent = ""; $("file").value = "";
const b = bubble("bot", ""); const det = document.createElement("details"); const sum = document.createElement("summary");
sum.textContent = "thinking"; det.appendChild(sum); const th = document.createElement("div"); det.appendChild(th);
const ans = document.createElement("div"); b.appendChild(ans);
$("send").disabled = true; $("stop").disabled = false; ctrl = new AbortController();
let answer = "";
try {
const headers = {"Content-Type": "application/json"};
if ($("key").value) headers["Authorization"] = "Bearer " + $("key").value;
const body = {model: "strata", messages: history, stream: true, max_tokens: 8192};
body.reasoning_effort = $("think").value;
const r = await fetch("/v1/chat/completions", {method: "POST", headers, body: JSON.stringify(body), signal: ctrl.signal});
if (!r.ok) { ans.textContent = "error: " + ((await r.json()).error || {}).message; throw 0; }
const rd = r.body.getReader(), dec = new TextDecoder(); let buf = "";
for (;;) {
const {value, done} = await rd.read(); if (done) break;
buf += dec.decode(value, {stream: true});
let i;
while ((i = buf.indexOf("\n\n")) >= 0) {
const line = buf.slice(0, i).trim(); buf = buf.slice(i + 2);
if (!line.startsWith("data: ") || line === "data: [DONE]") continue;
const d = JSON.parse(line.slice(6)).choices[0].delta;
if (d.reasoning_content) { if (!det.parentNode) b.insertBefore(det, ans); th.textContent += d.reasoning_content; }
if (d.content) { answer += d.content; ans.textContent = answer; }
$("log").scrollTop = $("log").scrollHeight;
}
}
} catch (err) { if (err && err.name === "AbortError") ans.textContent = answer + " [stopped]"; }
history.push({role: "assistant", content: answer});
$("send").disabled = false; $("stop").disabled = true; ctrl = null;
};
</script>
</body>
</html>
serve/server.py
+606
View File
@@ -0,0 +1,606 @@
"""serve/server.py - plan v0.3 P8: OpenAI and Anthropic endpoints over any engine that maps token ids to tokens.
python -m serve.server --engine mock --port 8095 (a scripted engine, for clients and tests)
python -m serve.server --engine strata --config strata.json --port 8080 (the real engine, resident)
Endpoints: POST /v1/chat/completions (OpenAI, stream and non-stream), POST /v1/messages (Anthropic, stream and
non-stream), GET /v1/models, GET /health. One sequence at a time behind a FIFO (plan: one resident sequence).
Images (optional, when the config has a "vision" entry): OpenAI image_url parts and Anthropic image blocks (base64
data, http(s) URLs or local file paths; JPEG/PNG/BMP/GIF) go through `strata-vision` (the model's mmproj file) and
reach the engine as embeddings (`GENI`).
Requests whose prompt plus max tokens exceed the engine's context are REJECTED with 400, never truncated.
The engine boundary is `Engine.generate(prompt_ids, max_new, sampling, cancel) -> iterator of token ids`.
`StrataEngine` keeps one `strata --serve` process resident (weights, expert arena and VRAM tier load once) and
talks to it over stdin/stdout; `MockEngine` is a scripted stand-in that makes every API path testable without a GPU.
"""
from __future__ import annotations
import argparse
import base64
import hashlib
import json
import os
import subprocess
import sys
import tempfile
import threading
import time
import urllib.request
import uuid
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from typing import Iterator, Protocol
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "tools"))
sys.path.insert(0, str(ROOT)) # run as a script (run-<model>.bat) as well as a module
from serve.frontend import (ChatTemplate, Event, OutputParser, anthropic_to_messages, # noqa: E402
images_of, openai_to_messages)
IM_END = "<|im_end|>"
IMAGE_PAD = "<|image_pad|>"
# ------------------------------------------------------------------------------------------------ engines
class Engine(Protocol):
max_context: int
def generate(self, ids: list[int], max_new: int, sampling: dict, cancel: threading.Event) -> Iterator[int]: ...
class MockEngine:
"""Replays a scripted completion (text) as token ids, one per step, then the end-of-turn token."""
def __init__(self, tokenizer, script: str, max_context: int = 32768, delay_s: float = 0.0):
self.tok, self.max_context, self.delay = tokenizer, max_context, delay_s
self.script = tokenizer.encode(script, parse_special=True) + tokenizer.encode(IM_END, parse_special=True)
self.last_prompt: list[int] = []
def generate(self, ids, max_new, sampling, cancel, embeddings=None):
self.last_prompt = list(ids)
self.last_embeddings = embeddings
for t in self.script[:max_new]:
if cancel.is_set():
return
if self.delay:
time.sleep(self.delay)
yield t
class StrataEngine:
"""The resident engine: `strata --serve` reads `GEN <max_new> <ids>` lines and streams `T <id>` lines, then
`DONE ...`. Requests are serialized by the service's FIFO, so one pipe is enough."""
def __init__(self, exe: str, args: list[str], cwd: str | None = None, log: str | None = None,
env: dict | None = None):
self.log = open(log, "a", encoding="utf-8") if log else subprocess.DEVNULL
self.proc = subprocess.Popen([exe, "--serve", *args], cwd=cwd, stdin=subprocess.PIPE, stdout=subprocess.PIPE,
stderr=self.log, text=True, encoding="utf-8", bufsize=1, env=env)
self.max_context = 0
self.last = {}
for line in self.proc.stdout:
if line.startswith("READY"):
self.max_context = int(line.split()[1])
break
if self.max_context <= 0:
raise RuntimeError("the engine exited before it was ready" + (f" (see {log})" if log else ""))
def generate(self, ids, max_new, sampling, cancel, embeddings=None):
head = f"GENI {int(max_new)} {embeddings}" if embeddings else f"GEN {int(max_new)}"
self.proc.stdin.write(f"{head} {','.join(str(int(t)) for t in ids)}\n")
self.proc.stdin.flush()
done = False
try:
for line in self.proc.stdout:
if line.startswith("T "):
if not cancel.is_set():
yield int(line[2:])
elif line.startswith("DONE"):
f = line.split()
self.last = {"generated": int(f[1]), "prompt_tokens": int(f[2]), "prompt_ms": float(f[3]),
"decode_ms": float(f[4]), "finish": f[5]}
done = True
return
elif line.startswith("ERR"):
done = True
raise ValueError(line[4:].strip())
raise RuntimeError("the engine process ended")
finally:
if not done: # the consumer stopped early: drain to DONE
for line in self.proc.stdout:
if line.startswith("DONE") or line.startswith("ERR"):
break
def close(self):
try:
self.proc.stdin.write("QUIT\n")
self.proc.stdin.flush()
self.proc.wait(timeout=10)
except Exception:
self.proc.kill()
class Vision:
"""The resident image encoder: `strata-vision` (llama.cpp mtmd + the mmproj file) reads `ENC <image> <out>`
lines and writes each image's embeddings; results are cached by the image's hash, so a conversation that
sends the same picture again (every turn, with most clients) encodes it once."""
def __init__(self, cfg: dict, log=None, env: dict | None = None):
args = [cfg["exe"], "--mmproj", cfg["mmproj"], "--model", cfg["model"]]
if cfg.get("gpu"):
args.append("--gpu")
if cfg.get("threads"):
args += ["--threads", str(cfg["threads"])]
if cfg.get("max_tokens"):
args += ["--max-tokens", str(cfg["max_tokens"])]
self.dir = Path(tempfile.mkdtemp(prefix="strata-vision-"))
self.proc = subprocess.Popen(args, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=log or subprocess.DEVNULL,
text=True, encoding="utf-8", bufsize=1, env=env)
line = self.proc.stdout.readline()
if not line.startswith("READY"):
raise RuntimeError("the vision encoder did not start: " + line.strip())
self.lock = threading.Lock()
self.cache: dict[str, tuple[Path, int]] = {}
@staticmethod
def load(source: str) -> bytes:
if source.startswith("data:"):
return base64.b64decode(source.split(",", 1)[1])
if source.startswith(("http://", "https://")):
req = urllib.request.Request(source, headers={"User-Agent": "strata"})
with urllib.request.urlopen(req, timeout=60) as r:
return r.read()
path = source[7:] if source.startswith("file://") else source
if path and os.path.isfile(path):
return Path(path).read_bytes()
raise ValueError("an image must be a data: URL, an http(s) URL or a local file path")
def encode(self, source: str) -> tuple[Path, int]:
"""-> (embeddings file, number of image tokens)."""
data = self.load(source)
key = hashlib.sha256(data).hexdigest()[:32]
with self.lock:
if key in self.cache:
return self.cache[key]
img, out = self.dir / f"{key}.img", self.dir / f"{key}.sve"
img.write_bytes(data)
self.proc.stdin.write(f"ENC {img} {out}\n")
self.proc.stdin.flush()
line = self.proc.stdout.readline().strip()
img.unlink(missing_ok=True)
if not line.startswith("OK"):
raise ValueError("the image could not be read: " + (line[4:] if line.startswith("ERR") else
"the vision encoder stopped"))
self.cache[key] = (out, int(line.split()[1]))
if len(self.cache) > 64: # oldest first
old = next(iter(self.cache))
self.cache.pop(old)[0].unlink(missing_ok=True)
return self.cache[key]
def close(self):
try:
self.proc.stdin.write("QUIT\n")
self.proc.stdin.flush()
self.proc.wait(timeout=10)
except Exception:
self.proc.kill()
def child_env(cfg: dict) -> dict:
"""The engine's environment: the CUDA libraries setup installed (pip's nvidia packages, or the toolkit that
compiled it) first on the library search path."""
env = dict(os.environ)
dirs = [d for d in cfg.get("lib_dirs") or [] if Path(d).is_dir()]
if dirs:
var = "PATH" if os.name == "nt" else "LD_LIBRARY_PATH"
env[var] = os.pathsep.join(dirs + ([env[var]] if env.get(var) else []))
return env
class ByteTokenizer:
"""Tiny stand-in tokenizer for tests without the pack: one id per UTF-8 byte, specials as ids >= 256."""
SPECIALS = ["<|im_start|>", "<|im_end|>", "<|endoftext|>", "<|vision_start|>", "<|image_pad|>", "<|vision_end|>"]
def encode(self, text, parse_special=False):
out, i = [], 0
while i < len(text):
for k, s in enumerate(self.SPECIALS):
if parse_special and text.startswith(s, i):
out.append(256 + k)
i += len(s)
break
else:
out.extend(text[i].encode("utf-8"))
i += 1
return out
def decode(self, ids, errors="replace"):
raw = bytearray()
for t in ids:
raw += self.SPECIALS[t - 256].encode() if t >= 256 else bytes([t])
return raw.decode("utf-8", errors=errors)
# ------------------------------------------------------------------------------------------------ core
class Detokenizer:
"""Incremental decode: re-decode the generated ids and emit only the new, complete suffix (a multi-byte
character split across tokens is held until complete)."""
def __init__(self, tok):
self.tok, self.ids, self.sent = tok, [], 0
def push(self, t: int) -> str:
self.ids.append(t)
text = self.tok.decode(self.ids)
if text.endswith("�"):
return ""
delta, self.sent = text[self.sent:], len(text)
return delta
class Service:
def __init__(self, engine: Engine, tokenizer, template: ChatTemplate, model_name: str = "qwen3.8-flash-next",
vision: Vision | None = None):
self.engine, self.tok, self.template, self.model, self.vision = engine, tokenizer, template, model_name, vision
self.fifo = threading.Lock()
self.embeddings = threading.local() # the current request's image embeddings file (GENI)
self.api_key = "" # when set, /v1/* needs it (Bearer or x-api-key)
self.stop_ids = set(tokenizer.encode(IM_END, parse_special=True) +
tokenizer.encode("<|endoftext|>", parse_special=True))
def prepare(self, messages, tools, kwargs, max_new):
prompt = self.template.render(messages, tools=tools, **kwargs)
ids = self.tok.encode(prompt, parse_special=True)
self.embeddings.path = None
images = images_of(messages)
if images:
if self.vision is None:
raise ValueError("this server was started without the vision encoder (run setup again and choose "
"'vision'), so it cannot read images")
pad = self.tok.encode(IMAGE_PAD, parse_special=True)[0]
encoded = [self.vision.encode(src) for src in images]
out, k = [], 0
for t in ids: # one <|image_pad|> per image -> one per image token
if t == pad and k < len(encoded):
out += [pad] * encoded[k][1]
k += 1
else:
out.append(t)
if k != len(encoded):
raise ValueError("the prompt and its images do not match")
ids = out
combined = self.vision.dir / f"req-{uuid.uuid4().hex[:12]}.sve"
with open(combined, "wb") as f:
for path, _ in encoded:
f.write(path.read_bytes())
self.embeddings.path = combined
if len(ids) + max_new > self.engine.max_context:
raise ValueError(f"prompt ({len(ids)} tokens) + max tokens ({max_new}) exceeds the context "
f"({self.engine.max_context}); requests are never truncated")
return ids, kwargs.get("enable_thinking", True) is not False
def run(self, ids, thinking, tools, max_new, sampling, cancel) -> Iterator[tuple[str, object]]:
"""Yields ("event", Event) as text arrives, then ("done", {"finish": .., "completion_tokens": ..})."""
parser, detok, n, finish = OutputParser(thinking=thinking, tools=tools), Detokenizer(self.tok), 0, "length"
emb = getattr(self.embeddings, "path", None)
try:
with self.fifo:
gen = self.engine.generate(ids, max_new, sampling, cancel, embeddings=emb) if emb else \
self.engine.generate(ids, max_new, sampling, cancel)
for t in gen:
n += 1
if t in self.stop_ids:
finish = "stop"
break
for ev in parser.feed(detok.push(t)):
yield "event", ev
finally:
if emb:
Path(emb).unlink(missing_ok=True)
for ev in parser.finish():
yield "event", ev
yield "done", {"finish": finish, "completion_tokens": n}
# ------------------------------------------------------------------------------------------------ OpenAI
def openai_chunks(svc: Service, req: dict, ids, thinking, tools, max_new, cancel):
cid, created = "chatcmpl-" + uuid.uuid4().hex[:24], int(time.time())
def chunk(delta, finish=None):
return {"id": cid, "object": "chat.completion.chunk", "created": created, "model": svc.model,
"choices": [{"index": 0, "delta": delta, "finish_reason": finish}]}
yield chunk({"role": "assistant", "content": ""})
calls = 0
for kind, x in svc.run(ids, thinking, tools, max_new, req, cancel):
if kind == "event":
ev: Event = x
if ev.kind == "reasoning" and ev.text:
yield chunk({"reasoning_content": ev.text})
elif ev.kind == "content" and ev.text:
yield chunk({"content": ev.text})
elif ev.kind == "tool_call":
yield chunk({"tool_calls": [{"index": calls, "id": ev.call.id, "type": "function",
"function": {"name": ev.call.name,
"arguments": json.dumps(ev.call.arguments, ensure_ascii=False)}}]})
calls += 1
else:
finish = "tool_calls" if calls and x["finish"] == "stop" else x["finish"]
last = chunk({}, finish)
last["usage"] = {"prompt_tokens": len(ids), "completion_tokens": x["completion_tokens"],
"total_tokens": len(ids) + x["completion_tokens"]}
yield last
def openai_collect(chunks) -> dict:
content, reasoning, calls, last = [], [], [], None
for c in chunks:
d = c["choices"][0]["delta"]
content.append(d.get("content") or "")
reasoning.append(d.get("reasoning_content") or "")
calls += [{k: v for k, v in tc.items() if k != "index"} for tc in d.get("tool_calls") or []]
last = c
msg = {"role": "assistant", "content": "".join(content) or None}
if "".join(reasoning):
msg["reasoning_content"] = "".join(reasoning)
if calls:
msg["tool_calls"] = calls
return {"id": last["id"], "object": "chat.completion", "created": last["created"], "model": last["model"],
"choices": [{"index": 0, "message": msg, "finish_reason": last["choices"][0]["finish_reason"]}],
"usage": last["usage"]}
# ------------------------------------------------------------------------------------------------ Anthropic
def anthropic_events(svc: Service, req: dict, ids, thinking, tools, max_new, cancel):
mid = "msg_" + uuid.uuid4().hex[:24]
yield "message_start", {"type": "message_start", "message": {
"id": mid, "type": "message", "role": "assistant", "model": svc.model, "content": [],
"stop_reason": None, "stop_sequence": None, "usage": {"input_tokens": len(ids), "output_tokens": 0}}}
index, open_kind, used_tool = -1, None, False
def close():
return ("content_block_stop", {"type": "content_block_stop", "index": index})
for kind, x in svc.run(ids, thinking, tools, max_new, req, cancel):
if kind == "event":
ev: Event = x
want = {"reasoning": "thinking", "content": "text", "tool_call": "tool_use"}[ev.kind]
if ev.kind != "tool_call" and not ev.text:
continue
if open_kind != want or want == "tool_use":
if open_kind is not None:
yield close()
index += 1
open_kind = want
block = {"thinking": {"type": "thinking", "thinking": "", "signature": ""},
"text": {"type": "text", "text": ""},
"tool_use": {"type": "tool_use", "id": ev.call.id if ev.call else "", "name":
ev.call.name if ev.call else "", "input": {}}}[want]
yield "content_block_start", {"type": "content_block_start", "index": index, "content_block": block}
if want == "thinking":
yield "content_block_delta", {"type": "content_block_delta", "index": index,
"delta": {"type": "thinking_delta", "thinking": ev.text}}
elif want == "text":
yield "content_block_delta", {"type": "content_block_delta", "index": index,
"delta": {"type": "text_delta", "text": ev.text}}
else:
used_tool = True
yield "content_block_delta", {"type": "content_block_delta", "index": index, "delta": {
"type": "input_json_delta", "partial_json": json.dumps(ev.call.arguments, ensure_ascii=False)}}
else:
if open_kind is not None:
yield close()
stop = "tool_use" if used_tool and x["finish"] == "stop" else \
{"stop": "end_turn", "length": "max_tokens"}[x["finish"]]
yield "message_delta", {"type": "message_delta", "delta": {"stop_reason": stop, "stop_sequence": None},
"usage": {"output_tokens": x["completion_tokens"]}}
yield "message_stop", {"type": "message_stop"}
def anthropic_collect(events) -> dict:
msg, blocks = None, []
for name, e in events:
if name == "message_start":
msg = e["message"]
elif name == "content_block_start":
blocks.append(dict(e["content_block"]))
elif name == "content_block_delta":
d, b = e["delta"], blocks[-1]
if d["type"] == "text_delta":
b["text"] += d["text"]
elif d["type"] == "thinking_delta":
b["thinking"] += d["thinking"]
else:
b["input"] = json.loads(d["partial_json"])
elif name == "message_delta":
msg["stop_reason"] = e["delta"]["stop_reason"]
msg["usage"]["output_tokens"] = e["usage"]["output_tokens"]
msg["content"] = blocks
return msg
# ------------------------------------------------------------------------------------------------ HTTP
def make_handler(svc: Service):
class Handler(BaseHTTPRequestHandler):
protocol_version = "HTTP/1.0" # SSE ends by closing the connection
def log_message(self, fmt, *args):
pass
def _json(self, code, obj):
body = json.dumps(obj, ensure_ascii=False).encode()
self.send_response(code)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(body)))
self.end_headers()
self.wfile.write(body)
def _authorized(self) -> bool:
if not svc.api_key:
return True
auth = self.headers.get("Authorization", "")
given = auth[7:].strip() if auth.lower().startswith("bearer ") else self.headers.get("x-api-key", "")
if given == svc.api_key:
return True
self._json(401, {"error": {"type": "authentication_error", "message": "missing or wrong API key"}})
return False
def do_GET(self):
path = self.path.split("?")[0].rstrip("/")
if path == "":
body = (ROOT / "serve" / "index.html").read_bytes()
self.send_response(200)
self.send_header("Content-Type", "text/html; charset=utf-8")
self.send_header("Content-Length", str(len(body)))
self.end_headers()
self.wfile.write(body)
elif path == "/health":
self._json(200, {"status": "ok", "max_context": svc.engine.max_context, "model": svc.model,
"images": svc.vision is not None, "api_key": bool(svc.api_key)})
elif path == "/v1/models":
if self._authorized():
self._json(200, {"object": "list", "data": [{"id": svc.model, "object": "model"}]})
else:
self._json(404, {"error": {"message": "not found"}})
def do_POST(self):
if not self._authorized():
return
try:
req = json.loads(self.rfile.read(int(self.headers.get("Content-Length", 0))) or b"{}")
if self.path.rstrip("/") == "/v1/chat/completions":
self._openai(req)
elif self.path.rstrip("/") == "/v1/messages":
self._anthropic(req)
else:
self._json(404, {"error": {"message": "not found"}})
except ValueError as e:
self._json(400, {"error": {"type": "invalid_request_error", "message": str(e)}})
def _sse(self):
self.send_response(200)
self.send_header("Content-Type", "text/event-stream")
self.send_header("Cache-Control", "no-cache")
self.end_headers()
def _openai(self, req):
messages, tools, kw = openai_to_messages(req)
max_new = int(req.get("max_completion_tokens") or req.get("max_tokens") or 1024)
ids, thinking = svc.prepare(messages, tools, kw, max_new)
cancel = threading.Event()
chunks = openai_chunks(svc, req, ids, thinking, tools, max_new, cancel)
if not req.get("stream"):
return self._json(200, openai_collect(chunks))
self._sse()
try:
for c in chunks:
self.wfile.write(b"data: " + json.dumps(c, ensure_ascii=False).encode() + b"\n\n")
self.wfile.flush()
self.wfile.write(b"data: [DONE]\n\n")
except OSError:
cancel.set() # client went away: stop at the next step
def _anthropic(self, req):
messages, tools, kw = anthropic_to_messages(req)
max_new = int(req.get("max_tokens") or 1024)
ids, thinking = svc.prepare(messages, tools, kw, max_new)
cancel = threading.Event()
events = anthropic_events(svc, req, ids, thinking, tools, max_new, cancel)
if not req.get("stream"):
return self._json(200, anthropic_collect(events))
self._sse()
try:
for name, e in events:
self.wfile.write(f"event: {name}\n".encode() + b"data: " +
json.dumps(e, ensure_ascii=False).encode() + b"\n\n")
self.wfile.flush()
except OSError:
cancel.set()
return Handler
class Server(ThreadingHTTPServer):
# On Windows SO_REUSEADDR lets a second server bind a port that is already serving, and requests then land on
# either one (a forgotten second start of run-<model>.bat). Without it the second start fails loudly instead.
allow_reuse_address = os.name != "nt"
def serve(svc: Service, host="127.0.0.1", port=8095) -> ThreadingHTTPServer:
httpd = Server((host, port), make_handler(svc))
threading.Thread(target=httpd.serve_forever, daemon=True).start()
return httpd
def main() -> int:
ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
ap.add_argument("--engine", choices=["mock", "strata"], default="mock")
ap.add_argument("--config", help="strata engine config (JSON: exe, args, cwd, tokenizer, model_name), "
"written by setup.py")
ap.add_argument("--host", default="127.0.0.1")
ap.add_argument("--script", default="Thinking about it.</think>\n\nHello from the mock engine.")
ap.add_argument("--port", type=int, default=8095)
ap.add_argument("--tokenizer", default=str(ROOT / "pack/full/tokenizer"),
help="pack tokenizer directory (falls back to a byte tokenizer if absent)")
ap.add_argument("--open", action="store_true", help="open the local page in the browser once the model is ready")
ap.add_argument("--api-key", default=os.environ.get("STRATA_API_KEY", ""),
help="require this key on /v1/* (Authorization: Bearer ... or x-api-key); also $STRATA_API_KEY")
a = ap.parse_args()
cfg = json.loads(Path(a.config).read_text(encoding="utf-8")) if a.config else {}
try: # before the minutes of loading: is the port free?
Server((a.host, a.port), BaseHTTPRequestHandler).server_close()
except OSError:
ap.error(f"port {a.port} is already in use - is Strata (or another server) already running? "
f"Close it, or start this one with a different --port")
if cfg.get("tokenizer"):
a.tokenizer = cfg["tokenizer"]
tok = ByteTokenizer()
tpath = Path(a.tokenizer)
if a.engine == "strata" and not (tpath / "vocab.json").exists():
ap.error(f"the model's tokenizer is missing ({tpath / 'vocab.json'}); run setup again")
if (tpath / "vocab.json").exists():
import strata_tokenizer as ST
vocab = json.loads((tpath / "vocab.json").read_text(encoding="utf-8"))
tokens = [None] * len(vocab)
for t, i in vocab.items():
tokens[i] = t
merges = (tpath / "merges.txt").read_text(encoding="utf-8").split("\n")
types = json.loads((tpath / "token_type.json").read_text())
tok = ST.Tokenizer(tokens, merges, types)
if a.engine == "strata":
if not cfg:
ap.error("--engine strata needs --config")
vision = None
env = child_env(cfg)
if cfg.get("vision"):
print("loading the vision encoder ...", flush=True)
vision = Vision(cfg["vision"], log=open(cfg["log"], "a", encoding="utf-8") if cfg.get("log") else None,
env=env)
print("loading the model (the first start takes a minute or two) ...", flush=True)
engine = StrataEngine(cfg["exe"], cfg["args"], cwd=cfg.get("cwd"), log=cfg.get("log"), env=env)
else:
engine, vision = MockEngine(tok, a.script), None
svc = Service(engine, tok, ChatTemplate(ROOT / "serve/chat_template.jinja"),
model_name=cfg.get("model_name", "qwen3.8-flash-next"), vision=vision)
svc.api_key = a.api_key or cfg.get("api_key", "")
httpd = serve(svc, host=a.host, port=a.port)
print(f"ready: http://{a.host}:{a.port}/v1 (OpenAI: /v1/chat/completions, Anthropic: /v1/messages, "
f"context {engine.max_context} tokens{', images on' if vision else ''}"
f"{', API key required' if svc.api_key else ''})", flush=True)
print(f" open http://{a.host}:{a.port}/ in a browser to chat; close this window to stop the model", flush=True)
if a.open:
import webbrowser
webbrowser.open(f"http://{'127.0.0.1' if a.host in ('0.0.0.0', '') else a.host}:{a.port}/")
try:
threading.Event().wait()
except KeyboardInterrupt:
httpd.shutdown()
if hasattr(engine, "close"):
engine.close()
if vision:
vision.close()
return 0
if __name__ == "__main__":
sys.exit(main())
setup.py
+806
View File
@@ -0,0 +1,806 @@
#!/usr/bin/env python3
"""Strata one-click setup and start (Windows and Linux, NVIDIA GPUs).
START-HERE.bat (Windows) / ./setup.sh (Linux) - they install Python if needed and run this file
The first time it asks three questions - which model, how much context, and whether the model should also read
images - then installs everything and starts the model on http://127.0.0.1:8080 (OpenAI- and Anthropic-compatible
API; a small page there shows that it runs). Every later start skips straight to running the model: nothing that
is already downloaded, installed or prepared is done again.
What the first run does (each step is skipped when it is already done):
1. checks your PC: NVIDIA GPU and driver, RAM, CPU, free disk space
2. asks the three questions
3. installs the Python packages it needs into .venv (numpy, jinja2, ..., and NVIDIA's CUDA libraries)
4. gets the Strata engine: a ready-made build for RTX 30/40/50 cards (no compiler needed); if none fits your PC,
it installs the build tools (asks first) and compiles the engine for your GPU
5. downloads the model from Hugging Face (resumable), and the vision encoder if you want images
6. prepares the model for Strata and fetches the MTP draft layer (~5 GB, from the original Qwen checkpoint)
7. writes run-<model>.bat / run-<model>.sh and starts the model
Options: --model Q2_0|IQ2_XS|IQ3_XXS, --context 32768, --vision yes|no|gpu|cpu, --port 8080, --yes (recommended
answers, no questions), --setup (install another model / change settings instead of starting), --no-start,
--models-dir DIR, --gguf-dir DIR (use GGUF files you already have), --build (compile instead of the ready-made
engine), --check (only check this PC).
"""
from __future__ import annotations
import argparse
import ctypes
import json
import os
import platform
import re
import shutil
import subprocess
import sys
import time
import urllib.request
import zipfile
from pathlib import Path
ROOT = Path(__file__).resolve().parent
WIN = os.name == "nt"
HF = "https://huggingface.co/ISTA-DASLab/Qwen3.8-Flash-Next-GSQ-RCO-GGUF/resolve/main/"
LLAMA_CPP_COMMIT = "3cf03257f219afbe7334045ff7c6a06ac68c627d"
LLAMA_CPP_ZIP = f"https://github.com/ggml-org/llama.cpp/archive/{LLAMA_CPP_COMMIT}.zip"
# The ready-made engine: <PREBUILT_URL><asset>, a zip with strata(.exe), strata-vision(.exe) and BUILD.json, built
# by tools/make_release.py. Set this to the GitHub release download folder when publishing, e.g.
# "https://github.com/<you>/Strata/releases/latest/download/" (or pass --prebuilt / set STRATA_PREBUILT_URL).
PREBUILT_URL = "https://github.com/Niko1221/Strata/releases/latest/download/"
PREBUILT_ASSET = "strata-windows-x64.zip" if WIN else "strata-linux-x64.zip"
# the CUDA libraries the ready-made engine loads (the same CUDA 13.0 it is built with), from NVIDIA's pip packages
CUDA_WHEELS = ["nvidia-cublas==13.0.2.14", "nvidia-cuda-runtime==13.0.96"]
MIN_DRIVER = 580 # CUDA 13.0
PY_PACKAGES = ["numpy", "jinja2", "regex", "pyyaml", "tqdm", "requests", "cmake", "ninja"]
MODELS = {
"Q2_0": {"about": "2-bit, the fastest", "download_gb": 66.4, "ram_gb": 48, "arena_gb": 34.0},
"IQ2_XS": {"about": "2-bit i-quant, a little better quality, close in speed", "download_gb": 68.0, "ram_gb": 48,
"arena_gb": 35.5},
"IQ3_XXS": {"about": "3-bit i-quant, the best quality, slower (more CPU work per token)", "download_gb": 75.8,
"ram_gb": 60, "arena_gb": 42.9},
}
CONTEXTS = [8192, 32768, 65536, 131072, 262144]
MMPROJ = "mmproj-Qwen3.8-Flash-Next-BF16.gguf"
# the image encoder on the GPU: ~0.9 GB of weights + ~0.3 GB of work buffers at 1024 image tokens, kept free of
# expert slots (the engine's default reserve is 700 MiB)
VISION = {"gpu": {"max_tokens": 1024, "reserve_mib": 2100},
"cpu": {"max_tokens": 300, "reserve_mib": 700}}
EXE = "strata.exe" if WIN else "strata"
VEXE = "strata-vision.exe" if WIN else "strata-vision"
# ------------------------------------------------------------------------------------------------ output
def say(msg=""):
print(msg, flush=True)
def step(n, title):
say()
say(f"=== Step {n}: {title} ===")
def ok(msg):
say(f" [ok] {msg}")
def warn(msg):
say(f" [!] {msg}")
def fail(msg, hint=None):
say(f"\n [X] {msg}")
if hint:
say(f" {hint}")
say("\nSetup stopped. Fix the item above and run it again - everything already done is kept and skipped.")
sys.exit(1)
def ask(question, choices, default, yes):
if yes:
return default
while True:
try:
a = input(f"{question} [{default}]: ").strip()
except EOFError:
return default
if not a:
return default
if a.lower() in [c.lower() for c in choices]:
return next(c for c in choices if c.lower() == a.lower())
say(f" please answer one of: {', '.join(choices)}")
def run(cmd, cwd=None, env=None, check=True, quiet=False):
say(" > " + " ".join(str(c) for c in cmd))
r = subprocess.run([str(c) for c in cmd], cwd=cwd, env=env,
stdout=subprocess.PIPE if quiet else None, stderr=subprocess.STDOUT if quiet else None,
text=True)
if check and r.returncode != 0:
if quiet and r.stdout:
say(r.stdout[-4000:])
fail(f"command failed (exit {r.returncode}): {Path(str(cmd[0])).name}")
return r
def out(cmd):
try:
return subprocess.run(cmd, capture_output=True, text=True, timeout=60).stdout
except (OSError, subprocess.TimeoutExpired):
return ""
def done(path: Path) -> bool:
"""A step's finish mark: <path>.done exists (written only after the step completed)."""
return path.with_name(path.name + ".done").exists()
def mark(path: Path, text=""):
path.with_name(path.name + ".done").write_text(text or time.strftime("%Y-%m-%d %H:%M"), encoding="utf-8")
# ------------------------------------------------------------------------------------------------ the PC
def ram_gb():
if WIN:
class MS(ctypes.Structure):
_fields_ = [("dwLength", ctypes.c_ulong), ("dwMemoryLoad", ctypes.c_ulong),
("ullTotalPhys", ctypes.c_ulonglong), ("ullAvailPhys", ctypes.c_ulonglong),
("ullTotalPageFile", ctypes.c_ulonglong), ("ullAvailPageFile", ctypes.c_ulonglong),
("ullTotalVirtual", ctypes.c_ulonglong), ("ullAvailVirtual", ctypes.c_ulonglong),
("ullAvailExtendedVirtual", ctypes.c_ulonglong)]
m = MS()
m.dwLength = ctypes.sizeof(MS)
ctypes.windll.kernel32.GlobalMemoryStatusEx(ctypes.byref(m))
return m.ullTotalPhys / 2**30
for line in open("/proc/meminfo"):
if line.startswith("MemTotal"):
return int(line.split()[1]) * 1024 / 2**30
return 0.0
def cpu_info():
"""(name, avx2, avx512): avx512 means everything Strata's fast AVX-512 kernels use (F, BW, VL, VNNI, VBMI),
the same test the engine makes (cpu_avx512_ok), not just AVX-512F."""
name, avx2, avx512 = platform.processor() or "unknown CPU", False, False
if WIN:
pf = ctypes.windll.kernel32.IsProcessorFeaturePresent
avx2 = bool(pf(40)) # PF_AVX2_INSTRUCTIONS_AVAILABLE
n = out(["powershell", "-NoProfile", "-Command", "(Get-CimInstance Win32_Processor).Name"]).strip()
name = n or name
avx512 = bool(pf(41)) and _cpuid_avx512_full()
else:
try:
txt = open("/proc/cpuinfo").read()
flags = set(re.search(r"^flags\s*:\s*(.*)$", txt, re.M).group(1).split())
avx2 = "avx2" in flags
avx512 = {"avx512f", "avx512bw", "avx512vl", "avx512_vnni", "avx512vbmi"} <= flags
m = re.search(r"^model name\s*:\s*(.*)$", txt, re.M)
name = m.group(1) if m else name
except OSError:
pass
return name, avx2, avx512
def _cpuid_avx512_full() -> bool:
"""Windows has no feature bit for VNNI / VBMI: ask the CPU (CPUID leaf 7) through a tiny machine-code stub."""
try:
code = bytes([0x53, 0x49, 0x89, 0xC8, 0xB8, 0x07, 0x00, 0x00, 0x00, 0x31, 0xC9, 0x0F, 0xA2, # push rbx; r8=rcx; cpuid(7,0)
0x41, 0x89, 0x18, 0x41, 0x89, 0x48, 0x04, 0x5B, 0xC3]) # [r8]=ebx,[r8+4]=ecx; pop rbx
k32 = ctypes.windll.kernel32
k32.VirtualAlloc.restype = ctypes.c_void_p
buf = k32.VirtualAlloc(None, len(code), 0x3000, 0x40)
if not buf:
return False
ctypes.memmove(buf, code, len(code))
regs = (ctypes.c_uint32 * 2)()
ctypes.CFUNCTYPE(None, ctypes.c_void_p)(buf)(ctypes.addressof(regs))
ebx, ecx = regs[0], regs[1]
need_ebx = (1 << 16) | (1 << 30) | (1 << 31) # F, BW, VL
need_ecx = (1 << 1) | (1 << 11) # VBMI, VNNI
return (ebx & need_ebx) == need_ebx and (ecx & need_ecx) == need_ecx
except Exception:
return False
def gpu_info():
s = out(["nvidia-smi", "--query-gpu=name,memory.total,compute_cap,driver_version", "--format=csv,noheader,nounits"])
if not s.strip():
return None
name, mem, cc, drv = [x.strip() for x in s.strip().splitlines()[0].split(",")]
return {"name": name, "vram_gb": float(mem) / 1024.0, "arch": cc.replace(".", ""), "driver": drv}
def find_nvcc():
cands = [shutil.which("nvcc")]
if os.environ.get("CUDA_PATH"):
cands.append(str(Path(os.environ["CUDA_PATH"]) / "bin" / ("nvcc.exe" if WIN else "nvcc")))
if WIN:
base = Path(r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA")
if base.exists():
cands += [str(p / "bin" / "nvcc.exe") for p in sorted(base.iterdir(), reverse=True)]
else:
cands += [str(p / "bin" / "nvcc") for p in sorted(Path("/usr/local").glob("cuda*"), reverse=True)]
best = (None, None)
for c in dict.fromkeys(cands): # every toolkit found; the newest wins
if c and Path(c).exists():
v = re.search(r"release (\d+)\.(\d+)", out([c, "--version"]))
if v and (best[1] is None or (int(v.group(1)), int(v.group(2))) > best[1]):
best = (c, (int(v.group(1)), int(v.group(2))))
return best
def find_vcvars():
vswhere = Path(os.environ.get("ProgramFiles(x86)", r"C:\Program Files (x86)")) / "Microsoft Visual Studio/Installer/vswhere.exe"
if not vswhere.exists():
return None
p = out([str(vswhere), "-latest", "-products", "*", "-requires",
"Microsoft.VisualStudio.Component.VC.Tools.x86.x64", "-property", "installationPath"]).strip()
v = Path(p) / "VC/Auxiliary/Build/vcvars64.bat" if p else None
return v if v and v.exists() else None
def find_tool(name):
"""A tool on PATH, or the one pip installed next to this Python (cmake, ninja)."""
p = shutil.which(name)
if p:
return p
for d in (Path(sys.executable).parent / "Scripts", Path(sys.executable).parent,
Path.home() / ".local" / "bin"):
c = d / (name + (".exe" if WIN else ""))
if c.exists():
return str(c)
return None
def free_gb(path):
path.mkdir(parents=True, exist_ok=True)
return shutil.disk_usage(path).free / 1e9
# ------------------------------------------------------------------------------------------------ downloads
def download(url, dst: Path, what=None):
"""Resumable HTTP(S) download with a progress line; `file://` and plain paths are copied (tests, mirrors).
A finished file gets a <name>.done mark, so a later run skips it without asking the server."""
dst.parent.mkdir(parents=True, exist_ok=True)
if dst.exists() and done(dst):
ok(f"{what or dst.name} already downloaded")
return
if not url.startswith(("http://", "https://")):
src = Path(url[7:] if url.startswith("file://") else url)
if not src.exists():
fail(f"not found: {src}")
shutil.copyfile(src, dst)
mark(dst)
ok(f"{what or dst.name} copied")
return
part = dst.with_name(dst.name + ".part")
total = 0
for attempt in range(5):
try:
req = urllib.request.Request(url, method="HEAD", headers={"User-Agent": "strata-setup"})
total = int(urllib.request.urlopen(req, timeout=60).headers.get("Content-Length", 0))
break
except OSError as e:
if attempt == 4:
fail(f"cannot reach {url.split('/')[2]} ({e})", "check your internet connection and run it again")
time.sleep(5)
if dst.exists() and total and dst.stat().st_size == total: # finished by an older setup (no mark yet)
mark(dst)
ok(f"{what or dst.name} already downloaded")
return
have = part.stat().st_size if part.exists() else 0
for attempt in range(30):
try:
req = urllib.request.Request(url, headers={"User-Agent": "strata-setup", "Range": f"bytes={have}-"})
with urllib.request.urlopen(req, timeout=60) as r, open(part, "ab" if have else "wb") as f:
if have and r.status != 206: # the server ignored the range: start over
f.seek(0)
f.truncate()
have = 0
last = 0.0
while True:
b = r.read(8 << 20)
if not b:
break
f.write(b)
have += len(b)
if time.time() - last > 2:
last = time.time()
size = f"{have / 1e9:6.2f} / {total / 1e9:.2f} GB ({100 * have / total:.0f}%)" if total \
else f"{have / 1e6:7.1f} MB"
print(f"\r {what or dst.name}: {size} ", end="", flush=True)
print()
if not total or have >= total:
break
except OSError as e:
print()
warn(f"download interrupted ({e}); retrying in 10 s ...")
time.sleep(10)
if total and part.stat().st_size != total:
fail(f"could not finish downloading {dst.name}", "check your internet connection and run it again")
part.replace(dst)
mark(dst)
ok(f"{what or dst.name} downloaded")
def get_llama_cpp():
"""llama.cpp at the pinned commit (ggml for the build, gguf-py for the tools, mtmd for images), as a zip: no git."""
llama = ROOT / "third_party" / "llama.cpp"
if (llama / "ggml" / "CMakeLists.txt").exists() and (llama / "gguf-py").is_dir():
return llama
z = ROOT / "third_party" / f"llama.cpp-{LLAMA_CPP_COMMIT[:7]}.zip"
download(LLAMA_CPP_ZIP, z, "llama.cpp source")
tmp = ROOT / "third_party" / "_unpack"
shutil.rmtree(tmp, ignore_errors=True)
with zipfile.ZipFile(z) as f:
f.extractall(tmp)
top = next(tmp.iterdir())
shutil.rmtree(llama, ignore_errors=True)
top.replace(llama)
shutil.rmtree(tmp, ignore_errors=True)
z.unlink(missing_ok=True)
z.with_name(z.name + ".done").unlink(missing_ok=True)
return llama
def pip_install(packages, what):
"""pip install into .venv, skipped when the same list was installed before."""
stamp = Path(sys.prefix) / ".strata-pip.json"
have = json.loads(stamp.read_text()) if stamp.exists() else []
need = [p for p in packages if p not in have]
if not need:
ok(f"{what} already installed")
return
say(f" Installing {what} ...")
run([sys.executable, "-m", "pip", "install", "--quiet", "--disable-pip-version-check", *need])
stamp.write_text(json.dumps(sorted(set(have) | set(need)), indent=0))
ok(f"{what} installed")
def cuda_lib_dirs():
"""Where pip put NVIDIA's CUDA libraries (nvidia/cu13/bin/x86_64 on Windows, nvidia/cu13/lib on Linux)."""
pattern = "cublas64_13.dll" if WIN else "libcublas.so.13*"
dirs = []
for sp in {Path(p) for p in sys.path if p.endswith("site-packages")}:
for hit in (sp / "nvidia").rglob(pattern) if (sp / "nvidia").is_dir() else []:
if hit.parent not in dirs:
dirs.append(hit.parent)
return [str(d) for d in dirs]
# ------------------------------------------------------------------------------------------------ the engine
def driver_major(gpu):
try:
return int(gpu["driver"].split(".")[0])
except (ValueError, KeyError):
return 0
def get_prebuilt(url_base, gpu, vision) -> Path | None:
"""The ready-made engine in engine/ (kept between runs), or None when there is none for this PC."""
eng = ROOT / "engine"
info = eng / "BUILD.json"
if info.exists() and (eng / EXE).exists():
ok("ready-made engine already installed")
return eng
if not url_base:
return None
z = ROOT / "engine" / PREBUILT_ASSET
base = url_base if url_base.endswith(("/", "\\")) else url_base + "/"
if base.startswith(("http://", "https://")):
try: # not published (yet), or no internet: compile instead
req = urllib.request.Request(base + PREBUILT_ASSET, method="HEAD", headers={"User-Agent": "strata-setup"})
urllib.request.urlopen(req, timeout=60).close()
except OSError as e:
warn(f"no ready-made engine at {base} ({e}): compiling instead")
return None
say(" Downloading the ready-made Strata engine ...")
download(base + PREBUILT_ASSET, z, "Strata engine")
tmp = ROOT / "engine" / "_unpack"
shutil.rmtree(tmp, ignore_errors=True)
with zipfile.ZipFile(z) as f:
f.extractall(tmp)
meta = json.loads((tmp / "BUILD.json").read_text())
archs = [int(a) for a in meta.get("archs", [])]
arch = int(gpu["arch"])
if arch not in archs and not (meta.get("ptx") and arch > max(archs)):
warn(f"the ready-made engine is built for {', '.join(str(a) for a in archs)}; your GPU is {arch}: compiling instead")
shutil.rmtree(tmp, ignore_errors=True)
return None
for p in tmp.iterdir():
dst = eng / p.name
if dst.exists():
shutil.rmtree(dst) if dst.is_dir() else dst.unlink()
p.replace(dst)
shutil.rmtree(tmp, ignore_errors=True)
z.unlink(missing_ok=True)
z.with_name(z.name + ".done").unlink(missing_ok=True)
if not (eng / EXE).exists():
fail("the ready-made engine archive has no " + EXE)
if not WIN:
for x in (EXE, VEXE):
if (eng / x).exists():
(eng / x).chmod(0o755)
ok(f"ready-made engine {meta.get('version', '')} for {', '.join('sm_' + str(a) for a in archs)} (CUDA "
f"{meta.get('cuda', '?')})")
return eng
def install_build_tools(gpu, yes):
"""The compiler and the CUDA toolkit, installed for the user (asks once). Returns (nvcc, vcvars)."""
nvcc, cuda_v = find_nvcc()
need_cuda = (12, 8) if int(gpu["arch"]) >= 120 else (12, 0)
vcvars = find_vcvars() if WIN else None
have_cc = vcvars is not None if WIN else shutil.which("g++") is not None
missing = []
if not have_cc:
missing.append("Visual Studio 2022 Build Tools (C++)" if WIN else "the C++ compiler (build-essential)")
if nvcc is None or cuda_v < need_cuda:
missing.append("the NVIDIA CUDA Toolkit 13.0")
if not missing:
ok(f"build tools present (CUDA {cuda_v[0]}.{cuda_v[1]})")
return nvcc, vcvars
say(" The engine has to be compiled for your PC, which needs: " + " and ".join(missing) + ".")
say(" They can be installed now (about 8-10 GB, 15-40 minutes" + (", Windows will ask for permission" if WIN else
", sudo will ask for your password") + ").")
if ask(" Install them now?", ["y", "n"], "y", yes) != "y":
fail("the build tools are needed", "install them yourself (see README.md) and run it again")
if WIN:
if shutil.which("winget") is None:
fail("winget (Windows package manager) is not available",
"install 'App Installer' from the Microsoft Store, or install the tools by hand (README.md)")
wg = ["winget", "install", "-e", "--source", "winget", "--accept-package-agreements",
"--accept-source-agreements", "--disable-interactivity"]
if not have_cc:
run([*wg, "--id", "Microsoft.VisualStudio.2022.BuildTools", "--override",
"--quiet --wait --norestart --nocache --add Microsoft.VisualStudio.Workload.VCTools --includeRecommended"],
check=False)
if nvcc is None or cuda_v < need_cuda:
run([*wg, "--id", "Nvidia.CUDA", "--version", "13.0"], check=False)
vcvars = find_vcvars()
else:
apt = shutil.which("apt-get")
if apt is None:
fail("automatic install is only done on Ubuntu/Debian",
"install g++ and the CUDA Toolkit 13 (https://developer.nvidia.com/cuda-downloads), then run it again")
if not have_cc:
run(["sudo", "apt-get", "install", "-y", "build-essential"])
if nvcc is None or cuda_v < need_cuda:
osr = dict(line.split("=", 1) for line in open("/etc/os-release").read().splitlines() if "=" in line)
ver = osr.get("VERSION_ID", "").strip('"').replace(".", "")
if osr.get("ID") != "ubuntu" or ver not in ("2204", "2404"):
fail("the CUDA Toolkit can be installed automatically on Ubuntu 22.04 / 24.04 only",
"install it from https://developer.nvidia.com/cuda-downloads and run it again")
deb = Path("/tmp/cuda-keyring.deb")
download(f"https://developer.download.nvidia.com/compute/cuda/repos/ubuntu{ver}/x86_64/cuda-keyring_1.1-1_all.deb",
deb, "CUDA repository key")
run(["sudo", "dpkg", "-i", str(deb)])
run(["sudo", "apt-get", "update"])
run(["sudo", "apt-get", "install", "-y", "cuda-toolkit-13-0"])
nvcc, cuda_v = find_nvcc()
if (WIN and find_vcvars() is None) or (not WIN and shutil.which("g++") is None):
fail("the C++ build tools did not install", "install them by hand (README.md) and run it again")
if nvcc is None or cuda_v < need_cuda:
fail("the CUDA Toolkit did not install", "install it from https://developer.nvidia.com/cuda-downloads, then run it again")
ok(f"build tools installed (CUDA {cuda_v[0]}.{cuda_v[1]})")
return nvcc, find_vcvars() if WIN else None
def cmake_build(src, bdir, target, defs, vcvars, bat_name):
cmake, ninja = find_tool("cmake"), find_tool("ninja")
if cmake is None or ninja is None:
fail("cmake / ninja not found after installing them", "run: .venv python -m pip install cmake ninja")
conf = [cmake, "-G", "Ninja", f"-DCMAKE_MAKE_PROGRAM={ninja}", "-S", str(src), "-B", str(bdir),
"-DCMAKE_BUILD_TYPE=Release", *defs]
build = [cmake, "--build", str(bdir), "--target", target, "-j", str(max(2, (os.cpu_count() or 4) // 2))]
if WIN:
bat = ROOT / bat_name
q = lambda c: " ".join(f'"{x}"' if " " in str(x) else str(x) for x in c) # noqa: E731
bat.write_text(f'@echo off\r\ncall "{vcvars}" >nul\r\n{q(conf)} || exit /b 1\r\n{q(build)} || exit /b 1\r\n',
encoding="utf-8")
run(["cmd", "/c", str(bat)])
else:
run(conf)
run(build)
def build_engine(gpu, vision, yes, llama) -> Path:
"""Compile the engine (and, for images, the encoder) for this GPU; the results go to engine/."""
eng = ROOT / "engine"
eng.mkdir(exist_ok=True)
stamp = eng / "BUILD.json"
meta = json.loads(stamp.read_text()) if stamp.exists() else {}
want_vision = vision != "none"
if meta.get("source") == "local" and (eng / EXE).exists() and (not want_vision or (eng / VEXE).exists()):
ok("engine already built for this PC")
return eng
nvcc, vcvars = install_build_tools(gpu, yes)
if not (eng / EXE).exists() or meta.get("source") != "local":
say(" Compiling the Strata engine for your GPU (10-20 minutes, once) ...")
cmake_build(ROOT, ROOT / "build", "strata",
["-DSTRATA_ENABLE_CUDA=ON", "-DSTRATA_BUILD_TESTS=OFF", f"-DCMAKE_CUDA_ARCHITECTURES={gpu['arch']}",
f"-DCMAKE_CUDA_COMPILER={nvcc}", f"-DSTRATA_GGML_DIR={llama}"], vcvars, "build-strata.bat")
shutil.copy2(ROOT / "build" / EXE, eng / EXE)
if want_vision and not (eng / VEXE).exists():
say(" Compiling the image encoder" + (" with CUDA (10-20 minutes, once) ..." if vision == "gpu" else " ..."))
defs = [f"-DLLAMA_DIR={llama}", f"-DSTRATA_VISION_CUDA={'ON' if vision == 'gpu' else 'OFF'}"]
if vision == "gpu":
defs += [f"-DCMAKE_CUDA_ARCHITECTURES={gpu['arch']}", f"-DCMAKE_CUDA_COMPILER={nvcc}"]
cmake_build(ROOT / "tools" / "vision", ROOT / "build-vision", "strata-vision", defs, vcvars, "build-vision.bat")
shutil.copy2(ROOT / "build-vision" / "bin" / VEXE, eng / VEXE)
bindir = Path(nvcc).parent # the toolkit's own libraries (bin, bin/x64, lib64)
dirs = [str(d) for d in (bindir, bindir / "x64", bindir.parent / "lib64") if d.is_dir()]
stamp.write_text(json.dumps({"source": "local", "archs": [int(gpu["arch"])], "vision": vision,
"cuda_dirs": dirs}, indent=1))
ok(f"engine compiled: {eng / EXE}")
return eng
# ------------------------------------------------------------------------------------------------ start
def installed_configs():
return sorted(ROOT.glob("strata-*.json"), key=lambda p: p.stat().st_mtime, reverse=True)
def start(cfg_path: Path, port: int | None, open_browser=True) -> int:
cfg = json.loads(cfg_path.read_text(encoding="utf-8"))
missing = [p for p in [cfg["exe"], *[a for a in cfg["args"] if a.endswith(".gguf")]] if not Path(p).exists()]
if missing:
fail(f"{cfg_path.name} refers to missing files: {missing[0]}", "run it again with --setup to repair")
cfg_path.touch() # the most recently used model
cmd = [sys.executable, str(ROOT / "serve" / "server.py"), "--engine", "strata", "--config", str(cfg_path),
"--port", str(port or cfg.get("port", 8080))]
if open_browser:
cmd.append("--open")
say()
say(f"Starting {cfg.get('model_name', 'the model')} (loads 34-43 GB into RAM: 30-90 s). Close this window to stop it.")
return subprocess.call(cmd)
def write_run_script(model, cfg_path, port):
serve = [sys.executable, str(ROOT / "serve" / "server.py"), "--engine", "strata", "--config", str(cfg_path),
"--port", str(port), "--open"]
if WIN:
script = ROOT / f"run-{model.lower()}.bat"
script.write_text("@echo off\r\ntitle Strata " + model + "\r\ncd /d \"" + str(ROOT) + "\"\r\n" +
" ".join(f'"{x}"' for x in serve) + "\r\npause\r\n", encoding="utf-8")
else:
script = ROOT / f"run-{model.lower()}.sh"
script.write_text("#!/bin/sh\ncd \"" + str(ROOT) + "\"\nexec " + " ".join(f'"{x}"' for x in serve) + "\n",
encoding="utf-8")
script.chmod(0o755)
return script
# ------------------------------------------------------------------------------------------------ main
def main() -> int:
ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
ap.add_argument("--model", choices=list(MODELS))
ap.add_argument("--context", type=int)
ap.add_argument("--vision", choices=["yes", "no", "none", "gpu", "cpu"],
help="let the model read images (yes = the encoder on the GPU)")
ap.add_argument("--port", type=int, default=8080)
ap.add_argument("--models-dir", default=str(ROOT / "models"), help="where the GGUF files go (~70 GB)")
ap.add_argument("--gguf-dir", help="use GGUF files you already have (a folder with the two shards)")
ap.add_argument("--yes", action="store_true", help="accept the recommended answers")
ap.add_argument("--setup", action="store_true", help="install another model or change settings")
ap.add_argument("--no-start", action="store_true", help="install only, do not start the model")
ap.add_argument("--build", action="store_true", help="compile the engine instead of using the ready-made one")
ap.add_argument("--prebuilt", default=os.environ.get("STRATA_PREBUILT_URL", PREBUILT_URL),
help="where the ready-made engine is (a URL folder or a local folder)")
ap.add_argument("--check", action="store_true", help="only check this PC and exit")
ap.add_argument("--skip-build", action="store_true", help=argparse.SUPPRESS)
a = ap.parse_args()
say("Strata - Qwen3.8-Flash-Next on a normal PC (NVIDIA GPU + system RAM + CPU)")
# ---- 0. already installed: just start it
have = installed_configs()
if have and not (a.setup or a.model or a.check or a.no_start):
if len(have) == 1:
return start(have[0], None)
say()
for i, c in enumerate(have, 1):
say(f" {i}) {c.stem[len('strata-'):].upper()}")
say(f" {len(have) + 1}) install another model / change settings")
pick = int(ask("Which one?", [str(i) for i in range(1, len(have) + 2)], "1", a.yes))
if pick <= len(have):
return start(have[pick - 1], None)
# ---- 1. the PC
step(1, "checking your PC")
gpu = gpu_info()
if gpu is None:
fail("no NVIDIA GPU found (nvidia-smi did not answer)",
"install the NVIDIA driver from https://www.nvidia.com/drivers and restart the PC")
ok(f"GPU: {gpu['name']}, {gpu['vram_gb']:.1f} GB VRAM, compute capability {gpu['arch'][:-1]}.{gpu['arch'][-1]}, "
f"driver {gpu['driver']}")
if int(gpu["arch"]) < 80:
fail("this GPU is older than the RTX 30 series (compute capability 8.0 is required)")
if driver_major(gpu) < MIN_DRIVER:
fail(f"the NVIDIA driver is too old ({gpu['driver']}; {MIN_DRIVER} or newer is needed)",
"update it with the NVIDIA App or from https://www.nvidia.com/drivers, restart, and run this again")
if gpu["vram_gb"] < 11:
warn("less than 12 GB of VRAM: Strata will run, but most experts stay on the CPU and it will be slow")
ram = ram_gb()
cpu, avx2, avx512 = cpu_info()
ok(f"RAM: {ram:.0f} GB")
ok(f"CPU: {cpu} ({'AVX-512' if avx512 else 'AVX2' if avx2 else 'no AVX2'})")
if not avx2:
fail("this CPU has no AVX2; Strata needs at least AVX2")
if a.check:
say()
for m, d in MODELS.items():
verdict = "fits" if ram >= d["ram_gb"] else "tight" if ram >= d["ram_gb"] - 8 else "does not fit"
say(f" {m:8s} needs ~{d['ram_gb']} GB RAM: {verdict}")
say("\nThis PC can run Strata. Run it again without --check to install.")
return 0
# ---- 2. the questions
step(2, "your choices")
names = list(MODELS)
for i, m in enumerate(names, 1):
d = MODELS[m]
fit = "" if ram >= d["ram_gb"] else f" <- needs {d['ram_gb']} GB RAM, you have {ram:.0f}"
say(f" {i}) {m:8s} {d['about']}; download {d['download_gb']:.0f} GB, uses ~{d['arena_gb']:.0f} GB of RAM{fit}")
rec = "3" if ram >= 60 else "1"
model = a.model or names[int(ask("Which model?", ["1", "2", "3"], rec, a.yes)) - 1]
if ram < MODELS[model]["ram_gb"] - 4:
fail(f"{model} needs about {MODELS[model]['ram_gb']} GB of RAM; this PC has {ram:.0f} GB",
"choose Q2_0 or IQ2_XS, or add RAM")
ok(f"model: {model}")
rec_ctx = 32768 if gpu["vram_gb"] < 14 else 65536 if gpu["vram_gb"] < 20 else 131072
if a.context:
ctx = a.context
else:
say()
say(" Context length = how much text the model can see at once (your chat, files, tool output).")
say(" Longer needs more VRAM for it, so fewer experts fit on the GPU:")
for i, c in enumerate(CONTEXTS, 1):
say(f" {i}) {c // 1024}K tokens" + (" (recommended for your GPU)" if c == rec_ctx else ""))
ctx = CONTEXTS[int(ask("Context?", [str(i) for i in range(1, 6)], str(CONTEXTS.index(rec_ctx) + 1), a.yes)) - 1]
if model == "IQ3_XXS" and ram < 90 and ctx > 131072:
warn("IQ3_XXS with a 262K context needs more than 64 GB of RAM (43 GB of experts + the context): using 128K")
ctx = 131072
ok(f"context: {ctx} tokens")
if a.vision:
vision = {"yes": "gpu", "no": "none"}.get(a.vision, a.vision)
else:
say()
say(" Images: the model can also read pictures (screenshots, photos, scanned pages). This adds a 0.9 GB")
say(" download and keeps ~1.4 GB of VRAM free for the image encoder, so text is a few % slower.")
vision = "gpu" if ask("Do you want images?", ["y", "n"], "n", a.yes) == "y" else "none"
ok("images: " + {"none": "off", "gpu": "on", "cpu": "on (encoder on the CPU)"}[vision])
models_dir = Path(a.gguf_dir) if a.gguf_dir else Path(a.models_dir) / model
shards = [models_dir / f"Qwen3.8-Flash-Next-GSQ-RCO-{model}-0000{i}-of-00002.gguf" for i in (1, 2)]
have_model = all(s.exists() and (done(s) or a.gguf_dir) for s in shards)
need = (0 if a.gguf_dir or have_model else MODELS[model]["download_gb"]) + 8 + \
(40 if model == "Q2_0" and avx512 else 0) + (1 if vision != "none" else 0)
if free_gb(models_dir) < need:
fail(f"not enough free disk space in {models_dir}: need ~{need:.0f} GB", "use --models-dir on a bigger drive")
# ---- 3. python packages
step(3, "Python packages")
pip_install(PY_PACKAGES, "numpy, jinja2, regex, pyyaml, tqdm, requests, cmake, ninja")
# ---- 4. the engine
step(4, "the Strata engine")
llama = get_llama_cpp()
ok(f"llama.cpp {LLAMA_CPP_COMMIT[:7]} (gguf-py, ggml, mtmd)")
eng = None if a.build else get_prebuilt(a.prebuilt, gpu, vision)
if eng is not None and json.loads((eng / "BUILD.json").read_text()).get("source") != "local":
pip_install(CUDA_WHEELS, "NVIDIA CUDA libraries (cuBLAS, CUDA runtime; ~0.4 GB)")
if vision != "none" and not (eng / VEXE).exists():
warn("the ready-made engine has no image encoder: compiling it")
eng = None
if eng is None:
eng = build_engine(gpu, vision, a.yes, llama)
meta = json.loads((eng / "BUILD.json").read_text())
lib_dirs = meta.get("cuda_dirs") or cuda_lib_dirs()
ok(f"engine: {eng / EXE}")
# ---- 5. the model files
step(5, f"downloading {model}")
if not a.gguf_dir:
for s in shards:
if s.exists() and done(s):
ok(f"{s.name} already downloaded")
continue
# shard 2 is the same file for all three models: reuse one that is already here
other = [p for p in Path(a.models_dir).glob("*/Qwen3.8-Flash-Next-GSQ-RCO-*-00002-of-00002.gguf") if done(p)]
if s.name.endswith("00002-of-00002.gguf") and other and not s.exists():
try:
os.link(other[0], s)
mark(s)
ok(f"{s.name} shared with {other[0].parent.name} (identical file)")
continue
except OSError:
pass
download(HF + f"{model}/{s.name}", s)
for s in shards:
if not s.exists():
fail(f"missing {s}")
ok("model files present")
mmproj = Path(a.models_dir) / MMPROJ
if vision != "none":
if not mmproj.exists() and a.gguf_dir and (Path(a.gguf_dir) / MMPROJ).exists():
mmproj = Path(a.gguf_dir) / MMPROJ
else:
download(HF + MMPROJ, mmproj, "vision encoder")
ok(f"vision encoder: {mmproj}")
# ---- 6. the pack and the MTP draft layer
step(6, "preparing the model for Strata")
pack = ROOT / "packs" / model.lower()
env = dict(os.environ, STRATA_GGUF_PY=str(llama / "gguf-py"))
if model == "Q2_0" and avx512:
# the Q2_0 experts repacked for the AVX-512 kernel (the measured speed): a one-time ~40 GB conversion
if not (pack / "index.txt").exists() or not (pack / "experts.bin").exists(): # index.txt is written last
say(" Converting the Q2_0 experts for the AVX-512 kernel (one time, ~40 GB written, 2-5 min) ...")
run([sys.executable, str(ROOT / "tools" / "strata_pack.py"), "build", "--gguf", str(shards[0]),
"--out", str(pack), "--skip-hash"], env=env)
run([sys.executable, str(ROOT / "tools" / "pack_index.py"), "--pack", str(pack)], env=env)
if not (pack / "tokenizer" / "vocab.json").exists():
run([sys.executable, str(ROOT / "tools" / "strata_tokenizer.py"), "--gguf", str(shards[0]),
"--out", str(pack)], env=env) # writes <pack>/tokenizer/
elif not (pack / "native_experts.txt").exists() or not (pack / "tokenizer" / "vocab.json").exists():
# every tensor as the GGUF stores it; the experts are read from the GGUF at start (seconds to build)
run([sys.executable, str(ROOT / "tools" / "iq_pack.py"), "--gguf", str(shards[0]), "--out", str(pack)], env=env)
ok(f"model prepared: {pack}")
mtp = ROOT / "mtp"
rt = mtp / "rt"
if not (rt / "experts.bin").exists():
say(" The MTP draft layer (speculative decoding, ~2x faster output) comes from the original Qwen checkpoint:")
say(" only its ~5 GB of MTP tensors are downloaded.")
run([sys.executable, str(ROOT / "tools" / "mtp_fetch.py"), "fetch", "--out", str(mtp)], env=env)
run([sys.executable, str(ROOT / "tools" / "mtp_pack.py"), "--src", str(mtp), "--experts", "q2_0",
"--out", str(mtp / "mtp-q2_0.gguf")], env=env)
run([sys.executable, str(ROOT / "tools" / "mtp_rt.py"), "--gguf", str(mtp / "mtp-q2_0.gguf"), "--out", str(rt)],
env=env)
if not (rt / "draft_vocab.bin").exists():
shutil.copyfile(ROOT / "data" / "draft_vocab.bin", rt / "draft_vocab.bin")
ok(f"MTP draft layer: {rt}")
# ---- 7. the start script
step(7, "writing the start script")
args = ["--pack", str(pack), "--native", str(shards[0]), "--ple-gguf", str(shards[1]),
"--expert-profile", str(ROOT / "data" / "expert-profile.bin"), "--expert-cache", "auto",
"--prefill", "2048", "--spec", "4", "--spec-min-p", "0.5", "--mtp", str(rt),
"--max-context", str(ctx)]
if ctx > 8192:
args += ["--kv", "int8"]
if vision != "none":
args += ["--vision", "--vram-reserve-mib", str(VISION[vision]["reserve_mib"])]
cfg = {"exe": str(eng / EXE), "args": args, "cwd": str(ROOT), "tokenizer": str(pack / "tokenizer"),
"model_name": f"qwen3.8-flash-next-{model.lower()}", "log": str(ROOT / f"strata-{model.lower()}.log"),
"lib_dirs": lib_dirs, "port": a.port}
if vision != "none":
cfg["vision"] = {"exe": str(eng / VEXE), "mmproj": str(mmproj), "model": str(shards[0]),
"gpu": vision == "gpu", "max_tokens": VISION[vision]["max_tokens"]}
if vision == "cpu":
cfg["vision"]["threads"] = max(1, (os.cpu_count() or 8) // 2)
cfg_path = ROOT / f"strata-{model.lower()}.json"
cfg_path.write_text(json.dumps(cfg, indent=1), encoding="utf-8")
script = write_run_script(model, cfg_path, a.port)
ok(f"start script: {script.name}")
say()
say("All set.")
say(f" API (OpenAI): http://127.0.0.1:{a.port}/v1 (any API key; model name: anything)")
say(f" API (Anthropic): http://127.0.0.1:{a.port}/v1/messages")
say(f" Next time: just run {'START-HERE.bat' if WIN else './setup.sh'} (or {script.name}) - it starts right away")
if vision != "none":
say(" Images: send them in the chat page, in chat.py (/image <path>) or over the API")
if a.no_start:
return 0
return start(cfg_path, a.port)
if __name__ == "__main__":
try:
sys.exit(main())
except KeyboardInterrupt:
say("\nstopped.")
sys.exit(1)
setup.sh
+30
View File
@@ -0,0 +1,30 @@
#!/bin/sh
# Strata for Linux: the first run installs everything and starts the model; later runs just start it.
# Needs only an NVIDIA driver. Python (with venv) is installed through apt/dnf if it is missing (asks for sudo).
cd "$(dirname "$0")" || exit 1
if [ ! -x .venv/bin/python ]; then
PY=""
for c in python3 python; do
if command -v $c >/dev/null 2>&1 && $c -c 'import sys, venv; sys.exit(0 if sys.version_info >= (3, 10) else 1)' 2>/dev/null; then
PY=$c; break
fi
done
if [ -z "$PY" ]; then
echo "Python 3.10+ with venv is needed; installing it (sudo will ask for your password) ..."
if command -v apt-get >/dev/null 2>&1; then
sudo apt-get update && sudo apt-get install -y python3 python3-venv python3-pip
elif command -v dnf >/dev/null 2>&1; then
sudo dnf install -y python3 python3-pip
elif command -v pacman >/dev/null 2>&1; then
sudo pacman -S --noconfirm python python-pip
fi
PY=python3
if ! $PY -c 'import sys, venv; sys.exit(0 if sys.version_info >= (3, 10) else 1)' 2>/dev/null; then
echo "Please install Python 3.10 or newer (with venv), then run ./setup.sh again."
exit 1
fi
fi
# a private environment inside this folder (system Python stays untouched; newer distros refuse global pip)
$PY -m venv .venv || { echo "could not create .venv: sudo apt install python3-venv"; exit 1; }
fi
exec .venv/bin/python setup.py "$@"
src/artifact/dequant.cpp
+123
View File
@@ -0,0 +1,123 @@
// src/artifact/dequant.cpp - the `strata-dequant` CLI (dequantize a GGUF tensor and dump F32).
// The dequantizers themselves are header-only in include/strata/artifact/dequant.hpp.
#include "strata/artifact/dequant.hpp"
#ifdef STRATA_DEQUANT_SELFTEST
int main(int argc, char** argv) {
if (argc < 2) {
std::printf("usage: dequant <file.gguf>\n");
return 2;
}
try {
strata::GgufFile g(argv[1]);
int checked = 0, bad = 0;
for (const auto& t : g.tensors()) {
if (t.type != 42) continue; // Q2_0 only for now
if (checked >= 8) break;
const uint8_t* base = g.tensor_data(t);
const uint64_t nblocks = t.elements() / 64;
for (uint64_t b = 0; b < nblocks && b < 64; ++b) {
const uint8_t* blk = base + b * 18;
float out[64];
strata::dequantize_q2_0(blk, out);
const float d = strata::fp16_to_fp32(strata::read_u16(blk));
// STRUCTURAL INVARIANT: every value must be exactly one of {-d, 0, d, 2d}.
for (int j = 0; j < 64; ++j) {
const float v = out[j];
const bool ok = (v == -d) || (v == 0.0f) || (v == d) || (v == 2.0f * d) || (d == 0.0f);
if (!ok) {
++bad;
if (bad < 4)
std::printf(" VIOLATION %s block %llu elem %d: v=%g d=%g\n", t.name.c_str(),
(unsigned long long)b, j, v, d);
}
}
}
if (checked == 0) {
// A folded checksum over the FIRST 64 blocks, for cross-checking against the Python
// implementation on identical bytes. A single implementation being self-consistent
// proves little - two disagreeing is what finds bugs (round 67).
double sum = 0, sumabs = 0;
const uint8_t* base = g.tensor_data(t);
for (uint64_t b = 0; b < nblocks && b < 64; ++b) {
float out[64];
strata::dequantize_q2_0(base + b * 18, out);
for (int j = 0; j < 64; ++j) {
sum += out[j];
sumabs += std::fabs(out[j]);
}
}
std::printf(" CHECKSUM %s sum=%.10g sumabs=%.10g\n", t.name.c_str(), sum, sumabs);
}
++checked;
}
// Q5_0 cross-check on the same 64-block window gguf-py is given, using the first Q5_0 tensor.
for (const auto& t : g.tensors()) {
if (t.type != 6 && t.type != 20 && t.type != 14 && t.type != 12 && t.type != 13 && t.type != 11 &&
t.type != 23)
continue;
const uint8_t* base = g.tensor_data(t);
const int bsz =
(t.type == 20)
? 18
: (t.type == 14
? 210
: (t.type == 12
? 144
: (t.type == 13 ? 176 : (t.type == 11 ? 110 : (t.type == 23 ? 136 : 22)))));
const int nel =
(t.type == 14 || t.type == 12 || t.type == 13 || t.type == 11 || t.type == 23) ? 256 : 32;
const uint64_t nb = t.elements() / (uint64_t)nel;
double sum = 0, sumabs = 0;
for (uint64_t b = 0; b < nb && b < 64; ++b) {
float out[256];
if (t.type == 20)
strata::dequantize_iq4_nl(base + b * bsz, out);
else if (t.type == 14)
strata::dequantize_q6_K(base + b * bsz, out);
else if (t.type == 12)
strata::dequantize_q4_K(base + b * bsz, out);
else if (t.type == 13)
strata::dequantize_q5_K(base + b * bsz, out);
else if (t.type == 11)
strata::dequantize_q3_K(base + b * bsz, out);
else if (t.type == 23)
strata::dequantize_iq4_xs(base + b * bsz, out);
else
strata::dequantize_q5_0(base + b * bsz, out);
for (int j = 0; j < nel; ++j) {
sum += out[j];
sumabs += std::fabs(out[j]);
}
}
std::printf(" CHECKSUM %s %s sum=%.10g sumabs=%.10g\n", t.type_name(), t.name.c_str(), sum,
sumabs);
float f8[256];
if (t.type == 20)
strata::dequantize_iq4_nl(base, f8);
else if (t.type == 14)
strata::dequantize_q6_K(base, f8);
else if (t.type == 12)
strata::dequantize_q4_K(base, f8);
else if (t.type == 13)
strata::dequantize_q5_K(base, f8);
else if (t.type == 11)
strata::dequantize_q3_K(base, f8);
else if (t.type == 23)
strata::dequantize_iq4_xs(base, f8);
else
strata::dequantize_q5_0(base, f8);
std::printf(" REF first8:");
for (int j = 0; j < 8; ++j) std::printf(" %.6g", f8[j]);
std::printf("\n");
if (t.type == 23) break; // emit Q5_0 first, then IQ4_NL
}
std::printf("%s\n Q2_0 tensors spot-checked: %d\n symbolic invariant violations: %d -> %s\n",
argv[1], checked, bad, bad ? "FAIL" : "PASS");
return bad ? 1 : 0;
} catch (const std::exception& e) {
std::printf("ERROR: %s\n", e.what());
return 1;
}
}
#endif
src/artifact/gguf_reader.cpp
+67
View File
@@ -0,0 +1,67 @@
// src/artifact/gguf_reader.cpp - the `strata-gguf` CLI. The reader itself is header-only in
// include/strata/artifact/gguf_reader.hpp; this file is just the entry point, so it cannot drift
// from the library it exercises.
#include "strata/artifact/gguf_reader.hpp"
#ifndef STRATA_GGUF_MAIN_DISABLED
int main(int argc, char** argv) {
if (argc < 2) {
std::printf("usage: gguf_reader <file.gguf> [--check]\n");
return 2;
}
const bool check = (argc > 2 && std::string(argv[2]) == "--check");
try {
strata::GgufFile g(argv[1]);
std::printf("%s\n", argv[1]);
std::printf(" version %u tensors %zu metadata %zu data_start %llu size %llu\n", g.version(),
g.tensors().size(), g.metadata().size(), (unsigned long long)g.data_start(),
(unsigned long long)g.file_size());
if (check) {
// Independent expectations, all established by earlier phases.
size_t n_blk = 0;
for (const auto& t : g.tensors())
if (t.name.rfind("blk.", 0) == 0) ++n_blk;
std::printf(" blk.* tensors %zu globals %zu\n", n_blk, g.tensors().size() - n_blk);
std::printf(" Q2_0 %llu BF16 %llu F32 %llu\n", (unsigned long long)g.count_type("Q2_0"),
(unsigned long long)g.count_type("BF16"), (unsigned long long)g.count_type("F32"));
// Every tensor's byte size must land inside the gap to the next tensor's offset, since
// GGUF aligns every tensor. That is the same bracket test tools/verify_q2_0_geometry.py
// uses, and it validates block_geometry() against the file rather than asserting it.
std::vector<const strata::TensorInfo*> ordered;
for (const auto& t : g.tensors()) ordered.push_back(&t);
std::sort(ordered.begin(), ordered.end(), [](auto a, auto b) { return a->offset < b->offset; });
size_t bad = 0, unknown = 0;
for (size_t i = 0; i < ordered.size(); ++i) {
const auto* t = ordered[i];
int el = 0, by = 0;
if (!strata::block_geometry(t->type, el, by)) {
++unknown;
continue;
}
if (t->elements() % (uint64_t)el) {
++bad;
continue;
}
const uint64_t nb = t->elements() / (uint64_t)el * (uint64_t)by;
// The bracket is on the GAP to the next tensor's offset, not on that offset itself -
// offsets are relative to data_start, so the gap is a difference. Comparing against
// the raw offset flagged almost every tensor (213/214) while the same test in Python
// (tools/verify_q2_0_geometry.py) passed.
const uint64_t end = (i + 1 < ordered.size()) ? ordered[i + 1]->offset - t->offset
: g.file_size() - g.data_start() - t->offset;
if (!(nb <= end && end < nb + 32)) ++bad;
}
std::printf(" geometry bracket: %zu out of range, %zu types unknown\n", bad, unknown);
const std::string err = strata::check_architecture(g);
std::printf(" architecture guard: %s\n", err.empty() ? "PASS" : ("FAIL - " + err).c_str());
if (bad) return 1;
}
return 0;
} catch (const std::exception& e) {
std::printf("ERROR: %s\n", e.what());
return 1;
}
}
#endif // STRATA_GGUF_MAIN_DISABLED
src/core/concurrent_main.cpp
+197
View File
@@ -0,0 +1,197 @@
// src/core/concurrent_main.cpp - do the CPU expert path and the GPU stream path actually run at once?
//
// L9 concluded that the two miss paths are INDEPENDENT RESOURCES - 42.55 GB/s of DRAM bandwidth for the CPU
// expert kernel and 28.9 GB/s of PCIe for the stream - and that running them together would cover a token's
// 552.7 MB of misses in 7.74 ms, which is what puts 100 tps inside reach on this Gen4 machine.
//
// THAT CLAIM HAS A HOLE. An H2D copy from PINNED host memory READS HOST DRAM to do it, so the bus path is not
// free of the resource the CPU path is bound by. If the two contend, the combined rate is not 42.55 + 28.9 and
// the 129 tok/s figure is optimistic; if they overlap cleanly, it stands. This measures it three ways in one
// process - CPU alone, GPU alone, both - because comparing two separately-measured rates is exactly the
// inference that has been wrong twice already in this ledger (L2's dependency chain, L8's bus-only ceiling).
//
// The CPU side is a READ LOOP over the same pinned arena rather than the real VNNI kernel, and that is
// deliberate: the question is whether the two paths contend for DRAM BANDWIDTH, and a read loop is a pure
// measurement of that demand, with no compute that could mask it.
#include "strata/core/pinned.hpp"
#include "strata/kernels/s_gemv.hpp"
#include <cuda_runtime.h>
#include <atomic>
#include <chrono>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <fstream>
#include <string>
#include <thread>
#include <vector>
namespace {
void check(cudaError_t e, const char* what) {
if (e != cudaSuccess) {
std::fprintf(stderr, "%s: %s\n", what, cudaGetErrorString(e));
std::exit(1);
}
}
// Volatile sink so the read loop cannot be optimised away; summed across threads and printed, so a compiler
// that removed the loads would produce a suspiciously round zero rather than a plausible number.
std::atomic<uint64_t> g_sink{0};
struct Run {
double seconds = 0;
double cpu_gb = 0; // GB/s read by the CPU threads
double gpu_experts = 0; // experts streamed and computed by the GPU
};
} // namespace
int main(int argc, char** argv) {
std::string path;
// 4096 planes = 2 GiB, which is 64x the 32 MB L3 on this CPU. At 0.1 GiB the read loop measured 206 GB/s -
// that is CACHE bandwidth, not DRAM, and it would have answered the contention question with a number about
// a different resource entirely.
int experts = 4096;
int cpu_threads = 6;
int gpu_roles = 256;
for (int i = 1; i < argc; ++i) {
const std::string a = argv[i];
auto val = [&]() -> const char* {
if (i + 1 >= argc) { std::fprintf(stderr, "%s needs a value\n", a.c_str()); std::exit(2); }
return argv[++i];
};
if (a == "--file") path = val();
else if (a == "--experts") experts = std::atoi(val());
else if (a == "--cpu-threads") cpu_threads = std::atoi(val());
else if (a == "--gpu-roles") gpu_roles = std::atoi(val());
else {
std::fprintf(stderr, "usage: strata-concurrent --file experts.bin [--experts N] "
"[--cpu-threads N] [--gpu-roles N]\n");
return 2;
}
}
if (path.empty()) { std::fprintf(stderr, "--file is required\n"); return 2; }
const long long n_in = 2560, n_out = 640;
const uint64_t role_codes = (uint64_t) n_out * n_in / 4;
const uint64_t plane_bytes = role_codes; // the CPU loop reads the same planes
strata::core::PinnedArena arena((uint64_t) experts * role_codes);
if (!arena.valid()) { std::fprintf(stderr, "arena allocation failed\n"); return 1; }
{
std::ifstream f(path, std::ios::binary);
if (!f) { std::fprintf(stderr, "cannot open %s\n", path.c_str()); return 1; }
for (int i = 0; i < experts; ++i) {
f.seekg((std::streamoff) i * 1382400);
f.read((char*) arena.data() + (uint64_t) i * role_codes, (std::streamsize) role_codes);
if ((uint64_t) f.gcount() != role_codes) { std::fprintf(stderr, "short read\n"); return 1; }
}
}
std::printf("arena %.2f GiB pinned; %d CPU threads; %d GPU roles\n",
(double) arena.capacity / (1024.0 * 1024 * 1024), cpu_threads, gpu_roles);
uint8_t* d_codes = nullptr;
check(cudaMalloc(&d_codes, role_codes), "cudaMalloc codes");
uint16_t* d_x = nullptr;
float *d_scales = nullptr, *d_y = nullptr;
check(cudaMalloc(&d_x, n_in * sizeof(uint16_t)), "cudaMalloc x");
check(cudaMalloc(&d_scales, (size_t) n_out * (n_in / 64) * sizeof(float)), "cudaMalloc scales");
check(cudaMalloc(&d_y, n_out * sizeof(float)), "cudaMalloc y");
std::vector<uint16_t> hx((size_t) n_in, 0x3C00);
std::vector<float> hs((size_t) n_out * (size_t) (n_in / 64), 0.001f);
check(cudaMemcpy(d_x, hx.data(), hx.size() * sizeof(uint16_t), cudaMemcpyHostToDevice), "copy x");
check(cudaMemcpy(d_scales, hs.data(), hs.size() * sizeof(float), cudaMemcpyHostToDevice), "copy scales");
cudaStream_t s_copy{}, s_comp{};
check(cudaStreamCreate(&s_copy), "stream copy");
check(cudaStreamCreate(&s_comp), "stream compute");
const strata::kernels::SForm form{2, -1, 64, strata::kernels::Codebook::Affine, false};
// The CPU loop: read the whole arena, Nthreads ways, and sum bytes so nothing is elided.
auto cpu_work = [&](std::atomic<bool>& stop, std::atomic<uint64_t>& bytes) {
uint64_t acc = 0, n = 0;
const uint8_t* p = arena.data();
while (!stop.load(std::memory_order_relaxed)) {
// touch one byte per 64 - a stride that defeats prefetch into a pure streaming read
for (uint64_t i = 0; i < plane_bytes * (uint64_t) experts; i += 64) acc += p[i];
n += plane_bytes * (uint64_t) experts;
}
bytes.fetch_add(n);
g_sink.fetch_add(acc);
};
// The GPU loop: stream each role plane on the copy stream while computing on the compute stream.
auto gpu_work = [&](int roles) {
cudaEvent_t ev[2];
for (int b = 0; b < 2; ++b) check(cudaEventCreateWithFlags(&ev[b], cudaEventDisableTiming), "event");
uint8_t* buf[2] = {d_codes, d_codes}; // one buffer is enough: the timing is what is measured
check(cudaMemcpyAsync(buf[0], arena.data(), role_codes, cudaMemcpyHostToDevice, s_copy), "prologue");
check(cudaEventRecord(ev[0], s_copy), "prologue event");
for (int i = 0; i < roles; ++i) {
const int cur = i & 1, nxt = (i + 1) & 1;
if (i + 1 < roles) {
check(cudaMemcpyAsync(buf[nxt], arena.data() + (uint64_t) ((i + 1) % experts) * role_codes,
role_codes, cudaMemcpyHostToDevice, s_copy),
"copy");
check(cudaEventRecord(ev[nxt], s_copy), "event");
}
check(cudaStreamWaitEvent(s_comp, ev[cur], 0), "wait");
strata::kernels::s_gemv_split_async(d_x, buf[cur], d_scales, nullptr, d_y, n_in, n_out, form, 32,
(void*) s_comp);
}
check(cudaStreamSynchronize(s_comp), "sync comp");
check(cudaStreamSynchronize(s_copy), "sync copy");
for (int b = 0; b < 2; ++b) cudaEventDestroy(ev[b]);
};
auto run = [&](bool do_cpu, bool do_gpu, const char* name) {
std::atomic<bool> stop{false};
std::atomic<uint64_t> bytes{0};
std::vector<std::thread> cpu;
check(cudaStreamSynchronize(0), "pre");
const auto t0 = std::chrono::steady_clock::now();
if (do_cpu) for (int t = 0; t < cpu_threads; ++t) cpu.emplace_back(cpu_work, std::ref(stop), std::ref(bytes));
if (do_gpu) {
gpu_work(gpu_roles);
} else {
// CPU-only needs a REAL window. The first version let t0..t1 span nothing when there was no GPU
// work, so the threads were stopped immediately and the run reported 0.000 s and 0.00 GB/s.
std::this_thread::sleep_for(std::chrono::milliseconds(3000));
}
const auto t1 = std::chrono::steady_clock::now();
stop = true;
for (auto& t : cpu) t.join();
Run r;
r.seconds = std::chrono::duration<double>(t1 - t0).count();
r.cpu_gb = (double) bytes.load() / 1e9 / r.seconds;
r.gpu_experts = do_gpu ? (double) gpu_roles / r.seconds : 0.0;
std::printf(" %-18s %7.3f s CPU read %6.2f GB/s GPU %7.1f roles/s\n", name, r.seconds, r.cpu_gb,
r.gpu_experts);
return r;
};
std::printf("\n mode wall CPU read GPU\n");
const Run cpu_only = run(true, false, "CPU only");
const Run gpu_only = run(false, true, "GPU only");
const Run both = run(true, true, "BOTH");
std::printf("\n independence check (1.00 = perfectly shared, 0.00 = fully independent):\n");
const double cpu_interference = cpu_only.cpu_gb > 0 ? 1.0 - both.cpu_gb / cpu_only.cpu_gb : 0.0;
const double gpu_interference = gpu_only.gpu_experts > 0 ? 1.0 - both.gpu_experts / gpu_only.gpu_experts : 0.0;
std::printf(" CPU throughput lost when the GPU also runs : %5.1f%%\n", cpu_interference * 100);
std::printf(" GPU throughput lost when the CPU also runs : %5.1f%%\n", gpu_interference * 100);
std::printf(" combined vs the sum of the parts : CPU %6.2f + GPU-equivalent vs %6.2f measured\n",
cpu_only.cpu_gb, both.cpu_gb);
std::printf("\n NOTE: the CPU side is a READ LOOP, not the VNNI kernel - it measures DRAM DEMAND, which is the\n"
" resource the question is about. L9's 129 tok/s assumes these two do not contend.\n");
cudaFree(d_codes);
cudaFree(d_x);
cudaFree(d_scales);
cudaFree(d_y);
cudaStreamDestroy(s_copy);
cudaStreamDestroy(s_comp);
return 0;
}
Loaded 100 of 219 files, more files were not shown because too many files have changed in this diff. Show more